From ea47b7f5ab5c1d16d41a843e357103429398f4b0 Mon Sep 17 00:00:00 2001 From: Alex Date: Sat, 18 Jul 2026 19:09:50 -0400 Subject: [PATCH 1/2] fix: hash UNet engine cache key so ONNX export path stays under Windows MAX_PATH EngineManager.get_engine_path encoded every UNet build flag verbatim into the on-disk directory name. With a realistic engine_dir (fp8 + static batch + pin_cache_frames + optlvl + resolution), the folder grew to 246 chars, and the derived unet.engine.onnx / unet.engine.opt.onnx paths hit 263/267 chars -- over Windows' 260-char MAX_PATH. Because the *directory* fit but the *file* didn't, mkdir succeeded while torch.onnx.export's open(onnx_path, "wb") raised FileNotFoundError, which wrapper.py's OOM-only fallback doesn't catch, so the build failed with "Acceleration has failed: [Errno 2] No such file or directory: ...unet.engine.onnx". Fix: for EngineType.UNET only, hash the fully-assembled flag prefix (sha1, first 12 hex chars) using the existing _lora_signature hashing idiom, keeping model name + fp8 + resolution human-readable in the folder name. Cache-key precision is unchanged (identical configs still hash identically, differing configs still diverge) -- only the on-disk encoding is compacted, from ~267 chars down to ~132. VAE/ControlNet directory naming is untouched, so existing VAE engines are not invalidated; UNet engines rebuild once under their new hashed folder name. Also adds a belt-and-suspenders warning if a derived path would still approach MAX_PATH (e.g. an unusually deep user engine_dir), and casts the "filename" config lookup to str to keep the type checker precise about self._configs' heterogeneous per-key value types. Adds a regression test (test_engine_path_length.py) that reproduces the exact crash-log config and asserts the derived .onnx / .opt.onnx paths stay under MAX_PATH, plus determinism and no-collision checks. Co-Authored-By: Claude Opus 4.8 --- .../acceleration/tensorrt/engine_manager.py | 41 +++++++- tests/unit/test_engine_path_length.py | 98 +++++++++++++++++++ 2 files changed, 136 insertions(+), 3 deletions(-) create mode 100644 tests/unit/test_engine_path_length.py diff --git a/src/streamdiffusion/acceleration/tensorrt/engine_manager.py b/src/streamdiffusion/acceleration/tensorrt/engine_manager.py index 3135756f..fa2c2c41 100644 --- a/src/streamdiffusion/acceleration/tensorrt/engine_manager.py +++ b/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.""" diff --git a/tests/unit/test_engine_path_length.py b/tests/unit/test_engine_path_length.py new file mode 100644 index 00000000..04f6205e --- /dev/null +++ b/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 From 3a3d4b8d85614c9849c9710f95674d89b24023b7 Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 19 Jul 2026 08:15:24 -0400 Subject: [PATCH 2/2] ci: trigger review workflow now that SDTD_040_beta_release has the CI gate (see #11)