From 46c4ce9bdd6e60fa501f36021f6694994fc4aece Mon Sep 17 00:00:00 2001
From: tastelikefeet <58414341+tastelikefeet@users.noreply.github.com>
Date: Sun, 1 Dec 2024 15:17:25 +0800
Subject: [PATCH] support tag ci_image (#1112)

---
 .github/workflows/docker-image.yml |  6 +-
 docker/build_image.py              | 88 ++++++++++++++++--------------
 2 files changed, 52 insertions(+), 42 deletions(-)

diff --git a/.github/workflows/docker-image.yml b/.github/workflows/docker-image.yml
index 6d20383d..13f61ff3 100644
--- a/.github/workflows/docker-image.yml
+++ b/.github/workflows/docker-image.yml
@@ -19,6 +19,10 @@ on:
       swift_branch:
         description: 'SWIFT branch to use(release/x.xx)'
         required: true
+      ci_image:
+        description: 'Set as the CI image'
+        default: '0'
+        required: false
       other_params:
         description: 'Other params in --xxx xxx'
         required: false
@@ -47,4 +51,4 @@ jobs:
       run: |
         set -e
         source ~/.bashrc
-        python docker/build_image.py --image_type ${{ github.event.inputs.image_type }} --modelscope_branch ${{ github.event.inputs.modelscope_branch }} --modelscope_version ${{ github.event.inputs.modelscope_version }} --swift_branch ${{ github.event.inputs.swift_branch }} ${{ github.event.inputs.other_params }}
+        python docker/build_image.py --image_type ${{ github.event.inputs.image_type }} --modelscope_branch ${{ github.event.inputs.modelscope_branch }} --modelscope_version ${{ github.event.inputs.modelscope_version }} --swift_branch ${{ github.event.inputs.swift_branch }} --ci_image ${{ github.event.inputs.ci_image }} ${{ github.event.inputs.other_params }}
diff --git a/docker/build_image.py b/docker/build_image.py
index ddc67581..344fc9d3 100644
--- a/docker/build_image.py
+++ b/docker/build_image.py
@@ -61,6 +61,9 @@ class Builder:
     def push(self) -> int:
         pass
 
+    def image(self) -> str:
+        pass
+
     def __call__(self):
         content = self.generate_dockerfile()
         self._save_dockerfile(content)
@@ -68,10 +71,18 @@ class Builder:
             ret = self.build()
             if ret != 0:
                 raise RuntimeError(f'Docker build error with errno: {ret}')
+
             ret = self.push()
             if ret != 0:
                 raise RuntimeError(f'Docker push error with errno: {ret}')
 
+            if self.args.ci_image != 0:
+                ret = os.system(
+                    f'docker tag {self.image()} {docker_registry}:ci_image')
+                if ret != 0:
+                    raise RuntimeError(
+                        f'Docker tag ci_image error with errno: {ret}')
+
 
 class BaseCPUImageBuilder(Builder):
 
@@ -87,18 +98,18 @@ class BaseCPUImageBuilder(Builder):
         content = content.replace('{tf_version}', self.args.tf_version)
         return content
 
-    def build(self):
-        image_tag = (
+    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 {image_tag} -f Dockerfile .')
+            f'DOCKER_BUILDKIT=0 docker build -t {self.image()} -f Dockerfile .'
+        )
 
     def push(self):
-        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}')
+        return os.system(f'docker push {self.image()}')
 
 
 class BaseGPUImageBuilder(Builder):
@@ -115,18 +126,18 @@ class BaseGPUImageBuilder(Builder):
         content = content.replace('{tf_version}', self.args.tf_version)
         return content
 
-    def build(self) -> int:
-        image_tag = (
+    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')
+
+    def build(self) -> int:
         return os.system(
-            f'DOCKER_BUILDKIT=0 docker build -t {image_tag} -f Dockerfile .')
+            f'DOCKER_BUILDKIT=0 docker build -t {self.image()} -f Dockerfile .'
+        )
 
     def push(self):
-        image_tag = (
-            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}')
+        return os.system(f'docker push {self.image()}')
 
 
 class CPUImageBuilder(Builder):
@@ -162,26 +173,24 @@ class CPUImageBuilder(Builder):
             content = content.replace('{swift_branch}', self.args.swift_branch)
         return content
 
-    def build(self) -> int:
-        image_tag = (
+    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'
         )
-        return os.system(f'docker build -t {image_tag} -f Dockerfile .')
+
+    def build(self) -> int:
+        return os.system(f'docker build -t {self.image()} -f Dockerfile .')
 
     def push(self):
-        image_tag = (
-            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}')
+        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 {image_tag} {image_tag2}')
+        ret = os.system(f'docker tag {self.image()} {image_tag2}')
         if ret != 0:
             return ret
         return os.system(f'docker push {image_tag2}')
@@ -224,26 +233,24 @@ RUN pip install tf-keras==2.16.0 --no-dependencies && \
             content = content.replace('{swift_branch}', self.args.swift_branch)
         return content
 
-    def build(self) -> int:
-        image_tag = (
+    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')
-        return os.system(f'docker build -t {image_tag} -f Dockerfile .')
+
+    def build(self) -> int:
+        return os.system(f'docker build -t {self.image()} -f Dockerfile .')
 
     def push(self):
-        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}-test')
-        ret = os.system(f'docker push {image_tag}')
+        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 {image_tag} {image_tag2}')
+        ret = os.system(f'docker tag {self.image()} {image_tag2}')
         if ret != 0:
             return ret
         return os.system(f'docker push {image_tag2}')
@@ -298,26 +305,24 @@ class LLMImageBuilder(Builder):
             content = content.replace('{swift_branch}', self.args.swift_branch)
         return content
 
-    def build(self) -> int:
-        image_tag = (
+    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'
         )
-        return os.system(f'docker build -t {image_tag} -f Dockerfile .')
+
+    def build(self) -> int:
+        return os.system(f'docker build -t {self.image()} -f Dockerfile .')
 
     def push(self):
-        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}-{self.args.modelscope_version}-LLM-test'
-        )
-        ret = os.system(f'docker push {image_tag}')
+        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}-'
             f'{self.args.modelscope_version}-LLM-{formatted_time}-test')
-        ret = os.system(f'docker tag {image_tag} {image_tag2}')
+        ret = os.system(f'docker tag {self.image()} {image_tag2}')
         if ret != 0:
             return ret
         return os.system(f'docker push {image_tag2}')
@@ -331,6 +336,7 @@ parser.add_argument('--ubuntu_version', type=str, default='22.04')
 parser.add_argument('--torch_version', type=str, default=None)
 parser.add_argument('--torchvision_version', type=str, default=None)
 parser.add_argument('--cuda_version', type=str, default=None)
+parser.add_argument('--ci_image', type=int, default=0)
 parser.add_argument('--torchaudio_version', type=str, default=None)
 parser.add_argument('--tf_version', type=str, default=None)
 parser.add_argument('--vllm_version', type=str, default=None)