fix(tensorrt): hash UNet engine cache key so ONNX export path stays under Windows MAX_PATH

src/streamdiffusion/acceleration/tensorrt/engine_manager.py

@@ -2,7 +2,7 @@
 import logging
 from enum import Enum
 from pathlib import Path
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, cast
 
 logger = logging.getLogger(__name__)
 
@@ -125,7 +125,12 @@ def get_engine_path(
         Moves and consolidates create_prefix() function from wrapper.py lines 995-1014.
         Special handling for ControlNet engines which use model_id-based directories.
         """
-        filename = self._configs[engine_type]["filename"]
+        # self._configs' inner dicts also hold compile_fn/loader callables under other
+        # keys, so the type checker infers "filename" as str | Callable | Callable across
+        # the whole dict shape rather than narrowing per-key. The "filename" value is
+        # always a plain str at runtime (see __init__); cast pins that down for the
+        # `self.engine_dir / prefix / filename` join below.
+        filename = cast(str, self._configs[engine_type]["filename"])
         optlvl_suffix = f"--optlvl{builder_optimization_level}" if builder_optimization_level is not None else ""
 
         if engine_type == EngineType.CONTROLNET:
@@ -216,7 +221,37 @@ def get_engine_path(
             if resolution is not None:
                 prefix += f"--res-{resolution[0]}x{resolution[1]}"
 
-            return self.engine_dir / prefix / filename
+            if engine_type == EngineType.UNET:
+                # The UNet prefix above encodes every build flag verbatim and can exceed
+                # ~170 chars. Combined with a real-world engine_dir root, that pushes
+                # unet.engine.onnx / .opt.onnx past Windows' 260-char MAX_PATH — the
+                # directory itself fits (mkdir succeeds) but torch.onnx.export's
+                # open(path, "wb") then fails with a bare FileNotFoundError. Compact the
+                # directory name to a short hash of the full prefix (same idea as
+                # _lora_signature) so every distinct config still maps to a distinct,
+                # stable directory, just without spelling out every flag on disk.
+                canonical = prefix
+                short_hash = hashlib.sha1(canonical.encode("utf-8")).hexdigest()[:12]
+                fp8_tag = "--fp8v3" if fp8 else ""
+                res_tag = f"--res-{resolution[0]}x{resolution[1]}" if resolution is not None else ""
+                prefix = f"{base_name}{fp8_tag}--h{short_hash}{res_tag}"
+                logger.debug(f"EngineManager: UNet cache key h{short_hash} = {canonical}")
+
+            engine_path = self.engine_dir / prefix / filename
+
+            # Belt-and-suspenders: warn (rather than silently mkdir-then-fail-on-write)
+            # if the longest artifact derived from this path would still exceed Windows'
+            # 260-char MAX_PATH — e.g. an unusually deep user-configured engine_dir.
+            # ".opt.onnx" is the longest suffix appended to engine_path (see builder.py).
+            longest_derived = len(str(engine_path)) + len(".opt.onnx")
+            if longest_derived >= 250:
+                logger.warning(
+                    f"EngineManager: engine path is {longest_derived} chars once .opt.onnx is "
+                    f"appended (Windows MAX_PATH=260). Move engine_dir closer to the drive root "
+                    f"if TensorRT export fails with FileNotFoundError: {engine_path}"
+                )
+
+            return engine_path
 
     def _get_embedding_dim_for_model_type(self, model_type: str) -> int:
         """Get embedding dimension based on model type."""

tests/unit/test_engine_path_length.py

@@ -0,0 +1,98 @@
+"""
+Regression test for the UNet TensorRT engine path exceeding Windows MAX_PATH (260).
+
+``EngineManager.get_engine_path`` (acceleration/tensorrt/engine_manager.py) encodes every
+UNet build flag into the on-disk directory name. With a realistic ``engine_dir`` and a
+config using fp8 + static batch + pin_cache_frames + optlvl + resolution, the generated
+directory was 246 chars; the derived ``unet.engine.onnx`` path was 263 and
+``unet.engine.opt.onnx`` was 267 — both over Windows' 260-char MAX_PATH. Because the
+*directory* fit but the *file* didn't, ``mkdir`` silently succeeded while
+``torch.onnx.export`` -> ``open(onnx_path, "wb")`` raised ``FileNotFoundError``, which
+wrapper.py's OOM-only fallback does not catch — see
+"Acceleration has failed: [Errno 2] No such file or directory: ...unet.engine.onnx".
+
+This test constructs ``EngineManager`` via ``__new__`` (bypassing ``__init__``'s compile-fn
+imports, which need TensorRT/onnx/polygraphy installed) so it only exercises the pure-path
+logic in ``get_engine_path``.
+
+Run with: pytest tests/unit/test_engine_path_length.py -v
+"""
+
+from pathlib import Path
+
+import pytest
+
+try:
+    from streamdiffusion.acceleration.tensorrt.engine_manager import EngineManager, EngineType
+
+    IMPORT_OK = True
+except ImportError:
+    IMPORT_OK = False
+
+pytestmark = pytest.mark.skipif(
+    not IMPORT_OK,
+    reason="acceleration.tensorrt.engine_manager not importable",
+)
+
+# Mirrors the crash report: a nested Documents path, matching the real-world depth
+# that pushed the UNet engine path over MAX_PATH.
+_CRASH_ENGINE_DIR = r"C:\Users\deswh\Documents\sdtd040\StreamDiffusion\engines\td"
+
+# The exact flag combination from the crash log's UNet directory name.
+_CRASH_UNET_KWARGS = {
+    "engine_type": EngineType.UNET,
+    "model_id_or_path": "stabilityai/sd-turbo",
+    "max_batch_size": 4,
+    "min_batch_size": 1,
+    "mode": "img2img",
+    "use_tiny_vae": True,
+    "fp8": True,
+    "use_cached_attn": False,
+    "use_feature_injection": False,
+    "build_static_batch": True,
+    "static_batch_size": 2,
+    "pin_cache_frames": True,
+    "cache_maxframes": 4,
+    "builder_optimization_level": 4,
+    "resolution": (512, 512),
+}
+
+
+def _make_engine_manager(engine_dir: str) -> EngineManager:
+    """Build an EngineManager without running __init__'s heavy compile-fn imports."""
+    em = EngineManager.__new__(EngineManager)
+    em.engine_dir = Path(engine_dir)
+    em._configs = {EngineType.UNET: {"filename": "unet.engine"}}
+    return em
+
+
+class TestUnetEnginePathLength:
+    def test_onnx_export_paths_stay_under_max_path(self):
+        """RED today (267 > 260 for .opt.onnx); GREEN once the dir name is hashed."""
+        em = _make_engine_manager(_CRASH_ENGINE_DIR)
+        engine_path = em.get_engine_path(**_CRASH_UNET_KWARGS)
+
+        onnx_path = str(engine_path) + ".onnx"
+        opt_onnx_path = str(engine_path) + ".opt.onnx"
+
+        assert len(onnx_path) < 260, f"onnx path is {len(onnx_path)} chars (MAX_PATH=260): {onnx_path}"
+        assert len(opt_onnx_path) < 260, f"opt.onnx path is {len(opt_onnx_path)} chars (MAX_PATH=260): {opt_onnx_path}"
+
+    def test_engine_path_is_deterministic(self):
+        """Same config -> same path, so a previously-built engine is still found on rebuild."""
+        em = _make_engine_manager(_CRASH_ENGINE_DIR)
+        path_a = em.get_engine_path(**_CRASH_UNET_KWARGS)
+        path_b = em.get_engine_path(**_CRASH_UNET_KWARGS)
+
+        assert path_a == path_b
+
+    def test_distinct_configs_do_not_collide(self):
+        """A differing flag (static_batch_size) must still produce a distinct directory."""
+        em = _make_engine_manager(_CRASH_ENGINE_DIR)
+        kwargs_b = dict(_CRASH_UNET_KWARGS, static_batch_size=4)
+
+        path_a = em.get_engine_path(**_CRASH_UNET_KWARGS)
+        path_b = em.get_engine_path(**kwargs_b)
+
+        assert path_a != path_b
+        assert path_a.parent != path_b.parent