Fix tests, adjust release workflow

.github/workflows/release_pypi.yaml

@@ -2,36 +2,143 @@ name: Release PyPi
 
 on:
   workflow_dispatch:
-  release:
-    types:
-      - created
+# release:
+#   types:
+#     - created
+  pull_request:
+    branches: [main]
 
 jobs:
-  build_wheels:
-    name: Build wheels on ${{ matrix.os }}
-    runs-on: ${{ matrix.os }}
+  linux:
+    name: Build ${{ matrix.platform.runner}} ${{ matrix.platform.target }} ${{ matrix.python-version }}
+    runs-on: ${{ matrix.platform.runner }}
     strategy:
       matrix:
-        # FIXME: rather do one job by OS
-        os: [ubuntu-latest, macos-latest, windows-latest]
+        platform:
+          - runner: ubuntu-latest
+            target: x86_64
+          - runner: ubuntu-latest
+            target: aarch64
+        python-version: ["3.9", "3.10", "3.11", "3.12"]
+
     steps:
       - uses: actions/checkout@v3
+
       - name: Set up Python
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
-          python-version: '3.10'
-      # - name: Set up Rust
-      #   uses: dtolnay/rust-toolchain@stable
-      # - uses: Swatinem/rust-cache@v2
-      - name: Build wheels with Maturin
+          python-version: ${{ matrix.python-version }}
+
+      - name: Build wheels
+        uses: PyO3/maturin-action@v1
+        with:
+          target: ${{ matrix.platform.target }}
+          args: --release --out dist --find-interpreter
+          sccache: ${{ !startsWith(github.ref, 'refs/tags/') }}
+          manylinux: auto
+
+      - uses: actions/upload-artifact@v4
+        with:
+          path: dist/*.whl
+          name: wheels-linux-${{ matrix.platform.target }}
+
+      - uses: actions/download-artifact@v4
+        with:
+          name: sdist
+          path: dist
+
+      - run: ls -R dist
+
+  windows:
+    name: Build ${{ matrix.platform.runner}} ${{ matrix.platform.target }} ${{ matrix.python-version }}
+    runs-on: ${{ matrix.platform.runner }}
+    strategy:
+      matrix:
+        platform:
+          - runner: windows-latest
+            target: x64
+          - runner: windows-latest
+            target: x86
+        python-version: ["3.9", "3.10", "3.11", "3.12"]
+
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Build wheels
+        uses: PyO3/maturin-action@v1
+        with:
+          target: ${{ matrix.platform.target }}
+          args: --release --out dist --find-interpreter
+          sccache: ${{ !startsWith(github.ref, 'refs/tags/') }}
+          manylinux: auto
+
+      - uses: actions/upload-artifact@v4
+        with:
+          path: dist/*.whl
+          name: wheels-linux-${{ matrix.platform.target }}
+
+      - uses: actions/download-artifact@v4
+        with:
+          name: sdist
+          path: dist
+
+      - run: ls -R dist
+
+  macos:
+    name: Build ${{ matrix.platform.runner}} ${{ matrix.platform.target }} ${{ matrix.python-version }}
+    runs-on: ${{ matrix.platform.runner }}
+    strategy:
+      matrix:
+        platform:
+          - runner: macos-14
+            target: x86_64
+            macos_version: "14.0"
+          - runner: macos-14
+            target: aarch64
+            macos_version: "14.0"
+          - runner: macos-15
+            target: x86_64
+            macos_version: "15.0"
+          - runner: macos-15
+            target: aarch64
+            macos_version: "15.0"
+        python-version: ["3.9", "3.10", "3.11", "3.12"]
+
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Set macOS version
+        run: echo "MACOSX_DEPLOYMENT_TARGET=${{ matrix.platform.macos_version }}" >> $GITHUB_ENV
+
+      - name: Build wheels
         uses: PyO3/maturin-action@v1
         with:
-          command: build
-          args: --release --out dist # --features python-bindings
+          target: ${{ matrix.platform.target }}
+          args: --release --out dist --find-interpreter
+          sccache: ${{ !startsWith(github.ref, 'refs/tags/') }}
+          manylinux: auto
+
       - uses: actions/upload-artifact@v4
         with:
           path: dist/*.whl
-          name: wheels-${{ matrix.os }}
+          name: wheels-linux-${{ matrix.platform.target }}
+
+      - uses: actions/download-artifact@v4
+        with:
+          name: sdist
+          path: dist
+
+      - run: ls -R dist
 
   build_sdist:
     name: Build source distribution
@@ -52,6 +159,13 @@ jobs:
           path: dist/*.tar.gz
           name: sdist
 
+      - uses: actions/download-artifact@v4
+        with:
+          name: sdist
+          path: dist
+
+      - run: ls -R dist
+
   # release:
   #   name: Release to PyPI
   #   needs: [build_wheels, build_sdist]
@@ -66,6 +180,7 @@ jobs:
   #       with:
   #         name: sdist
   #         path: dist
+
   #     - name: Publish to PyPI
   #       uses: pypa/[email protected]
   #       with:

.github/workflows/tests.yml

@@ -37,10 +37,9 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Set up test environment
       run: |
-        # python -m venv .venv
-        # source .venv/bin/activate
-        # echo "VIRTUAL_ENV=$(pwd)/.venv" >> $GITHUB_ENV
-        # echo "$(pwd)/.venv/bin" >> $GITHUB_PATH
+        python -m venv .venv
+        source .venv/bin/activate
+        python -m pip install --upgrade pip
         pip install .[test]
         maturin develop
     - name: Create matrix id
@@ -54,7 +53,7 @@ jobs:
         echo "::set-output name=id::$MATRIX_ID"
     - name: Run tests
       run: |
-        pytest --cov=outlines_core -vv
+        .venv/bin/python -m pytest --cov=outlines_core -vv
       env:
         COVERAGE_FILE: .coverage.${{ steps.matrix-id.outputs.id }}
     - name: Upload coverage data

pyproject.toml

@@ -42,7 +42,7 @@ test = [
     "scipy",
     "asv",
     "psutil",
-    "maturin",
+    "maturin[patchelf]",
 ]
 
 [project.urls]

src/python_bindings/mod.rs

@@ -338,6 +338,8 @@ pub fn build_regex_from_schema_py(
 
 fn register_child_module(parent_module: &Bound<'_, PyModule>) -> PyResult<()> {
     let m = PyModule::new_bound(parent_module.py(), "json_schema")?;
+    parent_module.add_submodule(&m)?;
+
     m.add("BOOLEAN", json_schema::BOOLEAN)?;
     m.add("DATE", json_schema::DATE)?;
     m.add("DATE_TIME", json_schema::DATE_TIME)?;
@@ -352,16 +354,21 @@ fn register_child_module(parent_module: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add("EMAIL", json_schema::EMAIL)?;
     m.add("URI", json_schema::URI)?;
     m.add_function(wrap_pyfunction!(build_regex_from_schema_py, &m)?)?;
-    parent_module.add_submodule(&m)?;
+
+    let sys = PyModule::import_bound(m.py(), "sys")?;
+    let sys_modules_bind = sys.as_ref().getattr("modules")?;
+    let sys_modules = sys_modules_bind.downcast::<PyDict>()?;
+    sys_modules.set_item("outlines_core.json_schema", &m)?;
+
     Ok(())
 }
 
-// FIXME: there's an unwanted outlines_core.outlines_core module generated by PyO3...
 #[pymodule]
 fn outlines_core(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<PyIndex>()?;
     m.add_class::<PyVocabulary>()?;
     m.add_class::<PyGuide>()?;
     register_child_module(m)?;
+
     Ok(())
 }

tests/test_guide.py

@@ -73,9 +73,9 @@ def test_token_trans_keys_identical():
     guide1 = Guide(index)
     guide2 = Guide(index)
 
-    assert guide1.advance(3) == guide2.advance(3)
+    assert sorted(guide1.advance(3)) == sorted(guide2.advance(3))
     # `a` and `b` have similar transitions to `z`
-    assert guide1.advance(1) == guide2.advance(2)
+    assert sorted(guide1.advance(1)) == sorted(guide2.advance(2))
     assert guide1.advance(3) == guide2.advance(3) == [eos_token_id]
     assert guide1.is_finished()
     assert guide2.is_finished()
@@ -94,9 +94,9 @@ def test_str_and_bytes_produce_the_same():
     guide1 = Guide(index1)
     guide2 = Guide(index2)
 
-    assert guide1.advance(3) == guide2.advance(3)
+    assert sorted(guide1.advance(3)) == sorted(guide2.advance(3))
     # `a` and `b` have similar transitions to `z`
-    assert guide1.advance(1) == guide2.advance(2)
+    assert sorted(guide1.advance(1)) == sorted(guide2.advance(2))
     assert guide1.advance(3) == guide2.advance(3) == [eos_token_id]
     assert guide1.is_finished()
     assert guide2.is_finished()