From 1f8413fa21984e211978d9c5f3344c4e3bbe4779 Mon Sep 17 00:00:00 2001 From: Johan de Ruiter Date: Fri, 3 Jul 2026 05:49:30 +0000 Subject: [PATCH] refactor: rename LTX_MEGACACHE_* env to generic DI_MEGACACHE_* (backward-compat) The compile-cache mechanism is torch.compiler Mega-Cache (save/load_cache_artifacts), a PyTorch-level portable-compile-cache that is NOT LTX-specific -- it is wired at the FastVideo gpu_worker level and applies to any video family (fastwan, ltx2, ltx23, ...). As we make it the default for our self-run video models, the LTX_-prefixed env name is misleading, so rename before it proliferates into more baked images: LTX_MEGACACHE_DIR -> DI_MEGACACHE_DIR LTX_MEGACACHE_BLOB -> DI_MEGACACHE_BLOB LTX_MEGACACHE_SAVE_EVERY -> DI_MEGACACHE_SAVE_EVERY Also renames patches/ltx23_gpu_worker_megacache.patch -> gpu_worker_megacache.patch and the private helper methods _ltx_megacache_* -> _megacache_*. Backward-compatible so nothing breaks mid-transition: - lib/pool.py reads `DI_MEGACACHE_DIR or LTX_MEGACACHE_DIR` and exports BOTH DI_MEGACACHE_BLOB and LTX_MEGACACHE_BLOB to the worker child. - the gpu_worker patch reads `DI_MEGACACHE_BLOB or LTX_MEGACACHE_BLOB` (and same for SAVE_EVERY). So already-baked images and any in-flight bakes that set LTX_MEGACACHE_DIR keep working. Drop the LTX_ fallbacks once every video image has been rebuilt on this change. Co-Authored-By: Claude Opus 4.8 --- examples/diffusers/Dockerfile | 2 +- examples/diffusers/lib/pool.py | 14 ++++++++---- examples/diffusers/ltx23/CACHING.md | 13 ++++++----- ...cache.patch => gpu_worker_megacache.patch} | 22 +++++++++---------- 4 files changed, 29 insertions(+), 22 deletions(-) rename examples/diffusers/patches/{ltx23_gpu_worker_megacache.patch => gpu_worker_megacache.patch} (83%) diff --git a/examples/diffusers/Dockerfile b/examples/diffusers/Dockerfile index b557bfc6157a..c5438d453ef3 100644 --- a/examples/diffusers/Dockerfile +++ b/examples/diffusers/Dockerfile @@ -91,7 +91,7 @@ ENV FASTVIDEO_SHA=${FASTVIDEO_SHA} # reuses baked compile artifacts instead of recompiling. COPY'd here (before # the FastVideo build) so `git apply` below picks it up; the full app tree is # only COPYd to /opt/app later. See examples/diffusers/ltx23/CACHING.md. -COPY patches/ltx23_gpu_worker_megacache.patch /tmp/megacache.patch +COPY patches/gpu_worker_megacache.patch /tmp/megacache.patch # deepinfra encode patch: cap libx264 -threads so the CPU mp4 encode can't grab # the whole shared node (no NVENC on B200 -> CPU libx264). git apply'd below; diff --git a/examples/diffusers/lib/pool.py b/examples/diffusers/lib/pool.py index 1528e94be9d6..d2e2a93f2e23 100644 --- a/examples/diffusers/lib/pool.py +++ b/examples/diffusers/lib/pool.py @@ -611,7 +611,9 @@ def _set_parent_death_signal(sig: int = signal.SIGTERM) -> None: def _megacache_blob_path(shape_key: str) -> str | None: """Per-shape Mega-Cache blob path, or None if the feature is off. - Opt-in via ``LTX_MEGACACHE_DIR``. When set, the dir holds one + Opt-in via ``DI_MEGACACHE_DIR`` (legacy alias ``LTX_MEGACACHE_DIR`` still + honored during the transition -- the mechanism is torch.compiler Mega-Cache + and not LTX-specific). When set, the dir holds one ``.megacache.bin`` per shape: a portable ``torch.compiler`` artifact bundle (FX graphs + Triton + autotune best_configs). At bake time the dir is writable and starts empty (worker compiles cold, then @@ -619,7 +621,7 @@ def _megacache_blob_path(shape_key: str) -> str | None: time the worker LOADs them so a fresh pod skips the autotune recompile. Unset => no-op (legacy behavior: compile cold per process). """ - base = os.environ.get("LTX_MEGACACHE_DIR") + base = os.environ.get("DI_MEGACACHE_DIR") or os.environ.get("LTX_MEGACACHE_DIR") if not base: return None return os.path.join(base, f"{shape_key}.megacache.bin") @@ -631,13 +633,17 @@ def _export_megacache_env(shape_key: str) -> None: The actual ``torch.compiler.save/load_cache_artifacts`` calls MUST run in the process that compiles -- and FastVideo's MultiprocExecutor compiles in a spawned worker child, NOT here. So this pool-worker only resolves the - per-shape path and exports it as ``LTX_MEGACACHE_BLOB`` BEFORE the factory + per-shape path and exports it as ``DI_MEGACACHE_BLOB`` BEFORE the factory builds the generator (which spawns that child). The child inherits the env and does the load (before first forward) / save (after first forward). See - fastvideo/worker/gpu_worker.py. No-op when ``LTX_MEGACACHE_DIR`` is unset. + fastvideo/worker/gpu_worker.py. No-op when ``DI_MEGACACHE_DIR`` is unset. """ path = _megacache_blob_path(shape_key) if path: + os.environ["DI_MEGACACHE_BLOB"] = path + # Transition: also export the legacy name so a pre-rename gpu_worker + # patch (older baked image) still resolves the blob. Drop once every + # video image has been rebuilt with the renamed patch below. os.environ["LTX_MEGACACHE_BLOB"] = path print( f"[pool-worker/{shape_key}] megacache blob path -> {path} " diff --git a/examples/diffusers/ltx23/CACHING.md b/examples/diffusers/ltx23/CACHING.md index dd3ec3d39e82..2bdb16943606 100644 --- a/examples/diffusers/ltx23/CACHING.md +++ b/examples/diffusers/ltx23/CACHING.md @@ -34,8 +34,9 @@ One resident process, both per-mode blobs loaded (glob-load), sequence t2v→i2v wrong" risk from the cache. ## What's wired (code) -- **Mega-Cache** — `lib/pool.py` exports `LTX_MEGACACHE_BLOB=/.megacache.bin` (from - `LTX_MEGACACHE_DIR`) before building the generator; `fastvideo/worker/gpu_worker.py` (fork patch) LOADS +- **Mega-Cache** — `lib/pool.py` exports `DI_MEGACACHE_BLOB=/.megacache.bin` (from + `DI_MEGACACHE_DIR`; legacy alias `LTX_MEGACACHE_DIR` still honored during transition) before building the + generator; `fastvideo/worker/gpu_worker.py` (fork patch) LOADS the blob before the first forward and SAVES it after the first forward. Per-shape (one blob per shape). Best-effort: a missing/incompatible blob silently falls back to a cold compile (never wedges a pod). - **CuTe / CUDA / QuACK env caches** — `lib/pool.py` sets `CUTE_DSL_CACHE_DIR=/cache/cutedsl`, @@ -47,11 +48,11 @@ One resident process, both per-mode blobs loaded (glob-load), sequence t2v→i2v The blob is keyed by (profile + shape + torch version + GPU arch + code/image), so **bake per profile** and **rebake on any image change** (a mismatch just falls back to cold compile, so it's safe, just slow). -1. **Bake** on a Blackwell box, `LTX_MEGACACHE_DIR` pointed at a writable scratch (one run does all shapes +1. **Bake** on a Blackwell box, `DI_MEGACACHE_DIR` pointed at a writable scratch (one run does all shapes in `shapes.json`; the worker saves a blob per shape): ``` docker run --gpus '"device="' \ - -e LTX23_PROFILE=quality -e VIDEO_MODEL_FAMILY=ltx23 -e LTX_MEGACACHE_DIR=/cache/megacache \ + -e LTX23_PROFILE=quality -e VIDEO_MODEL_FAMILY=ltx23 -e DI_MEGACACHE_DIR=/cache/megacache \ -v :/data/default:ro -v /scratch/ltx23-bake:/cache \ python3 ltx23/warmup.py --shapes ltx23/shapes.json --output-dir /cache/out --per-shape-timeout 3600 # -> /scratch/ltx23-bake/megacache/.megacache.bin (one per shape, ~300 MB each) @@ -61,7 +62,7 @@ The blob is keyed by (profile + shape + torch version + GPU arch + code/image), tar -cf megacache.tar -C /scratch/ltx23-bake megacache # Dockerfile: ADD megacache.tar /opt/app/ -> /opt/app/megacache/.megacache.bin ``` -3. **Serve**: set `LTX_MEGACACHE_DIR=/opt/app/megacache` in the LTX-2.3 model's `extra_env` (per-model, so +3. **Serve**: set `DI_MEGACACHE_DIR=/opt/app/megacache` in the LTX-2.3 model's `extra_env` (per-model, so LTX-2 is unaffected). The per-shape pool worker loads its shape's blob before the first compile. 4. **Boot-warm gate**: have the pod route one dummy request per shape at startup (the ~560 s cost) before marking ready, so no real request ever hits a cold compile. The resident pool then serves warm. @@ -71,7 +72,7 @@ The blob is keyed by (profile + shape + torch version + GPU arch + code/image), bites and it stops recompiling that location. Fine for our 2-shape menu; bump `torch._dynamo.config.cache_size_limit` if a single process must serve more. - **One blob per shape — NOT one multi-shape blob.** A blob saved from a process that compiled multiple - shapes (via `LTX_MEGACACHE_SAVE_EVERY`) fails to hit the *first* shape on reload (its key gets polluted by + shapes (via `DI_MEGACACHE_SAVE_EVERY`) fails to hit the *first* shape on reload (its key gets polluted by later compilation). Keep the per-shape pool + per-shape blobs (the proven path). - **Implicit inductor cache (`TORCHINDUCTOR_CACHE_DIR`) does NOT port** across processes — don't ship/rely on it for cold-start (LTX-2's tar+ADD bake only "worked" because of the di-slc-35 host-pin). diff --git a/examples/diffusers/patches/ltx23_gpu_worker_megacache.patch b/examples/diffusers/patches/gpu_worker_megacache.patch similarity index 83% rename from examples/diffusers/patches/ltx23_gpu_worker_megacache.patch rename to examples/diffusers/patches/gpu_worker_megacache.patch index 4d93afae6f01..0c385dcd88d6 100644 --- a/examples/diffusers/patches/ltx23_gpu_worker_megacache.patch +++ b/examples/diffusers/patches/gpu_worker_megacache.patch @@ -7,30 +7,30 @@ + # deepinfra: load baked torch.compile artifacts (Mega-Cache) BEFORE the + # first forward triggers lazy compilation, so a fresh pod reuses them + # instead of recompiling (~halves cold start). Per-mode blobs (t2v/i2v) -+ # are glob-loaded from the dir of LTX_MEGACACHE_BLOB. Best-effort: a ++ # are glob-loaded from the dir of DI_MEGACACHE_BLOB. Best-effort: a + # missing/incompatible blob just falls back to a cold compile and never + # wedges boot. This MUST live here (the worker child is the process that + # compiles). See examples/diffusers/ltx23/CACHING.md. + self._megacache_saved = False -+ self._megacache_loaded = self._ltx_megacache_load() ++ self._megacache_loaded = self._megacache_load() + def execute_forward(self, forward_batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: output_batch = self.pipeline.forward(forward_batch, self.fastvideo_args) + # deepinfra: after the first (cold) forward compiled all graphs, persist + # them once so a future fresh process can skip the recompile (bake path). -+ self._ltx_megacache_save() ++ self._megacache_save() return cast(ForwardBatch, output_batch) + # ---- deepinfra Mega-Cache helpers (torch.compiler save/load_cache_artifacts) ---- + @staticmethod -+ def _ltx_megacache_path() -> str | None: -+ return os.environ.get("LTX_MEGACACHE_BLOB") ++ def _megacache_path() -> str | None: ++ return os.environ.get("DI_MEGACACHE_BLOB") or os.environ.get("LTX_MEGACACHE_BLOB") + -+ def _ltx_megacache_load(self) -> bool: ++ def _megacache_load(self) -> bool: + """Load every ``*.megacache.bin`` in the blob's directory (one per + shape*mode, e.g. t2v + i2v) before the first compile. Repeated + ``load_cache_artifacts`` calls are additive (validated). Best-effort.""" -+ path = self._ltx_megacache_path() ++ path = self._megacache_path() + if not path: + return False + import glob @@ -49,14 +49,14 @@ + logger.warning("[megacache] load failed for %s (%r)", bp, exc) + return loaded_any + -+ def _ltx_megacache_save(self) -> None: ++ def _megacache_save(self) -> None: + """Persist compile artifacts after the first (cold) compile. No-op when + disabled, when we booted from a blob, or already saved. With -+ LTX_MEGACACHE_SAVE_EVERY set, re-save after every forward (bake helper).""" -+ path = self._ltx_megacache_path() ++ DI_MEGACACHE_SAVE_EVERY set, re-save after every forward (bake helper).""" ++ path = self._megacache_path() + if not path or getattr(self, "_megacache_loaded", False): + return -+ if getattr(self, "_megacache_saved", False) and not os.environ.get("LTX_MEGACACHE_SAVE_EVERY"): ++ if getattr(self, "_megacache_saved", False) and not (os.environ.get("DI_MEGACACHE_SAVE_EVERY") or os.environ.get("LTX_MEGACACHE_SAVE_EVERY")): + return + try: + os.makedirs(os.path.dirname(path), exist_ok=True)