diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py index 09fd5eec7..02860daf3 100644 --- a/src/dstack/_internal/core/models/configurations.py +++ b/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: diff --git a/src/dstack/_internal/server/compatibility/runs.py b/src/dstack/_internal/server/compatibility/runs.py index dd1510f67..688581f02 100644 --- a/src/dstack/_internal/server/compatibility/runs.py +++ b/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) diff --git a/src/tests/_internal/core/models/test_configurations.py b/src/tests/_internal/core/models/test_configurations.py index e0d91c718..d72b5ba9d 100644 --- a/src/tests/_internal/core/models/test_configurations.py +++ b/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,7 +1130,7 @@ 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", @@ -1162,15 +1138,16 @@ def test_dump_is_legacy_canonical(self): "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): diff --git a/src/tests/_internal/server/compatibility/test_runs.py b/src/tests/_internal/server/compatibility/test_runs.py index 9cdcd074a..2a02c66db 100644 --- a/src/tests/_internal/server/compatibility/test_runs.py +++ b/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")