fixed mlflow tests

.github/workflows/pypi_publish.yml

@@ -1,39 +1,63 @@
 name: Build and Publish Python Package
 
 on:
+  pull_request:
+    paths:
+      - 'src/python/**'
+      - '.github/workflows/pypi_publish.yml'
   release:
     types: [created]
-  workflow_dispatch:  
+  workflow_dispatch:
 
 jobs:
   test:
     if: |
+      github.event_name == 'pull_request' ||
       github.event_name == 'workflow_dispatch' ||
       (github.event_name == 'release' &&
        startsWith(github.ref, 'refs/tags/') &&
        contains(github.ref, 'py_v'))
     runs-on: ubuntu-latest
     strategy:
+      fail-fast: false
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        # 3.11-3.13 run the full suite. 3.10 and 3.14 install only the lean
+        # essential deps and skip the extra tests (`test_extra_*`, which cover
+        # the optional mlflow converter and ML frameworks).
+        include:
+          - python-version: "3.10"
+            extras: test-essential
+            pytest-args: "--ignore-glob='*test_extra_*'"
+          - python-version: "3.11"
+            extras: test
+            pytest-args: ""
+          - python-version: "3.12"
+            extras: test
+            pytest-args: ""
+          - python-version: "3.13"
+            extras: test
+            pytest-args: ""
+          - python-version: "3.14"
+            extras: test-essential
+            pytest-args: "--ignore-glob='*test_extra_*'"
     steps:
       - uses: actions/checkout@v4
-      
+
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
           cache: pip
           cache-dependency-path: '**/pyproject.toml'
-          
+
       - name: Install dependencies
         run: |
-          pip install -e 'src/python[core,test]'
-          
+          pip install -e "src/python[core,${{ matrix.extras }}]"
+
       - name: Run tests
         working-directory: src/python
         run: |
-          python -m unittest discover tests
+          pytest tests ${{ matrix.pytest-args }}
 
   deploy:
     runs-on: ubuntu-latest

.gitignore

@@ -161,6 +161,10 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
+# MLflow local tracking artifacts
+mlruns/
+mlflow.db
+
 spec/schemas/*.json
 !spec/schemas/combined.json
 src/python/tmp.py

src/python/pyproject.toml

@@ -17,10 +17,19 @@ classifiers = [
 core = ["numpy>=1.0.0,<3.0.0", "onnxruntime>=1.0.0,<2.0.0"]
 extras = ["pydantic>=2.0.0,<3.0.0", "pyyaml>=5.1"]
 mlflow = ["mlflow>=3.4"]
-test = [
+# Essential test deps: only the pure FNNX core (format/runtime/validators),
+# no mlflow converter and no ML frameworks. Drives the 3.10 / 3.14 CI legs,
+# which skip the extra suite (the `test_extra_*` files); see the test workflow.
+test-essential = [
     "jsonschema>=4.0.0,<5.0.0",
     "pydantic>=2.0.0,<3.0.0",
     "pyyaml>=5.1",
+    "pytest>=9.1.1",
+]
+# Full test deps: essential + the mlflow converter (extra) and every ML
+# framework it integrates with. Drives the 3.11-3.13 CI legs.
+test = [
+    "fnnx[test-essential]",
     "mlflow>=3.4",
     "scikit-learn>=1.0.0",
     "pandas>=1.0.0",
@@ -33,5 +42,7 @@ test = [
     "langgraph>=0.2.0",
     "prophet>=1.1.0",
     "tensorflow>=2.12.0",
-    "pytest>=9.1.1",
 ]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

src/python/tests/test_mlflow_torch.py → src/python/tests/test_extra_mlflow_torch.py

@@ -73,11 +73,12 @@ def _save_pytorch_model(tmp: str, *, with_code_paths: bool):
     )
 
     model_dir = os.path.join(tmp, "torch_model")
-    # mlflow.pytorch always serializes via torch.save (pickle, by-reference),
-    # which is exactly what exercises the by-reference self-containment
-    # heuristic. (There is no `serialization_format` arg on this flavor — it
-    # belongs to mlflow.sklearn — and passing it raises in torch.save.)
-    kwargs: dict = {"signature": signature}
+    # Force the pickle (torch.save, by-reference) serialization format, which
+    # is exactly what exercises the by-reference self-containment heuristic.
+    # mlflow >= 3 defaults this flavor to the self-contained 'pt2' traced-graph
+    # format (which additionally requires an input_example), so the default
+    # would bypass the path under test.
+    kwargs: dict = {"signature": signature, "serialization_format": "pickle"}
     if with_code_paths:
         kwargs["code_paths"] = [os.path.join(_FIXTURES_DIR, "torch_net.py")]
     mlflow.pytorch.save_model(model, model_dir, **kwargs)  # type: ignore[attr-defined]