add version and date to nightly container

neuralmagic · Mar 29, 2024 · 4bea3af · 4bea3af
1 parent 2d3de11
commit 4bea3af
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/.github/workflows/build-wheel-and-container.yml b/.github/workflows/build-wheel-and-container.yml
@@ -74,7 +74,6 @@ jobs:
         id: date
         run: echo "::set-output name=date::$(date +'%Y%m%d')"
       - name: Get the current version
-        if: ${{ env.RELEASE == 'true' }}
         id: version
         run: echo "::set-output name=version::$(echo ${{ github.base_ref }} | cut -c 9-15)"
       - name: Login to Github Packages
@@ -106,5 +105,8 @@ jobs:
         uses: docker/build-push-action@v4
         with:
           context: ./docker/containers/docker_nightly
+           build-args: |
+            VERSION=${{ steps.version.outputs.version }}
+            DATE=${{ steps.date.outputs.date }}
           push: true
           tags: ghcr.io/neuralmagic/test-sparseml-nightly:latest, ghcr.io/neuralmagic/test-sparseml-nightly:${{ steps.date.outputs.date }}
diff --git a/docker/containers/docker_nightly/Dockerfile b/docker/containers/docker_nightly/Dockerfile
@@ -3,6 +3,8 @@ ARG SOURCE=ghcr.io/neuralmagic/cuda-python3.10
 ARG TORCH_VERSION=2.1.2
 ARG TORCHVISION_VERSION=0.16.2
 ARG CUDA=121
+ARG VERSION
+ARG DATE
 
 FROM $SOURCE
 
@@ -12,9 +14,11 @@ RUN python3.10 -m pip install --upgrade pip \
 ARG CUDA 
 ARG TORCH_VERSION
 ARG TORCHVISION_VERSION
+ARG VERSION
+ARG DATE
 
 RUN python3.10 -m pip install torch==${TORCH_VERSION}+cu${CUDA} torchvision==${TORCHVISION_VERSION}+cu${CUDA} -f https://download.pytorch.org/whl/torch_stable.html \
-    && python3.10 -m pip install --no-cache-dir "sparseml-nightly[onnxruntime,torchvision,transformers,yolov5,ultralytics]" 
+    && python3.10 -m pip install --no-cache-dir "sparseml-nightly[onnxruntime,torchvision,transformers,yolov5,ultralytics]==$VERSION.$DATE" 
 
 HEALTHCHECK CMD python3.10 -c 'import sparseml'
 RUN python3.10 -m pip list | grep sparseml