Build Primus Docker (Private) (#247)

Author: wenxie-amd

.github/workflows/ci.yaml

@@ -9,6 +9,9 @@ on:
       - "v*"
   pull_request:
 
+env:
+  PRIMUS_TURBO_COMMIT: a97e986f01893fec5789243400c7046d18591ba9 # chore: update aiter to fix pybind11 issue (#132)
+
 jobs:
   code-lint:
     runs-on: ubuntu-latest
@@ -47,20 +50,83 @@ jobs:
         if: always()
         with:
           options: "--check --diff --color --verbose --line-length=110"
+
+  build-docker:
+    needs: [code-lint]
+    runs-on: build-docker
+    strategy:
+      matrix:
+        python-version: ["3.12"]
+    steps:
+      - run: echo "🎉 Begin Build Primus Docker Image."
+      - uses: actions/checkout@v4
+        with:
+          submodules: recursive
+      - name: Show Environment Info
+        run: |
+          echo "Hostname: $(hostname)"
+          echo "PWD: $(pwd)"
+          echo "HOME: $HOME"
+          echo "GITHUB_WORKSPACE: $GITHUB_WORKSPACE"
+          echo "Runner Temp Dir: $RUNNER_TEMP"
+          echo "Runner Tool Cache: $RUNNER_TOOL_CACHE"
+      - name: Parse Commit Info
+        run: |
+          if [[ "${{ github.event_name }}" == "pull_request" ]]; then
+            echo "IMAGE_TAG=pr-${{ github.event.pull_request.number }}" >> $GITHUB_ENV
+          elif [[ "${{ github.event_name }}" == "push" && "${{ github.ref }}" == "refs/heads/main" ]]; then
+            echo "IMAGE_TAG=latest" >> $GITHUB_ENV
+          elif [[ "${{ github.event_name }}" == "release" ]]; then
+            TAG_NAME="${{ github.ref }}"
+            TAG="${TAG_NAME#refs/tags/}"
+            echo "IMAGE_TAG=$TAG" >> $GITHUB_ENV
+          else
+            echo "IMAGE_TAG=others" >> $GITHUB_ENV
+          fi
+      - name: Build and Push Docker Image
+        run: |
+          echo "> Login to Docker Hub"
+          docker login -u tasimage -p ${{ secrets.PRIMUS_DOCKER_HUB_TOKEN }}
+
+          echo "> Build Docker Image with tag: ${{ env.IMAGE_TAG }}"
+          start_time=$(date +%s)
+          docker build -f $GITHUB_WORKSPACE/.github/workflows/docker/Dockerfile \
+               -t tasimage/primus:${{env.IMAGE_TAG}} \
+               --build-arg PRIMUS_TURBO_COMMIT=${PRIMUS_TURBO_COMMIT} .
+          end_time=$(date +%s)
+          elapsed=$((end_time - start_time))
+          echo "⏱️ [build primus docker] Total elapsed time: ${elapsed} seconds"
+
+          echo "> Docker tag image for Docker Hub"
+          docker tag tasimage/primus:${{env.IMAGE_TAG}} docker.io/tasimage/primus:${{env.IMAGE_TAG}}
+
+          echo "> Docker push to Docker Hub"
+          start_time=$(date +%s)
+          docker push docker.io/tasimage/primus:${{env.IMAGE_TAG}}
+          end_time=$(date +%s)
+          elapsed=$((end_time - start_time))
+          echo "⏱️ [push primus docker] Total elapsed time: ${elapsed} seconds"
+
+          # echo "> Docker cleanup local images"
+          # docker rmi tasimage/primus:${{env.IMAGE_TAG}}
+          # echo "> build-docker success"
+
   run-unittest:
     env:
       PRIMUS_WORKDIR: /apps/tas/0_public/primus_k8s_ci
     needs: [code-lint]
     runs-on: [tas-k8s]
     steps:
       - run: echo "🎉 Begin Primus-Turbo Checkout."
+      - name: Set commit hash to env
+        run: echo "PRIMUS_TURBO_COMMIT=${PRIMUS_TURBO_COMMIT}" >> $GITHUB_ENV
       - name: Checkout Repo Primus-Turbo
         uses: actions/checkout@v4
         with:
           repository: AMD-AIG-AIMA/Primus-Turbo
           submodules: "recursive"
           path: Primus-Turbo
-          ref: a97e986f01893fec5789243400c7046d18591ba9 # chore: update aiter to fix pybind11 issue (#132)
+          ref: ${{ env.PRIMUS_TURBO_COMMIT }}
       - run: echo "Begin Primus-Turbo Install."
       - name: Install Primus-Turbo
         run: |

.github/workflows/docker/Dockerfile

@@ -0,0 +1,25 @@
+# Base image
+FROM docker.io/rocm/megatron-lm:v25.9_gfx942
+
+# Specify the commit of Primus-Turbo when building: docker build --build-arg PRIMUS_TURBO_COMMIT=xxx .)
+ARG PRIMUS_TURBO_COMMIT
+
+# Install basic dependencies
+RUN apt-get update
+
+# Clone and install the Primus-Turbo
+WORKDIR /opt
+RUN mkdir -p /opt && cd /opt && \
+    git clone https://github.com/AMD-AGI/Primus-Turbo.git && \
+    cd Primus-Turbo && \
+    git checkout ${PRIMUS_TURBO_COMMIT} && \
+    git submodule update --init --recursive && \
+    pip3 install -r requirements.txt && \
+    GPU_ARCHS="gfx942;gfx950" pip3 install --no-build-isolation .
+
+
+# Set the default working directory
+WORKDIR /opt
+
+# check the installed Primus-Turbo package
+RUN python3 -m pip show primus-turbo || true