Update docker scripts (#1044)

.github/workflows/docker-image.yml

@@ -8,16 +8,16 @@ on:
         required: true
         default: 'build'
       modelscope_branch:
-        description: 'ModelScope branch to build from'
+        description: 'ModelScope branch to build from(release/x.xx)'
         required: true
       image_type:
-        description: 'The image type to build'
+        description: 'The image type to build(cpu/gpu/llm)'
         required: true
       modelscope_version:
-        description: 'ModelScope version to use'
+        description: 'ModelScope version to use(x.xx.x)'
         required: true
       swift_branch:
-        description: 'SWIFT branch to use'
+        description: 'SWIFT branch to use(release/x.xx)'
         required: true
       other_params:
         description: 'Other params in --xxx xxx'

docker/build_image.py

@@ -1,9 +1,12 @@
 import argparse
 import os
+from datetime import datetime
 from typing import Any
 
 docker_registry = os.environ['DOCKER_REGISTRY']
 assert docker_registry, 'You must pass a valid DOCKER_REGISTRY'
+timestamp = datetime.now()
+formatted_time = timestamp.strftime('%Y-%m-%d-%H:%M:%S')
 
 
 class Builder:
@@ -85,12 +88,16 @@ class BaseCPUImageBuilder(Builder):
         return content
 
     def build(self):
-        image_tag = f'{docker_registry}:ubuntu{self.args.ubuntu_version}-torch{self.args.torch_version}-base'
+        image_tag = (
+            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-{self.args.python_tag}-'
+            f'torch{self.args.torch_version}-base')
         return os.system(
             f'DOCKER_BUILDKIT=0 docker build -t {image_tag} -f Dockerfile .')
 
     def push(self):
-        image_tag = f'{docker_registry}:ubuntu{self.args.ubuntu_version}-torch{self.args.torch_version}-base'
+        image_tag = (
+            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-{self.args.python_tag}-'
+            f'torch{self.args.torch_version}-base')
         return os.system(f'docker push {image_tag}')
 
 
@@ -110,14 +117,14 @@ class BaseGPUImageBuilder(Builder):
 
     def build(self) -> int:
         image_tag = (
-            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-cuda{self.args.cuda_version}-'
+            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')
         return os.system(
             f'DOCKER_BUILDKIT=0 docker build -t {image_tag} -f Dockerfile .')
 
     def push(self):
         image_tag = (
-            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-cuda{self.args.cuda_version}-'
+            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')
         return os.system(f'docker push {image_tag}')
 
@@ -129,7 +136,9 @@ class CPUImageBuilder(Builder):
         version_args = (
             f'{self.args.torch_version} {self.args.torchvision_version} '
             f'{self.args.torchaudio_version}')
-        base_image = f'{docker_registry}:ubuntu{self.args.ubuntu_version}-torch{self.args.torch_version}-base'
+        base_image = (
+            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-{self.args.python_tag}'
+            f'-torch{self.args.torch_version}-base')
         extra_content = """\nRUN pip install adaseq\nRUN pip install pai-easycv"""
 
         with open('docker/Dockerfile.ubuntu', 'r') as f:
@@ -155,6 +164,13 @@ class CPUImageBuilder(Builder):
             f'{docker_registry}:ubuntu{self.args.ubuntu_version}-{self.args.python_tag}-'
             f'torch{self.args.torch_version}-{self.args.modelscope_version}-test'
         )
+        ret = os.system(f'docker push {image_tag}')
+        if ret != 0:
+            return ret
+        image_tag = (
+            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'
+        )
         return os.system(f'docker push {image_tag}')
 
 
@@ -168,7 +184,7 @@ class GPUImageBuilder(Builder):
             f'{self.args.vllm_version} {self.args.lmdeploy_version} {self.args.autogptq_version}'
         )
         base_image = (
-            f'{docker_registry}:ubuntu{self.args.ubuntu_version}-cuda{self.args.cuda_version}-'
+            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()
@@ -193,6 +209,13 @@ class GPUImageBuilder(Builder):
             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')
+        ret = os.system(f'docker push {image_tag}')
+        if ret != 0:
+            return ret
+        image_tag = (
+            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')
         return os.system(f'docker push {image_tag}')
 
 
@@ -249,6 +272,13 @@ class LLMImageBuilder(Builder):
             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'
         )
+        ret = os.system(f'docker push {image_tag}')
+        if ret != 0:
+            return ret
+        image_tag = (
+            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')
         return os.system(f'docker push {image_tag}')