[Bugfix][CI/Build] Fix layerwise offload shutdown and enable it for the LTX-2 distilled T2V accuracy test (#8453)

Signed-off-by: nikiwang-125 <yixinyan79@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

tests/diffusion/offloader/test_layerwise_backend.py

@@ -542,6 +542,24 @@ class TestLayerwiseComponentSelection:
         for name, tensor in pipeline.state_dict().items():
             torch.testing.assert_close(tensor, expected[name])
 
+    def test_shutdown_skips_block_restore(self, patched_offload_runtime, monkeypatch):
+        # Restoring every block copies the whole DiT back onto the device, which
+        # outlived the executor's shutdown grace period (issue 7625).
+        pipeline = nn.Module()
+        pipeline.transformer = _SingleBlockModel(num_blocks=3)
+        backend = _layer_backend()
+        backend.enable(pipeline)
+        restores = Mock()
+        for hook in backend._dit_hooks:
+            monkeypatch.setattr(hook, "restore_next_block", restores)
+
+        backend.shutdown()
+
+        restores.assert_not_called()
+        assert not backend.enabled
+        assert not backend._dit_hooks
+        assert not backend._hooked_dit_blocks
+
     def test_encoder_disable_failure_keeps_host_masters_for_retry(
         self,
         patched_offload_runtime,

tests/e2e/accuracy/ltx/test_ltx_official_similarity.py

@@ -103,7 +103,7 @@ class LTXAccuracyCase:
     spatial_upsampler: LTXArtifact | None = None
     distilled_lora: LTXArtifact | None = None
     # Official uses block streaming; Omni uses layerwise offload. Both preserve
-    # bf16 arithmetic while keeping full-resolution two-stage cases on one H100.
+    # bf16 arithmetic while keeping full-resolution cases on one H100.
     enable_layerwise_offload: bool = False
 
 
@@ -182,6 +182,7 @@ TWO_STAGE_CASES = (
         gemma_model_id="Lightricks/LTX-2",
         gemma_model_revision=LTX2_REVISION,
         gemma_model_env="VLLM_TEST_LTX2_MODEL",
+        enable_layerwise_offload=True,
     ),
     LTXAccuracyCase(
         name="ltx23_distilled_i2v",

vllm_omni/diffusion/offloader/layerwise_backend.py

@@ -521,3 +521,9 @@ class LayerWiseOffloadBackend(OffloadBackend):
 
     def disable(self) -> None:
         self._disable(restore_weights=True)
+
+    def shutdown(self) -> None:
+        # Restoring every block copies the whole DiT back onto the device, which
+        # only matters for a later ``enable`` and can outlast the executor's
+        # shutdown grace period.
+        self._disable(restore_weights=False)