Skip to content

Commit 3690d74

Browse files
[perf]: test the real load and stage-construction paths
The existing tests build stages and pipelines by hand, so they would stay green through both defects review found: a load_modules that never reaches the deferral, and a stage constructor that reads a component's attributes and materializes it during post_init. Three tests now run the production paths. Two call the real ComposedPipelineBase.load_modules with the component loader stubbed and a counter on it, asserting that only opted-in names become proxies and that the loader is never asked for them. The third builds the real MiniMax-H3 stages over tracked proxies and asserts nothing materialized. The third one was mutation-checked: adding a single transformer.patch_size read to MiniMaxH3DenoisingStage.__init__ makes it fail, which is the habit that defeats deferral in the shared DenoisingStage today. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent 1b82f83 commit 3690d74

1 file changed

Lines changed: 98 additions & 0 deletions

File tree

fastvideo/tests/stages/test_lazy_module_load.py

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -510,3 +510,101 @@ def forward(self, batch, fastvideo_args):
510510

511511
assert not vae.is_materialized
512512
assert not transformer.is_materialized
513+
514+
515+
# ----------------------------------------------------------------------
516+
# Production wiring
517+
#
518+
# The tests above build stages and pipelines by hand. These two run the real
519+
# code paths where the two defects review found would live: a `load_modules`
520+
# that never reaches the deferral, and a stage constructor that reads a
521+
# component's attributes and materializes it before the first request.
522+
# ----------------------------------------------------------------------
523+
524+
525+
class _StubLoader:
526+
"""Stands in for PipelineComponentLoader and counts what it is asked for."""
527+
528+
def __init__(self):
529+
self.loaded: list[str] = []
530+
531+
def load_module(self, *, module_name, component_model_path, transformers_or_diffusers, fastvideo_args):
532+
self.loaded.append(module_name)
533+
return _Component(module_name)
534+
535+
536+
def _run_real_load_modules(monkeypatch, lazy_names, manifest_modules):
537+
from fastvideo.pipelines import composed_pipeline_base as cpb
538+
539+
stub = _StubLoader()
540+
monkeypatch.setattr(cpb.PipelineComponentLoader, "load_module", stub.load_module)
541+
542+
class _Pipeline(ComposedPipelineBase):
543+
_required_config_modules = list(manifest_modules)
544+
_lazy_module_names = lazy_names
545+
546+
def __init__(self): # deliberately does not call super()
547+
self.model_path = "/nowhere"
548+
self.fastvideo_args = None
549+
550+
def _load_config(self, model_path):
551+
index = {"_class_name": "X", "_diffusers_version": "0"}
552+
index.update({name: ["diffusers", "Cls", {}] for name in manifest_modules})
553+
return index
554+
555+
def create_pipeline_stages(self, fastvideo_args):
556+
raise NotImplementedError
557+
558+
args = SimpleNamespace(lazy_module_load=True, training_mode=False, revision=None)
559+
modules = _Pipeline().load_modules(args)
560+
return modules, stub.loaded
561+
562+
563+
def test_real_load_modules_defers_only_the_opted_in_components(monkeypatch):
564+
modules, loaded = _run_real_load_modules(monkeypatch, ("transformer", "vae"), ["transformer", "vae", "scheduler"])
565+
566+
assert is_lazy_module(modules["transformer"])
567+
assert is_lazy_module(modules["vae"])
568+
assert not is_lazy_module(modules["scheduler"])
569+
# The loader is asked only for what stays eager.
570+
assert loaded == ["scheduler"]
571+
572+
573+
def test_real_load_modules_defers_nothing_when_the_pipeline_opts_out(monkeypatch):
574+
# The base class ships an empty list, so an unchecked pipeline must load
575+
# everything eagerly even with the flag on.
576+
modules, loaded = _run_real_load_modules(monkeypatch, (), ["transformer", "vae", "scheduler"])
577+
578+
assert not any(is_lazy_module(m) for m in modules.values())
579+
assert sorted(loaded) == ["scheduler", "transformer", "vae"]
580+
581+
582+
def test_building_the_real_h3_stages_materializes_nothing():
583+
# `DenoisingStage.__init__` in the shared stage set reads
584+
# `transformer.hidden_size` to pick an attention backend, which would pull
585+
# the DiT in during post_init. H3's stages must not acquire that habit.
586+
from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3Pipeline
587+
588+
loaded: list[str] = []
589+
590+
def tracked(name):
591+
return LazyModule(name, lambda: loaded.append(name) or _Component(name))
592+
593+
pipeline = MiniMaxH3Pipeline.__new__(MiniMaxH3Pipeline)
594+
pipeline._stages = []
595+
pipeline._stage_name_mapping = {}
596+
pipeline.modules = {
597+
"text_encoder": tracked("text_encoder"),
598+
"transformer": tracked("transformer"),
599+
"vae": tracked("vae"),
600+
"audio_vae": tracked("audio_vae"),
601+
"tokenizer": object(),
602+
"processor": object(),
603+
"scheduler": object(),
604+
"audio_scheduler": object(),
605+
}
606+
607+
pipeline._add_stages(ref2va=False)
608+
609+
assert loaded == [], f"building stages materialized {loaded}"
610+
assert len(pipeline._stages) == 6

0 commit comments

Comments
 (0)