diff --git a/src/dstack/_internal/cli/services/configurators/run.py b/src/dstack/_internal/cli/services/configurators/run.py index e47c236df..75f4c6773 100644 --- a/src/dstack/_internal/cli/services/configurators/run.py +++ b/src/dstack/_internal/cli/services/configurators/run.py @@ -677,6 +677,9 @@ def register_args(cls, parser: argparse.ArgumentParser): def apply_args(self, conf: ServiceConfiguration, args: argparse.Namespace): super().apply_args(conf, args) self.apply_commands_args(conf, args) + if conf.groups is not None: + for group in conf.groups: + self._interpolate_commands(group.commands, args) def _get_ready_wait_interval(attempt: int) -> float: diff --git a/src/dstack/_internal/cli/services/presets/build.py b/src/dstack/_internal/cli/services/presets/build.py index f3fb37246..ab695710e 100644 --- a/src/dstack/_internal/cli/services/presets/build.py +++ b/src/dstack/_internal/cli/services/presets/build.py @@ -96,7 +96,7 @@ def set_service_gpu_vendor_from_verification( service: ServiceConfiguration, verified_on: list[PresetVerificationReplicaGroup], ) -> None: - for group_num, group in enumerate(service.replica_groups): + for group in service.replica_groups: resources = group.resources if resources is None or not _requires_gpu(resources): continue @@ -105,9 +105,7 @@ def set_service_gpu_vendor_from_verification( continue if resources.gpu.vendor is not None and resources.gpu.vendor != verification_vendor: raise ValueError("preset service GPU vendor does not match verification") - group_resources = _get_service_group_resources(service, group_num) - if group_resources.gpu is not None: - group_resources.gpu.vendor = verification_vendor + resources.gpu.vendor = verification_vendor def _without_excluded_fields(configuration: ConfigurationT) -> ConfigurationT: @@ -146,20 +144,6 @@ def _get_resources_gpu_vendor(resources: ResourcesSpec) -> gpuhunt.AcceleratorVe return next(iter(vendors), None) -def _get_service_group_resources( - service: ServiceConfiguration, - group_num: int, -) -> ResourcesSpec: - resources = ( - service.replicas[group_num].resources - if isinstance(service.replicas, list) - else service.resources - ) - if resources is None: - raise ValueError("preset service object must specify resources") - return resources - - def _requires_gpu(resources: ResourcesSpec) -> bool: gpu = resources.gpu if gpu is None or gpu.count.max == 0: diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py index 795bdbef8..61f7dd3a8 100644 --- a/src/dstack/_internal/core/models/configurations.py +++ b/src/dstack/_internal/core/models/configurations.py @@ -11,12 +11,13 @@ Field, GetCoreSchemaHandler, RootModel, - Tag, + SerializerFunctionWrapHandler, ValidationError, ValidationInfo, conint, constr, field_validator, + model_serializer, model_validator, ) from pydantic_core import CoreSchema, core_schema @@ -707,10 +708,12 @@ class ConfigurationWithCommandsParams(CoreModel): @model_validator(mode="after") def check_image_or_commands_present(self) -> Self: # If replicas is list, skip validation - commands come from replica groups + # (legacy in-memory shape; after parse, service groups live on `groups`). replicas = getattr(self, "replicas", None) if isinstance(replicas, list): return self # If groups is set, skip validation - commands come from node groups + # or service replica groups. if getattr(self, "groups", None) is not None: return self @@ -990,7 +993,7 @@ class ReplicaGroup(CoreModel): description="The name of the replica group. If not provided, defaults to '0', '1', etc. based on position." ), ] = None - count: Annotated[ + replicas: Annotated[ Range[int], Field( description="The number of replicas. Can be a number (e.g. `2`) or a range (`0..4` or `1..8`). " @@ -999,7 +1002,7 @@ class ReplicaGroup(CoreModel): ] scaling: Annotated[ Optional[ScalingSpec], - Field(description="The auto-scaling rules. Required if `count` is set to a range"), + Field(description="The auto-scaling rules. Required if `replicas` is set to a range"), ] = None resources: Annotated[ @@ -1069,6 +1072,19 @@ class ReplicaGroup(CoreModel): ), ] = None + @model_validator(mode="before") + @classmethod + def _alias_count_to_replicas(cls, data: Any) -> Any: + if not isinstance(data, dict): + return data + data = dict(data) + if "replicas" in data: + data.pop("count", None) + return data + if "count" in data: + data["replicas"] = data.pop("count") + return data + @field_validator("name") @classmethod def validate_name(cls, v: Optional[str]) -> Optional[str]: @@ -1077,9 +1093,9 @@ def validate_name(cls, v: Optional[str]) -> Optional[str]: raise ValueError("Resource name should match regex '^[a-z0-9][a-z0-9-]{0,39}$'") return v - @field_validator("count") + @field_validator("replicas") @classmethod - def convert_count(cls, v: Range[int]) -> Range[int]: + def convert_replicas(cls, v: Range[int]) -> Range[int]: return _validate_replica_range(v) @field_validator("python", mode="before") @@ -1124,11 +1140,11 @@ def _privileged(cls, v, info: ValidationInfo) -> Optional[bool]: @model_validator(mode="after") def validate_scaling(self) -> Self: scaling = self.scaling - count = self.count - if count and count.min != count.max and not scaling: - raise ValueError("When you set `count` to a range, ensure to specify `scaling`.") - if count and count.min == count.max and scaling: - raise ValueError("To use `scaling`, `count` must be set to a range.") + replicas = self.replicas + if replicas and replicas.min != replicas.max and not scaling: + raise ValueError("When you set `replicas` to a range, ensure to specify `scaling`.") + if replicas and replicas.min == replicas.max and scaling: + raise ValueError("To use `scaling`, `replicas` must be set to a range.") return self @@ -1206,27 +1222,53 @@ class ServiceConfigurationParams(CoreModel): ] = None # None = omitted (may get default when model is set); [] = explicit empty replicas: Annotated[ - Optional[ - # `Tag` only names the arm in validation errors. Without it the `loc` of a bad - # `replicas` spells out the whole wrapped schema — - # `service.replicas.list[function-after[validate_scaling(), ReplicaGroup]].0` — which - # is what `dstack apply` shows the user. - Union[ - Annotated[list[ReplicaGroup], Tag("ReplicaGroup")], - Annotated[Range[int], Tag("Range[int]")], - ] - ], + Optional[Range[int]], Field( description=( - "The number of replicas or a list of replica groups. " - "Can be an integer (e.g., `2`), a range (e.g., `0..4`), or a list of replica groups. " - "Each replica group defines replicas with shared configuration " - "(commands, resources, scaling). " - "When `replicas` is a list of replica groups, top-level `scaling`, `commands`, " - "and `resources` are not allowed and must be specified in each replica group instead. " + "The number of replicas for a homogeneous service. " + "Can be an integer (e.g. `2`) or a range (e.g. `0..4`). " + "Mutually exclusive with `groups`." ) ), ] = None + groups: Annotated[ + Optional[List[ReplicaGroup]], + Field( + description=( + "A list of replica groups for heterogeneous services. " + "Mutually exclusive with `replicas`. " + "When `groups` is set, top-level `scaling` and `commands` are not allowed." + ) + ), + ] = None + + @model_validator(mode="before") + @classmethod + def _normalize_legacy_replica_groups(cls, data: Any) -> Any: + if not isinstance(data, dict): + return data + data = dict(data) + replicas = data.get("replicas") + if data.get("groups") is None and isinstance(replicas, list): + data["groups"] = replicas + data.pop("replicas", None) + if data.get("groups") is not None and data.get("replicas") is not None: + 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 @@ -1280,31 +1322,28 @@ def validate_gateway( @field_validator("replicas") @classmethod - def validate_replicas( - cls, v: Optional[Union[Range[int], List[ReplicaGroup]]] - ) -> Optional[Union[Range[int], List[ReplicaGroup]]]: + def validate_replicas(cls, v: Optional[Range[int]]) -> Optional[Range[int]]: if v is None: return v - if isinstance(v, Range): - return _validate_replica_range(v) - - if isinstance(v, list): - if not v: - raise ValueError("`replicas` cannot be an empty list") - - # Assign default names to groups without names - for index, group in enumerate(v): - if group.name is None: - group.name = str(index) - - # Check for duplicate names - names = [group.name for group in v] - if len(names) != len(set(names)): - duplicates = [name for name in set(names) if names.count(name) > 1] - raise ValueError( - f"Duplicate replica group names found: {duplicates}. " - "Each replica group must have a unique name." - ) + return _validate_replica_range(v) + + @field_validator("groups") + @classmethod + def validate_groups(cls, v: Optional[List[ReplicaGroup]]) -> Optional[List[ReplicaGroup]]: + if v is None: + return v + if not v: + raise ValueError("`groups` cannot be an empty list") + for index, group in enumerate(v): + if group.name is None: + group.name = str(index) + counts = Counter(group.name for group in v) + duplicates = [name for name, count in counts.items() if count > 1] + if duplicates: + raise ValueError( + f"Duplicate replica group names found: {duplicates}. " + "Each replica group must have a unique name." + ) return v @model_validator(mode="after") @@ -1323,25 +1362,22 @@ def validate_scaling(self) -> Self: @model_validator(mode="after") def validate_top_level_properties_with_replica_groups(self) -> Self: - """ - When replicas is a list of ReplicaGroup, forbid top-level scaling and commands. - """ - replicas = self.replicas - - if not isinstance(replicas, list): + """When groups is set, forbid top-level scaling and commands.""" + groups = self.groups + if groups is None: return self scaling = self.scaling if scaling is not None: raise ValueError( - "Top-level `scaling` is not allowed when `replicas` is a list. " + "Top-level `scaling` is not allowed when `groups` is set. " "Specify `scaling` in each replica group instead." ) commands = getattr(self, "commands", None) if commands: raise ValueError( - "Top-level `commands` is not allowed when `replicas` is a list. " + "Top-level `commands` is not allowed when `groups` is set. " "Specify `commands` in each replica group instead." ) @@ -1350,13 +1386,13 @@ def validate_top_level_properties_with_replica_groups(self) -> Self: @model_validator(mode="after") def validate_no_mixed_service_and_group_container_fields(self) -> Self: """ - When replicas is a list, certain fields may be set + When `groups` is set, certain fields may be set at the service level OR in replica groups, never both. Mixing is rejected — including partial mixing, where only some groups set a field the service also sets — because it leaves precedence ambiguous. """ - replicas = self.replicas - if not isinstance(replicas, list): + groups = self.groups + if groups is None: return self checks = [ @@ -1399,7 +1435,7 @@ def validate_no_mixed_service_and_group_container_fields(self) -> Self: for field, service_set, group_set in checks: if service_set: - conflicting = [g.name for g in replicas if group_set(g)] + conflicting = [g.name for g in groups if group_set(g)] if conflicting: raise ValueError( f"`{field}` is set at both the service level and in " @@ -1415,8 +1451,8 @@ def validate_no_conflicting_image_sources_across_levels(self) -> Self: Image-source fields (`image`, `docker`, `python`, `nvcc`) cannot be mixed across service and group levels in conflicting ways. """ - replicas = self.replicas - if not isinstance(replicas, list): + groups = self.groups + if groups is None: return self forbidden = [ @@ -1479,7 +1515,7 @@ def validate_no_conflicting_image_sources_across_levels(self) -> Self: for s_field, s_set, g_field, g_pred in forbidden: if s_set: - conflicting = [g.name for g in replicas if g_pred(g)] + conflicting = [g.name for g in groups if g_pred(g)] if conflicting: raise ValueError( f"Service-level `{s_field}` conflicts with group-level " @@ -1491,19 +1527,18 @@ def validate_no_conflicting_image_sources_across_levels(self) -> Self: @model_validator(mode="after") def validate_replica_groups_have_commands_or_image(self) -> Self: """ - When replicas is a list, ensure each ReplicaGroup has something + When `groups` is set, ensure each ReplicaGroup has something to run. Mirrors the service-level rule: either explicit `commands` or an `image` (group-level or service-level) is required. """ - replicas = self.replicas - - if not isinstance(replicas, list): + groups = self.groups + if groups is None: return self service_has_image = getattr(self, "image", None) is not None - for group in replicas: + for group in groups: if not group.commands and group.image is None and not service_has_image: raise ValueError( f"Replica group '{group.name}': either `commands` or " @@ -1515,16 +1550,16 @@ def validate_replica_groups_have_commands_or_image(self) -> Self: @model_validator(mode="after") def validate_at_most_one_router_replica_group(self) -> Self: - replicas = self.replicas - if not isinstance(replicas, list): + groups = self.groups + if groups is None: return self - router_groups = [g for g in replicas if g.router is not None] + router_groups = [g for g in groups if g.router is not None] if len(router_groups) > 1: raise ValueError("At most one replica group may specify `router`.") if router_groups: router_group = router_groups[0] - if router_group.count.min != 1 or router_group.count.max != 1: - raise ValueError("For now replica group with `router` must have `count: 1`.") + if router_group.replicas.min != 1 or router_group.replicas.max != 1: + raise ValueError("For now replica group with `router` must have `replicas: 1`.") return self @@ -1538,31 +1573,18 @@ class ServiceConfiguration( @property def replica_groups(self) -> List[ReplicaGroup]: - if self.replicas is None: - return [ - ReplicaGroup( - name=DEFAULT_REPLICA_GROUP_NAME, - count=Range[int](min=1, max=1), - commands=self.commands, - resources=self.resources, - scaling=self.scaling, - ) - ] - if isinstance(self.replicas, list): - return self.replicas - if isinstance(self.replicas, Range): - return [ - ReplicaGroup( - name=DEFAULT_REPLICA_GROUP_NAME, - count=self.replicas, - commands=self.commands, - resources=self.resources, - scaling=self.scaling, - ) - ] - raise ValueError( - f"Invalid replicas type: {type(self.replicas)}. Expected None, Range[int], or List[ReplicaGroup]" - ) + if self.groups is not None: + return self.groups + replicas = self.replicas if self.replicas is not None else Range[int](min=1, max=1) + return [ + ReplicaGroup( + name=DEFAULT_REPLICA_GROUP_NAME, + replicas=replicas, + commands=self.commands, + resources=self.resources, + scaling=self.scaling, + ) + ] AnyRunConfiguration = Union[DevEnvironmentConfiguration, TaskConfiguration, ServiceConfiguration] diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py index a89ef1ede..0a4b90bac 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py @@ -18,6 +18,7 @@ ) from dstack._internal.core.models.configurations import ( DevEnvironmentConfiguration, + ServiceConfiguration, ) from dstack._internal.core.models.files import FileArchiveMapping from dstack._internal.core.models.gateways import GatewayReplicaStatus @@ -117,6 +118,7 @@ get_router_env_for_job, get_router_replica_group, ) +from dstack._internal.server.services.runs.spec import run_spec_has_replica_ip_refs from dstack._internal.server.services.secrets import get_project_secrets_mapping from dstack._internal.server.services.storage import get_default_storage from dstack._internal.server.utils import tracing @@ -124,9 +126,12 @@ from dstack._internal.utils.interpolator import InterpolatorError from dstack._internal.utils.logging import get_logger from dstack._internal.utils.nodes_interpolator import ( + GroupsIpMember, find_groups_ip_refs, interpolate_groups_ip_address, + interpolate_groups_replica_ip_address, validate_groups_ref_bounds, + validate_groups_ref_member, validate_groups_refs, ) @@ -640,6 +645,12 @@ async def _prepare_startup_context( try: for c in commands: validate_groups_refs(c) + if not _substitute_groups_ip_refs(commands, context): + logger.debug( + "%s: waiting for referenced group IPs", + fmt(context.job_model), + ) + return None except InterpolatorError as e: _terminate_job( job_model=context.job_model, @@ -649,31 +660,6 @@ async def _prepare_startup_context( ) return None - # Hetero commands may reference peer nodes via ${{ groups[i].nodes[j].IP_ADDRESS }}. - # Wait until every referenced node has an internal IP, then substitute placeholders - # in-place before the job is sent to the runner. Bad refs are rejected above; - # out-of-range or still-missing IPs terminate the job below. - if any(find_groups_ip_refs(c) for c in commands): - nodes_view = _build_nodes_ip_view(context.run.jobs, context.job.job_spec.replica_num) - try: - if not _referenced_ips_ready(commands, nodes_view): - logger.debug( - "%s: waiting for referenced node group IPs", - fmt(context.job_model), - ) - return None - context.job.job_spec.commands = [ - interpolate_groups_ip_address(c, nodes_view) for c in commands - ] - except InterpolatorError as e: - _terminate_job( - job_model=context.job_model, - job_update_map=result.job_update_map, - termination_reason=JobTerminationReason.TERMINATED_BY_SERVER, - termination_reason_message=f"Groups IP interpolation error: {e.args[0]}", - ) - return None - return _StartupContext( cluster_info=cluster_info, volumes=volumes, @@ -717,14 +703,11 @@ async def _fetch_run_model( replica_num: If None, skip loading jobs (for RUNNING jobs that don't need siblings). If set, load only latest-submission jobs for that replica (for PROVISIONING/PULLING jobs that need same-replica siblings for cluster coordination). When the run has - a Dynamo router replica group, all non-terminated latest-submission jobs for the - run are loaded so find_router_job can identify the router by replica-group - membership. - run_spec: Required whenever `replica_num` is set. Used only to detect - whether the run has a Dynamo router replica group. The caller is - expected to parse it once from the eager-loaded JobModel.run - (see _refetch_locked_job_model) so we don't issue a separate - query for it here. + a Dynamo router replica group, or any command contains + ${{ groups[i].replicas[j].IP_ADDRESS }}, all non-terminated latest-submission + jobs for the run are loaded so sibling replica IPs are visible. + run_spec: Required whenever `replica_num` is set. Used to detect whether the run + has a Dynamo router replica group or replica IP refs. """ query = ( select(RunModel) @@ -752,6 +735,7 @@ async def _fetch_run_model( and router_group.router is not None and router_group.router.type == RouterType.DYNAMO ) + load_all_replicas = is_dynamo or run_spec_has_replica_ip_refs(run_spec) latest_submissions_sq = ( select( @@ -762,9 +746,9 @@ async def _fetch_run_model( ) .where( JobModel.run_id == run_id, - # For Service with Dynamo router: load all replicas. For Non-Dynamo: only the worker's - # own replica. - true() if is_dynamo else JobModel.replica_num == replica_num, + # Dynamo and replica-IP interpolation need sibling replicas. + # Other services load only the worker's own replica. + true() if load_all_replicas else JobModel.replica_num == replica_num, ) .group_by(JobModel.run_id, JobModel.replica_num, JobModel.job_num) .subquery() @@ -779,12 +763,12 @@ async def _fetch_run_model( job_alias.replica_num == latest_submissions_sq.c.replica_num, job_alias.job_num == latest_submissions_sq.c.job_num, job_alias.submission_num == latest_submissions_sq.c.max_submission_num, - # For Dynamo runs, drop terminated rows so accumulated - # scale-down history doesn't bloat the load. Non-Dynamo - # runs are already restricted to the worker's own - # replica above, so this filter is a no-op for them. + # When loading all replicas, drop terminated rows so + # accumulated scale-down history doesn't bloat the load. + # Own-replica loads are already restricted above, so this + # filter is a no-op for them. or_( - false() if is_dynamo else true(), + false() if load_all_replicas else true(), ~job_alias.status.in_(JobStatus.finished_statuses()) & (job_alias.status != JobStatus.TERMINATING), ), @@ -1797,6 +1781,41 @@ def _reset_disconnected_at(job_model: JobModel, result: _ProcessResult) -> None: result.job_update_map["disconnected_at"] = None +def _substitute_groups_ip_refs(commands: list[str], context: _ProcessContext) -> bool: + """Wait for / substitute groups IP refs. Returns False if a min-slot IP is not ready.""" + configuration = context.run.run_spec.configuration + has_replica_refs = any( + member == "replicas" for c in commands for _, member, _ in find_groups_ip_refs(c) + ) + has_node_refs = any( + member == "nodes" for c in commands for _, member, _ in find_groups_ip_refs(c) + ) + if isinstance(configuration, ServiceConfiguration): + for command in commands: + validate_groups_ref_member(command, "replicas") + if not has_replica_refs: + return True + replica_view = _build_replica_groups_ip_view(context.run.jobs, configuration) + if not _referenced_ips_ready(commands, replica_view, member="replicas"): + return False + context.job.job_spec.commands = [ + interpolate_groups_replica_ip_address(c, replica_view) for c in commands + ] + return True + if has_replica_refs: + for command in commands: + validate_groups_ref_member(command, "nodes") + if not has_node_refs: + return True + nodes_view = _build_nodes_ip_view(context.run.jobs, context.job.job_spec.replica_num) + if not _referenced_ips_ready(commands, nodes_view): + return False + context.job.job_spec.commands = [ + interpolate_groups_ip_address(c, nodes_view) for c in commands + ] + return True + + def _build_nodes_ip_view(jobs: list[Job], replica_num: int) -> list[list[str]]: replica_jobs = [job for job in jobs if job.job_spec.replica_num == replica_num] if not replica_jobs: @@ -1808,22 +1827,60 @@ def _build_nodes_ip_view(jobs: list[Job], replica_num: int) -> list[list[str]]: local_index = job.job_spec.node_group_job_index while len(nodes[group_index]) <= local_index: nodes[group_index].append("") - ip = "" - if job.job_submissions: - jpd = job.job_submissions[-1].job_provisioning_data - if jpd is not None: - ip = jpd.internal_ip or "" - nodes[group_index][local_index] = ip + nodes[group_index][local_index] = _job_internal_ip(job) return nodes -def _referenced_ips_ready(commands: list[str], nodes_view: list[list[str]]) -> bool: +def _build_replica_groups_ip_view( + jobs: list[Job], configuration: ServiceConfiguration +) -> list[list[str]]: + """Fixed-length rows of replicas.min; slot j is the j-th live job in that group.""" + view: list[list[str]] = [] + for group_index, group in enumerate(configuration.replica_groups): + size = group.replicas.min or 0 + row = [""] * size + group_name = group.name if group.name is not None else str(group_index) + group_jobs = [ + job + for job in jobs + if job.job_spec.replica_group == group_name and not _job_is_finished(job) + ] + group_jobs.sort(key=lambda job: job.job_spec.replica_num) + for slot, job in enumerate(group_jobs[:size]): + row[slot] = _job_internal_ip(job) + view.append(row) + return view + + +def _job_internal_ip(job: Job) -> str: + if job.job_submissions: + jpd = job.job_submissions[-1].job_provisioning_data + if jpd is not None: + return jpd.internal_ip or "" + return "" + + +def _job_is_finished(job: Job) -> bool: + if not job.job_submissions: + return False + status = job.job_submissions[-1].status + return status.is_finished() or status == JobStatus.TERMINATING + + +def _referenced_ips_ready( + commands: list[str], + nodes_view: list[list[str]], + *, + member: GroupsIpMember = "nodes", +) -> bool: group_sizes = [len(g) for g in nodes_view] for command in commands: - validate_groups_ref_bounds(command, group_sizes) - for group_index, node_index in find_groups_ip_refs(command): + validate_groups_ref_bounds(command, group_sizes, member=member) + for group_index, ref_member, index in find_groups_ip_refs(command): + if ref_member != member: + continue # Wait until every referenced slot has a non-empty internal IP. - if not nodes_view[group_index][node_index]: + if not nodes_view[group_index][index]: return False return True diff --git a/src/dstack/_internal/server/background/pipeline_tasks/runs/common.py b/src/dstack/_internal/server/background/pipeline_tasks/runs/common.py index 3f27e3b2a..03c528c9c 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/runs/common.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/runs/common.py @@ -38,17 +38,17 @@ def compute_desired_replica_counts( and len(replica_groups) == 1 and replica_groups[0].name == DEFAULT_REPLICA_GROUP_NAME ): - # Special case to avoid dropping the replica count to group.count.min + # Special case to avoid dropping the replica count to group.replicas.min # when a 0.20.7+ server first processes a service created by a pre-0.20.7 server. # TODO: remove once most users upgrade to 0.20.7+. prev_counts = {DEFAULT_REPLICA_GROUP_NAME: run_model.desired_replica_count} desired_counts: PerGroupDesiredCounts = {} total = 0 for group in replica_groups: - scaler = get_service_scaler(group.count, group.scaling) + scaler = get_service_scaler(group.replicas, group.scaling) assert group.name is not None, "Group name is always set" group_desired = scaler.get_desired_count( - current_desired_count=prev_counts.get(group.name, group.count.min or 0), + current_desired_count=prev_counts.get(group.name, group.replicas.min or 0), stats=gateway_stats, last_scaled_at=last_scaled_at, ) diff --git a/src/dstack/_internal/server/services/runs/__init__.py b/src/dstack/_internal/server/services/runs/__init__.py index a541b3904..bf3e1c575 100644 --- a/src/dstack/_internal/server/services/runs/__init__.py +++ b/src/dstack/_internal/server/services/runs/__init__.py @@ -793,7 +793,7 @@ async def submit_run( if run_spec.merged_profile.schedule is not None: group_initial_replicas = 0 else: - group_initial_replicas = replica_group.count.min or 0 + group_initial_replicas = replica_group.replicas.min or 0 # Each replica in this group gets the same group-specific configuration for group_replica_num in range(group_initial_replicas): diff --git a/src/dstack/_internal/server/services/runs/spec.py b/src/dstack/_internal/server/services/runs/spec.py index ec132ff7f..0b5c506f5 100644 --- a/src/dstack/_internal/server/services/runs/spec.py +++ b/src/dstack/_internal/server/services/runs/spec.py @@ -30,7 +30,9 @@ from dstack._internal.utils.logging import get_logger from dstack._internal.utils.nodes_interpolator import ( contains_groups_ref, + find_groups_ip_refs, validate_groups_ref_bounds, + validate_groups_ref_member, validate_groups_refs, ) @@ -53,6 +55,7 @@ "service": [ # in-place "replicas", + "groups", "scaling", # rolling deployment # NOTE: keep this list in sync with the "Rolling deployment" section in services.md @@ -120,7 +123,7 @@ def validate_run_spec_and_set_defaults( ) if isinstance(run_spec.configuration, ServiceConfiguration): if run_spec.merged_profile.schedule and all( - group.count.min == 0 for group in run_spec.configuration.replica_groups + group.replicas.min == 0 for group in run_spec.configuration.replica_groups ): raise ServerClientError( "Scheduled services with autoscaling to zero are not supported" @@ -146,7 +149,7 @@ def validate_run_spec_and_set_defaults( and run_spec.configuration.nodes is None ): run_spec.configuration.nodes = 1 - # We do not reject top-level `resources` when `replicas` is a list. Adding strict checks + # We do not reject top-level `resources` when `groups` is set. Adding strict checks # would be fragile because the spec may be changed later (for example by plugins). # Same for task `groups`: provisioning uses each group's resources; top-level is not banned. set_run_spec_resources_defaults(run_spec) @@ -169,8 +172,8 @@ def set_run_spec_resources_defaults(run_spec: RunSpec) -> None: image=configuration.image, docker=configuration.docker, ) - if configuration.type == "service" and isinstance(configuration.replicas, list): - for replica_group in configuration.replicas: + if configuration.type == "service" and configuration.groups is not None: + for replica_group in configuration.groups: image, docker = _get_replica_group_image_and_docker(replica_group, configuration) _set_resources_defaults( resources_spec=replica_group.resources, @@ -214,9 +217,19 @@ def _validate_groups_ip_refs(run_spec: RunSpec) -> None: for command in _iter_configuration_commands(run_spec.configuration): validate_groups_refs(command) if isinstance(run_spec.configuration, TaskConfiguration): + for command in _iter_configuration_commands(run_spec.configuration): + validate_groups_ref_member(command, "nodes") group_sizes = [g.nodes for g in run_spec.configuration.node_groups] for command in _iter_configuration_commands(run_spec.configuration): validate_groups_ref_bounds(command, group_sizes) + elif isinstance(run_spec.configuration, ServiceConfiguration): + for command in _iter_configuration_commands(run_spec.configuration): + validate_groups_ref_member(command, "replicas") + group_sizes = [ + group.replicas.min or 0 for group in run_spec.configuration.replica_groups + ] + for command in _iter_configuration_commands(run_spec.configuration): + validate_groups_ref_bounds(command, group_sizes, member="replicas") except InterpolatorError as e: raise ServerClientError(e.args[0]) from e @@ -233,13 +246,23 @@ def _iter_configuration_commands(configuration: AnyRunConfiguration): yield from group.commands +def run_spec_has_replica_ip_refs(run_spec: RunSpec) -> bool: + """True if any service command contains ${{ groups[i].replicas[j].IP_ADDRESS }}.""" + if not isinstance(run_spec.configuration, ServiceConfiguration): + return False + for command in _iter_configuration_commands(run_spec.configuration): + if any(member == "replicas" for _, member, _ in find_groups_ip_refs(command)): + return True + return False + + def _validate_gpu_vendor_and_image(run_spec: RunSpec) -> None: configuration = run_spec.configuration vendors: set[gpuhunt.AcceleratorVendor] = set() invalid_replicas: list[int] = [] invalid_groups: list[int] = [] - if configuration.type == "service" and isinstance(configuration.replicas, list): - for idx, replica_group in enumerate(configuration.replicas): + if configuration.type == "service" and configuration.groups is not None: + for idx, replica_group in enumerate(configuration.groups): image, docker = _get_replica_group_image_and_docker(replica_group, configuration) _vendors = _detect_gpu_vendors_requiring_image( gpu_spec=replica_group.resources.gpu, @@ -273,7 +296,7 @@ def _validate_gpu_vendor_and_image(run_spec: RunSpec) -> None: f" the default image: {sorted_vendors}" ) if invalid_replicas: - msg = f"replicas{invalid_replicas}: {msg}" + msg = f"groups{invalid_replicas}: {msg}" elif invalid_groups: msg = f"groups{invalid_groups}: {msg}" raise ServerClientError(msg) @@ -310,10 +333,10 @@ def _validate_cpu_arch_and_image(run_spec: RunSpec) -> None: image_msg = "`image` must be set when ARM CPU requested" docker_msg = "`docker: true` is not supported on ARM CPU" configuration = run_spec.configuration - if configuration.type == "service" and isinstance(configuration.replicas, list): + if configuration.type == "service" and configuration.groups is not None: invalid_replicas_without_image: list[int] = [] invalid_replicas_with_docker: list[int] = [] - for idx, replica_group in enumerate(configuration.replicas): + for idx, replica_group in enumerate(configuration.groups): image, docker = _get_replica_group_image_and_docker(replica_group, configuration) if replica_group.resources.cpu.arch == gpuhunt.CPUArchitecture.ARM: if docker: @@ -322,9 +345,9 @@ def _validate_cpu_arch_and_image(run_spec: RunSpec) -> None: invalid_replicas_without_image.append(idx) errors: list[str] = [] if invalid_replicas_without_image: - errors.append(f"replicas{invalid_replicas_without_image}: {image_msg}") + errors.append(f"groups{invalid_replicas_without_image}: {image_msg}") if invalid_replicas_with_docker: - errors.append(f"replicas{invalid_replicas_with_docker}: {docker_msg}") + errors.append(f"groups{invalid_replicas_with_docker}: {docker_msg}") if errors: raise ServerClientError("\n".join(errors)) elif isinstance(configuration, TaskConfiguration) and configuration.groups is not None: @@ -410,7 +433,7 @@ def _check_dynamo_in_place_update_compatibility( _router_affecting_top_level_fields = tuple( f for f in _TYPE_SPECIFIC_CONF_UPDATABLE_FIELDS.get("service", []) - if f not in ("replicas", "scaling") + if f not in ("replicas", "groups", "scaling") ) for field in _router_affecting_top_level_fields: if getattr(current_cfg, field, None) != getattr(new_cfg, field, None): @@ -476,7 +499,7 @@ def get_nodes_required_num(run_spec: RunSpec) -> int: nodes_required_num = run_spec.configuration.nodes_num elif run_spec.configuration.type == "service": nodes_required_num = sum( - group.count.min or 0 for group in run_spec.configuration.replica_groups + group.replicas.min or 0 for group in run_spec.configuration.replica_groups ) return nodes_required_num diff --git a/src/dstack/_internal/server/services/services/__init__.py b/src/dstack/_internal/server/services/services/__init__.py index ff6ce177c..8edb29efa 100644 --- a/src/dstack/_internal/server/services/services/__init__.py +++ b/src/dstack/_internal/server/services/services/__init__.py @@ -165,7 +165,7 @@ def _register_service_in_server( ) # Check if any group has autoscaling (min != max) has_autoscaling = any( - group.count.min != group.count.max for group in run_spec.configuration.replica_groups + group.replicas.min != group.replicas.max for group in run_spec.configuration.replica_groups ) if has_autoscaling: raise ServerClientError( diff --git a/src/dstack/_internal/utils/nodes_interpolator.py b/src/dstack/_internal/utils/nodes_interpolator.py index 86ec1ca40..11032a358 100644 --- a/src/dstack/_internal/utils/nodes_interpolator.py +++ b/src/dstack/_internal/utils/nodes_interpolator.py @@ -1,12 +1,19 @@ import re +from typing import Literal from dstack._internal.utils.interpolator import InterpolatorError, namespace_root -# Shared grammar for groups[i].nodes[j].IP_ADDRESS refs. -_GROUPS_IP_INNER = r"groups\[(\d+)\]\.nodes\[(\d+)\]\.IP_ADDRESS" +GroupsIpMember = Literal["nodes", "replicas"] + +# Shared grammar for groups[i].(nodes|replicas)[j].IP_ADDRESS refs. +_GROUPS_IP_INNER = r"groups\[(\d+)\]\.(nodes|replicas)\[(\d+)\]\.IP_ADDRESS" _GROUPS_IP_REF_NAME = re.compile(rf"^{_GROUPS_IP_INNER}$") # (? bool: return False -def find_groups_ip_refs(s: str) -> list[tuple[int, int]]: - return [(int(m.group(1)), int(m.group(2))) for m in _GROUPS_IP_REF.finditer(s)] +def find_groups_ip_refs(s: str) -> list[tuple[int, GroupsIpMember, int]]: + refs: list[tuple[int, GroupsIpMember, int]] = [] + for m in _GROUPS_IP_REF.finditer(s): + member = m.group(2) + if member != "nodes" and member != "replicas": + continue + refs.append((int(m.group(1)), member, int(m.group(3)))) + return refs + + +def validate_groups_ref_member(s: str, expected: GroupsIpMember) -> None: + """Reject groups refs whose member does not match the run type.""" + for group_index, member, index in find_groups_ip_refs(s): + if member != expected: + raise InterpolatorError( + f"Illegal reference name: groups[{group_index}].{member}[{index}].IP_ADDRESS" + ) -def validate_groups_ref_bounds(s: str, group_sizes: list[int]) -> None: - """Reject groups[i].nodes[j] refs that exceed configured group/node counts.""" - for group_index, node_index in find_groups_ip_refs(s): - if group_index >= len(group_sizes) or node_index >= group_sizes[group_index]: +def validate_groups_ref_bounds( + s: str, + group_sizes: list[int], + *, + member: GroupsIpMember = "nodes", +) -> None: + """Reject groups[i].{member}[j] refs that exceed configured group/slot counts.""" + for group_index, ref_member, index in find_groups_ip_refs(s): + if ref_member != member: + continue + if group_index >= len(group_sizes) or index >= group_sizes[group_index]: + if ( + member == "replicas" + and group_index < len(group_sizes) + and group_sizes[group_index] == 0 + ): + raise InterpolatorError( + f"Invalid reference groups[{group_index}].replicas[{index}].IP_ADDRESS: " + "this group scales to zero, so no replica is guaranteed at start" + ) raise InterpolatorError( - f"Invalid reference groups[{group_index}].nodes[{node_index}].IP_ADDRESS: " + f"Invalid reference groups[{group_index}].{member}[{index}].IP_ADDRESS: " "out of range" ) def interpolate_groups_ip_address(s: str, nodes: list[list[str]]) -> str: + return _interpolate_groups_member_ip_address(s, nodes, "nodes", _NODES_IP_REF) + + +def interpolate_groups_replica_ip_address(s: str, replicas: list[list[str]]) -> str: + return _interpolate_groups_member_ip_address(s, replicas, "replicas", _REPLICAS_IP_REF) + + +def _interpolate_groups_member_ip_address( + s: str, + view: list[list[str]], + member: GroupsIpMember, + pattern: re.Pattern[str], +) -> str: validate_groups_refs(s) - validate_groups_ref_bounds(s, [len(g) for g in nodes]) + validate_groups_ref_member(s, member) + validate_groups_ref_bounds(s, [len(g) for g in view], member=member) def repl(m: re.Match) -> str: - gi, ni = int(m.group(1)), int(m.group(2)) - ip = nodes[gi][ni] + gi, idx = int(m.group(1)), int(m.group(2)) + ip = view[gi][idx] if not ip: - raise InterpolatorError(f"IP not available for groups[{gi}].nodes[{ni}].IP_ADDRESS") + raise InterpolatorError( + f"IP not available for groups[{gi}].{member}[{idx}].IP_ADDRESS" + ) return ip - return _GROUPS_IP_REF.sub(repl, s) + return pattern.sub(repl, s) diff --git a/src/tests/_internal/cli/services/configurators/test_run.py b/src/tests/_internal/cli/services/configurators/test_run.py index 1e86759a5..aa77dec68 100644 --- a/src/tests/_internal/cli/services/configurators/test_run.py +++ b/src/tests/_internal/cli/services/configurators/test_run.py @@ -18,6 +18,7 @@ BaseRunConfiguration, DevEnvironmentConfiguration, PortMapping, + ServiceConfiguration, TaskConfiguration, ) from dstack._internal.core.models.envs import Env @@ -108,6 +109,57 @@ def test_interpolates_env(self): password="test_password", ) + def test_interpolates_service_group_commands(self): + conf = ServiceConfiguration.model_validate( + { + "type": "service", + "port": 8000, + "image": "debian", + "groups": [ + { + "replicas": 1, + "commands": [ + "echo ${{ run.args }} ${{ groups[0].replicas[0].IP_ADDRESS }}" + ], + } + ], + } + ) + modified, _ = self.apply_args(conf, ["hello"]) + assert modified.groups is not None + assert modified.groups[0].commands == [ + "echo hello ${{ groups[0].replicas[0].IP_ADDRESS }}" + ] + + def test_rejects_unknown_var_in_service_group_commands(self): + conf = ServiceConfiguration.model_validate( + { + "type": "service", + "port": 8000, + "image": "debian", + "groups": [ + { + "replicas": 1, + "commands": ["echo ${{ env.MISSING }}"], + } + ], + } + ) + with pytest.raises(ConfigurationError, match="missing vars"): + self.apply_args(conf, []) + + def test_homogeneous_service_interpolates_commands_once(self): + conf = ServiceConfiguration.model_validate( + { + "type": "service", + "port": 8000, + "image": "debian", + "commands": ["echo $${{ env.FOO }}"], + } + ) + modified, _ = self.apply_args(conf, []) + assert modified.commands == ["echo ${{ env.FOO }}"] + class TestApplyConfiguration: def test_composes_get_plan_and_apply_plan(self, monkeypatch): diff --git a/src/tests/_internal/cli/services/presets/test_build.py b/src/tests/_internal/cli/services/presets/test_build.py new file mode 100644 index 000000000..f215ae315 --- /dev/null +++ b/src/tests/_internal/cli/services/presets/test_build.py @@ -0,0 +1,100 @@ +import gpuhunt +import pytest + +from dstack._internal.cli.models.presets import PresetVerificationReplicaGroup +from dstack._internal.cli.services.presets.build import set_service_gpu_vendor_from_verification +from dstack._internal.core.models.configurations import ( + DEFAULT_REPLICA_GROUP_NAME, + ServiceConfiguration, +) +from dstack._internal.core.models.resources import ResourcesSpec + +pytestmark = pytest.mark.windows + + +def _gpu_resources(name: str, *, vendor: str | None = None) -> ResourcesSpec: + gpu = {"name": name, "memory": "48GB", "count": 1} + if vendor is not None: + gpu["vendor"] = vendor + return ResourcesSpec.model_validate({"gpu": gpu}) + + +class TestSetServiceGpuVendorFromVerification: + def test_stamps_vendor_on_homogeneous_service_resources(self): + service = ServiceConfiguration.model_validate( + { + "image": "vllm/vllm-openai:v0.11.0", + "commands": ["vllm serve model"], + "port": 8000, + "model": "m", + "resources": {"gpu": "40GB..48GB:1"}, + } + ) + verified_on = [ + PresetVerificationReplicaGroup( + name=DEFAULT_REPLICA_GROUP_NAME, + replicas=[_gpu_resources("A6000")], + ) + ] + + set_service_gpu_vendor_from_verification(service, verified_on) + + assert service.resources is not None + assert service.resources.gpu is not None + assert service.resources.gpu.vendor == gpuhunt.AcceleratorVendor.NVIDIA + + def test_stamps_vendor_on_each_replica_group_resources(self): + service = ServiceConfiguration.model_validate( + { + "image": "vllm/vllm-openai:v0.11.0", + "port": 8000, + "model": "m", + "groups": [ + { + "name": "prefill", + "replicas": 1, + "commands": ["prefill"], + "resources": {"gpu": "40GB..48GB:1"}, + }, + { + "name": "decode", + "replicas": 1, + "commands": ["decode"], + "resources": {"gpu": "40GB..48GB:1"}, + }, + ], + } + ) + verified_on = [ + PresetVerificationReplicaGroup(name="prefill", replicas=[_gpu_resources("H100")]), + PresetVerificationReplicaGroup(name="decode", replicas=[_gpu_resources("A100")]), + ] + + set_service_gpu_vendor_from_verification(service, verified_on) + + assert service.groups is not None + assert service.groups[0].resources.gpu is not None + assert service.groups[1].resources.gpu is not None + assert service.groups[0].resources.gpu.vendor == gpuhunt.AcceleratorVendor.NVIDIA + assert service.groups[1].resources.gpu.vendor == gpuhunt.AcceleratorVendor.NVIDIA + assert service.resources.gpu is None or service.resources.gpu.vendor is None + + def test_rejects_service_vendor_that_does_not_match_verification(self): + service = ServiceConfiguration.model_validate( + { + "image": "vllm/vllm-openai:v0.11.0", + "commands": ["vllm serve model"], + "port": 8000, + "model": "m", + "resources": {"gpu": "nvidia:40GB..48GB:1"}, + } + ) + verified_on = [ + PresetVerificationReplicaGroup( + name=DEFAULT_REPLICA_GROUP_NAME, + replicas=[_gpu_resources("MI300X", vendor="amd")], + ) + ] + + with pytest.raises(ValueError, match="GPU vendor does not match verification"): + set_service_gpu_vendor_from_verification(service, verified_on) diff --git a/src/tests/_internal/core/models/test_configurations.py b/src/tests/_internal/core/models/test_configurations.py index 80303e69e..77caaaa09 100644 --- a/src/tests/_internal/core/models/test_configurations.py +++ b/src/tests/_internal/core/models/test_configurations.py @@ -1,9 +1,12 @@ -from typing import Any, Optional +from copy import deepcopy +from typing import Any, Optional, Union import pytest +from pydantic import model_validator +from typing_extensions import Self from dstack._internal.core.errors import ConfigurationError -from dstack._internal.core.models.common import RegistryAuth +from dstack._internal.core.models.common import CoreModel, RegistryAuth, validate_extra_ignore from dstack._internal.core.models.configurations import ( DevEnvironmentConfigurationParams, PythonVersion, @@ -156,9 +159,8 @@ def test_replica_group_router(self): } parsed = parse_run_configuration(conf) assert isinstance(parsed, ServiceConfiguration) - assert parsed.replicas is not None - assert isinstance(parsed.replicas, list) - router_g = next(g for g in parsed.replicas if g.name == "router") + assert parsed.groups is not None + router_g = next(g for g in parsed.groups if g.name == "router") assert isinstance(router_g.router, ReplicaGroupRouterConfig) assert router_g.router.type == "sglang" @@ -242,7 +244,8 @@ def test_replica_group_accepts_image_python_nvcc_docker(self): } parsed = parse_run_configuration(conf) assert isinstance(parsed, ServiceConfiguration) - groups = {g.name: g for g in parsed.replicas} + assert parsed.groups is not None + groups = {g.name: g for g in parsed.groups} assert groups["a"].image == "nginx:latest" assert groups["b"].python == PythonVersion.PY312 assert groups["c"].nvcc is True @@ -263,7 +266,9 @@ def test_replica_group_accepts_privileged(self): ], } parsed = parse_run_configuration(conf) - assert parsed.replicas[0].privileged is True + assert isinstance(parsed, ServiceConfiguration) + assert parsed.groups is not None + assert parsed.groups[0].privileged is True @pytest.mark.parametrize( "yaml_value,expected", @@ -282,7 +287,9 @@ def test_replica_group_python_yaml_coercion(self, yaml_value, expected): "replicas": [{"count": 1, "python": yaml_value, "commands": ["x"]}], } parsed = parse_run_configuration(conf) - assert parsed.replicas[0].python == expected + assert isinstance(parsed, ServiceConfiguration) + assert parsed.groups is not None + assert parsed.groups[0].python == expected def test_replica_group_image_python_mutex(self): with pytest.raises( @@ -354,8 +361,10 @@ def test_replica_group_python_nvcc_allowed_together(self): ], } parsed = parse_run_configuration(conf) - assert parsed.replicas[0].python == PythonVersion.PY312 - assert parsed.replicas[0].nvcc is True + assert isinstance(parsed, ServiceConfiguration) + assert parsed.groups is not None + assert parsed.groups[0].python == PythonVersion.PY312 + assert parsed.groups[0].nvcc is True def test_replica_group_docker_with_privileged_false_rejected(self): with pytest.raises( @@ -1082,3 +1091,143 @@ def test_accepts_top_level_resources_with_groups(self): assert parsed.groups[0].resources == ResourcesSpec() assert parsed.resources.gpu is not None 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( + { + "type": "service", + "port": 8000, + "replicas": [{"count": 1, "commands": ["x"]}], + } + ) + assert isinstance(parsed, ServiceConfiguration) + assert parsed.replicas is None + assert parsed.groups is not None + assert parsed.groups[0].replicas == Range(min=1, max=1) + + def test_new_groups_syntax_parses_identically(self): + legacy = parse_run_configuration( + { + "type": "service", + "port": 8000, + "replicas": [{"count": 1, "commands": ["x"]}], + } + ) + new = parse_run_configuration( + { + "type": "service", + "port": 8000, + "groups": [{"replicas": 1, "commands": ["x"]}], + } + ) + assert isinstance(legacy, ServiceConfiguration) + assert isinstance(new, ServiceConfiguration) + assert legacy.replicas is None + assert new.replicas is None + assert legacy.groups == new.groups + + def test_dump_is_legacy_canonical(self): + parsed = parse_run_configuration( + { + "type": "service", + "port": 8000, + "groups": [{"replicas": 1, "commands": ["x"]}], + } + ) + 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] + 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] + + def test_dump_validate_is_fixed_point(self): + parsed = parse_run_configuration( + { + "type": "service", + "port": 8000, + "groups": [{"replicas": 1, "commands": ["x"]}], + } + ) + assert isinstance(parsed, ServiceConfiguration) + once = ServiceConfiguration.model_validate(parsed.model_dump()) + 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): + parsed = parse_run_configuration( + { + "type": "service", + "port": 8000, + "commands": ["x"], + "replicas": 2, + } + ) + dumped = parsed.model_dump() + assert "groups" not in dumped + assert dumped["replicas"] == {"min": 2, "max": 2} + + def test_replicas_and_groups_rejected(self): + with pytest.raises(ConfigurationError, match="mutually exclusive"): + parse_run_configuration( + { + "type": "service", + "port": 8000, + "replicas": 2, + "groups": [{"replicas": 1, "commands": ["x"]}], + } + ) + + def test_empty_groups_rejected(self): + with pytest.raises(ConfigurationError, match="empty"): + parse_run_configuration({"type": "service", "port": 8000, "groups": []}) + + def test_parse_does_not_mutate_caller_dict(self): + conf = { + "type": "service", + "port": 8000, + "replicas": [{"count": 1, "commands": ["x"]}], + } + original = deepcopy(conf) + parse_run_configuration(conf) + assert conf == original + assert conf["replicas"][0]["count"] == 1 diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py index 560e3c201..2b4b5a206 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py @@ -50,6 +50,7 @@ JobRunningPipelineItem, JobRunningWorker, _build_nodes_ip_view, + _build_replica_groups_ip_view, _fetch_run_model, _get_cluster_info, _prepare_startup_context, @@ -2959,6 +2960,7 @@ def _router_service_configuration( *, gateway: Optional[str] = None, probes: Optional[list[ProbeConfig]] = None, + router_commands: Optional[list[str]] = None, ) -> ServiceConfiguration: return ServiceConfiguration.model_validate( { @@ -2972,7 +2974,7 @@ def _router_service_configuration( { "name": "router", "router": {"type": router_type}, - "commands": ["echo router"], + "commands": router_commands or ["echo router"], "count": 1, }, ], @@ -3339,6 +3341,136 @@ async def test_does_not_defer_on_job_in_other_replica(self): self._assert_gate_passed(result) +@pytest.mark.asyncio +class TestPrepareStartupContextReplicaIpWait: + def _service_configuration(self) -> ServiceConfiguration: + return ServiceConfiguration.model_validate( + { + "type": "service", + "port": 8000, + "image": "debian", + "groups": [ + {"name": "router", "replicas": 1, "commands": ["echo router"]}, + { + "name": "worker", + "replicas": 1, + "commands": ["echo ${{ groups[0].replicas[0].IP_ADDRESS }}"], + }, + ], + } + ) + + def _make_context(self, *, router_ip: str, commands: Optional[list[str]] = None): + job_model = MagicMock() + job_model.submitted_at = datetime(2023, 1, 1, 10, 0, 0, tzinfo=timezone.utc) + router = _replica_group_job( + replica_num=0, + replica_group="router", + internal_ip=router_ip, + ) + job = MagicMock() + job.job_spec.replica_num = 1 + job.job_spec.replica_group = "worker" + job.job_spec.commands = commands or ["echo ${{ groups[0].replicas[0].IP_ADDRESS }}"] + run = MagicMock() + run.jobs = [router] + run.run_spec.configuration = self._service_configuration() + return _ProcessContext( + job_model=job_model, + run_model=MagicMock(), + run=run, + job=job, + job_submission=MagicMock(job_runtime_data=None), + job_provisioning_data=MagicMock(), + instance_access_revoked=False, + ) + + def _patches(self): + @asynccontextmanager + async def _fake_session_ctx(): + yield MagicMock() + + return ( + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_running.get_router_env_for_job", + return_value=None, + ), + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_running.get_session_ctx", + _fake_session_ctx, + ), + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_running.get_job_attached_volumes", + new_callable=AsyncMock, + return_value=[], + ), + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_running.get_repo_creds", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_running.get_project_secrets_mapping", + new_callable=AsyncMock, + return_value={}, + ), + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_running.repo_model_to_repo_head_with_creds", + return_value=MagicMock(repo_creds=None), + ), + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_running.interpolate_job_spec_secrets", + ), + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_running._get_cluster_info", + return_value=ClusterInfo( + job_ips=["10.0.0.1"], master_job_ip="10.0.0.1", gpus_per_job=0 + ), + ), + ) + + @freeze_time("2023-01-01 12:00:00Z") + async def test_replica_ip_not_ready_defers(self): + context = self._make_context(router_ip="") + result = _ProcessResult() + with ExitStack() as stack: + for p in self._patches(): + stack.enter_context(p) + out = await _prepare_startup_context(context=context, result=result) + assert out is None + assert result.job_update_map == {} + assert context.job.job_spec.commands == ["echo ${{ groups[0].replicas[0].IP_ADDRESS }}"] + + @freeze_time("2023-01-01 12:00:00Z") + async def test_replica_ip_ready_substitutes(self): + context = self._make_context(router_ip="10.0.0.5") + result = _ProcessResult() + with ExitStack() as stack: + for p in self._patches(): + stack.enter_context(p) + out = await _prepare_startup_context(context=context, result=result) + assert out is not None + assert context.job.job_spec.commands == ["echo 10.0.0.5"] + + @freeze_time("2023-01-01 12:00:00Z") + async def test_nodes_ref_in_service_terminates(self): + context = self._make_context( + router_ip="10.0.0.5", + commands=["echo ${{ groups[0].nodes[0].IP_ADDRESS }}"], + ) + result = _ProcessResult() + with ExitStack() as stack: + for p in self._patches(): + stack.enter_context(p) + out = await _prepare_startup_context(context=context, result=result) + assert out is None + assert result.job_update_map.get("status") == JobStatus.TERMINATING + assert "Illegal reference name" in ( + result.job_update_map.get("termination_reason_message") or "" + ) + assert context.job.job_spec.commands == ["echo ${{ groups[0].nodes[0].IP_ADDRESS }}"] + + @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) class TestFetchRunModelDynamoBranch: @@ -3415,6 +3547,129 @@ async def test_non_dynamo_loads_only_own_replica(self, test_db, session: AsyncSe ) assert {j.replica_num for j in run_model.jobs} == {0} + async def test_sglang_with_replica_ip_refs_loads_all_replicas( + self, test_db, session: AsyncSession + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run_spec = get_run_spec( + repo_id=repo.name, + configuration=_router_service_configuration( + "sglang", + router_commands=["echo ${{ groups[0].replicas[0].IP_ADDRESS }}"], + ), + ) + run = await create_run( + session=session, project=project, repo=repo, user=user, run_spec=run_spec + ) + await create_job( + session=session, + run=run, + replica_num=0, + status=JobStatus.PROVISIONING, + ) + await create_job( + session=session, + run=run, + replica_num=1, + status=JobStatus.PROVISIONING, + ) + run_id = run.id + parsed = validate_json_extra_ignore(RunSpec, get_or_error(run.run_spec)) + await session.commit() + session.expire_all() + run_model = await _fetch_run_model( + session=session, + run_id=run_id, + replica_num=1, + run_spec=parsed, + ) + assert {j.replica_num for j in run_model.jobs} == {0, 1} + + async def test_sglang_without_replica_ip_refs_loads_only_own_replica( + self, test_db, session: AsyncSession + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run_spec = get_run_spec( + repo_id=repo.name, + configuration=_router_service_configuration("sglang"), + ) + run = await create_run( + session=session, project=project, repo=repo, user=user, run_spec=run_spec + ) + await create_job( + session=session, + run=run, + replica_num=0, + status=JobStatus.PROVISIONING, + ) + await create_job( + session=session, + run=run, + replica_num=1, + status=JobStatus.PROVISIONING, + ) + run_id = run.id + parsed = validate_json_extra_ignore(RunSpec, get_or_error(run.run_spec)) + await session.commit() + session.expire_all() + run_model = await _fetch_run_model( + session=session, + run_id=run_id, + replica_num=1, + run_spec=parsed, + ) + assert {j.replica_num for j in run_model.jobs} == {1} + + async def test_sglang_with_replica_ip_refs_drops_terminated_replicas( + self, test_db, session: AsyncSession + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run_spec = get_run_spec( + repo_id=repo.name, + configuration=_router_service_configuration( + "sglang", + router_commands=["echo ${{ groups[0].replicas[0].IP_ADDRESS }}"], + ), + ) + run = await create_run( + session=session, project=project, repo=repo, user=user, run_spec=run_spec + ) + await create_job( + session=session, + run=run, + replica_num=0, + status=JobStatus.PROVISIONING, + ) + await create_job( + session=session, + run=run, + replica_num=1, + status=JobStatus.PROVISIONING, + ) + await create_job( + session=session, + run=run, + replica_num=2, + status=JobStatus.TERMINATED, + ) + run_id = run.id + parsed = validate_json_extra_ignore(RunSpec, get_or_error(run.run_spec)) + await session.commit() + session.expire_all() + run_model = await _fetch_run_model( + session=session, + run_id=run_id, + replica_num=1, + run_spec=parsed, + ) + assert {j.replica_num for j in run_model.jobs} == {0, 1} + def _node_group_job( *, @@ -3446,6 +3701,35 @@ def _node_group_job( ) +def _replica_group_job( + *, + replica_num: int, + replica_group: str, + internal_ip: str, + status: JobStatus = JobStatus.PROVISIONING, +) -> Job: + return Job.model_construct( + job_spec=JobSpec.model_construct( + replica_num=replica_num, + job_num=0, + replica_group=replica_group, + commands=[], + ), + job_submissions=[ + JobSubmission.model_construct( + id=uuid.uuid4(), + submitted_at=datetime.now(timezone.utc), + status=status, + job_provisioning_data=get_job_provisioning_data( + internal_ip=internal_ip, + gpu_count=1, + ), + job_runtime_data=None, + ) + ], + ) + + class TestGetClusterInfo: def test_fills_gpus_per_node(self): jobs = [ @@ -3550,3 +3834,83 @@ def test_referenced_ips_out_of_range(self): nodes_view = [["10.0.0.1"]] with pytest.raises(InterpolatorError, match="out of range"): _referenced_ips_ready(["echo ${{ groups[1].nodes[0].IP_ADDRESS }}"], nodes_view) + + +class TestReplicaGroupsIpView: + def _configuration(self, groups: list[dict]) -> ServiceConfiguration: + return ServiceConfiguration.model_validate( + { + "type": "service", + "port": 8000, + "image": "debian", + "groups": groups, + } + ) + + def test_builds_min_length_rows_by_group_order(self): + configuration = self._configuration( + [ + {"name": "router", "replicas": 1, "commands": ["echo router"]}, + {"name": "worker", "replicas": 1, "commands": ["echo worker"]}, + ] + ) + jobs = [ + _replica_group_job(replica_num=0, replica_group="router", internal_ip="10.0.0.1"), + _replica_group_job(replica_num=1, replica_group="worker", internal_ip="10.0.0.2"), + ] + assert _build_replica_groups_ip_view(jobs, configuration) == [ + ["10.0.0.1"], + ["10.0.0.2"], + ] + + def test_row_length_is_min_not_live_count(self): + configuration = self._configuration( + [ + { + "name": "workers", + "replicas": "1..4", + "scaling": {"metric": "rps", "target": 10}, + "commands": ["echo worker"], + } + ] + ) + jobs = [ + _replica_group_job(replica_num=0, replica_group="workers", internal_ip="10.0.0.1"), + _replica_group_job(replica_num=1, replica_group="workers", internal_ip="10.0.0.2"), + ] + assert _build_replica_groups_ip_view(jobs, configuration) == [["10.0.0.1"]] + + def test_skips_terminated_jobs(self): + configuration = self._configuration( + [{"name": "router", "replicas": 1, "commands": ["echo router"]}] + ) + jobs = [ + _replica_group_job( + replica_num=0, + replica_group="router", + internal_ip="10.0.0.9", + status=JobStatus.TERMINATED, + ), + _replica_group_job(replica_num=1, replica_group="router", internal_ip="10.0.0.1"), + ] + assert _build_replica_groups_ip_view(jobs, configuration) == [["10.0.0.1"]] + + def test_empty_ip_when_not_provisioned(self): + configuration = self._configuration( + [{"name": "router", "replicas": 1, "commands": ["echo router"]}] + ) + jobs = [_replica_group_job(replica_num=0, replica_group="router", internal_ip="")] + assert _build_replica_groups_ip_view(jobs, configuration) == [[""]] + + def test_referenced_replica_ips_ready(self): + replica_view = [["10.0.0.1"], [""]] + assert _referenced_ips_ready( + ["echo ${{ groups[0].replicas[0].IP_ADDRESS }}"], + replica_view, + member="replicas", + ) + assert not _referenced_ips_ready( + ["echo ${{ groups[1].replicas[0].IP_ADDRESS }}"], + replica_view, + member="replicas", + ) diff --git a/src/tests/_internal/server/services/runs/test_spec.py b/src/tests/_internal/server/services/runs/test_spec.py index 4e9974453..419d95220 100644 --- a/src/tests/_internal/server/services/runs/test_spec.py +++ b/src/tests/_internal/server/services/runs/test_spec.py @@ -19,6 +19,7 @@ from dstack._internal.server.services.runs.spec import ( _check_can_update_configuration, check_can_update_run_spec, + run_spec_has_replica_ip_refs, set_run_spec_resources_defaults, validate_run_spec_and_set_defaults, ) @@ -66,6 +67,17 @@ def _service_configuration( return ServiceConfiguration.model_validate(data) +def _service_with_groups(groups: list[dict]) -> ServiceConfiguration: + return ServiceConfiguration.model_validate( + { + "type": "service", + "port": 8000, + "image": "debian", + "groups": groups, + } + ) + + def _run_spec(configuration: ServiceConfiguration, **kwargs): return get_run_spec( repo_id="test-repo", run_name="test-run", configuration=configuration, **kwargs @@ -249,6 +261,225 @@ def test_rejects_out_of_range_node_index(self): SimpleNamespace(ssh_public_key="ssh-rsa test"), run_spec ) + def test_rejects_replicas_member_in_task_commands(self): + run_spec = get_run_spec( + repo_id="test-repo", + configuration=TaskConfiguration( + image="debian", + groups=[ + NodeGroup( + name="head", + nodes=1, + commands=["echo ${{ groups[0].replicas[0].IP_ADDRESS }}"], + ), + ], + ), + ) + + with pytest.raises(ServerClientError, match="Illegal reference name"): + validate_run_spec_and_set_defaults( + SimpleNamespace(ssh_public_key="ssh-rsa test"), run_spec + ) + + def test_rejects_nodes_member_in_service_commands(self): + run_spec = get_run_spec( + repo_id="test-repo", + configuration=ServiceConfiguration.model_validate( + { + "type": "service", + "port": 8000, + "image": "debian", + "groups": [ + { + "replicas": 1, + "commands": ["echo ${{ groups[0].nodes[0].IP_ADDRESS }}"], + } + ], + } + ), + ) + + with pytest.raises(ServerClientError, match="Illegal reference name"): + validate_run_spec_and_set_defaults( + SimpleNamespace(ssh_public_key="ssh-rsa test"), run_spec + ) + + def test_accepts_service_replicas_ref_on_fixed_and_min_slot(self): + for replicas in (1, "1..4"): + group = { + "replicas": replicas, + "commands": ["echo ${{ groups[0].replicas[0].IP_ADDRESS }}"], + } + if replicas == "1..4": + group["scaling"] = {"metric": "rps", "target": 10} + run_spec = get_run_spec( + repo_id="test-repo", + configuration=_service_with_groups([group]), + ) + validate_run_spec_and_set_defaults( + SimpleNamespace(ssh_public_key="ssh-rsa test"), run_spec + ) + + def test_accepts_service_replicas_indexes_for_fixed_count(self): + run_spec = get_run_spec( + repo_id="test-repo", + configuration=_service_with_groups( + [ + { + "replicas": 2, + "commands": [ + "echo ${{ groups[0].replicas[0].IP_ADDRESS }} " + "${{ groups[0].replicas[1].IP_ADDRESS }}" + ], + } + ] + ), + ) + validate_run_spec_and_set_defaults( + SimpleNamespace(ssh_public_key="ssh-rsa test"), run_spec + ) + + def test_rejects_service_replicas_index_above_min(self): + run_spec = get_run_spec( + repo_id="test-repo", + configuration=_service_with_groups( + [ + { + "replicas": "1..4", + "scaling": {"metric": "rps", "target": 10}, + "commands": ["echo ${{ groups[0].replicas[1].IP_ADDRESS }}"], + } + ] + ), + ) + + with pytest.raises(ServerClientError, match="out of range"): + validate_run_spec_and_set_defaults( + SimpleNamespace(ssh_public_key="ssh-rsa test"), run_spec + ) + + def test_rejects_service_replicas_ref_into_scale_to_zero_group(self): + run_spec = get_run_spec( + repo_id="test-repo", + configuration=_service_with_groups( + [ + { + "replicas": "0..4", + "scaling": {"metric": "rps", "target": 10}, + "commands": ["echo ${{ groups[0].replicas[0].IP_ADDRESS }}"], + } + ] + ), + ) + + with pytest.raises(ServerClientError, match="scales to zero"): + validate_run_spec_and_set_defaults( + SimpleNamespace(ssh_public_key="ssh-rsa test"), run_spec + ) + + def test_rejects_service_replicas_group_out_of_range(self): + run_spec = get_run_spec( + repo_id="test-repo", + configuration=_service_with_groups( + [ + { + "replicas": 1, + "commands": ["echo ${{ groups[7].replicas[9].IP_ADDRESS }}"], + } + ] + ), + ) + + with pytest.raises(ServerClientError, match="out of range"): + validate_run_spec_and_set_defaults( + SimpleNamespace(ssh_public_key="ssh-rsa test"), run_spec + ) + + def test_accepts_service_ref_to_another_group_min_slot(self): + run_spec = get_run_spec( + repo_id="test-repo", + configuration=_service_with_groups( + [ + { + "replicas": 1, + "commands": [ + "smg --prefill http://${{ groups[1].replicas[0].IP_ADDRESS }}:8000" + ], + }, + { + "replicas": 1, + "commands": ["echo prefill"], + }, + ] + ), + ) + validate_run_spec_and_set_defaults( + SimpleNamespace(ssh_public_key="ssh-rsa test"), run_spec + ) + + def test_rejects_service_groups_ref_in_env(self): + run_spec = get_run_spec( + repo_id="test-repo", + configuration=ServiceConfiguration.model_validate( + { + "type": "service", + "port": 8000, + "image": "debian", + "commands": ["echo ok"], + "env": { + "PREFILL_URL": "http://${{ groups[1].replicas[0].IP_ADDRESS }}", + }, + } + ), + ) + + with pytest.raises(ServerClientError, match="only supported in commands, not in `env`"): + validate_run_spec_and_set_defaults( + SimpleNamespace(ssh_public_key="ssh-rsa test"), run_spec + ) + + +class TestRunSpecHasReplicaIpRefs: + def test_true_when_service_group_command_has_replica_ref(self): + run_spec = get_run_spec( + repo_id="test-repo", + configuration=_service_with_groups( + [ + {"replicas": 1, "commands": ["echo router"]}, + { + "replicas": 1, + "commands": ["echo ${{ groups[0].replicas[0].IP_ADDRESS }}"], + }, + ] + ), + ) + assert run_spec_has_replica_ip_refs(run_spec) + + def test_false_when_service_has_no_replica_refs(self): + run_spec = get_run_spec( + repo_id="test-repo", + configuration=_service_with_groups( + [{"replicas": 1, "commands": ["echo ok"]}], + ), + ) + assert not run_spec_has_replica_ip_refs(run_spec) + + def test_false_for_task_node_refs(self): + run_spec = get_run_spec( + repo_id="test-repo", + configuration=TaskConfiguration( + image="debian", + groups=[ + NodeGroup( + name="head", + nodes=1, + commands=["echo ${{ groups[0].nodes[0].IP_ADDRESS }}"], + ), + ], + ), + ) + assert not run_spec_has_replica_ip_refs(run_spec) + class TestCheckCanUpdateConfigurationRouterType: def test_sglang_to_dynamo_router_type_change_is_rejected(self): @@ -529,7 +760,8 @@ def test_sets_defaults_for_every_replica_group(self): set_run_spec_resources_defaults(run_spec) - groups = run_spec.configuration.replicas + groups = run_spec.configuration.groups + assert groups is not None assert [g.resources.gpu.vendor for g in groups] == [ gpuhunt.AcceleratorVendor.AMD, gpuhunt.AcceleratorVendor.NVIDIA, @@ -557,7 +789,8 @@ def test_infers_vendor_from_the_image_used_by_the_group( set_run_spec_resources_defaults(run_spec) - assert run_spec.configuration.replicas[0].resources.gpu.vendor == expected_vendor + assert run_spec.configuration.groups is not None + assert run_spec.configuration.groups[0].resources.gpu.vendor == expected_vendor def test_sets_defaults_for_top_level_resources(self): # The top-level resources are ignored when replica groups are set, but they are still @@ -649,7 +882,7 @@ def test_reports_replica_groups_requiring_image(self): ], ) - with pytest.raises(ServerClientError, match=re.escape("replicas[0, 2]")): + with pytest.raises(ServerClientError, match=re.escape("groups[0, 2]")): _validate(run_spec) def test_allows_replica_group_with_its_own_image(self): @@ -714,7 +947,7 @@ def test_reports_replica_groups_requiring_image(self): ], ) - with pytest.raises(ServerClientError, match=re.escape("replicas[0, 2]")): + with pytest.raises(ServerClientError, match=re.escape("groups[0, 2]")): _validate(run_spec) def test_allows_replica_group_with_its_own_image(self): diff --git a/src/tests/_internal/utils/test_interpolator.py b/src/tests/_internal/utils/test_interpolator.py index 3b545ecd0..0491e7c9e 100644 --- a/src/tests/_internal/utils/test_interpolator.py +++ b/src/tests/_internal/utils/test_interpolator.py @@ -52,12 +52,14 @@ def test_illegal_name(self): get_interpolator().interpolate("${{ secrets.007 }}") def test_skips_groups_refs(self): - s = "ray start --address=${{ groups[0].nodes[0].IP_ADDRESS }}:6379" interpolator = VariablesInterpolator( {"run": {"args": "x"}}, skip={"groups": is_valid_groups_ip_ref}, ) - assert interpolator.interpolate(s) == s + nodes = "ray start --address=${{ groups[0].nodes[0].IP_ADDRESS }}:6379" + replicas = "python --host ${{ groups[0].replicas[0].IP_ADDRESS }}" + assert interpolator.interpolate(nodes) == nodes + assert interpolator.interpolate(replicas) == replicas def test_rejects_invalid_groups_refs(self): interpolator = VariablesInterpolator( diff --git a/src/tests/_internal/utils/test_nodes_interpolator.py b/src/tests/_internal/utils/test_nodes_interpolator.py index e4adfa2ef..e92f02ccc 100644 --- a/src/tests/_internal/utils/test_nodes_interpolator.py +++ b/src/tests/_internal/utils/test_nodes_interpolator.py @@ -5,7 +5,10 @@ contains_groups_ref, find_groups_ip_refs, interpolate_groups_ip_address, + interpolate_groups_replica_ip_address, + is_valid_groups_ip_ref, validate_groups_ref_bounds, + validate_groups_ref_member, validate_groups_refs, ) @@ -13,17 +16,50 @@ class TestFindGroupsIpRefs: def test_finds_refs(self): s = "ray start --address=${{ groups[0].nodes[0].IP_ADDRESS }}:6379" - assert find_groups_ip_refs(s) == [(0, 0)] + assert find_groups_ip_refs(s) == [(0, "nodes", 0)] + + def test_finds_replica_refs(self): + s = "python --host ${{ groups[0].replicas[0].IP_ADDRESS }}" + assert find_groups_ip_refs(s) == [(0, "replicas", 0)] def test_finds_multiple_refs(self): s = "${{ groups[0].nodes[1].IP_ADDRESS }} ${{groups[2].nodes[0].IP_ADDRESS}}" - assert find_groups_ip_refs(s) == [(0, 1), (2, 0)] + assert find_groups_ip_refs(s) == [(0, "nodes", 1), (2, "nodes", 0)] def test_no_refs(self): assert find_groups_ip_refs("echo hello") == [] def test_ignores_escaped_refs(self): assert find_groups_ip_refs("$${{ groups[0].nodes[0].IP_ADDRESS }}") == [] + assert find_groups_ip_refs("$${{ groups[0].replicas[0].IP_ADDRESS }}") == [] + + +class TestIsValidGroupsIpRef: + def test_accepts_nodes_and_replicas(self): + assert is_valid_groups_ip_ref("groups[0].nodes[0].IP_ADDRESS") + assert is_valid_groups_ip_ref("groups[1].replicas[2].IP_ADDRESS") + + def test_rejects_typos(self): + assert not is_valid_groups_ip_ref("groups[0].nodes[0].IP") + assert not is_valid_groups_ip_ref("groups[0].node[0].IP_ADDRESS") + assert not is_valid_groups_ip_ref("groups.prefill.nodes[0].IP_ADDRESS") + + +class TestValidateGroupsRefMember: + def test_accepts_expected_member(self): + validate_groups_ref_member("${{ groups[0].nodes[0].IP_ADDRESS }}", "nodes") + validate_groups_ref_member("${{ groups[0].replicas[0].IP_ADDRESS }}", "replicas") + + def test_rejects_replicas_when_nodes_expected(self): + with pytest.raises(InterpolatorError, match="Illegal reference name"): + validate_groups_ref_member("${{ groups[0].replicas[0].IP_ADDRESS }}", "nodes") + + def test_rejects_nodes_when_replicas_expected(self): + with pytest.raises(InterpolatorError, match="Illegal reference name"): + validate_groups_ref_member("${{ groups[0].nodes[0].IP_ADDRESS }}", "replicas") + + def test_ignores_escaped_refs(self): + validate_groups_ref_member("echo $${{ groups[0].replicas[0].IP_ADDRESS }}", "nodes") class TestValidateGroupsRefBounds: @@ -38,10 +74,42 @@ def test_rejects_node_out_of_range(self): with pytest.raises(InterpolatorError, match="out of range"): validate_groups_ref_bounds("${{ groups[0].nodes[1].IP_ADDRESS }}", [1, 2]) + def test_ignores_replica_refs_when_bounding_nodes(self): + validate_groups_ref_bounds("${{ groups[9].replicas[9].IP_ADDRESS }}", [1]) + + def test_accepts_replica_index_below_min(self): + # replicas: 1 or 1..4 → only replicas[0] is a guaranteed slot + validate_groups_ref_bounds( + "${{ groups[0].replicas[0].IP_ADDRESS }}", [1], member="replicas" + ) + + def test_rejects_replica_index_at_or_above_min(self): + # replicas: 1..4 → replicas[1] is a scale-up slot, not legal at submit + with pytest.raises(InterpolatorError, match="out of range"): + validate_groups_ref_bounds( + "${{ groups[0].replicas[1].IP_ADDRESS }}", [1], member="replicas" + ) + + def test_rejects_any_replica_ref_when_min_is_zero(self): + with pytest.raises(InterpolatorError, match="scales to zero"): + validate_groups_ref_bounds( + "${{ groups[0].replicas[0].IP_ADDRESS }}", [0], member="replicas" + ) + + def test_rejects_replica_group_out_of_range(self): + with pytest.raises(InterpolatorError, match="out of range"): + validate_groups_ref_bounds( + "${{ groups[7].replicas[0].IP_ADDRESS }}", [1], member="replicas" + ) + + def test_ignores_node_refs_when_bounding_replicas(self): + validate_groups_ref_bounds("${{ groups[9].nodes[9].IP_ADDRESS }}", [1], member="replicas") + class TestValidateGroupsRefs: def test_accepts_valid_ref(self): validate_groups_refs("ray start --address=${{ groups[0].nodes[0].IP_ADDRESS }}:6379") + validate_groups_refs("python --host ${{ groups[0].replicas[0].IP_ADDRESS }}") def test_rejects_typo_field(self): with pytest.raises(InterpolatorError, match="Illegal reference name"): @@ -62,6 +130,7 @@ def test_ignores_escaped_refs(self): class TestContainsGroupsRef: def test_detects_valid_and_invalid_refs(self): assert contains_groups_ref("http://${{ groups[1].nodes[0].IP_ADDRESS }}") + assert contains_groups_ref("${{ groups[0].replicas[0].IP_ADDRESS }}") assert contains_groups_ref("${{ groups[0].nodes[0].IP }}") assert not contains_groups_ref("${{ secrets.token }}") assert not contains_groups_ref("${{ groups_config.x }}") @@ -81,6 +150,12 @@ def test_replaces_nested_node(self): result = interpolate_groups_ip_address(s, [["10.0.0.1"], ["10.0.0.2"]]) assert result == "10.0.0.2" + def test_does_not_substitute_replica_refs(self): + with pytest.raises(InterpolatorError, match="Illegal reference name"): + interpolate_groups_ip_address( + "${{ groups[0].replicas[0].IP_ADDRESS }}", [["10.0.0.1"]] + ) + def test_raises_when_ip_missing(self): with pytest.raises(InterpolatorError, match="IP not available"): interpolate_groups_ip_address("${{ groups[0].nodes[0].IP_ADDRESS }}", [[""]]) @@ -92,3 +167,28 @@ def test_raises_when_out_of_range(self): def test_raises_on_invalid_ref(self): with pytest.raises(InterpolatorError, match="Illegal reference name"): interpolate_groups_ip_address("${{ groups[0].nodes[0].IP }}", [["10.0.0.1"]]) + + +class TestInterpolateGroupsReplicaIpAddress: + def test_replaces_ip(self): + s = "python --host ${{ groups[0].replicas[0].IP_ADDRESS }}" + result = interpolate_groups_replica_ip_address(s, [["10.0.0.1"], ["10.0.0.2"]]) + assert result == "python --host 10.0.0.1" + + def test_replaces_nested_group(self): + result = interpolate_groups_replica_ip_address( + "${{ groups[1].replicas[0].IP_ADDRESS }}", [["10.0.0.1"], ["10.0.0.2"]] + ) + assert result == "10.0.0.2" + + def test_does_not_substitute_node_refs(self): + with pytest.raises(InterpolatorError, match="Illegal reference name"): + interpolate_groups_replica_ip_address( + "${{ groups[0].nodes[0].IP_ADDRESS }}", [["10.0.0.1"]] + ) + + def test_raises_when_ip_missing(self): + with pytest.raises(InterpolatorError, match="IP not available"): + interpolate_groups_replica_ip_address( + "${{ groups[0].replicas[0].IP_ADDRESS }}", [[""]] + )