anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Model] Fold the DLO warmup recipe into _make_dummy_request

Browse Source 
One request builder instead of two: when the generic dummy run opts out
(num_frames <= 0), _make_dummy_request itself falls back to the model's
declared dlo_dummy_run_recipe under distributed layerwise offload, and
_dummy_run keeps its original shape. The recipe request is text-only even
for image-capable models — the recipe names its own task, and t2va rejects
image conditions; caught by a boot smoke after folding (the standalone
builder had built the prompt fresh) and pinned by a regression test.

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
This commit is contained in:
hsliuustc0106 committed 2026-08-26 15:35:13 +08:00
1 parent f4394f1d3d
commit 33297bdb9d
2 files changed
+92 -65
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/test_diffusion_engine_dummy_run.py
+57 -21
View File
@@ -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:
vllm_omni/diffusion/diffusion_engine.py
+35 -44
View File
@@ -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,