From 01881c3113980038bea43f95a37dc7e352d8a47e Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Wed, 13 Nov 2024 13:15:49 -0500
Subject: [PATCH] make sure to tag images in docker for tagged releases (#2051)
 [skip ci]

* make sure to tag images in docker for tagged releases

* fix tag event
---
 .github/workflows/main.yml | 14 +++++++++++---
 .github/workflows/pypi.yml |  3 ++-
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index e0443e912..8005f6683 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -4,6 +4,8 @@ on:
   push:
     branches:
       - "main"
+    tags:
+      - "v*"
   workflow_dispatch:
 
 jobs:
@@ -45,6 +47,9 @@ jobs:
           images: |
             winglian/axolotl
             axolotlai/axolotl
+          tags: |
+            type=ref,event=tag
+            type=ref,event=push
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
       - name: Login to Docker Hub
@@ -58,7 +63,7 @@ jobs:
         with:
           context: .
           build-args: |
-            BASE_TAG=${{ github.ref_name }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
+            BASE_TAG=${{ github.ref_type == 'tag' && 'main' || github.ref_name }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
             CUDA=${{ matrix.cuda }}
             PYTORCH_VERSION=${{ matrix.pytorch }}
             AXOLOTL_ARGS=${{ matrix.axolotl_args }}
@@ -121,7 +126,7 @@ jobs:
         with:
           context: .
           build-args: |
-            BASE_TAG=${{ github.ref_name }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+            BASE_TAG=${{ github.ref_type == 'tag' && 'main' || github.ref_name }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
             CUDA=${{ matrix.cuda }}
           file: ./docker/Dockerfile-cloud
           push: ${{ github.event_name != 'pull_request' }}
@@ -153,6 +158,9 @@ jobs:
           images: |
             winglian/axolotl-cloud-term
             axolotlai/axolotl-cloud-term
+          tags: |
+            type=ref,event=tag
+            type=ref,event=push
       - name: Login to Docker Hub
         uses: docker/login-action@v3
         with:
@@ -165,7 +173,7 @@ jobs:
         with:
           context: .
           build-args: |
-            BASE_TAG=${{ github.ref_name }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+            BASE_TAG=${{ github.ref_type == 'tag' && 'main' || github.ref_name }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
             CUDA=${{ matrix.cuda }}
           file: ./docker/Dockerfile-cloud-no-tmux
           push: ${{ github.event_name != 'pull_request' }}
diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml
index a8131fa54..ae0a77e2f 100644
--- a/.github/workflows/pypi.yml
+++ b/.github/workflows/pypi.yml
@@ -1,7 +1,7 @@
 name: publish pypi
 
 on:
-  create:
+  push:
     tags:
       - 'v*'
   workflow_dispatch:
@@ -27,6 +27,7 @@ jobs:
   pypi-publish:
     name: Upload release to PyPI
     runs-on: ubuntu-latest
+    needs: [setup_release]
     environment:
       name: pypi
       url: https://pypi.org/p/axolotl