wip

docker/Dockerfile.ascend

@@ -60,31 +60,21 @@ RUN pip uninstall ms-swift modelscope -y && pip --no-cache-dir install pip==23.*
 if [ "$INSTALL_MS_DEPS" = "True" ]; then \
     pip --no-cache-dir install omegaconf==2.0.6 && \
     pip install 'editdistance==0.8.1' && \
-    pip install --no-cache-dir 'cython<=0.29.36' versioneer 'numpy<2.0' -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
-    pip install --no-cache-dir -r /var/modelscope/framework.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
+    pip install --no-cache-dir 'cython<=0.29.36' versioneer 'numpy<2.0' && \
+    pip install --no-cache-dir -r /var/modelscope/framework.txt && \
     pip install --no-cache-dir -r /var/modelscope/audio.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
-    pip install --no-cache-dir -r /var/modelscope/cv.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
-    pip install --no-cache-dir -r /var/modelscope/multi-modal.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
-    pip install --no-cache-dir -r /var/modelscope/nlp.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
-    pip install --no-cache-dir -r /var/modelscope/science.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
-    pip install --no-cache-dir -r /var/modelscope/tests.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
+    pip install --no-cache-dir -r /var/modelscope/tests.txt && \
     pip install --no-cache-dir -r /var/modelscope/server.txt && \
     pip install --no-cache-dir https://modelscope.oss-cn-beijing.aliyuncs.com/packages/imageio_ffmpeg-0.4.9-py3-none-any.whl --no-dependencies --force && \
-    pip install adaseq pai-easycv && \
     pip install --no-cache-dir 'scipy<1.13.0' && \
-    pip install --no-cache-dir funtextprocessing typeguard==2.13.3 scikit-learn -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
-    pip install --no-cache-dir text2sql_lgesql==1.3.0 git+https://github.com/jin-s13/xtcocoapi.git@v1.14 git+https://github.com/gatagat/lap.git@v0.4.0 -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html --force --no-deps && \
-    pip install --no-cache-dir mmcls>=0.21.0 mmdet>=2.25.0 decord>=0.6.0 mpi4py paint_ldm ipykernel fasttext -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
-    pip uninstall ddpm_guided_diffusion -y && \
+    pip install --no-cache-dir funtextprocessing typeguard==2.13.3 scikit-learn && \
+    pip install --no-cache-dir decord>=0.6.0 mpi4py paint_ldm ipykernel fasttext && \
     pip install --no-cache-dir 'blobfile>=1.0.5' && \
-    pip install 'ddpm_guided_diffusion' -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html --no-index && \
-    pip uninstall shotdetect_scenedetect_lgss -y && \
-    pip install 'shotdetect_scenedetect_lgss' -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html --no-index && \
     pip uninstall MinDAEC -y && \
     pip install https://modelscope.oss-cn-beijing.aliyuncs.com/releases/dependencies/MinDAEC-0.0.2-py3-none-any.whl && \
     pip cache purge; \
 else \
-    pip install --no-cache-dir -r /var/modelscope/framework.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
+    pip install --no-cache-dir -r /var/modelscope/framework.txt && \
     pip cache purge; \
 fi
 
@@ -104,11 +94,6 @@ RUN curl -fsSL https://ollama.com/install.sh | sh && \
     pip install --no-cache-dir transformers diffusers timm>=0.9.0 && pip cache purge; \
     pip install --no-cache-dir omegaconf==2.3.0 && pip cache purge;
 
-
-RUN if [ "$INSTALL_MS_DEPS" = "True" ]; then \
-    pip install --no-cache-dir huggingface-hub transformers peft -U; \
-fi;
-
 ENV SETUPTOOLS_USE_DISTUTILS=stdlib
 ENV VLLM_USE_MODELSCOPE=True
 ENV LMDEPLOY_USE_MODELSCOPE=True

docker/Dockerfile.ubuntu

@@ -25,31 +25,21 @@ RUN pip uninstall ms-swift modelscope -y && pip --no-cache-dir install pip==23.*
 if [ "$INSTALL_MS_DEPS" = "True" ]; then \
     pip --no-cache-dir install omegaconf==2.0.6 && \
     pip install 'editdistance==0.8.1' && \
-    pip install --no-cache-dir 'cython<=0.29.36' versioneer 'numpy<2.0' -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
-    pip install --no-cache-dir -r /var/modelscope/framework.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
+    pip install --no-cache-dir 'cython<=0.29.36' versioneer 'numpy<2.0' && \
+    pip install --no-cache-dir -r /var/modelscope/framework.txt && \
     pip install --no-cache-dir -r /var/modelscope/audio.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
-    pip install --no-cache-dir -r /var/modelscope/cv.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
-    pip install --no-cache-dir -r /var/modelscope/multi-modal.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
-    pip install --no-cache-dir -r /var/modelscope/nlp.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
-    pip install --no-cache-dir -r /var/modelscope/science.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
-    pip install --no-cache-dir -r /var/modelscope/tests.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
+    pip install --no-cache-dir -r /var/modelscope/tests.txt && \
     pip install --no-cache-dir -r /var/modelscope/server.txt && \
     pip install --no-cache-dir https://modelscope.oss-cn-beijing.aliyuncs.com/packages/imageio_ffmpeg-0.4.9-py3-none-any.whl --no-dependencies --force && \
-    pip install adaseq pai-easycv && \
     pip install --no-cache-dir 'scipy<1.13.0' && \
-    pip install --no-cache-dir funtextprocessing typeguard==2.13.3 scikit-learn -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
-    pip install --no-cache-dir text2sql_lgesql==1.3.0 git+https://github.com/jin-s13/xtcocoapi.git@v1.14 git+https://github.com/gatagat/lap.git@v0.4.0 -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html --force --no-deps && \
-    pip install --no-cache-dir mmcls>=0.21.0 mmdet>=2.25.0 decord>=0.6.0 mpi4py paint_ldm ipykernel fasttext -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
-    pip uninstall ddpm_guided_diffusion -y && \
+    pip install --no-cache-dir funtextprocessing typeguard==2.13.3 scikit-learn && \
+    pip install --no-cache-dir decord>=0.6.0 mpi4py paint_ldm ipykernel fasttext && \
     pip install --no-cache-dir 'blobfile>=1.0.5' && \
-    pip install 'ddpm_guided_diffusion' -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html --no-index && \
-    pip uninstall shotdetect_scenedetect_lgss -y && \
-    pip install 'shotdetect_scenedetect_lgss' -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html --no-index && \
     pip uninstall MinDAEC -y && \
     pip install https://modelscope.oss-cn-beijing.aliyuncs.com/releases/dependencies/MinDAEC-0.0.2-py3-none-any.whl && \
     pip cache purge; \
 else \
-    pip install --no-cache-dir -r /var/modelscope/framework.txt -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html && \
+    pip install --no-cache-dir -r /var/modelscope/framework.txt && \
     pip cache purge; \
 fi
 
@@ -57,7 +47,6 @@ ARG CUR_TIME={cur_time}
 RUN echo $CUR_TIME
 
 RUN bash /tmp/install.sh {version_args} && \
-    curl -fsSL https://ollama.com/install.sh | sh && \
     pip install --no-cache-dir -U funasr scikit-learn && \
     pip install --no-cache-dir -U qwen_vl_utils qwen_omni_utils librosa timm transformers accelerate peft trl safetensors && \
     cd /tmp && GIT_LFS_SKIP_SMUDGE=1 git clone -b {swift_branch}  --single-branch https://github.com/modelscope/ms-swift.git && \
@@ -75,21 +64,21 @@ RUN bash /tmp/install.sh {version_args} && \
     cp /tmp/resources/ubuntu2204.aliyun /etc/apt/sources.list
 
 
-RUN if [ "$IMAGE_TYPE" = "swift" ]; then \
-    pip install "sglang[all]<0.5" "math_verify==0.5.2" "gradio<5.33" -U && \
-    pip install liger_kernel wandb swanlab nvitop pre-commit "transformers<4.57" "trl<0.21" huggingface-hub -U && \
-    SITE_PACKAGES=$(python -c "import site; print(site.getsitepackages()[0])") && echo $SITE_PACKAGES && \
-    CUDNN_PATH=$SITE_PACKAGES/nvidia/cudnn CPLUS_INCLUDE_PATH=$SITE_PACKAGES/nvidia/cudnn/include \
+RUN if [ "$IMAGE_TYPE" = "gpu" ]; then \
+    pip install math_verify "gradio<5.33" "deepspeed<0.18" ray -U && \
+    pip install liger_kernel wandb swanlab nvitop pre-commit "transformers" "trl<0.25" "peft<0.18" huggingface-hub -U && \
     pip install --no-build-isolation transformer_engine[pytorch]; \
     cd /tmp && GIT_LFS_SKIP_SMUDGE=1 git clone https://github.com/NVIDIA/apex && \
-    cd apex && git checkout e13873debc4699d39c6861074b9a3b2a02327f92 && pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./ && \
+    cd apex && pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./ && \
     cd / && rm -fr /tmp/apex && pip cache purge; \
-    pip install git+https://github.com/NVIDIA/Megatron-LM.git@core_r0.13.0; \
-elif [ "$IMAGE_TYPE" = "llm" ]; then \
-    pip install --no-cache-dir huggingface-hub transformers peft diffusers -U; \
+    pip install git+https://github.com/NVIDIA/Megatron-LM.git@core_r0.15.0; \
     pip uninstall autoawq -y; \
+    mkdir -p /root/.cache/modelscope/_github; \
+    git -C /root/.cache/modelscope/_github clone https://github.com/NVIDIA/Megatron-LM.git Megatron-LM --branch core_r0.15.0; \
+elif [ "$IMAGE_TYPE" = "cpu" ]; then \
+    pip install --no-cache-dir huggingface-hub transformers peft diffusers -U; \
 else \
-    pip install "transformers<4.56" "tokenizers<0.22" "trl<0.23" "diffusers<0.35" --no-dependencies; \
+    pip install "transformers<5.0" "tokenizers<0.22" "trl<0.23" "diffusers<0.35" --no-dependencies; \
 fi
 
 # install nvm and set node version to 18
@@ -103,5 +92,7 @@ RUN rm -f /etc/apt/sources.list.d/cuda-*.list && apt-get update
 
 ENV VLLM_USE_MODELSCOPE=True
 ENV LMDEPLOY_USE_MODELSCOPE=True
-ENV MODELSCOPE_CACHE=/mnt/workspace/.cache/modelscope/hub
+ENV MODELSCOPE_CACHE=/mnt/workspace/.cache/modelscope
+ENV MEGATRON_LM_PATH=/root/.cache/modelscope/_github/Megatron-LM
+ENV SWIFT_PATCH_CONV3D=1
 SHELL ["/bin/bash", "-c"]

docker/Dockerfile.ubuntu.old

@@ -0,0 +1,5 @@
+FROM {base_image}
+
+RUN cd modelscope && pip install . -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html
+
+SHELL ["/bin/bash", "-c"]

docker/Dockerfile.ubuntu_base

@@ -151,44 +151,6 @@ RUN if [ "$USE_GPU" = "True" ] ; then \
     fi
 
 
-# default shell bash
-ENV SHELL=/bin/bash
-# install special package
-RUN if [ "$USE_GPU" = "True" ] ; then \
-        pip install --no-cache-dir dgl -f https://data.dgl.ai/wheels/$CUDATOOLKIT_VERSION/repo.html; \
-    else \
-        pip install --no-cache-dir dgl dglgo -f https://data.dgl.ai/wheels/repo.html; \
-    fi
-
-# copy install scripts
-COPY docker/scripts/install_unifold.sh docker/scripts/install_colmap.sh docker/scripts/install_pytorch3d_nvdiffrast.sh docker/scripts/install_tiny_cuda_nn.sh docker/scripts/install_apex.sh /tmp/
-
-# 3d supports
-RUN if [ "$USE_GPU" = "True" ] ; then \
-        bash /tmp/install_colmap.sh; \
-    else \
-     echo 'cpu unsupport colmap'; \
-    fi
-# install pytorch3d
-RUN if [ "$USE_GPU" = "True" ] ; then \
-        bash /tmp/install_pytorch3d_nvdiffrast.sh; \
-    else \
-     echo 'cpu unsupport pytorch3d nvdiffrast'; \
-    fi
-
-# for uniford
-RUN if [ "$USE_GPU" = "True" ] ; then \
-        bash /tmp/install_unifold.sh; \
-    else \
-     echo 'cpu unsupport uniford'; \
-    fi
-
-RUN if [ "$USE_GPU" = "True" ] ; then \
-       export TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.9;9.0;8.6+PTX" && pip install --no-cache-dir git+https://github.com/gxd1994/Pointnet2.PyTorch.git@master#subdirectory=pointnet2; \
-    else \
-     echo 'CPU env does not support Pointnet2'; \
-    fi
-
 # NVIDIA apex building
 RUN if [ "$USE_GPU" = "True" ] ; then \
         pip install --no-cache-dir packaging &&\
@@ -198,19 +160,6 @@ RUN if [ "$USE_GPU" = "True" ] ; then \
     fi
 
 
-ARG TENSORFLOW_VERSION={tf_version}
- RUN if [ "$USE_GPU" = "True" ] ; then \
-            pip install --no-cache-dir tensorflow==$TENSORFLOW_VERSION; \
-        else \
-            echo 'cpu not install tensorflow'; \
-    fi
-
- RUN if [ "$USE_GPU" = "True" ] ; then \
-            cd /tmp && git clone -b ms_build --single-branch https://github.com/tastelikefeet/mmcv.git && cd mmcv && TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.9;9.0;8.6+PTX" MMCV_WITH_OPS=1 MAX_JOBS=32 FORCE_CUDA=1 pip install . && cd / && rm -fr /tmp/mmcv && pip cache purge; \
-        else \
-            cd /tmp && git clone -b ms_build --single-branch https://github.com/tastelikefeet/mmcv.git && cd mmcv && MMCV_WITH_OPS=1 MAX_JOBS=32 pip install . && cd / && rm -fr /tmp/mmcv && pip cache purge; \
-    fi
-
  RUN pip install --no-cache-dir https://github.com/liyaodev/fairseq/releases/download/v0.12.3.1/fairseq-0.12.3.1-cp311-cp311-linux_x86_64.whl  # for py311
 
 ENTRYPOINT []

docker/build_image.py

@@ -21,23 +21,23 @@ class Builder:
     def init_args(self, args: Any) -> Any:
         if not args.base_image:
             # A mirrored image of nvidia/cuda:12.4.0-devel-ubuntu22.04
-            args.base_image = 'nvidia/cuda:12.1.0-devel-ubuntu22.04'
+            args.base_image = 'nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04'
         if not args.torch_version:
-            args.torch_version = '2.3.1'
-            args.torchaudio_version = '2.3.1'
-            args.torchvision_version = '0.18.1'
+            args.torch_version = '2.9.1'
+            args.torchaudio_version = '2.9.1'
+            args.torchvision_version = '0.24.1'
         if not args.tf_version:
             args.tf_version = '2.16.1'
         if not args.cuda_version:
-            args.cuda_version = '12.1.0'
+            args.cuda_version = '12.8.1'
         if not args.vllm_version:
-            args.vllm_version = '0.5.3'
+            args.vllm_version = '0.15.1'
         if not args.lmdeploy_version:
-            args.lmdeploy_version = '0.6.2'
+            args.lmdeploy_version = '0.10.1'
         if not args.autogptq_version:
             args.autogptq_version = '0.7.1'
         if not args.flashattn_version:
-            args.flashattn_version = '2.7.1.post4'
+            args.flashattn_version = '2.8.3'
         return args
 
     def _generate_cudatoolkit_version(self, cuda_version: str) -> str:
@@ -86,6 +86,110 @@ class Builder:
                         f'Docker tag ci_image error with errno: {ret}')
 
 
+class OldCPUImageBuilder(Builder):
+
+    def init_args(self, args: Any) -> Any:
+        if not args.torch_version:
+            args.torch_version = '2.3.1'
+            args.torchaudio_version = '2.3.1'
+            args.torchvision_version = '0.18.1'
+        if not args.tf_version:
+            args.tf_version = '2.16.1'
+        if not args.cuda_version:
+            args.cuda_version = '12.1.0'
+        if not args.vllm_version:
+            args.vllm_version = '0.5.3'
+        if not args.lmdeploy_version:
+            args.lmdeploy_version = '0.6.2'
+        if not args.autogptq_version:
+            args.autogptq_version = '0.7.1'
+        if not args.flashattn_version:
+            args.flashattn_version = '2.7.1.post4'
+        return args
+
+    def generate_dockerfile(self) -> str:
+        with open('docker/Dockerfile.ubuntu.old', 'r') as f:
+            content = f.read()
+        content = content.replace('{base_image}', '')
+        return content
+
+    def image(self) -> str:
+        return (
+            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-{self.args.python_tag}-'
+            f'torch{self.args.torch_version}-{self.args.modelscope_version}-test'
+        )
+
+    def build(self):
+        return os.system(
+            f'DOCKER_BUILDKIT=0 docker build -t {self.image()} -f Dockerfile .'
+        )
+
+    def push(self):
+        ret = os.system(f'docker push {self.image()}')
+        if ret != 0:
+            return ret
+        image_tag2 = (
+            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-{self.args.python_tag}-'
+            f'torch{self.args.torch_version}-{self.args.modelscope_version}-{formatted_time}-test'
+        )
+        ret = os.system(f'docker tag {self.image()} {image_tag2}')
+        if ret != 0:
+            return ret
+        return os.system(f'docker push {image_tag2}')
+
+
+
+class OldGPUImageBuilder(Builder):
+
+    def init_args(self, args: Any) -> Any:
+        if not args.torch_version:
+            args.torch_version = '2.3.1'
+            args.torchaudio_version = '2.3.1'
+            args.torchvision_version = '0.18.1'
+        if not args.tf_version:
+            args.tf_version = '2.16.1'
+        if not args.cuda_version:
+            args.cuda_version = '12.1.0'
+        if not args.vllm_version:
+            args.vllm_version = '0.5.3'
+        if not args.lmdeploy_version:
+            args.lmdeploy_version = '0.6.2'
+        if not args.autogptq_version:
+            args.autogptq_version = '0.7.1'
+        if not args.flashattn_version:
+            args.flashattn_version = '2.7.1.post4'
+        return args
+
+    def generate_dockerfile(self) -> str:
+        with open('docker/Dockerfile.ubuntu.old', 'r') as f:
+            content = f.read()
+        content = content.replace('{base_image}', '')
+        return content
+
+    def image(self) -> str:
+        return (
+            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-{self.args.python_tag}-'
+            f'torch{self.args.torch_version}-base')
+
+    def build(self):
+        return os.system(
+            f'DOCKER_BUILDKIT=0 docker build -t {self.image()} -f Dockerfile .'
+        )
+
+    def push(self):
+        ret = os.system(f'docker push {self.image()}')
+        if ret != 0:
+            return ret
+        image_tag2 = (
+            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-cuda{self.args.cuda_version}-'
+            f'{self.args.python_tag}-torch{self.args.torch_version}-tf{self.args.tf_version}-'
+            f'{self.args.modelscope_version}-{formatted_time}-test')
+        ret = os.system(f'docker tag {self.image()} {image_tag2}')
+        if ret != 0:
+            return ret
+        return os.system(f'docker push {image_tag2}')
+
+
 class BaseCPUImageBuilder(Builder):
 
     def generate_dockerfile(self) -> str:
@@ -97,7 +201,6 @@ class BaseCPUImageBuilder(Builder):
         content = content.replace('{torch_version}', self.args.torch_version)
         content = content.replace('{cudatoolkit_version}',
                                   self.args.cudatoolkit_version)
-        content = content.replace('{tf_version}', self.args.tf_version)
         return content
 
     def image(self) -> str:
@@ -125,13 +228,12 @@ class BaseGPUImageBuilder(Builder):
         content = content.replace('{torch_version}', self.args.torch_version)
         content = content.replace('{cudatoolkit_version}',
                                   self.args.cudatoolkit_version)
-        content = content.replace('{tf_version}', self.args.tf_version)
         return content
 
     def image(self) -> str:
         return (
-            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-cuda{self.args.cuda_version}-{self.args.python_tag}-'
-            f'torch{self.args.torch_version}-tf{self.args.tf_version}-base')
+            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-cuda{self.args.cuda_version}-'
+            f'{self.args.python_tag}-torch{self.args.torch_version}-test')
 
     def build(self) -> int:
         return os.system(
@@ -142,7 +244,7 @@ class BaseGPUImageBuilder(Builder):
         return os.system(f'docker push {self.image()}')
 
 
-class CPUImageBuilder(Builder):
+class StableCPUImageBuilder(Builder):
 
     def generate_dockerfile(self) -> str:
         meta_file = './docker/install_cpu.sh'
@@ -200,33 +302,22 @@ class CPUImageBuilder(Builder):
         return os.system(f'docker push {image_tag2}')
 
 
-class GPUImageBuilder(Builder):
+class StableGPUImageBuilder(Builder):
+    """Dependencies will be stable versions"""
 
     def generate_dockerfile(self) -> str:
         meta_file = './docker/install.sh'
-        # pushd ~ popd is to solve the tf cannot use gpu problem.
-        extra_content = """
-RUN pip install tf-keras==2.16.0 --no-dependencies && \
-    pip install onnx==1.18.0 --no-dependencies && \
-    pip install deepspeed==0.17.4 --no-dependencies && \
-    pip install --no-cache-dir torchsde jupyterlab torchmetrics==0.11.4 basicsr pynvml shortuuid && \
-    CUDA_HOME=/usr/local/cuda TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0" \
-        pip install --no-cache-dir  'git+https://github.com/facebookresearch/detectron2.git'
-RUN pushd $(dirname $(python -c 'print(__import__("tensorflow").__file__)'))  && \
-    ln -svf ../nvidia/*/lib/*.so* .  && \
-    popd
-"""
-
+        with open('docker/Dockerfile.extra_install', 'r') as f:
+            extra_content = f.read()
+            extra_content = extra_content.replace('{python_version}',
+                                                  self.args.python_version)
         version_args = (
             f'{self.args.torch_version} {self.args.torchvision_version} {self.args.torchaudio_version} '
             f'{self.args.vllm_version} {self.args.lmdeploy_version} {self.args.autogptq_version} '
             f'{self.args.flashattn_version}')
-        base_image = (
-            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-cuda{self.args.cuda_version}-{self.args.python_tag}-'
-            f'torch{self.args.torch_version}-tf{self.args.tf_version}-base')
         with open('docker/Dockerfile.ubuntu', 'r') as f:
             content = f.read()
-            content = content.replace('{base_image}', base_image)
+            content = content.replace('{base_image}', self.args.base_image)
             content = content.replace('{extra_content}', extra_content)
             content = content.replace('{meta_file}', meta_file)
             content = content.replace('{version_args}', version_args)
@@ -248,83 +339,7 @@ RUN pushd $(dirname $(python -c 'print(__import__("tensorflow").__file__)'))  &&
     def image(self) -> str:
         return (
             f'{docker_registry}:ubuntu{self.args.ubuntu_version}-cuda{self.args.cuda_version}-'
-            f'{self.args.python_tag}-torch{self.args.torch_version}-tf{self.args.tf_version}-'
-            f'{self.args.modelscope_version}-test')
-
-    def build(self) -> int:
-        return os.system(f'docker build -t {self.image()} -f Dockerfile .')
-
-    def push(self):
-        ret = os.system(f'docker push {self.image()}')
-        if ret != 0:
-            return ret
-        image_tag2 = (
-            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-cuda{self.args.cuda_version}-'
-            f'{self.args.python_tag}-torch{self.args.torch_version}-tf{self.args.tf_version}-'
-            f'{self.args.modelscope_version}-{formatted_time}-test')
-        ret = os.system(f'docker tag {self.image()} {image_tag2}')
-        if ret != 0:
-            return ret
-        return os.system(f'docker push {image_tag2}')
-
-
-class LLMImageBuilder(Builder):
-
-    def init_args(self, args) -> Any:
-        if not args.base_image:
-            # A mirrored image of nvidia/cuda:12.4.0-devel-ubuntu22.04
-            args.base_image = 'nvidia/cuda:12.4.0-devel-ubuntu22.04'
-        if not args.torch_version:
-            args.torch_version = '2.8.0'
-            args.torchaudio_version = '2.8.0'
-            args.torchvision_version = '0.23.0'
-        if not args.cuda_version:
-            args.cuda_version = '12.4.0'
-        if not args.vllm_version:
-            args.vllm_version = '0.11.0'
-        if not args.lmdeploy_version:
-            args.lmdeploy_version = '0.10.1'
-        if not args.autogptq_version:
-            args.autogptq_version = '0.7.1'
-        if not args.flashattn_version:
-            args.flashattn_version = '2.7.4.post1'
-        return args
-
-    def generate_dockerfile(self) -> str:
-        meta_file = './docker/install.sh'
-        with open('docker/Dockerfile.extra_install', 'r') as f:
-            extra_content = f.read()
-            extra_content = extra_content.replace('{python_version}',
-                                                  self.args.python_version)
-        version_args = (
-            f'{self.args.torch_version} {self.args.torchvision_version} {self.args.torchaudio_version} '
-            f'{self.args.vllm_version} {self.args.lmdeploy_version} {self.args.autogptq_version} '
-            f'{self.args.flashattn_version}')
-        with open('docker/Dockerfile.ubuntu', 'r') as f:
-            content = f.read()
-            content = content.replace('{base_image}', self.args.base_image)
-            content = content.replace('{extra_content}', extra_content)
-            content = content.replace('{meta_file}', meta_file)
-            content = content.replace('{version_args}', version_args)
-            content = content.replace('{cur_time}', formatted_time)
-            content = content.replace('{install_ms_deps}', 'False')
-            content = content.replace('{image_type}', 'llm')
-            content = content.replace('{torch_version}',
-                                      self.args.torch_version)
-            content = content.replace('{torchvision_version}',
-                                      self.args.torchvision_version)
-            content = content.replace('{torchaudio_version}',
-                                      self.args.torchaudio_version)
-            content = content.replace('{index_url}', '')
-            content = content.replace('{modelscope_branch}',
-                                      self.args.modelscope_branch)
-            content = content.replace('{swift_branch}', self.args.swift_branch)
-        return content
-
-    def image(self) -> str:
-        return (
-            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-cuda{self.args.cuda_version}-'
-            f'{self.args.python_tag}-torch{self.args.torch_version}-{self.args.modelscope_version}-LLM-test'
+            f'{self.args.python_tag}-torch{self.args.torch_version}-{self.args.modelscope_version}-test'
         )
 
     def build(self) -> int:
@@ -337,30 +352,19 @@ class LLMImageBuilder(Builder):
         image_tag2 = (
             f'{docker_registry}:ubuntu{self.args.ubuntu_version}-cuda{self.args.cuda_version}-'
             f'{self.args.python_tag}-torch{self.args.torch_version}-'
-            f'{self.args.modelscope_version}-LLM-{formatted_time}-test')
+            f'{self.args.modelscope_version}-{formatted_time}-test')
         ret = os.system(f'docker tag {self.image()} {image_tag2}')
         if ret != 0:
             return ret
         return os.system(f'docker push {image_tag2}')
 
 
-class SwiftImageBuilder(LLMImageBuilder):
+class LatestGPUImageBuilder(StableGPUImageBuilder):
+    """Dependencies will be latest versions"""
 
-    def init_args(self, args) -> Any:
-        if not args.base_image:
-            args.base_image = 'nvidia/cuda:12.6.3-devel-ubuntu22.04'
-        if not args.cuda_version:
-            args.cuda_version = '12.6.3'
-        if not args.torch_version:
-            args.torch_version = '2.7.1'
-            args.torchaudio_version = '2.7.1'
-            args.torchvision_version = '0.22.1'
+    def init_args(self, args: Any) -> Any:
         if not args.vllm_version:
-            args.vllm_version = '0.10.1.1'
-        if not args.lmdeploy_version:
-            args.lmdeploy_version = '0.9.2.post1'
-        if not args.flashattn_version:
-            args.flashattn_version = '2.7.4.post1'
+            args.vllm_version = '0.16.0'
         return super().init_args(args)
 
     def generate_dockerfile(self) -> str:
@@ -383,8 +387,8 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
             content = content.replace('{meta_file}', meta_file)
             content = content.replace('{version_args}', version_args)
             content = content.replace('{cur_time}', formatted_time)
-            content = content.replace('{install_ms_deps}', 'False')
-            content = content.replace('{image_type}', 'swift')
+            content = content.replace('{install_ms_deps}', 'True')
+            content = content.replace('{image_type}', 'gpu')
             content = content.replace('{torch_version}',
                                       self.args.torch_version)
             content = content.replace('{torchvision_version}',
@@ -400,7 +404,7 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
     def image(self) -> str:
         return (
             f'{docker_registry}:ubuntu{self.args.ubuntu_version}-cuda{self.args.cuda_version}-'
-            f'{self.args.python_tag}-torch{self.args.torch_version}-{self.args.modelscope_version}-swift-test'
+            f'{self.args.python_tag}-torch{self.args.torch_version}-{self.args.modelscope_version}-latest-test'
         )
 
     def push(self):
@@ -410,14 +414,14 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
         image_tag2 = (
             f'{docker_registry}:ubuntu{self.args.ubuntu_version}-cuda{self.args.cuda_version}-'
             f'{self.args.python_tag}-torch{self.args.torch_version}-'
-            f'{self.args.modelscope_version}-swift-{formatted_time}-test')
+            f'{self.args.modelscope_version}-latest-{formatted_time}-test')
         ret = os.system(f'docker tag {self.image()} {image_tag2}')
         if ret != 0:
             return ret
         return os.system(f'docker push {image_tag2}')
 
 
-class AscendSwiftImageBuilder(SwiftImageBuilder):
+class AscendSwiftImageBuilder(StableGPUImageBuilder):
 
     def init_args(self, args) -> Any:
         if not args.base_image:
@@ -443,7 +447,7 @@ RUN pip install --no-cache-dir -U icecream soundfile pybind11 py-spy
     def image(self) -> str:
         return (
             f'{docker_registry}:{self.args.base_image.split(":")[-1]}-torch2.7.1'
-            f'-{self.args.modelscope_version}-ascend-swift-test')
+            f'-{self.args.modelscope_version}-ascend-test')
 
     def push(self):
         return 0
@@ -470,21 +474,17 @@ parser.add_argument('--swift_branch', type=str, default='main')
 parser.add_argument('--dry_run', type=int, default=0)
 args = parser.parse_args()
 
-if args.image_type.lower() == 'base_cpu':
-    builder_cls = BaseCPUImageBuilder
-elif args.image_type.lower() == 'base_gpu':
-    builder_cls = BaseGPUImageBuilder
-elif args.image_type.lower() == 'cpu':
-    builder_cls = CPUImageBuilder
-elif args.image_type.lower() == 'gpu':
-    builder_cls = GPUImageBuilder
-elif args.image_type.lower() == 'llm':
-    builder_cls = LLMImageBuilder
-elif args.image_type.lower() == 'swift':
-    builder_cls = SwiftImageBuilder
-elif args.image_type.lower() == 'ascend_swift':
-    builder_cls = AscendSwiftImageBuilder
+if args.image_type.lower() == 'base':
+    builder_cls = [BaseCPUImageBuilder, BaseGPUImageBuilder]
+elif args.image_type.lower() == 'old':
+    builder_cls = [OldCPUImageBuilder, OldGPUImageBuilder]
+elif args.image_type.lower() == 'stable':
+    builder_cls = [StableCPUImageBuilder, StableGPUImageBuilder, AscendSwiftImageBuilder]
+elif args.image_type.lower() == 'latest':
+    builder_cls = [LatestGPUImageBuilder]
 else:
     raise ValueError(f'Unsupported image_type: {args.image_type}')
 
-builder_cls(args, args.dry_run)()
+
+for builder in builder_cls:
+    builder(args, args.dry_run)()

docker/install.sh

@@ -21,7 +21,7 @@ pip install --no-cache-dir tiktoken transformers_stream_generator bitsandbytes d
 # pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.6.3/flash_attn-2.6.3+cu123torch2.4cxx11abiTRUE-cp310-cp310-linux_x86_64.whl
 # find on: https://github.com/Dao-AILab/flash-attention/releases
 # cd /tmp && git clone https://github.com/Dao-AILab/flash-attention.git && cd flash-attention && python setup.py install && cd / && rm -fr /tmp/flash-attention && pip cache purge;
-pip install --no-cache-dir flash_attn==$flashattn_version
+MAX_JOBS=16 pip install --no-cache-dir flash_attn==$flashattn_version --no-build-isolation
 
 pip install --no-cache-dir triton auto-gptq==$autogptq_version -U && pip cache purge