From 07955f85f6f745a04e92a4a627d2026279b1838f Mon Sep 17 00:00:00 2001
From: Richard Edgar <riedgar@microsoft.com>
Date: Fri, 10 May 2024 13:06:24 -0400
Subject: [PATCH] [Release] Add smoke test of wheels (#813)

Add some basic testing of the wheels built during the release workflow
---
 .github/workflows/pypi_upload.yml | 43 ++++++++++++++++++++++++++++---
 1 file changed, 39 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/pypi_upload.yml b/.github/workflows/pypi_upload.yml
index 81f7023cc..02c63edb2 100644
--- a/.github/workflows/pypi_upload.yml
+++ b/.github/workflows/pypi_upload.yml
@@ -111,10 +111,9 @@ jobs:
           name: sdist_files
           path: dist/*.tar.gz
 
-  publish_wheels:
-    permissions:
-      id-token: write
-    name: Publish wheels on pypi
+
+  assemble_wheels:
+    name: Combine wheels
     needs: [build_wheels, build_sdist]
     runs-on: ubuntu-latest
     steps:
@@ -134,6 +133,42 @@ jobs:
           path: ./dist/*
           name: collected_dist_files
 
+  test_wheels:
+    name: Test Wheels
+    needs: [assemble_wheels]
+    strategy:
+      matrix:
+        os: [ubuntu-latest, windows-latest, macos-12, macos-latest]
+        python-version: ["3.11", "3.12"]
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: collected_dist_files
+          path: wheelhouse
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - run: pip install guidance -f ./wheelhouse/
+        name: Install guidance from wheels
+      - run: pip install transformers torch
+        name: Other installs
+      - run: python -c "import guidance; lm = guidance.models.Transformers('gpt2'); lm += '1,2,3,4,5,' + guidance.gen('num', max_tokens=5, temperature=0); print(f'\n{str(lm)=}\n'); assert lm['num'].startswith('6')"
+        name: Run smoke test
+
+  publish_wheels:
+    permissions:
+      id-token: write
+    name: Publish wheels on pypi
+    needs: [test_wheels]
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: collected_dist_files
+          path: dist
+
       - name: Publish package to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         if: startsWith(github.ref, 'refs/tags')