diff --git a/components/src/dynamo/sglang/args.py b/components/src/dynamo/sglang/args.py index 072dd450534e..e5e885caee91 100644 --- a/components/src/dynamo/sglang/args.py +++ b/components/src/dynamo/sglang/args.py @@ -11,7 +11,7 @@ import warnings from argparse import Namespace from pathlib import Path -from typing import Any, Dict, Generator, Optional +from typing import Any, Dict, Generator, List, Optional import yaml from sglang.srt.server_args import ServerArgs @@ -326,6 +326,18 @@ async def parse_args(args: list[str]) -> Config: Raises: SystemExit: If arguments are invalid or incompatible. """ + # Help must match the parser the worker will actually use. Diffusion + # workers parse engine args with sglang's diffusion ServerArgs, so route + # --help through that parser instead of letting the Dynamo/LLM parsers + # print LLM engine options that a diffusion worker would reject. Checked + # on the raw strings because argparse exits on -h in whichever parser + # sees it first. + if ("-h" in args or "--help" in args) and ( + "--image-diffusion-worker" in args or "--video-generation-worker" in args + ): + _print_diffusion_worker_help() + sys.exit(0) + runtime_argspec = DynamoRuntimeArgGroup() dynamo_sglang_argspec = DynamoSGLangArgGroup() @@ -369,6 +381,13 @@ async def parse_args(args: list[str]) -> Config: dynamo_config.router_advertisement, unknown = parse_worker_router_config(unknown) dynamo_config.validate() + # Image/video diffusion workers configure DiffGenerator, whose options + # live in SGLang's *diffusion* ServerArgs — a different dataclass from the + # LLM ServerArgs everything below parses against. Branch off before any + # LLM-specific processing so diffusion engine args are parsed natively. + if dynamo_config.image_diffusion_worker or dynamo_config.video_generation_worker: + return await _resolve_diffusion_worker_config(unknown, dynamo_config) + # Dealing with SGLang native configs temp_config_file = None if dynamo_config.disagg_config and dynamo_config.disagg_config_key: @@ -526,11 +545,6 @@ async def parse_args(args: list[str]) -> Config: # fetch_model (download the model) here, in `parse_args`. `parse_args` should not # contain code to download a model, it should only parse the args. - # For diffusion/video workers, create a minimal dummy ServerArgs since diffusion - # doesn't use transformer models or sglang Engine - it uses DiffGenerator directly - image_diffusion_worker = dynamo_config.image_diffusion_worker - video_generation_worker = dynamo_config.video_generation_worker - # ServerArgs is read-only after resolution, so apply Dynamo defaults first. fpm_source = _forward_pass_metrics_source(dynamo_config) if fpm_source and not getattr(parsed_args, "enable_forward_pass_metrics", False): @@ -544,43 +558,11 @@ async def parse_args(args: list[str]) -> Config: parsed_args.max_running_requests = 8 logging.info("Defaulting max_running_requests to 8 for diffusion worker") - if image_diffusion_worker or video_generation_worker: - worker_type = ( - "image diffusion" if image_diffusion_worker else "video generation" - ) - logging.info( - f"{worker_type.title()} worker detected with model: {model_path}, creating minimal ServerArgs stub" - ) - # Create a minimal ServerArgs-like object that bypasses model config loading - # Diffusion/video workers don't actually use ServerArgs - they use DiffGenerator - import types - - server_args = types.SimpleNamespace() - # Copy over any attrs that might be needed, but avoid triggering __post_init__ - server_args.model_path = model_path - server_args.served_model_name = parsed_args.served_model_name - server_args.enable_metrics = getattr(parsed_args, "enable_metrics", False) - server_args.log_level = getattr(parsed_args, "log_level", "info") - server_args.kv_events_config = getattr(parsed_args, "kv_events_config", None) - server_args.tp_size = getattr(parsed_args, "tp_size", 1) - server_args.dp_size = getattr(parsed_args, "dp_size", 1) - server_args.speculative_algorithm = None - server_args.disaggregation_mode = None - server_args.dllm_algorithm = False - server_args.load_format = None - server_args.enable_trace = getattr(parsed_args, "enable_trace", False) - server_args.enable_forward_pass_metrics = getattr( - parsed_args, "enable_forward_pass_metrics", False - ) - logging.info( - f"Created stub ServerArgs for {worker_type}: model_path={server_args.model_path}" - ) - else: - # Dynamo expects disjoint output_ids; ServerArgs is read-only after resolution. - parsed_args.incremental_streaming_output = True - server_args = ServerArgs.from_cli_args(parsed_args) - if server_args.get_model_config().is_multimodal: - ensure_sglang_tensor_image_size() + # Dynamo expects disjoint output_ids; ServerArgs is read-only after resolution. + parsed_args.incremental_streaming_output = True + server_args = ServerArgs.from_cli_args(parsed_args) + if server_args.get_model_config().is_multimodal: + ensure_sglang_tensor_image_size() if getattr(server_args, "schedule_low_priority_values_first", False): raise ValueError( @@ -632,6 +614,99 @@ async def parse_args(args: list[str]) -> Config: return Config(server_args, dynamo_config) +def _print_diffusion_worker_help() -> None: + """Print combined Dynamo + native diffusion engine options.""" + from dynamo.sglang.diffusion_args import build_diffusion_parser + + dynamo_parser = argparse.ArgumentParser( + prog="dynamo.sglang", + description="Dynamo SGLang diffusion worker configuration", + formatter_class=argparse.RawTextHelpFormatter, + add_help=False, + ) + DynamoRuntimeArgGroup().add_arguments(dynamo_parser) + DynamoSGLangArgGroup().add_arguments(dynamo_parser) + print(dynamo_parser.format_help()) + print( + "SGLang Diffusion Engine Options (native sglang diffusion ServerArgs;" + " every option below is forwarded to the engine):\n" + ) + diffusion_parser, _ = build_diffusion_parser() + print(diffusion_parser.format_help()) + + +async def _resolve_diffusion_worker_config( + unknown: List[str], dynamo_config: "DynamoConfig" +) -> Config: + """Resolve Config for image/video diffusion workers. + + Engine arguments are parsed with SGLang's native diffusion ServerArgs CLI + (see diffusion_args.py), so every native diffusion engine argument is + reachable — no hand-copied stub. The returned server_args is a + DiffusionWorkerArgs adapter: engine fields resolve from the natively + parsed diffusion ServerArgs, Dynamo-side settings live on the adapter. + """ + from dynamo.sglang.diffusion_args import parse_diffusion_args + + worker_type = ( + "image diffusion" + if dynamo_config.image_diffusion_worker + else "video generation" + ) + logging.info( + f"{worker_type.title()} worker detected: parsing engine args with " + "SGLang's native diffusion ServerArgs" + ) + + # Fetch the model before the full parse: sglang's diffusion ServerArgs + # resolves model info during argument resolution, so the weights must be + # available first (mirrors the LLM path, which fetches before + # ServerArgs.from_cli_args). + pre_parser = argparse.ArgumentParser(add_help=False) + pre_parser.add_argument("--model-path", type=str, default=None) + pre_args, _ = pre_parser.parse_known_args(unknown) + if not pre_args.model_path: + raise ValueError("--model-path is required for diffusion workers") + if should_fetch_model(argparse.Namespace(), pre_args.model_path): + await fetch_model(pre_args.model_path) + + _parsed_args, server_args = parse_diffusion_args(unknown) + + # --served-model-name may pack several names; first is primary. + served_names = split_served_model_names(server_args.served_model_name) + if served_names: + server_args.served_model_name = served_names[0] + dynamo_config.served_model_aliases = served_names[1:] + if served_names[1:]: + logging.info( + "Multi-name registration: primary=%r, aliases=%s", + served_names[0], + served_names[1:], + ) + + if is_snapshot_enabled(): + configure_snapshot_capture_env() + + endpoint = ( + dynamo_config.endpoint or f"dyn://{dynamo_config.namespace}.backend.generate" + ) + parsed_namespace, parsed_component_name, parsed_endpoint_name = parse_endpoint( + endpoint + ) + dynamo_config.namespace = parsed_namespace + dynamo_config.component = parsed_component_name + dynamo_config.endpoint = parsed_endpoint_name + # dllm (text diffusion LLM) does not apply to image/video workers. The + # old stub accidentally made this True (False is not None); it was never + # load-bearing because main.py dispatches image/video workers first. + dynamo_config.diffusion_worker = False + dynamo_config.use_kv_events = False + + logging.debug(f"Dynamo configs: {dynamo_config}") + + return Config(server_args, dynamo_config) + + @contextlib.contextmanager def reserve_free_port(host: str = "localhost") -> Generator[int, None, None]: """Find and reserve a free port until context exits. diff --git a/components/src/dynamo/sglang/diffusion_args.py b/components/src/dynamo/sglang/diffusion_args.py new file mode 100644 index 000000000000..110d4840c9f2 --- /dev/null +++ b/components/src/dynamo/sglang/diffusion_args.py @@ -0,0 +1,232 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Native argument parsing for diffusion (image/video) workers. + +Diffusion workers configure SGLang's DiffGenerator, whose options live in +``sglang.multimodal_gen.runtime.server_args.ServerArgs`` — a different +dataclass from the LLM ``ServerArgs``. Building the CLI from that class (the +same way the LLM path builds its CLI from the LLM ``ServerArgs``) makes every +native diffusion engine argument reachable from Dynamo without hand-copying +fields. + +``DiffusionWorkerArgs`` is the thin adapter the rest of the Dynamo worker code +sees: it delegates engine fields to the real diffusion ``ServerArgs`` and adds +the few Dynamo-side settings that class does not define. Fields Dynamo's +shared worker code probes on every ``server_args`` (speculative decoding, +disaggregation, load format) are pinned to their inert values here instead of +being hand-maintained in a SimpleNamespace stub. +""" + +import argparse +import logging +from typing import Any, List, Optional, Tuple + +logger = logging.getLogger(__name__) + +# Dynamo-side flags the diffusion ServerArgs may not define. Each entry is +# only added to the parser when the engine parser does not already provide it, +# so a future SGLang version adding the same flag natively wins automatically. +_DYNAMO_SIDE_FLAGS = ( + { + "flags": ("--served-model-name",), + "kwargs": { + "type": str, + "default": None, + "help": "Model name reported to the frontend and /v1/models. " + "Defaults to --model-path.", + }, + }, + { + "flags": ("--enable-metrics",), + "kwargs": { + "action": "store_true", + "default": False, + "help": "Enable Dynamo worker metrics publishing.", + }, + }, +) + + +class DiffusionWorkerArgs: + """Adapter combining SGLang's diffusion ServerArgs with Dynamo settings. + + Attribute access falls through to the wrapped diffusion ``ServerArgs``, + so shared worker code reading e.g. ``model_path``, ``tp_size`` or + ``log_level`` sees the natively parsed values. ``engine_args`` exposes the + wrapped object for handing to ``DiffGenerator.from_server_args``. + """ + + # Fields Dynamo's shared worker code reads on any server_args. Diffusion + # workers do not use these subsystems; pin them to inert values. + _INERT_DEFAULTS = { + "speculative_algorithm": None, + "disaggregation_mode": None, + "dllm_algorithm": False, + "load_format": None, + "kv_events_config": None, + "enable_forward_pass_metrics": False, + } + + def __init__( + self, + engine_args: Any, + served_model_name: Optional[str], + enable_metrics: bool, + ): + self.engine_args = engine_args + self.served_model_name = served_model_name or engine_args.model_path + self.enable_metrics = enable_metrics + for name, value in self._INERT_DEFAULTS.items(): + # Prefer a native field if the engine args grow one later. + setattr(self, name, getattr(engine_args, name, value)) + + def __getattr__(self, name: str) -> Any: + # Only called when normal lookup fails: delegate to the engine args. + return getattr(self.__dict__["engine_args"], name) + + +def _existing_option_strings(parser: argparse.ArgumentParser) -> set: + return {opt for action in parser._actions for opt in action.option_strings} + + +def _import_diffusion_server_args(): + """Import sglang's diffusion ServerArgs (lazily; layout varies by version).""" + # Imported lazily: only diffusion/video workers need the multimodal_gen + # extra, and importing it pulls in heavy dependencies. server_args became + # a package in newer SGLang releases; older ones expose it as a module. + try: + from sglang.multimodal_gen.runtime.server_args.server_args import ( + ServerArgs as DiffusionServerArgs, + ) + except (ImportError, ModuleNotFoundError): + from sglang.multimodal_gen.runtime.server_args import ( + ServerArgs as DiffusionServerArgs, # type: ignore[no-redef] + ) + return DiffusionServerArgs + + +def build_diffusion_parser() -> Tuple[argparse.ArgumentParser, List[str]]: + """Build the diffusion worker CLI: native engine args + Dynamo-side flags. + + Returns: + (parser, dynamo_side_dests) where dynamo_side_dests names the flags + registered by Dynamo rather than the engine. + """ + DiffusionServerArgs = _import_diffusion_server_args() + + try: + from sglang.multimodal_gen.utils import FlexibleArgumentParser + except ImportError: + FlexibleArgumentParser = argparse.ArgumentParser + + parser = FlexibleArgumentParser( + description="Dynamo SGLang diffusion worker configuration", + add_help=False, + ) + DiffusionServerArgs.add_cli_args(parser) + + # Register Dynamo-side flags the engine parser does not define. + existing = _existing_option_strings(parser) + dynamo_group = parser.add_argument_group("Dynamo Options") + dynamo_side_dests = [] + for spec in _DYNAMO_SIDE_FLAGS: + if any(flag in existing for flag in spec["flags"]): + continue + action = dynamo_group.add_argument(*spec["flags"], **spec["kwargs"]) + dynamo_side_dests.append(action.dest) + return parser, dynamo_side_dests + + +def parse_diffusion_args( + unknown_args: List[str], +) -> Tuple[argparse.Namespace, DiffusionWorkerArgs]: + """Parse worker args against SGLang's native diffusion ServerArgs CLI. + + Args: + unknown_args: Argument strings left over after Dynamo's own parser. + + Returns: + (parsed argparse namespace, DiffusionWorkerArgs adapter) + """ + DiffusionServerArgs = _import_diffusion_server_args() + parser, dynamo_side_dests = build_diffusion_parser() + + parsed, remaining = parser.parse_known_args(unknown_args) + + # Split Dynamo-side values out of the namespace before handing it to + # from_cli_args, so they never masquerade as engine arguments. + dynamo_values = {} + for dest in dynamo_side_dests: + dynamo_values[dest] = getattr(parsed, dest) + delattr(parsed, dest) + + # from_cli_args distinguishes explicitly-set flags from argparse defaults + # by scanning sys.argv — correct for a worker process, but wrong for any + # caller that passes an argument list (tests, embedding). Communicate the + # flags we parsed through the side channel the engine parser supports, so + # resolution never depends on process argv. Each raw flag is resolved to + # its parser destination (not its raw spelling): argparse accepts + # abbreviations and aliases (e.g. --tp for --tp-size), and recording the + # raw text would make the engine treat the real field as unspecified. + option_to_dest = { + opt: action.dest for action in parser._actions for opt in action.option_strings + } + + def _resolve_dest(raw_flag: str) -> Optional[str]: + if raw_flag in option_to_dest: + return option_to_dest[raw_flag] + # argparse allows unambiguous prefixes; mirror that resolution. + matches = {d for o, d in option_to_dest.items() if o.startswith(raw_flag)} + return matches.pop() if len(matches) == 1 else None + + explicit_names = set() + for arg in unknown_args: + if not arg.startswith("--"): + continue + raw = arg.split("=", 1)[0] + dest = _resolve_dest(raw) + if dest is None: + # Not a registered option (e.g. dynamic ---path flags + # the engine resolves itself): keep the normalized raw name. + dest = raw.replace("-", "_").lstrip("_") + explicit_names.add(dest) + explicit_names -= set(dynamo_side_dests) + + # The engine defaults num_gpus to 1 and does not derive it from + # parallelism degrees, so tp/dp without an explicit --num-gpus would + # under-allocate. Preserve the historical num_gpus = tp * dp behavior. + if "num_gpus" not in explicit_names: + tp = getattr(parsed, "tp_size", None) or 1 + dp = getattr(parsed, "dp_size", None) or 1 + if tp * dp > 1: + parsed.num_gpus = tp * dp + explicit_names.add("num_gpus") + logger.info( + "Derived num_gpus=%d from tp_size=%d * dp_size=%d", tp * dp, tp, dp + ) + + if hasattr(parsed, "_sglang_explicit_arg_names"): + explicit_names |= set(parsed._sglang_explicit_arg_names) + parsed._sglang_explicit_arg_names = tuple(sorted(explicit_names)) + + engine_args = DiffusionServerArgs.from_cli_args(parsed, remaining) + + adapter = DiffusionWorkerArgs( + engine_args=engine_args, + served_model_name=dynamo_values.get( + "served_model_name", getattr(parsed, "served_model_name", None) + ), + enable_metrics=dynamo_values.get( + "enable_metrics", getattr(parsed, "enable_metrics", False) + ), + ) + + logger.info( + "Diffusion worker args parsed natively: model_path=%s, " + "served_model_name=%s (SGLang diffusion ServerArgs, %d fields reachable)", + adapter.model_path, + adapter.served_model_name, + len(getattr(engine_args, "__dataclass_fields__", ())), + ) + return parsed, adapter diff --git a/components/src/dynamo/sglang/init_diffusion.py b/components/src/dynamo/sglang/init_diffusion.py index 26ca30588e07..ce20b3fe7bcb 100644 --- a/components/src/dynamo/sglang/init_diffusion.py +++ b/components/src/dynamo/sglang/init_diffusion.py @@ -141,19 +141,10 @@ async def init_image_diffusion( if not server_args.model_path: raise ValueError("--model is required for diffusion workers") - tp_size = getattr(server_args, "tp_size", 1) - dp_size = getattr(server_args, "dp_size", 1) - num_gpus = tp_size * dp_size - - dist_timeout = getattr(server_args, "dist_timeout", None) - - generator = DiffGenerator.from_pretrained( - model_path=server_args.model_path, - num_gpus=num_gpus, - tp_size=tp_size, - dp_size=dp_size, - dist_timeout=dist_timeout, - ) + # server_args is a DiffusionWorkerArgs adapter; engine_args is the + # natively parsed sglang diffusion ServerArgs, handed over wholesale so + # every engine option the user set on the CLI takes effect. + generator = DiffGenerator.from_server_args(server_args.engine_args) fs_url = dynamo_args.media_output_fs_url @@ -228,19 +219,10 @@ async def init_video_diffusion( if not server_args.model_path: raise ValueError("--model is required for video generation workers") - tp_size = getattr(server_args, "tp_size", 1) - dp_size = getattr(server_args, "dp_size", 1) - num_gpus = tp_size * dp_size - - dist_timeout = getattr(server_args, "dist_timeout", None) - - generator = DiffGenerator.from_pretrained( - model_path=server_args.model_path, - num_gpus=num_gpus, - tp_size=tp_size, - dp_size=dp_size, - dist_timeout=dist_timeout, - ) + # server_args is a DiffusionWorkerArgs adapter; engine_args is the + # natively parsed sglang diffusion ServerArgs, handed over wholesale so + # every engine option the user set on the CLI takes effect. + generator = DiffGenerator.from_server_args(server_args.engine_args) fs_url = dynamo_args.media_output_fs_url diff --git a/components/src/dynamo/sglang/tests/test_sglang_diffusion_args.py b/components/src/dynamo/sglang/tests/test_sglang_diffusion_args.py new file mode 100644 index 000000000000..9ec16f819d29 --- /dev/null +++ b/components/src/dynamo/sglang/tests/test_sglang_diffusion_args.py @@ -0,0 +1,290 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for native diffusion worker argument parsing. + +The real sglang.multimodal_gen ServerArgs resolves model info at parse time, +so these tests inject a lightweight fake through sys.modules. They cover the +Dynamo layer only: dynamo-side flag registration, adapter delegation, and the +handoff contract with the engine parser. +""" + +import argparse +import dataclasses +import sys +from types import ModuleType +from typing import Optional + +import pytest + +from dynamo.sglang.diffusion_args import DiffusionWorkerArgs + +pytestmark = [ + pytest.mark.unit, + pytest.mark.sglang, + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.parallel, +] + + +@dataclasses.dataclass +class FakeEngineArgs: + """Stand-in for sglang.multimodal_gen ServerArgs (0.5.17 field subset).""" + + model_path: str + tp_size: Optional[int] = None + dp_size: int = 1 + log_level: str = "info" + enable_trace: bool = False + attention_backend: Optional[str] = None + performance_mode: str = "auto" + num_gpus: int = 1 + + @staticmethod + def add_cli_args(parser): + parser.add_argument("--model-path", type=str) + parser.add_argument("--tp-size", type=int, default=None) + parser.add_argument("--dp-size", type=int, default=1) + parser.add_argument("--num-gpus", type=int, default=1) + parser.add_argument("--log-level", type=str, default="info") + parser.add_argument("--enable-trace", action="store_true") + parser.add_argument("--attention-backend", type=str, default=None) + parser.add_argument("--performance-mode", type=str, default="auto") + return parser + + @classmethod + def from_cli_args(cls, args, unknown_args=None): + if unknown_args: + raise SystemExit(f"error: unrecognized arguments: {' '.join(unknown_args)}") + field_names = {f.name for f in dataclasses.fields(cls)} + values = {k: v for k, v in vars(args).items() if k in field_names} + return cls(**values) + + +def _install_fake_sglang(monkeypatch, engine_cls=FakeEngineArgs): + """Override the sglang.multimodal_gen leaf modules with fakes. + + Only the leaves are faked: parent packages stay real (other dynamo + modules import e.g. sglang.srt at module level), and Python resolves a + fully dotted name from sys.modules directly, so leaf overrides suffice. + Missing parents (envs without the multimodal extra) get bare stubs. + """ + server_args_mod = ModuleType( + "sglang.multimodal_gen.runtime.server_args.server_args" + ) + server_args_mod.ServerArgs = engine_cls + + utils_mod = ModuleType("sglang.multimodal_gen.utils") + utils_mod.FlexibleArgumentParser = argparse.ArgumentParser + + import importlib + + for parent in ( + "sglang", + "sglang.multimodal_gen", + "sglang.multimodal_gen.runtime", + "sglang.multimodal_gen.runtime.server_args", + ): + if parent in sys.modules: + continue + try: + importlib.import_module(parent) + except ImportError: + monkeypatch.setitem(sys.modules, parent, ModuleType(parent)) + + monkeypatch.setitem( + sys.modules, + "sglang.multimodal_gen.runtime.server_args.server_args", + server_args_mod, + ) + monkeypatch.setitem(sys.modules, "sglang.multimodal_gen.utils", utils_mod) + + +class TestDiffusionWorkerArgs: + """Adapter behavior: delegation, fallbacks, inert defaults.""" + + def _engine(self, **overrides): + return FakeEngineArgs(model_path="org/model", **overrides) + + def test_engine_fields_delegate(self): + adapter = DiffusionWorkerArgs( + self._engine(tp_size=4, log_level="debug"), + served_model_name="my-model", + enable_metrics=True, + ) + assert adapter.model_path == "org/model" + assert adapter.tp_size == 4 + assert adapter.log_level == "debug" + assert adapter.served_model_name == "my-model" + assert adapter.enable_metrics is True + + def test_served_model_name_falls_back_to_model_path(self): + adapter = DiffusionWorkerArgs( + self._engine(), served_model_name=None, enable_metrics=False + ) + assert adapter.served_model_name == "org/model" + + def test_inert_defaults_present(self): + """Fields Dynamo's shared worker code probes must exist and be inert.""" + adapter = DiffusionWorkerArgs( + self._engine(), served_model_name=None, enable_metrics=False + ) + assert adapter.speculative_algorithm is None + assert adapter.disaggregation_mode is None + assert adapter.dllm_algorithm is False + assert adapter.load_format is None + assert adapter.kv_events_config is None + assert adapter.enable_forward_pass_metrics is False + + def test_missing_engine_field_raises(self): + adapter = DiffusionWorkerArgs( + self._engine(), served_model_name=None, enable_metrics=False + ) + with pytest.raises(AttributeError): + _ = adapter.not_a_real_field + + def test_native_field_wins_over_inert_default(self): + """If a future engine ServerArgs grows one of the pinned fields, + its native value must win over the inert default.""" + + @dataclasses.dataclass + class EngineWithLoadFormat(FakeEngineArgs): + load_format: Optional[str] = "auto" + + adapter = DiffusionWorkerArgs( + EngineWithLoadFormat(model_path="org/model"), + served_model_name=None, + enable_metrics=False, + ) + assert adapter.load_format == "auto" + + +class TestParseDiffusionArgs: + """Parsing behavior against a fake engine parser.""" + + def test_engine_args_parse_natively(self, monkeypatch): + _install_fake_sglang(monkeypatch) + from dynamo.sglang.diffusion_args import parse_diffusion_args + + _, adapter = parse_diffusion_args( + [ + "--model-path", + "org/model", + "--attention-backend", + "flash", + "--performance-mode", + "fast", + "--tp-size", + "2", + ] + ) + assert adapter.engine_args.attention_backend == "flash" + assert adapter.engine_args.performance_mode == "fast" + assert adapter.engine_args.tp_size == 2 + + def test_dynamo_side_flags_split_out(self, monkeypatch): + """--served-model-name / --enable-metrics must not leak into the + engine args when the engine parser does not define them.""" + _install_fake_sglang(monkeypatch) + from dynamo.sglang.diffusion_args import parse_diffusion_args + + _, adapter = parse_diffusion_args( + [ + "--model-path", + "org/model", + "--served-model-name", + "alias", + "--enable-metrics", + ] + ) + assert adapter.served_model_name == "alias" + assert adapter.enable_metrics is True + assert not hasattr(adapter.engine_args, "served_model_name") + assert not hasattr(adapter.engine_args, "enable_metrics") + + def test_native_served_model_name_preferred(self, monkeypatch): + """When the engine parser defines --served-model-name natively (sglang + main does), the Dynamo-side flag must not be registered twice.""" + + @dataclasses.dataclass + class EngineWithServedName(FakeEngineArgs): + served_model_name: Optional[str] = None + + @staticmethod + def add_cli_args(parser): + FakeEngineArgs.add_cli_args(parser) + parser.add_argument("--served-model-name", type=str, default=None) + return parser + + _install_fake_sglang(monkeypatch, engine_cls=EngineWithServedName) + from dynamo.sglang.diffusion_args import parse_diffusion_args + + _, adapter = parse_diffusion_args( + ["--model-path", "org/model", "--served-model-name", "native-name"] + ) + assert adapter.engine_args.served_model_name == "native-name" + assert adapter.served_model_name == "native-name" + + def test_abbreviated_flag_resolves_to_destination(self, monkeypatch): + """--tp (abbreviation of --tp-size) must be recorded under its parser + destination in the explicit-args side channel, not its raw spelling, + or the engine treats tp_size as unspecified and defaults to 1.""" + _install_fake_sglang(monkeypatch) + from dynamo.sglang.diffusion_args import parse_diffusion_args + + parsed, adapter = parse_diffusion_args( + ["--model-path", "org/model", "--tp", "2"] + ) + assert adapter.engine_args.tp_size == 2 + assert "tp_size" in parsed._sglang_explicit_arg_names + assert "tp" not in parsed._sglang_explicit_arg_names + + def test_num_gpus_derived_from_parallelism(self, monkeypatch): + """Without --num-gpus, tp*dp > 1 must derive num_gpus (the engine + defaults it to 1 and does not derive it itself).""" + _install_fake_sglang(monkeypatch) + from dynamo.sglang.diffusion_args import parse_diffusion_args + + parsed, adapter = parse_diffusion_args( + ["--model-path", "org/model", "--tp-size", "2", "--dp-size", "2"] + ) + assert adapter.engine_args.num_gpus == 4 + assert "num_gpus" in parsed._sglang_explicit_arg_names + + def test_explicit_num_gpus_not_overridden(self, monkeypatch): + """An explicit --num-gpus wins over the tp*dp derivation.""" + _install_fake_sglang(monkeypatch) + from dynamo.sglang.diffusion_args import parse_diffusion_args + + _, adapter = parse_diffusion_args( + ["--model-path", "org/model", "--tp-size", "2", "--num-gpus", "8"] + ) + assert adapter.engine_args.num_gpus == 8 + + def test_help_routed_to_diffusion_parser(self, monkeypatch, capsys): + """--help on a diffusion worker must show the diffusion engine's + options (and Dynamo's), not the LLM engine's.""" + import asyncio + + _install_fake_sglang(monkeypatch) + from dynamo.sglang.args import parse_args + + with pytest.raises(SystemExit) as excinfo: + asyncio.run(parse_args(["--image-diffusion-worker", "--help"])) + assert excinfo.value.code == 0 + out = capsys.readouterr().out + assert "--performance-mode" in out # engine option visible + assert "--served-model-name" in out # dynamo-side option visible + assert "--image-diffusion-worker" in out # dynamo worker flag visible + + def test_unknown_flag_rejected(self, monkeypatch): + """Typos and unsupported flags fail loudly instead of being silently + absorbed (the failure mode the old stub had).""" + _install_fake_sglang(monkeypatch) + from dynamo.sglang.diffusion_args import parse_diffusion_args + + with pytest.raises(SystemExit): + parse_diffusion_args( + ["--model-path", "org/model", "--not-a-real-flag", "value"] + ) diff --git a/examples/backends/sglang/launch/image_diffusion.sh b/examples/backends/sglang/launch/image_diffusion.sh index 0a70521535f8..a521328efb24 100755 --- a/examples/backends/sglang/launch/image_diffusion.sh +++ b/examples/backends/sglang/launch/image_diffusion.sh @@ -107,7 +107,6 @@ python3 -m dynamo.sglang \ --media-output-fs-url "$FS_URL" \ "${HTTP_URL_ARGS[@]}" \ --trust-remote-code \ - --skip-tokenizer-init \ --enable-metrics \ "${EXTRA_ARGS[@]}" & diff --git a/examples/backends/sglang/launch/text-to-video-diffusion.sh b/examples/backends/sglang/launch/text-to-video-diffusion.sh index 30ec9ef3b7e1..e19f2fd69203 100755 --- a/examples/backends/sglang/launch/text-to-video-diffusion.sh +++ b/examples/backends/sglang/launch/text-to-video-diffusion.sh @@ -127,11 +127,10 @@ echo "Starting T2V Worker ($WAN_SIZE)..." python3 -m dynamo.sglang \ --model-path "$MODEL_PATH" \ --served-model-name "$MODEL_PATH" \ - --tp "$TP_SIZE" \ + --tp-size "$TP_SIZE" \ --video-generation-worker \ --media-output-fs-url "$FS_URL" \ --trust-remote-code \ - --skip-tokenizer-init \ --enable-metrics \ "${EXTRA_ARGS[@]}" &