From 9765c8d017ab145b5ff622db7f8918d95589edca Mon Sep 17 00:00:00 2001 From: RomirJ Date: Sat, 20 Jun 2026 14:05:35 -0700 Subject: [PATCH] fix: add mypy hook and SmolVLA export guards --- .pre-commit-config.yaml | 5 +++ pyproject.toml | 21 ++++++++++ src/tether/exporters/monolithic.py | 61 ++++++++++++++++++++++++++--- tests/test_cli_export_end_to_end.py | 48 +++++++++++++++++++++++ 4 files changed, 130 insertions(+), 5 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1ab448dd..60b83bf7 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -12,3 +12,8 @@ repos: - id: end-of-file-fixer - id: check-yaml - id: check-added-large-files +- repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.10.0 + hooks: + - id: mypy + args: [ --config-file=pyproject.toml ] diff --git a/pyproject.toml b/pyproject.toml index d58bccd8..2d8c24c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -324,6 +324,27 @@ artifacts = ["*.cpp", "*.cu", "*.txt", "*.json"] target-version = "py310" line-length = 100 +[tool.mypy] +python_version = "3.10" +mypy_path = "src" +ignore_missing_imports = true +follow_imports = "silent" +check_untyped_defs = false +warn_unused_configs = true +exclude = [ + "src/tether/models/third_party/", +] + +[[tool.mypy.overrides]] +module = "tether.exporters.monolithic" +disable_error_code = [ + "assignment", + "attr-defined", + "method-assign", + "misc", + "no-redef", +] + [tool.pytest.ini_options] testpaths = ["tests"] # pytest-asyncio v1.x requires explicit mode declaration. Tests use diff --git a/src/tether/exporters/monolithic.py b/src/tether/exporters/monolithic.py index f1006704..052baf5a 100644 --- a/src/tether/exporters/monolithic.py +++ b/src/tether/exporters/monolithic.py @@ -26,6 +26,7 @@ import json import logging import time +from importlib.metadata import PackageNotFoundError, version as _dist_version from pathlib import Path from typing import Any @@ -64,6 +65,8 @@ "smolvla": "HuggingFaceTB/SmolLM2-135M", } +_SMOLVLA_EXPORT_PATCH_FAILURES: list[str] = [] + def _quiet_noisy_export_loggers() -> None: """Keep streamed export logs readable while preserving warnings/errors.""" @@ -120,7 +123,7 @@ def _require_monolithic_deps() -> None: """Check that the ``[monolithic]`` optional dep group is installed. Raises ImportError with a clean message if anything's missing or a - transformers version mismatch is detected (5.4+ has the q_length bug). + required package version mismatch is detected. """ _quiet_noisy_export_loggers() missing = [] @@ -136,8 +139,25 @@ def _require_monolithic_deps() -> None: except ImportError as e: missing.append(f"transformers==5.3.0 ({e})") + try: + __import__("lerobot") + try: + lerobot_version = _dist_version("lerobot") + except PackageNotFoundError as exc: + raise ImportError( + "lerobot imports but package metadata is missing; install with: " + "pip install 'fastcrest-tether[monolithic]'" + ) from exc + if lerobot_version != "0.5.1": + raise ImportError( + f"lerobot {lerobot_version} detected; monolithic export requires " + "exactly lerobot==0.5.1. Install with: " + "pip install 'fastcrest-tether[monolithic]'" + ) + except ImportError as e: + missing.append(f"lerobot==0.5.1 ({e})") + for mod_name, pip_name in [ - ("lerobot", "lerobot==0.5.1"), ("onnx_diagnostic", "onnx-diagnostic>=0.9"), ("onnxscript", "onnxscript>=0.1"), ("optree", "optree"), @@ -156,6 +176,29 @@ def _require_monolithic_deps() -> None: ) +def _record_smolvla_export_patch_failure(patch_name: str, exc: BaseException) -> None: + message = f"{patch_name}: {type(exc).__name__}: {exc}" + _SMOLVLA_EXPORT_PATCH_FAILURES.append(message) + logger.warning( + "[smolvla] export patch failed; SmolVLA monolithic export will stop " + "before torch.export instead of failing later: %s", + message, + ) + + +def _raise_if_smolvla_export_patches_failed() -> None: + if not _SMOLVLA_EXPORT_PATCH_FAILURES: + return + details = "\n - ".join(_SMOLVLA_EXPORT_PATCH_FAILURES) + raise RuntimeError( + "SmolVLA monolithic export patches failed to install. This usually " + "means the lerobot/transformers export stack does not match the pinned " + "monolithic environment; install with: pip install " + "'fastcrest-tether[monolithic]'.\n - " + + details + ) + + def apply_export_patches() -> None: """Install the full set of transformers + lerobot patches required for the monolithic `torch.export` path. Safe to call multiple times; later @@ -168,6 +211,8 @@ def apply_export_patches() -> None: import torch + _SMOLVLA_EXPORT_PATCH_FAILURES.clear() + # Stub GR00T imports to avoid Python 3.13 dataclass issue (harmless on 3.12 too) for _mod in ("lerobot.policies.groot.groot_n1", "lerobot.policies.groot.modeling_groot"): if _mod not in sys.modules: @@ -255,7 +300,10 @@ def _embed_image_with_explicit_patch_mask(self, image): _embed_image_with_explicit_patch_mask._tether_patched = True _smolvla.SmolVLMWithExpertModel.embed_image = _embed_image_with_explicit_patch_mask except Exception as exc: - logger.debug("SmolVLA explicit patch-mask export patch not installed: %s", exc) + _record_smolvla_export_patch_failure( + "SmolVLA explicit patch-mask export patch", + exc, + ) try: from transformers.models.smolvlm import modeling_smolvlm as _smolvlm @@ -294,7 +342,10 @@ def _vision_attention_forward_no_dense_mask(self, hidden_states, attention_mask= _vision_attention_forward_no_dense_mask._tether_patched = True _smolvlm.SmolVLMVisionAttention.forward = _vision_attention_forward_no_dense_mask except Exception as exc: - logger.debug("SmolVLM vision attention export patch not installed: %s", exc) + _record_smolvla_export_patch_failure( + "SmolVLM vision attention export patch", + exc, + ) # DynamicCache deepcopy bypass (FakeTensor can't be deepcopied) _orig_deepcopy = copy.deepcopy @@ -506,7 +557,6 @@ def _fix_onnx_where_dtype_mismatches(onnx_path: Path) -> int: for init in shape_info.graph.initializer: name_dtype[init.name] = init.data_type - FLOAT_TYPES = {TensorProto.FLOAT, TensorProto.FLOAT16, TensorProto.BFLOAT16, TensorProto.DOUBLE} INT_TYPES = {TensorProto.INT64, TensorProto.INT32, TensorProto.INT16, TensorProto.INT8} fixes = 0 @@ -589,6 +639,7 @@ def export_smolvla_monolithic( from onnx_diagnostic.torch_export_patches import torch_export_patches apply_export_patches() + _raise_if_smolvla_export_patches_failed() logger.info("[smolvla] Loading %s", model_id) from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy diff --git a/tests/test_cli_export_end_to_end.py b/tests/test_cli_export_end_to_end.py index 27cd9466..ee0b34a6 100644 --- a/tests/test_cli_export_end_to_end.py +++ b/tests/test_cli_export_end_to_end.py @@ -17,6 +17,10 @@ """ from __future__ import annotations +import logging +import sys +import types + import pytest @@ -63,3 +67,47 @@ def test_dep_check_catches_wrong_transformers(monkeypatch): monkeypatch.setattr(transformers, "__version__", "4.99.0") with pytest.raises(ImportError, match="5.3.0"): monolithic._require_monolithic_deps() + + +def test_dep_check_catches_wrong_lerobot(monkeypatch): + """_require_monolithic_deps() must fail early on the known-bad + lerobot 0.4.x stack instead of letting torch.export fail downstream.""" + from tether.exporters import monolithic + import transformers + + monkeypatch.setattr(transformers, "__version__", "5.3.0") + for mod_name in ("lerobot", "onnx_diagnostic", "onnxscript", "optree", "scipy"): + monkeypatch.setitem(sys.modules, mod_name, types.ModuleType(mod_name)) + + def _fake_dist_version(dist_name: str) -> str: + if dist_name == "lerobot": + return "0.4.4" + raise monolithic.PackageNotFoundError(dist_name) + + monkeypatch.setattr(monolithic, "_dist_version", _fake_dist_version) + + with pytest.raises(ImportError, match=r"lerobot 0\.4\.4.*lerobot==0\.5\.1"): + monolithic._require_monolithic_deps() + + +def test_smolvla_export_patch_failure_is_warning_then_fatal(caplog): + """SmolVLA-specific export patches should not hide at debug level. + + A failed patch is logged immediately and converted into a clear RuntimeError + before SmolVLA torch.export can hit a cryptic FakeTensor shape error. + """ + from tether.exporters import monolithic + + monolithic._SMOLVLA_EXPORT_PATCH_FAILURES.clear() + try: + with caplog.at_level(logging.WARNING, logger="tether.exporters.monolithic"): + monolithic._record_smolvla_export_patch_failure( + "SmolVLA explicit patch-mask export patch", + RuntimeError("boom"), + ) + + assert any("export patch failed" in rec.message for rec in caplog.records) + with pytest.raises(RuntimeError, match="SmolVLA monolithic export patches failed"): + monolithic._raise_if_smolvla_export_patches_failed() + finally: + monolithic._SMOLVLA_EXPORT_PATCH_FAILURES.clear()