diff --git a/docs/advanced/pd-disaggregation.md b/docs/advanced/pd-disaggregation.md index 1df3c7152c0..0e0c5574cc1 100644 --- a/docs/advanced/pd-disaggregation.md +++ b/docs/advanced/pd-disaggregation.md @@ -25,8 +25,8 @@ PD disaggregation splits them into two pools, each sized for its own workload. `--prefill-num-servers` is a Miles-native flag added by `add_prefill_decode_disaggregation_arguments` in `miles/utils/arguments.py`. -When set, `miles/ray/rollout/rollout_server.py` calls -`SglangConfig.from_prefill_num_servers(args)` to dedicate that many SGLang +When set, `resolve_sglang_config` in +`miles/backends/sglang_utils/sglang_config.py` dedicates that many SGLang servers to prefill, with the rest used for decode. `--prefill-num-servers` is mutually exclusive with the `sglang_config` diff --git a/miles/backends/sglang_utils/sglang_config.py b/miles/backends/sglang_utils/sglang_config.py index afffaf356b9..dda49284c30 100644 --- a/miles/backends/sglang_utils/sglang_config.py +++ b/miles/backends/sglang_utils/sglang_config.py @@ -2,16 +2,21 @@ import dataclasses import logging +from pathlib import Path +import pydantic import yaml from miles.backends.sglang_utils.arguments import collect_eval_sglang_overrides +from miles.utils.pydantic_utils import FrozenStrictBaseModel logger = logging.getLogger(__name__) -@dataclasses.dataclass -class ServerGroupConfig: +# ---------------------------- raw config ----------------------------- + + +class _RawServerGroupConfig(FrozenStrictBaseModel): """Configuration for a single server group. Attributes: @@ -28,18 +33,10 @@ class ServerGroupConfig: worker_type: str num_gpus: int num_gpus_per_engine: int | None = None - overrides: dict = dataclasses.field(default_factory=dict) - - def __post_init__(self): - valid_types = {"regular", "prefill", "decode", "placeholder"} - assert ( - self.worker_type in valid_types - ), f"Invalid worker_type '{self.worker_type}', must be one of {valid_types}" - assert self.num_gpus > 0, f"num_gpus must be > 0, got {self.num_gpus}" + overrides: dict = {} -@dataclasses.dataclass -class ModelConfig: +class _RawModelConfig(FrozenStrictBaseModel): """Configuration for a single model deployment. Attributes: @@ -59,51 +56,18 @@ class ModelConfig: name: str model_path: str | None = None num_gpus_per_engine: int | None = None - server_groups: list[ServerGroupConfig] = dataclasses.field(default_factory=list) + server_groups: list[_RawServerGroupConfig] = pydantic.Field( + default_factory=list, + validation_alias=pydantic.AliasChoices("server_groups", "engine_groups"), + ) update_weights: bool | None = None - def resolve(self, args) -> None: - """Resolve per-group defaults from model-level then args-level values.""" - default_gpus_per_engine = self.num_gpus_per_engine or args.rollout_num_gpus_per_engine - default_model_path = self.model_path or args.hf_checkpoint - for g in self.server_groups: - if g.num_gpus_per_engine is None: - g.num_gpus_per_engine = default_gpus_per_engine - if "model_path" not in g.overrides: - g.overrides["model_path"] = default_model_path - - if self.server_groups: - model_paths = {g.overrides["model_path"] for g in self.server_groups} - assert len(model_paths) == 1, ( - f"Model '{self.name}' has server groups with different model_path values: " - f"{model_paths}. All server groups within a model must use the same model_path." - ) - effective_model_path = model_paths.pop() - else: - effective_model_path = default_model_path - - if self.update_weights is None: - if effective_model_path != args.hf_checkpoint: - logger.warning( - f"Model '{self.name}' uses model_path='{effective_model_path}' which differs " - f"from hf_checkpoint='{args.hf_checkpoint}'. Defaulting update_weights to False. " - f"Set update_weights explicitly in the config to suppress this warning." - ) - self.update_weights = False - else: - self.update_weights = True - - @property - def has_pd_disaggregation(self) -> bool: - return any(g.worker_type in ("prefill", "decode") for g in self.server_groups) - @property def total_num_gpus(self) -> int: return sum(g.num_gpus for g in self.server_groups) -@dataclasses.dataclass -class SglangConfig: +class _RawSglangConfig(FrozenStrictBaseModel): """Configuration for SGLang engine deployment. Loaded from ``--sglang-config`` YAML file. @@ -142,73 +106,146 @@ class SglangConfig: ``server_groups`` in the YAML config. """ - models: list[ModelConfig] + models: list[_RawModelConfig] = pydantic.Field(validation_alias=pydantic.AliasChoices("models", "sglang")) - @staticmethod - def from_yaml(path: str) -> "SglangConfig": - with open(path) as f: - data = yaml.safe_load(f) - - assert "sglang" in data, ( - f"sglang config must have a 'sglang' key, got {list(data.keys())}. " - f"Wrap your server_groups inside a model entry under 'sglang'." - ) - models = [] - for m in data["sglang"]: - raw_groups = m.get("server_groups") or m.get("engine_groups") or [] - groups = [ServerGroupConfig(**g) for g in raw_groups] - models.append( - ModelConfig( - name=m["name"], - model_path=m.get("model_path"), - num_gpus_per_engine=m.get("num_gpus_per_engine"), - server_groups=groups, - update_weights=m.get("update_weights"), - ) - ) - return SglangConfig(models=models) + @classmethod + def from_yaml(cls, path: str) -> "_RawSglangConfig": + return cls.model_validate(yaml.safe_load(Path(path).read_text())) @staticmethod - def from_prefill_num_servers(args) -> "SglangConfig": + def from_prefill_num_servers(args) -> "_RawSglangConfig": """Build a config equivalent to the legacy --prefill-num-servers flag.""" total_gpus = args.rollout_num_gpus prefill_gpus = args.prefill_num_servers * args.rollout_num_gpus_per_engine decode_gpus = total_gpus - prefill_gpus assert decode_gpus > 0, f"No decode GPUs: total {total_gpus}, prefill {prefill_gpus}" - return SglangConfig( + return _RawSglangConfig( models=[ - ModelConfig( + _RawModelConfig( name="default", server_groups=[ - ServerGroupConfig(worker_type="prefill", num_gpus=prefill_gpus), - ServerGroupConfig(worker_type="decode", num_gpus=decode_gpus), + _RawServerGroupConfig(worker_type="prefill", num_gpus=prefill_gpus), + _RawServerGroupConfig(worker_type="decode", num_gpus=decode_gpus), ], ) ] ) + @property + def total_num_gpus(self) -> int: + return sum(m.total_num_gpus for m in self.models) + + +# ---------------------------- resolved config ----------------------------- + + +@dataclasses.dataclass +class ServerGroupConfig: + worker_type: str + num_gpus: int + num_gpus_per_engine: int | None = None + overrides: dict = dataclasses.field(default_factory=dict) + + def __post_init__(self): + valid_types = {"regular", "prefill", "decode", "placeholder"} + assert ( + self.worker_type in valid_types + ), f"Invalid worker_type '{self.worker_type}', must be one of {valid_types}" + assert self.num_gpus > 0, f"num_gpus must be > 0, got {self.num_gpus}" + + +@dataclasses.dataclass +class ModelConfig: + name: str + model_path: str | None = None + num_gpus_per_engine: int | None = None + server_groups: list[ServerGroupConfig] = dataclasses.field(default_factory=list) + update_weights: bool | None = None + + def resolve(self, args) -> None: + """Resolve per-group defaults from model-level then args-level values.""" + default_gpus_per_engine = self.num_gpus_per_engine or args.rollout_num_gpus_per_engine + default_model_path = self.model_path or args.hf_checkpoint + for g in self.server_groups: + if g.num_gpus_per_engine is None: + g.num_gpus_per_engine = default_gpus_per_engine + if "model_path" not in g.overrides: + g.overrides["model_path"] = default_model_path + + if self.server_groups: + model_paths = {g.overrides["model_path"] for g in self.server_groups} + assert len(model_paths) == 1, ( + f"Model '{self.name}' has server groups with different model_path values: " + f"{model_paths}. All server groups within a model must use the same model_path." + ) + effective_model_path = model_paths.pop() + else: + effective_model_path = default_model_path + + if self.update_weights is None: + if effective_model_path != args.hf_checkpoint: + logger.warning( + f"Model '{self.name}' uses model_path='{effective_model_path}' which differs " + f"from hf_checkpoint='{args.hf_checkpoint}'. Defaulting update_weights to False. " + f"Set update_weights explicitly in the config to suppress this warning." + ) + self.update_weights = False + else: + self.update_weights = True + @property def has_pd_disaggregation(self) -> bool: - return any(m.has_pd_disaggregation for m in self.models) + return any(g.worker_type in ("prefill", "decode") for g in self.server_groups) + + +@dataclasses.dataclass +class SglangConfig: + models: list[ModelConfig] @property - def total_num_gpus(self) -> int: - return sum(m.total_num_gpus for m in self.models) + def has_pd_disaggregation(self) -> bool: + return any(m.has_pd_disaggregation for m in self.models) def resolve_sglang_config(args) -> SglangConfig: """Build a SglangConfig from args, choosing the right source.""" - config = _compute_raw_sglang_config(args) + raw = _compute_raw_sglang_config(args) + + config = SglangConfig( + models=[ + ModelConfig( + name=m.name, + model_path=m.model_path, + num_gpus_per_engine=m.num_gpus_per_engine, + server_groups=[ServerGroupConfig(**g.model_dump()) for g in m.server_groups], + update_weights=m.update_weights, + ) + for m in raw.models + ] + ) + if args.eval_num_gpus > 0: + eval_models = [m for m in config.models if m.name == "eval"] + if not eval_models: + eval_models = [ + ModelConfig( + name="eval", + server_groups=[ServerGroupConfig(worker_type="regular", num_gpus=args.eval_num_gpus)], + ) + ] + config.models.append(eval_models[0]) + _apply_eval_model_config(eval_models[0], args) + for model in config.models: model.resolve(args) + return config -def _compute_raw_sglang_config(args) -> SglangConfig: +def _compute_raw_sglang_config(args) -> _RawSglangConfig: eval_num_gpus = args.eval_num_gpus if getattr(args, "sglang_config", None) is not None: - config = SglangConfig.from_yaml(args.sglang_config) + config = _RawSglangConfig.from_yaml(args.sglang_config) expected = args.rollout_num_gpus + eval_num_gpus actual = config.total_num_gpus assert ( @@ -224,25 +261,16 @@ def _compute_raw_sglang_config(args) -> SglangConfig: return config if args.prefill_num_servers is not None: - config = SglangConfig.from_prefill_num_servers(args) - else: - config = SglangConfig( - models=[ - ModelConfig( - name="default", - server_groups=[ServerGroupConfig(worker_type="regular", num_gpus=args.rollout_num_gpus)], - ) - ] - ) + return _RawSglangConfig.from_prefill_num_servers(args) - if eval_num_gpus > 0: - eval_model = ModelConfig( - name="eval", - server_groups=[ServerGroupConfig(worker_type="regular", num_gpus=eval_num_gpus)], - ) - _apply_eval_model_config(eval_model, args) - config.models.append(eval_model) - return config + return _RawSglangConfig( + models=[ + _RawModelConfig( + name="default", + server_groups=[_RawServerGroupConfig(worker_type="regular", num_gpus=args.rollout_num_gpus)], + ) + ] + ) def _eval_sglang_overrides(args) -> dict: @@ -263,7 +291,7 @@ def _eval_sglang_overrides(args) -> dict: return overrides | collect_eval_sglang_overrides(args) -def _apply_eval_model_config(model_cfg: ModelConfig, args) -> None: +def _apply_eval_model_config(model_cfg: "ModelConfig", args) -> None: """Fill the eval model from the ``--eval-*`` args: YAML > ``--eval-sglang-*`` > ``--sglang-*``.""" if model_cfg.update_weights is None: # Never joins the training broadcast group; the fleet is synced by snapshot only. diff --git a/tests/fast/backends/sglang_utils/test_sglang_config.py b/tests/fast/backends/sglang_utils/test_sglang_config.py index 78fdf486e91..84ccd2e408a 100644 --- a/tests/fast/backends/sglang_utils/test_sglang_config.py +++ b/tests/fast/backends/sglang_utils/test_sglang_config.py @@ -93,6 +93,36 @@ def test_a_yaml_without_the_sglang_key_is_rejected(self, tmp_path): with pytest.raises((AssertionError, ValueError), match="sglang|models"): _resolve_yaml(tmp_path, "other_key:\n - name: actor\n", rollout_num_gpus=8) + def test_an_unknown_model_key_is_rejected(self, tmp_path): + """Typos at the model level fail parsing instead of being silently dropped.""" + with pytest.raises(ValueError, match="typo_key"): + _resolve_yaml( + tmp_path, + "sglang:\n" + " - name: actor\n" + " typo_key: 1\n" + " server_groups:\n" + " - worker_type: regular\n" + " num_gpus: 8\n", + rollout_num_gpus=8, + ) + + def test_giving_both_group_spellings_is_rejected(self, tmp_path): + """server_groups plus engine_groups on one model is ambiguous and fails parsing.""" + with pytest.raises(ValueError, match="engine_groups"): + _resolve_yaml( + tmp_path, + "sglang:\n" + " - name: actor\n" + " server_groups:\n" + " - worker_type: regular\n" + " num_gpus: 8\n" + " engine_groups:\n" + " - worker_type: regular\n" + " num_gpus: 8\n", + rollout_num_gpus=8, + ) + class TestPrefillNumServersPath: def test_prefill_consuming_all_gpus_is_rejected(self):