Merge branch 'main' into export_instantiated_model

.github/workflows/build-container.yml

@@ -1,80 +1,56 @@
-name: Build Docker Container
+name: Build Container with wheel and push to GCR
 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
+  workflow_call:
+    inputs:
+      build-label:
+        description: "requested runner label"
+        type: string
+      dev:
+        type: string
+        required: true 
+      release:
+        type: string
+        required: true
+      name:
+        type: string
 
 jobs:
   build-container:
-    name: Build sparseml container
-    runs-on: ubuntu-20.04
+    runs-on: ${{ inputs.build-label }}
     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')"
+        run: echo "date=$(date +'%Y%m%d')" >> $GITHUB_OUTPUT
+
       - name: Get the current version
-        if: ${{ env.RELEASE == 'true' }}
+        if: ${{ inputs.release == 'true' }}
         id: version
-        run: echo "::set-output name=version::$(echo ${{ github.base_ref }} | cut -c 9-15)"
+        run: echo "version=$(echo ${{ github.base_ref }} | cut -c 9-15)" >> $GITHUB_OUTPUT
+
       - 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' }}
+        if: ${{ inputs.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 }}
+          tags: ghcr.io/neuralmagic/sparseml-dev:${{ inputs.name }}

.github/workflows/build-wheel-and-container.yml

@@ -0,0 +1,77 @@
+name: Build PyPi Wheel and Docker Container
+on:
+  pull_request:
+    types: [opened, synchronize, reopened]
+    branches:
+      - main
+      - 'release/[0-9]+.[0-9]+'
+  push:
+    branches:
+      - 'release/[0-9]+.[0-9]+'
+      - main
+  release:
+    types: [created, published]
+  schedule:
+    - cron: '0 0 * * *'
+
+permissions:
+  id-token: write
+  contents: read
+  packages: write
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+# if not dev or release, will create a nightly build
+env:
+  PRODUCTION: ${{ github.event_name == 'schedule' || github.event_name == 'release'}}
+  RELEASE: ${{ github.event_name =='release' || startsWith(github.base_ref, 'release/') }}
+  DEV: ${{ github.base_ref == 'main' && github.event_name == 'pull_request'}}
+
+jobs:
+  set-outputs:
+    runs-on: ubuntu-latest
+    outputs:
+      dev: ${{ steps.set-outputs.outputs.dev }}
+      release: ${{ steps.set-outputs.outputs.release }}
+    steps:
+      - name: Set variables for workflows
+        id: set-outputs
+        run: |
+          echo "dev=$DEV" >> $GITHUB_OUTPUT
+          echo "release=$RELEASE" >> $GITHUB_OUTPUT
+
+  build-wheel-and-push:
+    needs: set-outputs
+    uses: ./.github/workflows/build-wheel.yml
+    with:
+      build-label: ubuntu-20.04
+      dev: ${{ needs.set-outputs.outputs.dev }}
+      release: ${{ needs.set-outputs.outputs.release  }}
+      name: ${{ github.event.number }} 
+      filename: dist/*.whl
+      bucket_name: nm-actions-test
+      python: '3.10'
+    secrets: inherit
+
+  test-wheel-and-push-internal:
+    needs: build-wheel-and-push
+    uses: ./.github/workflows/test-wheel-push-to-internal.yml
+    with:
+      build-label: aws-avx2-64G
+      whl: ${{ needs.build-wheel-and-push.outputs.wheel }}
+      python: '3.10'
+    secrets: inherit
+
+  # TODO: add nightly and release container build steps once wheel build push
+  # to production is automated. Removed until then.
+  build-container-and-push:
+    needs: [set-outputs, test-wheel-and-push-internal]
+    uses: ./.github/workflows/build-container.yml
+    with:
+      build-label: aws-avx2-64G
+      dev: ${{ needs.set-outputs.outputs.dev }}
+      release: ${{ needs.set-outputs.outputs.release  }}
+      name: ${{ github.event.number }}
+    secrets: inherit

.github/workflows/build-wheel.yml

@@ -1,58 +1,67 @@
-name: Build PyPi Wheel
+name: Build Wheel and Push to s3
 on:
-  pull_request:
-    types: [opened, synchronize, reopened]
-    branches:
-      - main
-      - 'release/[0-9]+.[0-9]+'
-  push:
-    branches:
-      - main
-  release:
-    types: [created, published]
-  schedule:
-    - cron: '0 0 * * *'
-
-permissions:
-  id-token: write
-  contents: read
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
-  cancel-in-progress: true
-
-# if not dev or release, will create a nightly build
-# everything is pushed to internal unless created through a nightly scheduled cron job which creates the build or 
-# missing release tag workflow/needs to be added in
-env:
-  INTERNAL: ${{ github.event_name != 'schedule' && github.event_name != 'release'}}
-  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 }} 
+  workflow_call:
+    inputs:
+      build-label:
+        description: "requested runner label"
+        type: string
+        required: true
+      dev:
+        type: string
+        required: true 
+      release:
+        type: string
+        required: true
+      name:
+        type: string
+      filename:
+        type: string
+        required: true
+      bucket_name:
+        type: string
+        required: true
+      python:
+        type: string
+    outputs:
+      wheel:
+        value: ${{ jobs.build-wheel-and-push.outputs.wheel }}
 
 jobs:
-  build_and_push:
-    runs-on: ubuntu-latest
+  build-wheel-and-push:
+    runs-on: ${{ inputs.build-label }}
     outputs:
       wheel: ${{ steps.push-wheel.outputs.wheel }}
     steps:
-      - name: Checkout code
-        uses: actions/checkout@v3
-      - name: Login to s3
-        uses: aws-actions/configure-aws-credentials@v2
-        with:
-          role-to-assume: ${{ secrets.AWS_WEBIDENTITY_FOR_GITHUB_ACTIONS }}
-          aws-region: us-east-1 
-      - name: Build PyPi Wheel
-        id: build-wheel
-        uses: neuralmagic/nm-actions/actions/pypi_build@main
-        with:
-          dev: $DEV
-          release: $RELEASE
-          name: $NAME
-      - name: Push to s3 bucket
-        id: push-wheel
-        uses: neuralmagic/nm-actions/actions/s3_push@main
-        with:
-          filename: dist/*.whl
-          internal: $INTERNAL
+    - uses: actions/setup-python@v4
+      with:
+        python-version: ${{ inputs.python }}
+
+    - name: Checkout code
+      uses: actions/checkout@v3
+
+    - name: Login to s3
+      uses: aws-actions/configure-aws-credentials@v2
+      with:
+        role-to-assume: ${{ secrets.AWS_WEBIDENTITY_FOR_GITHUB_ACTIONS }}
+        aws-region: us-east-1 
+
+    - name: Set Env
+      run: |
+        pip3 install virtualenv
+        virtualenv venv
+        source venv/bin/activate
+
+    - name: Build PyPi Wheel
+      id: build-wheel
+      uses: neuralmagic/nm-actions/actions/pypi_build@main
+      with:
+        dev: ${{ inputs.dev }}
+        release: ${{ inputs.release }}
+        name: ${{ inputs.name }}
+
+    - name: Push to s3 bucket
+      id: push-wheel
+      uses: neuralmagic/nm-actions/actions/s3_push@main
+      with:
+        filename: ${{ inputs.filename }}
+        bucket_name: ${{ inputs.bucket_name }}

.github/workflows/test-wheel-push-to-internal.yml

@@ -0,0 +1,68 @@
+name: Test Wheel and Push to Internal PyPi
+on:
+  workflow_call:
+    inputs:
+      build-label:
+        description: "requested runner label"
+        type: string
+        required: true
+      whl:
+        type: string
+        required: true
+      python:
+        type: string
+
+jobs:
+  test-wheel-and-push-internal:
+    runs-on: ${{ inputs.build-label }}
+    steps:
+    - uses: actions/setup-python@v4
+      with:
+        python-version: ${{ inputs.python }}
+
+    - name: Login to s3
+      uses: aws-actions/configure-aws-credentials@v2
+      with:
+        role-to-assume: ${{ secrets.AWS_WEBIDENTITY_FOR_GITHUB_ACTIONS }}
+        aws-region: us-east-1 
+
+    - name: Make directory for wheel
+      run: |
+          mkdir dist_s3
+          
+    - name: Pull from s3
+      uses: neuralmagic/nm-actions/actions/s3_pull@main
+      with:
+        filename: ${{ inputs.whl }}
+        dst: dist_s3
+
+    - name: Set Env
+      run: |
+        pip3 install virtualenv
+        virtualenv venv
+        source venv/bin/activate
+  
+    - name: Fetch name of whl
+      run: |
+          echo "FILENAME=$(echo dist_s3/*.whl)" >> $GITHUB_ENV
+          
+    - name: Push to internal pypi
+      uses: neuralmagic/nm-actions/actions/nm-upload-whl@main
+      with:
+        server: ${{ secrets.NM_PRIVATE_PYPI_LOCATION }}
+        username: ${{ secrets.NM_PRIVATE_PYPI_USER }}
+        password: ${{ secrets.NM_PRIVATE_PYPI_AUTH }}
+        whl: ./$FILENAME
+        port: 8080
+
+    - name: Install whl
+      run: |
+          pip3 install $FILENAME[dev]
+
+    - name: Checkout code
+      uses: actions/checkout@v3
+
+    - name: Remove src files and run tests
+      run: |
+            rm -rf src
+            make test

integrations/huggingface-transformers/tutorials/sentiment-analysis/docs-sentiment-analysis-python-custom-teacher-rottentomatoes.ipynb

@@ -364,13 +364,11 @@
     "model_kwargs = {\"config\": model_config}\n",
     "model_kwargs[\"state_dict\"], s_delayed = SparseAutoModel._loadable_state_dict(model_path)\n",
     "model = AutoModelForSequenceClassification.from_pretrained(model_path,**model_kwargs,)\n",
-    "SparseAutoModel.log_model_load(model, model_path, \"student\", s_delayed) # prints metrics on sparsity profile\n",
     "\n",
     "# initialize teacher using familiar HF AutoModel\n",
     "teacher_kwargs = {\"config\": teacher_config}\n",
     "teacher_kwargs[\"state_dict\"], t_delayed = SparseAutoModel._loadable_state_dict(teacher_path)\n",
-    "teacher = AutoModelForSequenceClassification.from_pretrained(teacher_path,**teacher_kwargs,)\n",
-    "SparseAutoModel.log_model_load(teacher, teacher_path, \"teacher\", t_delayed)"
+    "teacher = AutoModelForSequenceClassification.from_pretrained(teacher_path,**teacher_kwargs,)"
    ]
   },
   {