Add workflows for dev, nightly, and release Docker containers

.github/workflows/build-container.yml

@@ -0,0 +1,80 @@
+name: Build Docker Container
+on:
+  pull_request:
+    types: [opened, synchronize, reopened]
+    branches:
+      - main
+      - 'release/[0-9]+.[0-9]+'
+  push:
+    branches:
+      - 'main'
+  release:
+    types: [created, published]
+  schedule:
+    - cron: '0 2 * * *'
+
+# TODO: docker containers created through a release cut vs PR to the release branch
+# will be pushed to different locations (i.e one will be sparseml the other will be test-sparseml).
+# These containers rely on the new internal pypi server being enabled. Once enabled,
+# this workflow can be expanded to make this distinction.
+env:
+  RELEASE: ${{ github.event_name =='release' || (startsWith(github.base_ref, 'release/') && github.event_name == 'pull_request')}}
+  DEV: ${{ github.base_ref == 'main' && github.event_name == 'pull_request'}}
+  NAME: ${{ github.event.number }} 
+
+permissions:
+  contents: read
+  packages: write
+
+jobs:
+  build-container:
+    name: Build sparseml container
+    runs-on: ubuntu-20.04
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 1
+      - name: Set up Docker Buildx
+        id: buildx
+        uses: docker/setup-buildx-action@v2
+        with:
+          buildkitd-flags: --debug
+      - name: Get current date
+        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
+        uses: docker/login-action@v2
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+      - name: Build Dev Docker Container
+        if: ${{ env.DEV == 'true' }}
+        uses: docker/build-push-action@v4
+        with: 
+          context: ./docker/containers/docker_dev
+          build-args: |
+            BRANCH=${{github.head_ref}}
+          push: true
+          tags: ghcr.io/neuralmagic/sparseml-dev:${{ env.NAME }}
+      - name: Build Release Docker Container
+        if: ${{ env.RELEASE == 'true' }}
+        uses: docker/build-push-action@v4
+        with: 
+          context: ./docker/containers/docker_release
+          build-args: |
+            VERSION=${{ steps.version.outputs.version }}
+          push: true
+          tags: ghcr.io/neuralmagic/test-sparseml:latest, ghcr.io/neuralmagic/test-sparseml:${{ steps.version.outputs.version }}
+      - name: Build Nightly Docker Container
+        if: ${{ env.DEV == 'false' && env.RELEASE == 'false'}}
+        uses: docker/build-push-action@v4
+        with:
+          context: ./docker/containers/docker_nightly
+          push: true
+          tags: ghcr.io/neuralmagic/test-sparseml-nightly:latest, ghcr.io/neuralmagic/test-sparseml-nightly:${{ steps.date.outputs.date }}

docker/containers/docker_dev/Dockerfile

@@ -0,0 +1,25 @@
+ARG SOURCE=ghcr.io/neuralmagic/cuda-python3.10
+
+ARG TORCH_VERSION=2.1.2
+ARG TORCHVISION_VERSION=0.16.2
+ARG CUDA=121
+ARG BRANCH
+
+FROM $SOURCE
+
+ARG BRANCH
+
+RUN python3.10 -m pip install --upgrade pip \
+    && python3.10 -m pip install --upgrade setuptools
+
+ARG CUDA
+ARG TORCH_VERSION
+ARG TORCHVISION_VERSION
+
+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 \
+    && git clone https://github.com/neuralmagic/sparseml.git --depth 1 --single-branch -b ${BRANCH} \
+    && python3.10 -m pip install -e "./sparseml[dev]"
+
+HEALTHCHECK CMD python3.10 -c 'import sparseml'
+RUN python3.10 -m pip list | grep sparseml
+CMD bash

docker/containers/docker_nightly/Dockerfile

@@ -0,0 +1,21 @@
+ARG SOURCE=ghcr.io/neuralmagic/cuda-python3.10
+
+ARG TORCH_VERSION=2.1.2
+ARG TORCHVISION_VERSION=0.16.2
+ARG CUDA=121
+
+FROM $SOURCE
+
+RUN python3.10 -m pip install --upgrade pip \
+    && python3.10 -m pip install --upgrade setuptools
+
+ARG CUDA 
+ARG TORCH_VERSION
+ARG TORCHVISION_VERSION
+
+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]" 
+
+HEALTHCHECK CMD python3.10 -c 'import sparseml'
+RUN python3.10 -m pip list | grep sparseml
+CMD bash

docker/containers/docker_release/Dockerfile

@@ -0,0 +1,24 @@
+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
+
+FROM $SOURCE
+
+ARG VERSION
+
+ARG CUDA
+ARG TORCH_VERSION
+ARG TORCHVISION_VERSION
+
+RUN python3.10 -m pip install --upgrade pip \
+    && python3.10 -m pip install --upgrade setuptools
+
+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[onnxruntime,torchvision,transformers,yolov5,ultralytics]==$VERSION" 
+
+HEALTHCHECK CMD python3.10 -c 'import sparseml'
+RUN python3.10 -m pip list | grep sparseml
+CMD bash