Replace the legacy replica groups serializer with a client patch

src/dstack/_internal/core/models/configurations.py

@@ -12,13 +12,11 @@
     GetCoreSchemaHandler,
     PositiveInt,
     RootModel,
-    SerializerFunctionWrapHandler,
     ValidationError,
     ValidationInfo,
     conint,
     constr,
     field_validator,
-    model_serializer,
     model_validator,
 )
 from pydantic_core import CoreSchema, core_schema
@@ -1257,20 +1255,6 @@ def _normalize_legacy_replica_groups(cls, data: Any) -> Any:
             raise ValueError("`replicas` and `groups` are mutually exclusive")
         return data
 
-    @model_serializer(mode="wrap")
-    def _serialize_legacy_replica_groups(
-        self, handler: SerializerFunctionWrapHandler
-    ) -> Dict[str, Any]:
-        res = handler(self)
-        groups = res.pop("groups", None)
-        if groups is None:
-            return res
-        for group in groups:
-            if "replicas" in group:
-                group["count"] = group.pop("replicas")
-        res["replicas"] = groups
-        return res
-
     @field_validator("port")
     @classmethod
     def convert_port(cls, v) -> PortMapping:

src/dstack/_internal/server/compatibility/runs.py

@@ -31,6 +31,20 @@ def patch_run(run: Run, client_version: Optional[Version]) -> None:
 def patch_run_spec(run_spec: RunSpec, client_version: Optional[Version]) -> None:
     if client_version is None:
         return
+    # Clients prior to 0.21.3 do not support `groups` on services
+    if (
+        client_version < Version("0.21.3")
+        and isinstance(run_spec.configuration, ServiceConfiguration)
+        and run_spec.configuration.groups is not None
+    ):
+        groups = [group.model_dump(mode="json") for group in run_spec.configuration.groups]
+        for group in groups:
+            group["count"] = group.pop("replicas")
+        # `replicas` was a list of replica groups in old clients but is now typed
+        # `Optional[Range[int]]`. Safe to ignore: pydantic validates on parse, not
+        # on assignment, and serializes the value as set.
+        run_spec.configuration.replicas = groups  # type: ignore[assignment]
+        run_spec.configuration.groups = None
     # Clients that type nodes as int reject null. Homogeneous default is 1.
     if (
         isinstance(run_spec.configuration, TaskConfiguration)

src/tests/_internal/core/models/test_configurations.py

@@ -1,12 +1,11 @@
 from copy import deepcopy
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 import pytest
-from pydantic import ValidationError, model_validator
-from typing_extensions import Self
+from pydantic import ValidationError
 
 from dstack._internal.core.errors import ConfigurationError
-from dstack._internal.core.models.common import CoreModel, RegistryAuth, validate_extra_ignore
+from dstack._internal.core.models.common import RegistryAuth
 from dstack._internal.core.models.configurations import (
     DevEnvironmentConfigurationParams,
     PresetConfiguration,
@@ -1096,29 +1095,6 @@ def test_accepts_top_level_resources_with_groups(self):
         assert parsed.resources.gpu.name == ["H100"]
 
 
-class _Legacy021ReplicaGroup(CoreModel):
-    """0.21-shaped group: size is `count`, no `groups` parent field."""
-
-    count: Range[int]
-    commands: list[str] = []
-
-
-class _Legacy021Service(CoreModel):
-    """Stand-in for a 0.21 client that does not know `groups`."""
-
-    commands: list[str] = []
-    image: Optional[str] = None
-    replicas: Optional[Union[list[_Legacy021ReplicaGroup], Range[int]]] = None
-
-    @model_validator(mode="after")
-    def check_image_or_commands_present(self) -> Self:
-        if isinstance(self.replicas, list):
-            return self
-        if not self.commands and self.image is None:
-            raise ValueError("Either `commands` or `image` must be set")
-        return self
-
-
 class TestServiceGroupsPhase1:
     def test_legacy_replicas_list_parses_to_groups(self):
         parsed = parse_run_configuration(
@@ -1154,23 +1130,24 @@ def test_new_groups_syntax_parses_identically(self):
         assert new.replicas is None
         assert legacy.groups == new.groups
 
-    def test_dump_is_legacy_canonical(self):
+    def test_dump_is_groups_canonical(self):
         parsed = parse_run_configuration(
             {
                 "type": "service",
                 "port": 8000,
                 "groups": [{"replicas": 1, "commands": ["x"]}],
             }
         )
+        # The legacy `replicas: [{count: ...}]` shape is produced only for old
+        # clients, by `server/compatibility/runs.py`, not by the model.
         dumped = parsed.model_dump()
-        assert "groups" not in dumped
-        assert isinstance(dumped["replicas"], list)
-        assert "count" in dumped["replicas"][0]
-        assert "replicas" not in dumped["replicas"][0]
+        assert dumped["replicas"] is None
+        assert isinstance(dumped["groups"], list)
+        assert "replicas" in dumped["groups"][0]
+        assert "count" not in dumped["groups"][0]
         dumped_json = parsed.model_dump(mode="json")
-        assert "groups" not in dumped_json
-        assert "count" in dumped_json["replicas"][0]
-        assert "replicas" not in dumped_json["replicas"][0]
+        assert dumped_json["replicas"] is None
+        assert "replicas" in dumped_json["groups"][0]
 
     def test_dump_validate_is_fixed_point(self):
         parsed = parse_run_configuration(
@@ -1185,18 +1162,7 @@ def test_dump_validate_is_fixed_point(self):
         twice = ServiceConfiguration.model_validate(once.model_dump())
         assert once.model_dump() == twice.model_dump() == parsed.model_dump()
 
-    def test_dumped_json_parses_as_0_21_client(self):
-        parsed = parse_run_configuration(
-            {
-                "type": "service",
-                "port": 8000,
-                "groups": [{"replicas": 1, "commands": ["x"]}],
-            }
-        )
-        dumped = parsed.model_dump()
-        validate_extra_ignore(_Legacy021Service, dumped)
-
-    def test_homogeneous_dump_has_no_groups_key(self):
+    def test_homogeneous_dump_has_null_groups(self):
         parsed = parse_run_configuration(
             {
                 "type": "service",
@@ -1205,8 +1171,9 @@ def test_homogeneous_dump_has_no_groups_key(self):
                 "replicas": 2,
             }
         )
+        # Nothing strips the key now that the model no longer rewrites groups.
         dumped = parsed.model_dump()
-        assert "groups" not in dumped
+        assert dumped["groups"] is None
         assert dumped["replicas"] == {"min": 2, "max": 2}
 
     def test_replicas_and_groups_rejected(self):

src/tests/_internal/server/compatibility/test_runs.py

@@ -1,16 +1,99 @@
-from typing import Optional
+from typing import Optional, Union
 
 import pytest
 from packaging.version import Version
+from pydantic_core import to_json
 
+from dstack._internal.core.models.common import CoreModel, validate_json_extra_ignore
 from dstack._internal.core.models.configurations import (
     AnyRunConfiguration,
     DevEnvironmentConfiguration,
+    ServiceConfiguration,
     TaskConfiguration,
+    parse_run_configuration,
+)
+from dstack._internal.core.models.resources import Range
+from dstack._internal.server.compatibility.runs import (
+    is_run_plan_for_offers_only,
+    patch_run_spec,
 )
-from dstack._internal.server.compatibility.runs import is_run_plan_for_offers_only
 from dstack._internal.server.testing.common import get_run_spec
 
+
+class _Legacy021ReplicaGroup(CoreModel):
+    """0.21-shaped group: size is `count`, no `groups` parent field."""
+
+    count: Range[int]
+    commands: list[str] = []
+
+
+class _Legacy021Service(CoreModel):
+    """Stand-in for a 0.21 client that does not know `groups`."""
+
+    commands: list[str] = []
+    image: Optional[str] = None
+    replicas: Optional[Union[list[_Legacy021ReplicaGroup], Range[int]]] = None
+
+
+def _grouped_service() -> ServiceConfiguration:
+    configuration = parse_run_configuration(
+        {
+            "type": "service",
+            "port": 8000,
+            "groups": [
+                {"replicas": 1, "commands": ["a"]},
+                {"replicas": 2, "commands": ["b"]},
+            ],
+        }
+    )
+    assert isinstance(configuration, ServiceConfiguration)
+    return configuration
+
+
+class TestPatchRunSpecReplicaGroups:
+    @pytest.mark.parametrize(
+        "client_version",
+        [Version("0.20.7"), Version("0.21.0"), Version("0.21.2")],
+    )
+    def test_downgrades_groups_for_clients_without_them(self, client_version):
+        run_spec = get_run_spec(repo_id="test", configuration=_grouped_service())
+
+        patch_run_spec(run_spec, client_version)
+
+        configuration = run_spec.configuration
+        assert configuration.groups is None
+        assert [group["count"] for group in configuration.replicas] == [
+            {"min": 1, "max": 1},
+            {"min": 2, "max": 2},
+        ]
+        # `to_json` is how the server actually renders a response, so this is
+        # exactly the payload an old client has to parse.
+        validate_json_extra_ignore(_Legacy021Service, to_json(configuration))
+
+    @pytest.mark.parametrize("client_version", [Version("0.21.3"), Version("0.22.0"), None])
+    def test_keeps_groups_for_clients_that_support_them(self, client_version):
+        run_spec = get_run_spec(repo_id="test", configuration=_grouped_service())
+
+        patch_run_spec(run_spec, client_version)
+
+        assert run_spec.configuration.replicas is None
+        assert [group.replicas for group in run_spec.configuration.groups] == [
+            Range[int](min=1, max=1),
+            Range[int](min=2, max=2),
+        ]
+
+    def test_leaves_a_service_without_groups_alone(self):
+        configuration = parse_run_configuration(
+            {"type": "service", "port": 8000, "commands": ["x"], "replicas": 2}
+        )
+        run_spec = get_run_spec(repo_id="test", configuration=configuration)
+
+        patch_run_spec(run_spec, Version("0.21.0"))
+
+        assert run_spec.configuration.groups is None
+        assert run_spec.configuration.replicas == Range[int](min=2, max=2)
+
+
 _OFFER_CLI_CONFIGURATION = TaskConfiguration(commands=[":"], image="scratch", user="root")