Refactor and further generalize GHA Docker workflows

mosaicml · Feb 15, 2023 · abda4e1 · abda4e1
1 parent d96bb11
commit abda4e1
Show file tree

Hide file tree

Showing 5 changed files with 178 additions and 186 deletions.
diff --git a/.github/bin/gen_docker_matrix.py b/.github/bin/gen_docker_matrix.py
@@ -23,10 +23,6 @@ def _parse_args() -> Namespace:
                       action='append',
                       required=False,
                       help='List of build args to override globally')
-    args.add_argument('--get_image_build_args',
-                      action='store',
-                      required=False,
-                      help='Return JSON output of build arguments for specified image')
 
     return args.parse_args()
 
@@ -36,26 +32,22 @@ def main(args: Namespace):
     image_configs = yaml.safe_load(args.yaml_file)
 
     for image_config in image_configs:
-        if args.get_image_build_args is None:
-            # Convert tags list to a CSV string
-            image_config['TAGS'] = ','.join(image_config['TAGS'])
-
-            # Generate a random UUID for staging
-            image_config['UUID'] = str(uuid4())
-
-            # Apply build args override
-            if args.build_args is not None:
-                for build_arg in args.build_args:
-                    arg, val = build_arg.split('=')
-                    if arg in image_config.keys():
-                        image_config[arg] = val
-        else:
-            if image_config['IMAGE_NAME'] == args.get_image_build_args:
-                print(json.dumps(image_config['BUILD_ARGS']))
-
-    if args.get_image_build_args is None:
-        json_string = json.dumps(image_configs)
-        print(f"""matrix={{"include": {json_string}}}""")
+
+        # Convert tags list to a CSV string
+        image_config['TAGS'] = ','.join(image_config['TAGS'])
+
+        # Generate a random UUID for staging
+        image_config['UUID'] = str(uuid4())
+
+        # Apply build args override
+        if args.build_args is not None:
+            for build_arg in args.build_args:
+                arg, val = build_arg.split('=')
+                if arg in image_config.keys():
+                    image_config[arg] = val
+
+    json_string = json.dumps(image_configs)
+    print(f"""matrix={{"include": {json_string}}}""")
 
 
 if __name__ == '__main__':

diff --git a/.github/workflows/docker-build.yaml b/.github/workflows/docker-build.yaml
@@ -1,13 +1,19 @@
-name: Docker Build
+name: Docker Image Configure-Build-Push
 on:
   workflow_call:
     inputs:
-      build-matrix:
-        required: true
+      build-args:
+        required: false
         type: string
       context:
         required: true
         type: string
+      image-name:
+        required: true
+        type: string
+      image-uuid:
+        required: false
+        type: string
       push:
         required: true
         type: boolean
@@ -17,18 +23,20 @@ on:
       staging-repo:
         required: false
         type: string
+      tags:
+        required: true
+        type: string
+      target:
+        required: false
+        type: string
     secrets:
       username:
         required: true
       password:
         required: true
 jobs:
-  build-images:
+  configure-build-push:
     runs-on: ubuntu-latest
-    strategy:
-      fail-fast: true
-      matrix: ${{ fromJSON(inputs.build-matrix) }}
-    name: "${{ matrix.IMAGE_NAME }}"
     steps:
       - name: Checkout
         uses: actions/checkout@v3
@@ -49,29 +57,15 @@ jobs:
         run: |
           set -euo pipefail
 
-          ###################
-          # Image config
-          ###################
-          echo "IMAGE_NAME: ${{ matrix.IMAGE_NAME }}"
-          echo "BASE_IMAGE: ${{ matrix.BASE_IMAGE }}"
-          echo "TARGET: ${{ matrix.TARGET }}"
-          echo "TAGS: ${{ matrix.TAGS }}"
-          echo "UUID: ${{ matrix.UUID }}"
-          echo "BUILD_ARGS: ${{ toJSON(matrix.BUILD_ARGS) }}"
-
-          BUILD_ARGS=$(python .github/bin/gen_docker_matrix.py docker/build_matrix.yaml --get_image_build_args ${{ matrix.IMAGE_NAME }} | jq -r 'to_entries | map("\(.key)=\(.value|tostring)") | join(" ")')
-          echo "BUILD_ARGS=$BUILD_ARGS" >> ${GITHUB_ENV}
-
-
           ###################
           # Calculate the tag
           ###################
           if [[ ${{ inputs.staging }} ]]; then
             STAGING_REPO=${{ inputs.staging-repo }}
-            IMAGE_TAG=${STAGING_REPO}:${{ matrix.UUID }}
-            IMAGE_CACHE="${STAGING_REPO}:${{ matrix.IMAGE_NAME }}-buildcache"
+            IMAGE_TAG=${STAGING_REPO}:${{ inputs.image-uuid }}
+            IMAGE_CACHE="${STAGING_REPO}:${{ inputs.image-name }}-buildcache"
           else
-            IMAGE_TAG=${{ matrix.TAGS }}
+            IMAGE_TAG=${{ inputs.tags }}
             IMAGE_CACHE="${IMAGE_TAG}-buildcache"
           fi
 
@@ -83,8 +77,8 @@ jobs:
         with:
           context: ${{ inputs.context }}
           tags: ${{ env.IMAGE_TAG }}
-          target: ${{ matrix.TARGET }}
+          target: ${{ inputs.target }}
           push: ${{ inputs.push }}
           cache-from: type=registry,ref=${{ env.IMAGE_CACHE }}
           cache-to: type=registry,ref=${{ env.IMAGE_CACHE }},mode=max
-          build-args: ${{ env.BUILD_ARGS }}
+          build-args: ${{ inputs.build-args }}
diff --git a/.github/workflows/pr-docker.yaml b/.github/workflows/pr-docker.yaml
@@ -7,7 +7,8 @@ on:
       - release/**
     paths:
       - .github/bin/gen_docker_matrix.py
-      - .github/workflows/docker**
+      - .github/workflows/docker-build.yaml
+      - .github/workflows/pr-docker.yaml
       - docker/**
 defaults:
   run:
@@ -35,15 +36,31 @@ jobs:
           # Generate build matrix
           BUILD_MATRIX=$(python .github/bin/gen_docker_matrix.py docker/build_matrix.yaml -b COMPOSER_INSTALL_COMMAND=$COMPOSER_INSTALL_COMMAND)
           echo $BUILD_MATRIX >> $GITHUB_OUTPUT
-  docker-build:
+  stage-docker-build:
     needs: build-image-matrix
     uses: ./.github/workflows/docker-build.yaml
+    strategy:
+      fail-fast: true
+      matrix: ${{ fromJSON(needs.build-image-matrix.outputs.matrix) }}
+    name: ${{ matrix.IMAGE_NAME }}
     with:
-      build-matrix: ${{ needs.build-image-matrix.outputs.matrix }}
+      build-args: |
+        AWS_OFI_NCCL_VERSION=${{ matrix.AWS_OFI_NCCL_VERSION }}
+        BASE_IMAGE=${{ matrix.BASE_IMAGE }}
+        CUDA_VERSION=${{ matrix.CUDA_VERSION }}
+        MOFED_VERSION=${{ matrix.MOFED_VERSION }}
+        PYTHON_VERSION=${{ matrix.PYTHON_VERSION }}
+        PYTORCH_VERSION=${{ matrix.PYTORCH_VERSION }}
+        TORCHTEXT_VERSION=${{ matrix.TORCHTEXT_VERSION }}
+        TORCHVISION_VERSION=${{ matrix.TORCHVISION_VERSION }}
       context: ./docker
+      image-name: ${{ matrix.IMAGE_NAME }}
+      image-uuid: ${{ matrix.UUID }}
       push: false
       staging: true
       staging-repo: mosaicml/ci-staging
+      tags: ${{ matrix.TAGS }}
+      target: ${{ matrix.TARGET }}
     secrets:
       username: ${{ secrets.DOCKER_HUB_USERNAME }}
       password: ${{ secrets.DOCKER_HUB_PASSWORD }}