diff --git a/tests/diffusion/test_diffusion_engine_dummy_run.py b/tests/diffusion/test_diffusion_engine_dummy_run.py index 6c648dc93..0e22d087a 100644 --- a/tests/diffusion/test_diffusion_engine_dummy_run.py +++ b/tests/diffusion/test_diffusion_engine_dummy_run.py @@ -125,19 +125,42 @@ def _dlo_recipe_engine(enable_dlo: bool) -> DiffusionEngine: return engine -def test_make_dlo_dummy_request_skipped_without_dlo(monkeypatch: pytest.MonkeyPatch) -> None: +def _opt_out_generic_warmup( + monkeypatch: pytest.MonkeyPatch, + recipe: dict | None, + *, + supports_image_input: bool = False, +) -> None: + """Make the generic dummy run opt out and stub the declared recipe.""" + + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.supports_multimodal_input", + lambda od_config: (supports_image_input, False), + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.get_dummy_run_num_frames", + lambda model_class_name, supports_audio_input: 0, + ) monkeypatch.setattr( "vllm_omni.diffusion.diffusion_engine.get_dlo_dummy_run_recipe", - lambda model_class_name: {"task": "t2va"}, + lambda model_class_name: recipe, ) - assert _dlo_recipe_engine(enable_dlo=False)._make_dlo_dummy_request() is None + +def test_make_dummy_request_skips_recipe_without_dlo(monkeypatch: pytest.MonkeyPatch) -> None: + _opt_out_generic_warmup(monkeypatch, {"task": "t2va"}) + + request = _dlo_recipe_engine(enable_dlo=False)._make_dummy_request(height=512, width=512, guidance_scale=0.0) + + assert request is None -def test_make_dlo_dummy_request_builds_recipe_request(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - "vllm_omni.diffusion.diffusion_engine.get_dlo_dummy_run_recipe", - lambda model_class_name: { +def test_make_dummy_request_builds_recipe_request_when_generic_opted_out( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _opt_out_generic_warmup( + monkeypatch, + { "task": "t2va", "duration": 4.0, "aspect_ratio": "16:9", @@ -146,7 +169,7 @@ def test_make_dlo_dummy_request_builds_recipe_request(monkeypatch: pytest.Monkey }, ) - request = _dlo_recipe_engine(enable_dlo=True)._make_dlo_dummy_request() + request = _dlo_recipe_engine(enable_dlo=True)._make_dummy_request(height=512, width=512, guidance_scale=0.0) assert request is not None assert request.is_dummy_run() @@ -159,31 +182,44 @@ def test_make_dlo_dummy_request_builds_recipe_request(monkeypatch: pytest.Monkey assert extra["cfg_text_scale"] == 1.0 -def test_dummy_run_falls_back_to_dlo_recipe_when_generic_opted_out(monkeypatch: pytest.MonkeyPatch) -> None: - engine = _dlo_recipe_engine(enable_dlo=True) - recipe_request = OmniDiffusionRequest( - prompt="dummy run", - request_id="dlo-dummy", - sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2), +def test_make_dummy_request_recipe_stays_text_only_with_image_capable_model( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # H3 advertises image input for fl2va, but the warmup recipe names t2va, + # which rejects image conditions; the recipe request must not inherit the + # generic dummy run's synthetic multimodal data. + _opt_out_generic_warmup( + monkeypatch, + {"task": "t2va", "duration": 4.0, "aspect_ratio": "16:9"}, + supports_image_input=True, ) - engine._make_dummy_request = Mock(return_value=None) - engine._make_dlo_dummy_request = Mock(return_value=recipe_request) + + request = _dlo_recipe_engine(enable_dlo=True)._make_dummy_request(height=512, width=512, guidance_scale=0.0) + + assert request is not None + assert not request.prompt.get("multi_modal_data") + + +def test_dummy_run_sends_recipe_request_when_generic_opted_out(monkeypatch: pytest.MonkeyPatch) -> None: + _opt_out_generic_warmup(monkeypatch, {"task": "t2va", "duration": 15.0, "aspect_ratio": "16:9"}) + engine = _dlo_recipe_engine(enable_dlo=True) engine.add_req_and_wait_for_response = Mock(return_value=SimpleNamespace(error=None, request_id="dlo-dummy")) engine._dummy_run() - engine._make_dlo_dummy_request.assert_called_once() - engine.add_req_and_wait_for_response.assert_called_once_with(recipe_request) + sent = engine.add_req_and_wait_for_response.call_args.args[0] + assert sent.is_dummy_run() + assert sent.sampling_params.extra_args["task"] == "t2va" def test_dummy_run_still_skips_without_recipe(monkeypatch: pytest.MonkeyPatch) -> None: + _opt_out_generic_warmup(monkeypatch, None) engine = _dlo_recipe_engine(enable_dlo=True) - engine._make_dummy_request = Mock(return_value=None) - engine._make_dlo_dummy_request = Mock(return_value=None) + engine.add_req_and_wait_for_response = Mock() engine._dummy_run() - engine._make_dlo_dummy_request.assert_called_once() + engine.add_req_and_wait_for_response.assert_not_called() def test_dense_mode_does_not_build_kv_profile_request() -> None: diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 380b45939..144dba961 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -65,6 +65,9 @@ logger = init_logger(__name__) _ASYNC_OUTPUT_TIMEOUT_ENV = "VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT" _ASYNC_OUTPUT_TIMEOUT_DEFAULT = 600.0 # seconds +# DLO warmup-recipe keys applied to the dummy request's sampling params; +# every other recipe entry is forwarded through the request's extra_args. +_DLO_DUMMY_RUN_SAMPLING_KEYS = frozenset({"height", "width", "num_inference_steps", "num_frames", "guidance_scale"}) def _async_output_timeout() -> float: @@ -1018,8 +1021,37 @@ class DiffusionEngine: ) -> OmniDiffusionRequest | None: """Build a one-step model request for startup profiling or warmup.""" - prompt: OmniTextPrompt = {"prompt": "dummy run"} supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config) + num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input) + if num_frames <= 0: + # The model opted out of the generic dummy run. Under distributed + # layerwise offload, fall back to its declared warmup recipe so one + # startup generation primes the component allocator-cache + # retention policy ahead of real traffic. The recipe names its own + # task, so the request stays text-only even when the model's other + # tasks accept image or audio conditions. + if not getattr(self.od_config, "enable_distributed_layerwise_offload", False): + return None + recipe = get_dlo_dummy_run_recipe(self.od_config.model_class_name) + if recipe is None: + return None + extra_args: dict[str, Any] = {"cfg_text_scale": 1.0, "cfg_img_scale": 1.0} + extra_args.update({k: v for k, v in recipe.items() if k not in _DLO_DUMMY_RUN_SAMPLING_KEYS}) + return OmniDiffusionRequest( + prompt={"prompt": "dummy run"}, + request_id=DUMMY_DIFFUSION_REQUEST_ID, + sampling_params=OmniDiffusionSamplingParams( + height=int(recipe.get("height", height)), + width=int(recipe.get("width", width)), + num_inference_steps=int(recipe.get("num_inference_steps", 1)), + num_frames=int(recipe.get("num_frames", 0)), + guidance_scale=float(recipe.get("guidance_scale", guidance_scale)), + num_outputs_per_prompt=1, + extra_args=extra_args, + ), + ) + + prompt: OmniTextPrompt = {"prompt": "dummy run"} if supports_image_input: color_format = image_color_format(self.od_config.model_class_name) images = [PIL.Image.new(color_format, (width, height)) for _ in range(num_image_inputs)] @@ -1029,9 +1061,6 @@ class DiffusionEngine: audio_sr = 16000 prompt.setdefault("multi_modal_data", {})["audio"] = np.random.randn(audio_sr * 2).astype(np.float32) - num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input) - if num_frames <= 0: - return None return OmniDiffusionRequest( prompt=prompt, request_id=DUMMY_DIFFUSION_REQUEST_ID, @@ -1103,10 +1132,6 @@ class DiffusionEngine: profile_requests.append(profile_request) return profile_requests - # Sampling-recipe keys consumed by _make_dlo_dummy_request directly; - # everything else is forwarded through the request's extra_args. - _DLO_DUMMY_RUN_SAMPLING_KEYS = frozenset({"height", "width", "num_inference_steps", "num_frames", "guidance_scale"}) - def _dummy_run(self): """A dummy run to warm up the model.""" req = self._make_dummy_request( @@ -1115,47 +1140,13 @@ class DiffusionEngine: guidance_scale=0.0, ) if req is None: - req = self._make_dlo_dummy_request() - if req is None: - logger.info("Skipping dummy warmup run (num_frames=0)") - return + logger.info("Skipping dummy warmup run (num_frames=0)") + return logger.info("dummy run to warm up the model") output = self.add_req_and_wait_for_response(req) if output.error: raise RuntimeError(f"Dummy run failed: {output.error}") - def _make_dlo_dummy_request(self) -> OmniDiffusionRequest | None: - """Build the model-declared startup warmup used under distributed offload. - - Models whose request geometry needs model-specific sampling arguments - (task, duration, aspect ratio, ...) opt out of the generic dummy run - and declare ``dlo_dummy_run_recipe`` instead. Under distributed - layerwise offload one startup generation also primes the component - allocator-cache retention policy's observed-peak budget, so the - cold-start flushes land before real traffic. - """ - - if not getattr(self.od_config, "enable_distributed_layerwise_offload", False): - return None - recipe = get_dlo_dummy_run_recipe(self.od_config.model_class_name) - if recipe is None: - return None - extra_args: dict[str, Any] = {"cfg_text_scale": 1.0, "cfg_img_scale": 1.0} - extra_args.update({k: v for k, v in recipe.items() if k not in self._DLO_DUMMY_RUN_SAMPLING_KEYS}) - return OmniDiffusionRequest( - prompt={"prompt": "dummy run"}, - request_id=DUMMY_DIFFUSION_REQUEST_ID, - sampling_params=OmniDiffusionSamplingParams( - height=int(recipe.get("height", 512)), - width=int(recipe.get("width", 512)), - num_inference_steps=int(recipe.get("num_inference_steps", 1)), - num_frames=int(recipe.get("num_frames", 0)), - guidance_scale=float(recipe.get("guidance_scale", 0.0)), - num_outputs_per_prompt=1, - extra_args=extra_args, - ), - ) - def _submit_rpc( self, method: str,