refactor: rename LTX_MEGACACHE_* env to generic DI_MEGACACHE_* (backward-compat)

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;

examples/diffusers/lib/pool.py

@@ -611,15 +611,17 @@ 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
     ``<shape_key>.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
     saves the blob); the blobs are then COPYd into the image, and at serve
     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} "

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=<dir>/<shape_key>.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=<dir>/<shape_key>.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=<UUID>"' \
-     -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 <weights>:/data/default:ro -v /scratch/ltx23-bake:/cache \
      <image> python3 ltx23/warmup.py --shapes ltx23/shapes.json --output-dir /cache/out --per-shape-timeout 3600
    # -> /scratch/ltx23-bake/megacache/<shape_key>.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/<shape_key>.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).

examples/diffusers/patches/ltx23_gpu_worker_megacache.patch → 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)