Centralize ONNX export warning suppression in export_pytorch()

src/winml/modelkit/_warnings.py

@@ -81,10 +81,9 @@ def filter(self, record: logging.LogRecord) -> bool:
     # NOTE: TracerWarning (from torch.jit) is intentionally NOT filtered here.
     # Importing torch.jit at startup would pull all of torch (~1.7s) into
     # `winml --help` and violate the CLI import budget (tests/cli/test_import_time.py).
-    # During ONNX export, build.py already wraps the export call in
+    # During ONNX export, export_pytorch() wraps torch.onnx.export in
     # `warnings.catch_warnings()` + `filterwarnings("ignore")`, which is strictly
-    # broader than a TracerWarning-only filter. Direct callers of export_pytorch()
-    # that want the same suppression can apply it locally at the call site.
+    # broader than a TracerWarning-only filter.
 
     # Diffusers
     warnings.filterwarnings(

src/winml/modelkit/commands/build.py

@@ -1047,27 +1047,21 @@ def _name(base: str) -> str:
     current_path = export_path
 
     # ── Export stage ──────────────────────────────────────────────
-    import warnings
-
     with StageLive("export", console) as sl:
         sl.set_status("Exporting to ONNX...")
 
         # Load + export (blocking)
-        # Suppress TracerWarning and other transformer warnings
-        # during export to keep Live display clean.
         pytorch_model = _load_model(config, model_id, trust_remote_code=False)
         t0 = time.monotonic()
-        with warnings.catch_warnings():
-            warnings.filterwarnings("ignore")
-            export_onnx(
-                model=pytorch_model,
-                output_path=export_path,
-                export_config=config.export,
-                model_id=model_label,
-                task=config.loader.task,
-                verbose=False,
-                use_external_data=True,
-            )
+        export_onnx(
+            model=pytorch_model,
+            output_path=export_path,
+            export_config=config.export,
+            model_id=model_label,
+            task=config.loader.task,
+            verbose=False,
+            use_external_data=True,
+        )
         _export_elapsed = time.monotonic() - t0
         sl.set_done(_export_elapsed)
         # Meta shown after export completes (avoids duplicate in Live frame)

src/winml/modelkit/export/pytorch.py

@@ -24,6 +24,7 @@
 from __future__ import annotations
 
 import logging
+import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
@@ -83,11 +84,13 @@ def export_pytorch(
         enable_reporting=enable_reporting,
         embed_hierarchy_attributes=export_config.enable_hierarchy_tags,
     )
-    return exporter.export(
-        model=model,
-        output_path=str(output_path),
-        export_config=export_config,
-        model_name_or_path=model_name_or_path,
-        task=task,
-        **kwargs,
-    )
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore")
+        return exporter.export(
+            model=model,
+            output_path=str(output_path),
+            export_config=export_config,
+            model_name_or_path=model_name_or_path,
+            task=task,
+            **kwargs,
+        )

tests/unit/test_pytorch_warning_suppression.py

@@ -0,0 +1,85 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Tests warning suppression behavior in export_pytorch()."""
+
+from __future__ import annotations
+
+import sys
+import types
+import warnings
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+from winml.modelkit.export.pytorch import export_pytorch
+
+
+def _install_fake_exporter(
+    monkeypatch,
+    export_return: dict[str, int],
+    *,
+    emit_warning: bool,
+) -> MagicMock:
+    exporter_cls = MagicMock()
+    exporter = exporter_cls.return_value
+
+    def _export(*args, **kwargs):
+        if emit_warning:
+            warnings.warn("tracer warning noise", UserWarning, stacklevel=2)
+        return export_return
+
+    exporter.export.side_effect = _export
+
+    fake_htp_pkg = types.ModuleType("winml.modelkit.export.htp")
+    fake_htp_pkg.__path__ = []  # mark as package
+    fake_exporter_module = types.ModuleType("winml.modelkit.export.htp.exporter")
+    fake_exporter_module.HTPExporter = exporter_cls
+
+    monkeypatch.setitem(sys.modules, "winml.modelkit.export.htp", fake_htp_pkg)
+    monkeypatch.setitem(sys.modules, "winml.modelkit.export.htp.exporter", fake_exporter_module)
+    return exporter
+
+
+def test_export_pytorch_uses_warning_context(tmp_path, monkeypatch) -> None:
+    """export_pytorch should use warnings.catch_warnings() context manager."""
+
+    class DummyModel:
+        pass
+
+    model = DummyModel()
+    config = SimpleNamespace(enable_hierarchy_tags=True)
+    expected = {"onnx_nodes": 1}
+    exporter = _install_fake_exporter(monkeypatch, expected, emit_warning=False)
+
+    with (
+        patch("winml.modelkit.export.pytorch.warnings.catch_warnings") as mock_catch_warnings,
+        patch("winml.modelkit.export.pytorch.warnings.filterwarnings") as mock_filterwarnings,
+    ):
+        result = export_pytorch(model, tmp_path / "model.onnx", config)
+
+    assert result == expected
+    exporter.export.assert_called_once()
+    mock_catch_warnings.assert_called_once_with()
+    mock_catch_warnings.return_value.__enter__.assert_called_once_with()
+    mock_filterwarnings.assert_called_once_with("ignore")
+
+
+def test_export_pytorch_suppresses_export_warnings(tmp_path, monkeypatch) -> None:
+    """Warnings emitted during export should not leak to callers."""
+
+    class DummyModel:
+        pass
+
+    model = DummyModel()
+    config = SimpleNamespace(enable_hierarchy_tags=True)
+    expected = {"onnx_nodes": 1}
+    exporter = _install_fake_exporter(monkeypatch, expected, emit_warning=True)
+
+    with warnings.catch_warnings(record=True) as captured:
+        warnings.simplefilter("always")
+        result = export_pytorch(model, tmp_path / "model.onnx", config)
+
+    assert result == expected
+    exporter.export.assert_called_once()
+    assert captured == []