Add conda support

Author: junwha

dockerfiles/conda/Dockerfile

@@ -15,5 +15,43 @@ RUN conda init && \
     conda install python=$PY_VERSION -y && \
     conda install -c "nvidia/label/cuda-$DOCKER_CUDA_VERSION" cuda -y
 
-# RUN conda init && \
-#     pip install torch==$TORCH_VERSION torchvision==$TORCH_VISION_VERSION torchaudio==$TORCH_VERSION --index-url https://download.pytorch.org/whl/cu$(cat /CUDA_VERSION | awk -F. '{print $1$2}')
+# torch installer
+COPY <<EOF /usr/local/bin/install_torch_env
+#!/bin/bash
+set -e
+TORCH_VER=\$1
+VISION_VER=\$2
+CUDA_VER=\$3
+FLASH_ATTN_LINK=\$4
+
+conda create -n torch\${TORCH_VER}
+
+export PATH_BCK=\$PATH
+export PATH=/root/anaconda/envs/torch\${TORCH_VER}:\$PATH
+
+pip install --no-cache-dir torch=="\${TORCH_VER}" torchvision=="\${VISION_VER}" torchaudio=="\${TORCH_VER}" --index-url "https://download.pytorch.org/whl/\${CUDA_VER}"
+pip install --no-cache-dir \$FLASH_ATTN_LINK
+
+"""conda activate torch\${TORCH_VER} && echo "export PATH=/root/anaconda/envs/torch\${TORCH_VER}:\\\$PATH && source activate torch\${TORCH_VER}" &>> ~/.bashrc""" > /usr/local/bin/init_torch\${TORCH_VER}
+chmod +x /usr/local/bin/init_torch\${TORCH_VER}
+
+export PATH=\$PATH_BCK
+
+EOF
+
+RUN chmod +x /usr/local/bin/install_torch_env
+
+# pytorch 2.4.1 (251209 Update)
+RUN install_torch_env 2.4.1 0.19.1 cu124 https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.3.12/flash_attn-2.8.0+cu124torch2.4-cp310-cp310-linux_x86_64.whl
+
+# pytorch 2.5.1 (251209 Update)
+RUN install_torch_env 2.5.1 0.20.1 cu124 https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.4.11/flash_attn-2.8.3+cu124torch2.5-cp310-cp310-linux_x86_64.whl
+
+# pytorch 2.6.0 (251209 Update)
+RUN install_torch_env 2.6.0 0.21.0 cu124 https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.4.11/flash_attn-2.8.3+cu124torch2.6-cp310-cp310-linux_x86_64.whl
+
+# pytorch 2.7.1 (251209 Update)
+RUN install_torch_env 2.7.1 0.22.1 cu126 https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.4.11/flash_attn-2.8.3+cu126torch2.7-cp310-cp310-linux_x86_64.whl
+
+# pytorch 2.9.0 (251209 Update)
+RUN install_torch_env 2.9.0 0.24.0 cu126 https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.4.17/flash_attn-2.6.3+cu126torch2.9-cp310-cp310-linux_x86_64.whl

dockerfiles/conda/build_conda.sh

@@ -2,5 +2,7 @@ PY_VERSION=$1
 DATE=251214 # $(date +"%y%m%d")
 TAG=junwha/ddiff-base:cu12.4.1-py${PY_VERSION}-conda-$DATE
 
-docker build -t $TAG --build-arg PY_VERSION=${PY_VERSION} ./ && \
-./test.sh $TAG && docker push $TAG
+docker build -t $TAG --build-arg PY_VERSION=${PY_VERSION} ./ 
+
+# && \
+# ./test.sh $TAG && docker push $TAG

dockerfiles/conda/deprecated/Dockerfile.torch

@@ -9,5 +9,5 @@ FROM junwha/dslice-base-py:${BASE_TAG}
 ARG TORCH_VERSION
 ARG TORCH_VISION_VERSION
 
-# RUN conda init && \
-#     pip install torch==$TORCH_VERSION torchvision==$TORCH_VISION_VERSION torchaudio==$TORCH_VERSION --index-url https://download.pytorch.org/whl/cu$(cat /CUDA_VERSION | awk -F. '{print $1$2}')
+RUN conda init && \
+    pip install torch==$TORCH_VERSION torchvision==$TORCH_VISION_VERSION torchaudio==$TORCH_VERSION --index-url https://download.pytorch.org/whl/cu$(cat /CUDA_VERSION | awk -F. '{print $1$2}')

dockerfiles/uv/Dockerfile

@@ -26,7 +26,7 @@ uv venv --python "\${PY_VERSION}" --system-site-packages --seed
 uv pip install --no-cache torch=="\${TORCH_VER}" torchvision=="\${VISION_VER}" torchaudio=="\${TORCH_VER}" --index-url "https://download.pytorch.org/whl/\${CUDA_VER}"
 uv pip install --no-cache \$FLASH_ATTN_LINK
 echo "ln -sfn \${BASE_DIR}/.venv ./.venv && [ -f pyproject.toml ] && uv add torch==\${TORCH_VER} torchvision==\${VISION_VER} torchaudio==\${TORCH_VER}; echo Created uv virtual environment with torch==\${TORCH_VER}" > /usr/local/bin/uv_init_torch\${TORCH_VER}
-chmod +x /usr/local/bin/uv_init_torch\${TORCH_VER}
+chmod +x /usr/local/bin/init_torch\${TORCH_VER}
 EOF
 RUN chmod +x /usr/local/bin/install_torch_env