Parse the raw sglang config into a dedicated pydantic model

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`

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.

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):