From 9ca358b6710b9b7576d2f5f40b75f47fde5559d5 Mon Sep 17 00:00:00 2001
From: Hamel Husain <hamel.husain@gmail.com>
Date: Sat, 6 Jan 2024 05:20:33 -0800
Subject: [PATCH] Simplify Docker Unit Test CI (#1055) [skip ci]

* Update tests-docker.yml

* Update tests-docker.yml

* run ci tests on ci yaml updates

---------

Co-authored-by: Wing Lian <wing.lian@gmail.com>
---
 .github/workflows/tests-docker.yml | 36 ++++++++++++------------------
 1 file changed, 14 insertions(+), 22 deletions(-)

diff --git a/.github/workflows/tests-docker.yml b/.github/workflows/tests-docker.yml
index e93884e64..380729637 100644
--- a/.github/workflows/tests-docker.yml
+++ b/.github/workflows/tests-docker.yml
@@ -5,6 +5,7 @@ on:
     paths:
       - '**.py'
       - 'requirements.txt'
+      - '.github/workflows/*.yml'
   workflow_dispatch:
 
 jobs:
@@ -19,36 +20,27 @@ jobs:
             cuda_version: 11.8.0
             python_version: "3.10"
             pytorch: 2.0.1
-            axolotl_extras:
           - cuda: 121
             cuda_version: 12.1.0
             python_version: "3.10"
             pytorch: 2.1.1
-            axolotl_extras:
     runs-on: [self-hosted, gpu, docker]
     steps:
       - name: Checkout
         uses: actions/checkout@v4
-      - name: Docker metadata
-        id: metadata
-        uses: docker/metadata-action@v5
-        with:
-          images: winglian/axolotl
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
       - name: Build Docker image
-        uses: docker/build-push-action@v5
-        with:
-          context: .
-          load: true
-          build-args: |
-            BASE_TAG=main-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
-            CUDA=${{ matrix.cuda }}
-            PYTORCH_VERSION=${{ matrix.pytorch }}
-          file: ./docker/Dockerfile
-          tags: |
-            ${{ steps.metadata.outputs.tags }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
-          labels: ${{ steps.metadata.outputs.labels }}
+        run: |
+          # Set up build arguments
+          BASE_TAG="main-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}"
+          CUDA="${{ matrix.cuda }}"
+          PYTORCH_VERSION="${{ matrix.pytorch }}"
+          # Build the Docker image
+          docker build . \
+            --file ./docker/Dockerfile \
+            --build-arg BASE_TAG=$BASE_TAG \
+            --build-arg CUDA=$CUDA \
+            --build-arg PYTORCH_VERSION=$PYTORCH_VERSION \
+            --tag test-axolotl
       - name: Unit Tests w docker image
         run: |
-          docker run --rm ${{ steps.metadata.outputs.tags }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }} pytest --ignore=tests/e2e/ /workspace/axolotl/tests/
+          docker run --rm test-axolotl pytest --ignore=tests/e2e/ /workspace/axolotl/tests/