Scan for CVEs in wheels, fix python versions

.github/infra_tests/test_parse_pr_config.py

@@ -20,7 +20,7 @@ def test_get_default_config(self):
         """Test default configuration."""
         config = parse_pr_config.get_default_config()
         assert config["build"] is True
-        assert config["test"] == ["ops", "benchmark"]
+        assert config["test"] == ["ops", "benchmark", "sanity"]
 
     def test_extract_yaml_from_pr_body(self):
         """Test YAML extraction from PR body."""
@@ -67,4 +67,4 @@ def test_resolve_config_empty_pr_body(self):
         """Test config resolution with empty PR body."""
         config = parse_pr_config.resolve_config("")
         assert config["build"] is True  # defaults
-        assert config["test"] == ["ops", "benchmark"]
+        assert config["test"] == ["ops", "benchmark", "sanity"]

.github/pull_request_template.md

@@ -10,7 +10,7 @@
 ```yaml
 config:
   build: true
-  # valid options are "ops" and "benchmark"
+  # valid options are "ops", "benchmark", and "sanity"
   test: []
 ```
 

.github/scripts/parse_pr_config.py

@@ -20,7 +20,7 @@
 
 def get_default_config():
     """Return default CI configuration."""
-    return {"build": True, "test": ["ops", "benchmark"]}
+    return {"build": True, "test": ["ops", "benchmark", "sanity"]}
 
 
 def extract_yaml_from_pr_body(pr_body):
@@ -77,6 +77,7 @@ def write_github_outputs(config):
         "build": str(config["build"]).lower(),
         "run_ops": str("ops" in test_list).lower(),
         "run_benchmark": str("benchmark" in test_list).lower(),
+        "run_sanity": str("sanity" in test_list).lower(),
     }
 
     # Write outputs
@@ -85,7 +86,7 @@ def write_github_outputs(config):
 
     # Log final config
     print(
-        f"Config: build={outputs['build']}, run_ops={outputs['run_ops']}, run_benchmark={outputs['run_benchmark']}",
+        f"Config: build={outputs['build']}, run_ops={outputs['run_ops']}, run_benchmark={outputs['run_benchmark']}, run_sanity={outputs['run_sanity']}",
         file=sys.stderr,
     )
 

.github/workflows/build-wheel.yml

@@ -51,6 +51,11 @@ on:
         required: false
         type: string
         default: ''
+      run-pip-audit:
+        description: 'Run pip-audit to check for known CVEs in dependencies (disabled by default)'
+        required: false
+        type: boolean
+        default: false
       pre-build-script:
         description: 'Optional path to a script (relative to repo root) to run before building the wheel (e.g. copy deps). Run with: python <path>.'
         required: false
@@ -103,6 +108,31 @@ jobs:
           fi
           twine check dist/*
 
+      - name: Audit dependencies for CVEs
+        if: ${{ inputs.run-pip-audit }}
+        run: |
+          pip install pip-audit
+          # Extract installable deps from wheel metadata.
+          # Skips git+ deps and extras-only deps which can't be resolved in a plain environment.
+          python - <<'EOF'
+          import zipfile, glob
+          whl = glob.glob('dist/*.whl')[0]
+          with zipfile.ZipFile(whl) as z:
+              meta = next(n for n in z.namelist() if n.endswith('/METADATA'))
+              lines = z.read(meta).decode().splitlines()
+          deps = [
+              l.split('Requires-Dist: ')[1].split(';')[0].strip()
+              for l in lines
+              if l.startswith('Requires-Dist:')
+              and 'extra ==' not in l
+              and '@' not in l.split('Requires-Dist: ')[1]
+          ]
+          with open('audit-requirements.txt', 'w') as f:
+              f.write('\n'.join(deps))
+          print('Auditing deps:\n' + '\n'.join(deps))
+          EOF
+          pip-audit -r audit-requirements.txt
+
       - name: Test import
         if: ${{ !inputs.skip-import-test }}
         run: |

.github/workflows/tilegym-ci.yml

@@ -32,6 +32,7 @@ jobs:
       build: ${{ steps.parse.outputs.build }}
       run_ops: ${{ steps.parse.outputs.run_ops }}
       run_benchmark: ${{ steps.parse.outputs.run_benchmark }}
+      run_sanity: ${{ steps.parse.outputs.run_sanity }}
       image_tag: ${{ steps.parse.outputs.image_tag }}
       image_name: ${{ steps.parse.outputs.image_name }}
       is_pr: ${{ steps.context.outputs.is_pr }}
@@ -144,11 +145,13 @@ jobs:
               echo "build=false" >> $GITHUB_OUTPUT
               echo "run_ops=false" >> $GITHUB_OUTPUT
               echo "run_benchmark=false" >> $GITHUB_OUTPUT
+              echo "run_sanity=false" >> $GITHUB_OUTPUT
             else
               echo "🔨 Building new image and running tests"
               echo "build=true" >> $GITHUB_OUTPUT
               echo "run_ops=true" >> $GITHUB_OUTPUT
               echo "run_benchmark=true" >> $GITHUB_OUTPUT
+              echo "run_sanity=true" >> $GITHUB_OUTPUT
             fi
           fi
 
@@ -168,6 +171,7 @@ jobs:
       architectures: '["x86_64", "arm64"]'  # Build for both architectures (6 wheels total)
       retention-days: 7  # All wheels kept for 7 days; only tested wheel gets -verified (30 days)
       skip-import-test: true  # TileGym requires CUDA, test in Docker instead
+      run-pip-audit: true
       # Optional: Override default runners (ubuntu-latest for x86_64, ubuntu-24.04-arm for arm64)
       # runner-x86-64: [self-hosted, linux, x64, gpu]
 
@@ -249,6 +253,32 @@ jobs:
             type=registry,ref=ghcr.io/${{ steps.vars.outputs.owner_lower }}/tilegym:latest
           cache-to: type=gha,mode=max
 
+  sanity-check:
+    name: sanity-check
+    needs: [config, build-wheel]
+    if: |
+      always() &&
+      needs.config.outputs.run_sanity == 'true' &&
+      (needs.build-wheel.result == 'success' || needs.build-wheel.result == 'skipped')
+    runs-on: ubuntu-latest
+    steps:
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+
+      - name: Download wheel (Python 3.10, x86_64)
+        uses: actions/download-artifact@v4
+        with:
+          name: ${{ needs.build-wheel.outputs.artifact-name }}-py310-x86_64
+          path: ./wheel
+
+      - name: pip check
+        run: |
+          python -m venv /tmp/sanity-env
+          /tmp/sanity-env/bin/pip install --quiet ./wheel/*.whl
+          /tmp/sanity-env/bin/pip check
+
   test-ops:
     name: test-ops
     needs: [config, build]
@@ -537,10 +567,11 @@ jobs:
 
   publish-wheel:
     name: publish-verified-wheel
-    needs: [config, build-wheel, test-ops, test-benchmark]
+    needs: [config, build-wheel, sanity-check, test-ops, test-benchmark]
     if: |
       always() &&
       needs.build-wheel.result == 'success' &&
+      needs.sanity-check.result == 'success' &&
       needs.test-ops.result == 'success' &&
       needs.test-benchmark.result == 'success'
     # Note: Only marks the py310-x86_64 wheel as "verified" because that's the wheel
@@ -555,11 +586,12 @@ jobs:
 
   promote-to-latest:
     name: promote-to-latest
-    needs: [config, build, test-ops, test-benchmark]
+    needs: [config, build, sanity-check, test-ops, test-benchmark]
     if: |
       always() &&
       needs.config.outputs.is_pr == 'false' &&
       needs.build.result == 'success' &&
+      needs.sanity-check.result == 'success' &&
       needs.test-ops.result == 'success' &&
       needs.test-benchmark.result == 'success'
     runs-on: ubuntu-latest

setup.py

@@ -16,17 +16,19 @@
     description="TileGym",
     long_description=README,
     long_description_content_type="text/markdown",
-    url="https://github.com/NVIDIA/tilegym",
+    url="https://github.com/NVIDIA/TileGym",
     packages=setuptools.find_packages(where="src"),
     package_dir={"": "src"},
     license="MIT",
     classifiers=[
         "Development Status :: 4 - Beta",
         "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
         "Operating System :: OS Independent",
     ],
-    python_requires=">=3.9",
+    python_requires=">=3.10",
     install_requires=[
         # Note: torch and triton should be pre-installed in your environment
         "transformers==4.56.2",
@@ -39,6 +41,8 @@
         "numpy",
         "cuda-tile",
         "cuda-tile-experimental @ git+https://github.com/NVIDIA/cutile-python.git#subdirectory=experimental",
+        "filelock>=3.20.3",  # CVE fix: GHSA-w853-jp5j-5j7f, GHSA-qmgc-5h2g-mvrw
+        "pillow>=12.1.1",  # CVE fix: GHSA-cfh3-3jmp-rvhc
         # 'nvidia-ml-py', # optional
     ],
     extras_require={