From fbc299ca6245f8b87fde19451167ba8cd525f075 Mon Sep 17 00:00:00 2001 From: Bihan Rana Date: Tue, 4 Aug 2026 22:01:12 +0545 Subject: [PATCH 1/4] Support Heterogenous Node Groups --- runner/internal/runner/executor/executor.go | 41 ++++-- runner/internal/runner/schemas/schemas.go | 1 + .../cli/services/configurators/run.py | 22 ++-- src/dstack/_internal/cli/utils/run.py | 2 +- .../_internal/core/backends/slurm/compute.py | 8 +- .../_internal/core/compatibility/runs.py | 11 ++ .../_internal/core/models/configurations.py | 100 ++++++++++++++- src/dstack/_internal/core/models/runs.py | 18 +++ .../background/pipeline_tasks/jobs_running.py | 84 +++++++++++-- .../pipeline_tasks/jobs_submitted.py | 119 +++++++++++++++--- .../services/jobs/configurators/base.py | 48 +++++-- .../server/services/jobs/configurators/dev.py | 10 +- .../services/jobs/configurators/service.py | 5 +- .../services/jobs/configurators/task.py | 44 +++++-- .../server/services/runs/__init__.py | 2 +- .../_internal/server/services/runs/spec.py | 2 +- src/dstack/_internal/server/testing/common.py | 14 ++- src/dstack/_internal/utils/interpolator.py | 11 +- .../_internal/utils/nodes_interpolator.py | 24 ++++ .../core/models/test_configurations.py | 117 +++++++++++++++++ .../pipeline_tasks/test_node_groups.py | 115 +++++++++++++++++ .../pipeline_tasks/test_submitted_jobs.py | 109 ++++++++++++++-- .../_internal/server/routers/test_runs.py | 6 + .../services/jobs/configurators/test_task.py | 62 ++++++++- .../_internal/utils/test_interpolator.py | 5 + .../utils/test_nodes_interpolator.py | 40 ++++++ 26 files changed, 929 insertions(+), 91 deletions(-) create mode 100644 src/dstack/_internal/utils/nodes_interpolator.py create mode 100644 src/tests/_internal/server/background/pipeline_tasks/test_node_groups.py create mode 100644 src/tests/_internal/utils/test_nodes_interpolator.py diff --git a/runner/internal/runner/executor/executor.go b/runner/internal/runner/executor/executor.go index 31f3d7fe92..bb86c991bd 100644 --- a/runner/internal/runner/executor/executor.go +++ b/runner/internal/runner/executor/executor.go @@ -473,7 +473,15 @@ func (ex *RunExecutor) execJob(ctx context.Context, jobLogFile io.Writer) error nodeRank := ex.jobSpec.JobNum nodesNum := ex.jobSpec.JobsPerReplica gpusPerNodeNum := ex.clusterInfo.GPUSPerJob - gpusNum := nodesNum * gpusPerNodeNum + gpusNum := 0 + if len(ex.clusterInfo.GPUSPerNode) > 0 { + for _, n := range ex.clusterInfo.GPUSPerNode { + gpusNum += n + } + } else { + // Old servers omit gpus_per_node; fall back to homogeneous math. + gpusNum = nodesNum * gpusPerNodeNum + } mpiHostfilePath := filepath.Join(ex.dstackDir, "mpi/hostfile") @@ -544,7 +552,7 @@ func (ex *RunExecutor) execJob(ctx context.Context, jobLogFile io.Writer) error log.Warning(ctx, "failed to include dstack_profile", "path", profilePath, "err", err) } - if err := writeMpiHostfile(ctx, ex.clusterInfo.JobIPs, gpusPerNodeNum, mpiHostfilePath); err != nil { + if err := writeMpiHostfile(ctx, ex.clusterInfo.JobIPs, ex.clusterInfo.GPUSPerNode, gpusPerNodeNum, mpiHostfilePath); err != nil { return fmt.Errorf("write MPI hostfile: %w", err) } @@ -759,7 +767,7 @@ func prepareUserSshDir(user *linuxuser.User) (string, error) { return sshDir, nil } -func writeMpiHostfile(ctx context.Context, ips []string, gpusPerNode int, path string) error { +func writeMpiHostfile(ctx context.Context, ips []string, gpusPerNode []int, fallbackGpusPerJob int, path string) error { if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { return fmt.Errorf("create MPI hostfile directory: %w", err) } @@ -775,16 +783,25 @@ func writeMpiHostfile(ctx context.Context, ips []string, gpusPerNode int, path s } } if len(nonEmptyIps) == len(ips) { - var template string - if gpusPerNode == 0 { - // CPU node: the number of slots defaults to the number of processor cores on that host - // See: https://docs.open-mpi.org/en/main/launching-apps/scheduling.html#calculating-the-number-of-slots - template = "%s\n" - } else { - template = fmt.Sprintf("%%s slots=%d\n", gpusPerNode) + if len(gpusPerNode) > 0 && len(gpusPerNode) != len(ips) { + return fmt.Errorf( + "gpus_per_node length %d != job_ips length %d", + len(gpusPerNode), len(ips), + ) } - for _, ip := range nonEmptyIps { - if _, err = fmt.Fprintf(file, template, ip); err != nil { + for i, ip := range nonEmptyIps { + n := fallbackGpusPerJob + if len(gpusPerNode) > 0 { + n = gpusPerNode[i] + } + if n == 0 { + // CPU node: the number of slots defaults to the number of processor cores on that host + // See: https://docs.open-mpi.org/en/main/launching-apps/scheduling.html#calculating-the-number-of-slots + _, err = fmt.Fprintf(file, "%s\n", ip) + } else { + _, err = fmt.Fprintf(file, "%s slots=%d\n", ip, n) + } + if err != nil { return fmt.Errorf("write MPI hostfile line: %w", err) } } diff --git a/runner/internal/runner/schemas/schemas.go b/runner/internal/runner/schemas/schemas.go index 47706228cd..c9102d732d 100644 --- a/runner/internal/runner/schemas/schemas.go +++ b/runner/internal/runner/schemas/schemas.go @@ -95,6 +95,7 @@ type ClusterInfo struct { JobIPs []string `json:"job_ips"` MasterJobIP string `json:"master_job_ip"` GPUSPerJob int `json:"gpus_per_job"` + GPUSPerNode []int `json:"gpus_per_node"` } type SSHKey struct { diff --git a/src/dstack/_internal/cli/services/configurators/run.py b/src/dstack/_internal/cli/services/configurators/run.py index df2d0b35d4..9ace4f474d 100644 --- a/src/dstack/_internal/cli/services/configurators/run.py +++ b/src/dstack/_internal/cli/services/configurators/run.py @@ -692,20 +692,25 @@ def register_commands_args(cls, parser: argparse.ArgumentParser): metavar="RUN_ARGS", ) - def apply_commands_args( - self, - conf: ConfigurationWithCommandsParams, - args: argparse.Namespace, - ): - commands = conf.commands + def _interpolate_commands(self, commands: list[str], args: argparse.Namespace) -> None: run_args = shlex.join(args.run_args) - interpolator = VariablesInterpolator({"run": {"args": run_args}}, skip=["secrets"]) + interpolator = VariablesInterpolator( + {"run": {"args": run_args}}, + skip=["secrets", "groups"], + ) try: for i, command in enumerate(commands): commands[i] = interpolator.interpolate_or_error(command) except InterpolatorError as e: raise ConfigurationError(e.args[0]) + def apply_commands_args( + self, + conf: ConfigurationWithCommandsParams, + args: argparse.Namespace, + ): + self._interpolate_commands(conf.commands, args) + class TaskConfigurator( RunWithPortsConfiguratorMixin, RunWithCommandsConfiguratorMixin, BaseRunConfigurator @@ -722,6 +727,9 @@ def apply_args(self, conf: TaskConfiguration, args: argparse.Namespace): super().apply_args(conf, args) self.apply_ports_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) class DevEnvironmentConfigurator(RunWithPortsConfiguratorMixin, BaseRunConfigurator): diff --git a/src/dstack/_internal/cli/utils/run.py b/src/dstack/_internal/cli/utils/run.py index 6c27f2aa6f..46dfba5803 100644 --- a/src/dstack/_internal/cli/utils/run.py +++ b/src/dstack/_internal/cli/utils/run.py @@ -122,7 +122,7 @@ def th(s: str) -> str: props.add_row(th("User"), run_plan.user) configuration_type = run_spec.configuration.type if run_spec.configuration.type == "task": - configuration_type += f" (nodes={run_spec.configuration.nodes})" + configuration_type += f" (nodes={run_spec.configuration.nodes_num})" props.add_row(th("Type"), configuration_type) props.add_row(th("Resources"), pretty_req) props.add_row(th("Spot policy"), spot_policy) diff --git a/src/dstack/_internal/core/backends/slurm/compute.py b/src/dstack/_internal/core/backends/slurm/compute.py index 3ae14d5d99..423de1f05d 100644 --- a/src/dstack/_internal/core/backends/slurm/compute.py +++ b/src/dstack/_internal/core/backends/slurm/compute.py @@ -164,6 +164,7 @@ def run_jobs( instance_offer=instance_offer, project_ssh_public_key=project_ssh_public_key, requirements=requirements, + node_count=len(job_configurations), ) def terminate_instance( @@ -186,6 +187,7 @@ def _run_slurm_job( instance_offer: InstanceOfferWithAvailability, project_ssh_public_key: str, requirements: Requirements, + node_count: Optional[int] = None, ) -> ComputeGroupProvisioningData: if job.job_spec.registry_auth is not None: self._skip_offer_cache.add(run, job, instance_offer) @@ -209,7 +211,11 @@ def _run_slurm_job( assert run.run_spec.ssh_key_pub is not None authorized_keys = [project_ssh_public_key.strip(), run.run_spec.ssh_key_pub.strip()] - node_count = job.job_spec.jobs_per_replica + # Heterogeneous groups provision one shape at a time; Slurm allocation + # size must match that batch. Fall back to jobs_per_replica for + # run_job / homogeneous single-call paths. + if node_count is None: + node_count = job.job_spec.jobs_per_replica resources_spec = requirements.resources requested_resources = get_requested_resources_from_resources_spec(resources_spec) diff --git a/src/dstack/_internal/core/compatibility/runs.py b/src/dstack/_internal/core/compatibility/runs.py index 847e7be303..fa95dcd689 100644 --- a/src/dstack/_internal/core/compatibility/runs.py +++ b/src/dstack/_internal/core/compatibility/runs.py @@ -8,6 +8,7 @@ ) from dstack._internal.core.models.configurations import ( ServiceConfiguration, + TaskConfiguration, ) from dstack._internal.core.models.routers import SGLangServiceRouterConfig from dstack._internal.core.models.runs import ( @@ -98,6 +99,10 @@ def get_run_spec_excludes(run_spec: RunSpec) -> IncludeExcludeDictType: if not run_spec.configuration.dstack: configuration_excludes["dstack"] = True + if isinstance(run_spec.configuration, TaskConfiguration): + if run_spec.configuration.groups is None: + configuration_excludes["groups"] = True + if isinstance(run_spec.configuration, ServiceConfiguration): if run_spec.configuration.probes: probe_excludes: IncludeExcludeDictType = {} @@ -160,6 +165,12 @@ def get_job_spec_excludes(job_specs: list[JobSpec]) -> IncludeExcludeDictType: spec_excludes: IncludeExcludeDictType = {} if all(s.replica_group == DEFAULT_REPLICA_GROUP_NAME for s in job_specs): spec_excludes["replica_group"] = True + if all(s.node_group_index == 0 for s in job_specs): + spec_excludes["node_group_index"] = True + if all(s.node_group_name == DEFAULT_REPLICA_GROUP_NAME for s in job_specs): + spec_excludes["node_group_name"] = True + if all(s.node_group_job_index == 0 for s in job_specs): + spec_excludes["node_group_job_index"] = True probe_excludes: IncludeExcludeDictType = {} spec_excludes["probes"] = {"__all__": probe_excludes} diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py index c05bd7f1db..04ea6a5385 100644 --- a/src/dstack/_internal/core/models/configurations.py +++ b/src/dstack/_internal/core/models/configurations.py @@ -708,6 +708,9 @@ def check_image_or_commands_present(self) -> Self: replicas = getattr(self, "replicas", None) if isinstance(replicas, list): return self + # If groups is set, skip validation - commands come from node groups + if getattr(self, "groups", None) is not None: + return self if not self.commands and not getattr(self, "image", None): raise ValueError("Either `commands` or `image` must be set") @@ -798,8 +801,85 @@ def validate_dstack_and_inactivity_duration(self) -> Self: return self +class NodeGroup(CoreModel): + name: Annotated[ + Optional[str], + Field( + description=( + "The name of the node group. If not provided, defaults to '0', '1', etc. " + "based on position." + ) + ), + ] = None + nodes: Annotated[int, Field(description="The number of nodes in this group", ge=1)] = 1 + resources: Annotated[ + ResourcesSpec, + Field(description="The resources requirements for nodes in this group"), + ] = ResourcesSpec() + commands: Annotated[ + CommandsList, + Field(description="The shell commands to run for nodes in this group"), + ] = [] + ports: Annotated[ + List[PortMappingOrShorthand], + Field(description="Port numbers/mapping to expose for nodes in this group"), + ] = [] + + @field_validator("name") + @classmethod + def validate_name(cls, v: Optional[str]) -> Optional[str]: + if v is not None: + if not is_valid_replica_group_name(v): + raise ValueError("Resource name should match regex '^[a-z0-9][a-z0-9-]{0,39}$'") + return v + + class TaskConfigurationParams(CoreModel): - nodes: Annotated[int, Field(description="Number of nodes", ge=1)] = 1 + nodes: Annotated[ + int, + Field(description="The number of nodes for homogeneous multi-node tasks", ge=1), + ] = 1 + groups: Annotated[ + Optional[List[NodeGroup]], + Field( + description=( + "A list of node groups for heterogeneous multi-node tasks. " + "Mutually exclusive with `nodes`." + ), + ), + ] = None + + @model_validator(mode="before") + @classmethod + def validate_nodes_xor_groups(cls, data): + if not isinstance(data, dict): + return data + # Allow groups with default nodes: 1 (serialized configs always include it). + # Reject nodes: N (N != 1) together with groups. + if data.get("groups") is not None and "nodes" in data: + nodes = data.get("nodes") + if nodes is not None and nodes != 1: + raise ValueError("`nodes` and `groups` are mutually exclusive") + return data + + @field_validator("groups") + @classmethod + def validate_groups(cls, v: Optional[List[NodeGroup]]) -> Optional[List[NodeGroup]]: + 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 node group names found: {duplicates}. " + "Each node group must have a unique name." + ) + return v class TaskConfiguration( @@ -811,6 +891,24 @@ class TaskConfiguration( ): type: Literal["task"] = "task" + @property + def node_groups(self) -> List[NodeGroup]: + if self.groups is not None: + return self.groups + return [ + NodeGroup( + name=DEFAULT_REPLICA_GROUP_NAME, + nodes=self.nodes, + commands=self.commands, + resources=self.resources, + ports=self.ports, + ) + ] + + @property + def nodes_num(self) -> int: + return sum(group.nodes for group in self.node_groups) + def _validate_replica_range(v: Range[int]) -> Range[int]: """Validate a Range[int] used for replica counts.""" diff --git a/src/dstack/_internal/core/models/runs.py b/src/dstack/_internal/core/models/runs.py index a292a928b8..c41c16084a 100644 --- a/src/dstack/_internal/core/models/runs.py +++ b/src/dstack/_internal/core/models/runs.py @@ -299,6 +299,18 @@ class JobSpec(CoreModel): service_port: Optional[int] = None """`service_port` is `None` for non-services and pre-0.19.19 services. See `get_service_port`.""" probes: list[ProbeSpec] = [] + node_group_index: int = 0 + """`node_group_index` uses a default value for backward compatibility.""" + node_group_name: str = DEFAULT_REPLICA_GROUP_NAME + """`node_group_name` uses a default value for backward compatibility.""" + node_group_job_index: int = 0 + """That node's index inside its group (0 .. group.nodes-1). + Example: + groups: + - nodes: 2 # jobs get node_group_job_index 0 and 1 + - nodes: 1 # job gets node_group_job_index 0 + Default for backward compatibility. + """ class JobProvisioningData(CoreModel): @@ -391,6 +403,12 @@ class ClusterInfo(CoreModel): job_ips: List[str] master_job_ip: str gpus_per_job: int + """GPU count on this node only.""" + gpus_per_node: List[int] = [] + """GPU count for each node in the run, in `job_ips` order. + Used for heterogeneous node groups where nodes can have different GPU + counts (e.g. `[2, 8]`). `0` means CPU-only. Empty for older servers. + """ class Probe(CoreModel): 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 720b0141ba..21d78da495 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py @@ -121,6 +121,10 @@ from dstack._internal.utils.common import get_current_datetime, get_or_error, run_async from dstack._internal.utils.interpolator import InterpolatorError from dstack._internal.utils.logging import get_logger +from dstack._internal.utils.nodes_interpolator import ( + find_groups_ip_refs, + interpolate_groups_ip_address, +) logger = get_logger(__name__) @@ -611,6 +615,28 @@ async def _prepare_startup_context( ) return None + commands = context.job.job_spec.commands + 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, @@ -1762,25 +1788,69 @@ def _reset_disconnected_at(job_model: JobModel, result: _ProcessResult) -> None: result.job_update_map["disconnected_at"] = None +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: + return [] + max_group_index = max(job.job_spec.node_group_index for job in replica_jobs) + nodes: list[list[str]] = [[] for _ in range(max_group_index + 1)] + for job in replica_jobs: + group_index = job.job_spec.node_group_index + 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 + return nodes + + +def _referenced_ips_ready(commands: list[str], nodes_view: list[list[str]]) -> bool: + for command in commands: + for group_index, node_index in find_groups_ip_refs(command): + if group_index >= len(nodes_view) or node_index >= len(nodes_view[group_index]): + raise InterpolatorError( + f"Invalid reference groups[{group_index}].nodes[{node_index}].IP_ADDRESS: " + "out of range" + ) + # Wait until every referenced slot has a non-empty internal IP. + if not nodes_view[group_index][node_index]: + return False + return True + + def _get_cluster_info( jobs: list[Job], replica_num: int, job_provisioning_data: JobProvisioningData, job_runtime_data: Optional[JobRuntimeData], ) -> ClusterInfo: - job_ips = [] - for job in jobs: - if job.job_spec.replica_num == replica_num: - job_ips.append( - get_or_error(job.job_submissions[-1].job_provisioning_data).internal_ip or "" - ) + job_ips: list[str] = [] + gpus_per_node: list[int] = [] + replica_jobs = sorted( + (job for job in jobs if job.job_spec.replica_num == replica_num), + key=lambda j: j.job_spec.job_num, + ) + for job in replica_jobs: + submission = job.job_submissions[-1] + jpd = get_or_error(submission.job_provisioning_data) + job_ips.append(jpd.internal_ip or "") + jrd = submission.job_runtime_data + if jrd is not None and jrd.offer is not None: + gpus_per_node.append(len(jrd.offer.instance.resources.gpus)) + else: + gpus_per_node.append(len(jpd.instance_type.resources.gpus)) gpus_per_job = len(job_provisioning_data.instance_type.resources.gpus) if job_runtime_data is not None and job_runtime_data.offer is not None: gpus_per_job = len(job_runtime_data.offer.instance.resources.gpus) return ClusterInfo( job_ips=job_ips, - master_job_ip=job_ips[0], + master_job_ip=job_ips[0] if job_ips else "", gpus_per_job=gpus_per_job, + gpus_per_node=gpus_per_node, ) diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py index 74e1031c5f..e7cf969c2a 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py @@ -49,6 +49,7 @@ JobTerminationReason, Requirements, Run, + RunSpec, ) from dstack._internal.core.models.volumes import Volume from dstack._internal.core.services.profiles import get_termination @@ -777,25 +778,45 @@ async def _fetch_run_model_for_submitted_job( ) -> RunModel: """Fetch run model with only the relevant latest-submission jobs. + Loading jobs is separate from provisioning them. job_num=0 may load all + siblings for coordination, but still provisions only its own node group. + Only a small subset is needed depending on the job type: - * Master multinode: all same-replica jobs (for cluster provisioning and releasing sibling waits). - * Non-master: master job + current job (for master provisioning data lookup). - * Master single-node: current job only (no siblings needed). + * Multinode master (job_num=0): all same-replica jobs. + * First job in a node group (not job 0): job 0 + jobs in its group + (same shape batch). + * Other multinode jobs: job 0 + current job. + * Single-node master: current job only. Only the latest submission per (replica_num, job_num) is loaded since historical submissions are never accessed in submitted job processing. """ + job_spec = get_job_spec(job_model) is_master = job_model.job_num == 0 - is_multinode = get_job_spec(job_model).jobs_per_replica > 1 + is_multinode = job_spec.jobs_per_replica > 1 job_num_filters: list = [] - if is_master and not is_multinode: - # Master single-node: only current job needed. - job_num_filters.append(JobModel.job_num == 0) - elif not is_master: - # Non-master: master job (for provisioning data) + current job. + if not is_multinode: + if is_master: + # Single-node master: only current job needed. + job_num_filters.append(JobModel.job_num == 0) + else: + # Non-master single-node should not happen; keep master + current for safety. + job_num_filters.append(JobModel.job_num.in_([0, job_model.job_num])) + elif is_master: + # Multinode master: load all jobs (fleet setup, release waiting_master_job). + # Provisioning still batches only this job's node group (same shape). + pass + elif job_spec.node_group_job_index == 0: + # First job in a node group (not job 0): job 0 + this group's jobs. + run_spec = await _get_run_spec(session, job_model.run_id) + group_job_nums = _job_nums_for_node_group( + run_spec.configuration, job_spec.node_group_index + ) + job_num_filters.append(JobModel.job_num.in_(sorted({0, *group_job_nums}))) + else: + # Other multinode jobs: job 0 + current job. job_num_filters.append(JobModel.job_num.in_([0, job_model.job_num])) - # else: master multinode — no job_num filter, load all jobs in replica. latest_submissions_sq = ( select( @@ -839,6 +860,21 @@ async def _fetch_run_model_for_submitted_job( return res.unique().scalar_one() +async def _get_run_spec(session: AsyncSession, run_id: uuid.UUID) -> RunSpec: + res = await session.execute(select(RunModel.run_spec).where(RunModel.id == run_id)) + return RunSpec.model_validate_json(res.scalar_one()) + + +def _job_nums_for_node_group(configuration, group_index: int) -> list[int]: + assert configuration.type == "task" + job_num = 0 + for index, group in enumerate(configuration.node_groups): + if index == group_index: + return list(range(job_num, job_num + group.nodes)) + job_num += group.nodes + raise ValueError(f"node_group_index {group_index} out of range") + + def _get_job_models_for_jobs( job_models: list[JobModel], jobs: list[Job], @@ -2142,15 +2178,56 @@ def _hint_pipelines_fetch( pipeline_hinter.hint_fetch(FleetModel.__name__) +def _is_node_group_master(job: Job, replica_jobs: list[Job]) -> bool: + """True if `job` has the lowest job_num among loaded jobs in its node group. + + `job` must be in `replica_jobs`. + """ + group_index = job.job_spec.node_group_index + group_job_nums = [ + j.job_spec.job_num for j in replica_jobs if j.job_spec.node_group_index == group_index + ] + return job.job_spec.job_num == min(group_job_nums) + + +def _job_needs_provisioning(job: Job) -> bool: + if not job.job_submissions: + return True + return job.job_submissions[-1].job_provisioning_data is None + + def _select_jobs_to_provision(job: Job, replica_jobs: list[Job], job_model: JobModel) -> list[Job]: - jobs_to_provision = [job] - if is_multinode_job(job) and is_master_job(job) and job_model.waiting_master_job is not None: - jobs_to_provision = replica_jobs - return jobs_to_provision + """Select jobs to launch in this provision attempt. + + Homogeneous multinode (`nodes: N` → one node group) still batches the whole + replica on the group master (rank 0). + + Heterogeneous node groups batch only jobs that share `node_group_index`, so + ComputeGroup backends (`run_jobs`) receive a single-shape offer set. + + Global `waiting_master_job` is unchanged: non-masters stay blocked until the + global master (job_num=0) finishes its provision attempt. + """ + if not is_multinode_job(job): + return [job] + # Legacy rows without the master-wait protocol: provision one-by-one only. + if job_model.waiting_master_job is None: + return [job] + if not _is_node_group_master(job, replica_jobs): + return [job] + + group_index = job.job_spec.node_group_index + group_jobs = [ + j + for j in replica_jobs + if j.job_spec.node_group_index == group_index and _job_needs_provisioning(j) + ] + return group_jobs if group_jobs else [job] def _get_required_targeted_instance_offers(context: _SubmittedJobContext) -> int: - if is_multinode_job(context.job) and is_master_job(context.job): + # Node-group masters (including non-zero groups) may batch multiple jobs. + if is_multinode_job(context.job) and len(context.jobs_to_provision) > 1: return len(context.jobs_to_provision) return 1 @@ -2160,9 +2237,15 @@ def _release_replica_jobs_from_master_wait( replica_job_models: list[JobModel], jobs_to_provision: list[Job], ) -> None: - if len(jobs_to_provision) > 1: - logger.debug("%s: allow replica jobs to be provisioned one-by-one", fmt(job_model)) - for replica_job_model in replica_job_models: + # Global master may only provision its own node group (len == 1). Still release + # waiting workers so other groups can provision on later ticks. + if job_model.job_num != 0: + return + if not any(m.waiting_master_job for m in replica_job_models): + return + logger.debug("%s: allow replica jobs to be provisioned one-by-one", fmt(job_model)) + for replica_job_model in replica_job_models: + if replica_job_model.waiting_master_job: replica_job_model.waiting_master_job = False diff --git a/src/dstack/_internal/server/services/jobs/configurators/base.py b/src/dstack/_internal/server/services/jobs/configurators/base.py index 5ec38790dd..9ae8b38b3d 100644 --- a/src/dstack/_internal/server/services/jobs/configurators/base.py +++ b/src/dstack/_internal/server/services/jobs/configurators/base.py @@ -3,6 +3,7 @@ import sys import threading from abc import ABC, abstractmethod +from dataclasses import dataclass from pathlib import PurePosixPath from typing import Dict, List, Optional @@ -26,6 +27,7 @@ LEGACY_REPO_DIR, OPENAI_MODEL_PROBE_TIMEOUT, HTTPHeaderSpec, + NodeGroup, PortMapping, ProbeConfig, PythonVersion, @@ -94,6 +96,13 @@ def get_default_image(nvcc: bool = False) -> str: return f"{settings.DSTACK_DOCKER_BASE_IMAGE}:{settings.DSTACK_DOCKER_BASE_IMAGE_VERSION}-{'devel' if nvcc else 'base'}-ubuntu{settings.DSTACK_DOCKER_BASE_IMAGE_UBUNTU_VERSION}" +@dataclass(frozen=True) +class NodeGroupJobContext: + group: NodeGroup + group_index: int + job_index: int + + class JobConfigurator(ABC): TYPE: RunConfigurationType @@ -116,7 +125,7 @@ async def get_job_specs(self, replica_num: int) -> List[JobSpec]: return [job_spec] @abstractmethod - def _shell_commands(self) -> List[str]: + def _shell_commands(self, node_group: Optional[NodeGroup] = None) -> List[str]: pass @abstractmethod @@ -135,7 +144,7 @@ def _reservation(self) -> Optional[str]: return self.run_spec.merged_profile.reservation @abstractmethod - def _ports(self) -> List[PortMapping]: + def _ports(self, node_group: Optional[NodeGroup] = None) -> List[PortMapping]: pass async def _get_image_config(self) -> ImageConfig: @@ -165,15 +174,17 @@ async def _get_job_spec( replica_num: int, job_num: int, jobs_per_replica: int, + node_group_context: Optional[NodeGroupJobContext] = None, ) -> JobSpec: + node_group = node_group_context.group if node_group_context is not None else None job_spec = JobSpec( replica_num=replica_num, # TODO(egor-s): add to env variables in the runner job_num=job_num, job_name=f"{self.run_spec.run_name}-{job_num}-{replica_num}", jobs_per_replica=jobs_per_replica, replica_group=self.replica_group_name or DEFAULT_REPLICA_GROUP_NAME, - app_specs=self._app_specs(), - commands=await self._commands(), + app_specs=self._app_specs(node_group), + commands=await self._commands(node_group), env=self._env(), home_dir=self._home_dir(), image_name=self._image_name(), @@ -184,7 +195,7 @@ async def _get_job_spec( stop_duration=self._stop_duration(), utilization_policy=self._utilization_policy(), registry_auth=self._registry_auth(), - requirements=self._requirements(jobs_per_replica), + requirements=self._requirements(jobs_per_replica, node_group), retry=self._retry(), working_dir=self._working_dir(), volumes=self._volumes(job_num), @@ -196,6 +207,17 @@ async def _get_job_spec( file_archives=self.run_spec.file_archives, service_port=self._service_port(), probes=self._probes(), + node_group_index=( + node_group_context.group_index if node_group_context is not None else 0 + ), + node_group_name=( + node_group.name + if node_group is not None and node_group.name is not None + else DEFAULT_REPLICA_GROUP_NAME + ), + node_group_job_index=( + node_group_context.job_index if node_group_context is not None else 0 + ), ) return job_spec @@ -210,12 +232,12 @@ def _shell(self) -> str: return "/bin/bash" return "/bin/sh" - async def _commands(self) -> List[str]: + async def _commands(self, node_group: Optional[NodeGroup] = None) -> List[str]: if self.run_spec.configuration.entrypoint is not None: # docker-like format assert self.run_spec.configuration.type != "dev-environment" entrypoint = shlex.split(self.run_spec.configuration.entrypoint) commands = self.run_spec.configuration.commands - elif shell_commands := self._shell_commands(): + elif shell_commands := self._shell_commands(node_group): entrypoint = [self._shell(), "-i", "-c"] dstack_image_commands = self._dstack_image_commands() commands = [_join_shell_commands(dstack_image_commands + shell_commands)] @@ -265,9 +287,9 @@ def _dstack_image_commands(self) -> List[str]: f"eval $(echo '. $DSTACK_VENV_DIR/bin/activate' | sudo tee -a {DSTACK_PROFILE_PATH})", ] - def _app_specs(self) -> List[AppSpec]: + def _app_specs(self, node_group: Optional[NodeGroup] = None) -> List[AppSpec]: specs = [] - for i, pm in enumerate(filter_reserved_ports(self._ports())): + for i, pm in enumerate(filter_reserved_ports(self._ports(node_group))): specs.append( AppSpec( port=pm.container_port, @@ -335,13 +357,19 @@ def _utilization_policy(self) -> Optional[UtilizationPolicy]: def _registry_auth(self) -> Optional[RegistryAuth]: return self.run_spec.configuration.registry_auth - def _requirements(self, jobs_per_replica: int) -> Requirements: + def _requirements( + self, + jobs_per_replica: int, + node_group: Optional[NodeGroup] = None, + ) -> Requirements: resources = self.run_spec.configuration.resources if self.run_spec.configuration.type == "service": for group in self.run_spec.configuration.replica_groups: if group.name == self.replica_group_name: resources = group.resources break + elif self.run_spec.configuration.type == "task" and node_group is not None: + resources = node_group.resources spot_policy = self._spot_policy() return Requirements( resources=resources, diff --git a/src/dstack/_internal/server/services/jobs/configurators/dev.py b/src/dstack/_internal/server/services/jobs/configurators/dev.py index e4ee0a2d56..39d77d63a2 100644 --- a/src/dstack/_internal/server/services/jobs/configurators/dev.py +++ b/src/dstack/_internal/server/services/jobs/configurators/dev.py @@ -1,7 +1,11 @@ from typing import Dict, List, Optional from dstack._internal.core.errors import ServerClientError -from dstack._internal.core.models.configurations import PortMapping, RunConfigurationType +from dstack._internal.core.models.configurations import ( + NodeGroup, + PortMapping, + RunConfigurationType, +) from dstack._internal.core.models.profiles import SpotPolicy from dstack._internal.core.models.runs import RunSpec from dstack._internal.server.services.ides import get_ide @@ -33,7 +37,7 @@ def __init__( self.ide = ide super().__init__(run_spec=run_spec, secrets=secrets, replica_group_name=replica_group_name) - def _shell_commands(self) -> List[str]: + def _shell_commands(self, node_group: Optional[NodeGroup] = None) -> List[str]: assert self.run_spec.configuration.type == "dev-environment" commands = [] @@ -65,6 +69,6 @@ def _default_max_duration(self) -> Optional[int]: def _spot_policy(self) -> SpotPolicy: return self.run_spec.merged_profile.spot_policy or SpotPolicy.ONDEMAND - def _ports(self) -> List[PortMapping]: + def _ports(self, node_group: Optional[NodeGroup] = None) -> List[PortMapping]: assert self.run_spec.configuration.type == "dev-environment" return self.run_spec.configuration.ports diff --git a/src/dstack/_internal/server/services/jobs/configurators/service.py b/src/dstack/_internal/server/services/jobs/configurators/service.py index 45bc4c8f72..9861e6fbf5 100644 --- a/src/dstack/_internal/server/services/jobs/configurators/service.py +++ b/src/dstack/_internal/server/services/jobs/configurators/service.py @@ -2,6 +2,7 @@ from dstack._internal import settings from dstack._internal.core.models.configurations import ( + NodeGroup, PortMapping, ReplicaGroup, RunConfigurationType, @@ -24,7 +25,7 @@ def _current_replica_group(self) -> Optional[ReplicaGroup]: return group return None - def _shell_commands(self) -> List[str]: + def _shell_commands(self, node_group: Optional[NodeGroup] = None) -> List[str]: assert self.run_spec.configuration.type == "service" group = self._current_replica_group() if group is not None: @@ -124,5 +125,5 @@ def _reservation(self) -> Optional[str]: return group.reservation return super()._reservation() - def _ports(self) -> List[PortMapping]: + def _ports(self, node_group: Optional[NodeGroup] = None) -> List[PortMapping]: return [] diff --git a/src/dstack/_internal/server/services/jobs/configurators/task.py b/src/dstack/_internal/server/services/jobs/configurators/task.py index 51c136dfe7..03bdf9e8cb 100644 --- a/src/dstack/_internal/server/services/jobs/configurators/task.py +++ b/src/dstack/_internal/server/services/jobs/configurators/task.py @@ -1,9 +1,16 @@ from typing import List, Optional -from dstack._internal.core.models.configurations import PortMapping, RunConfigurationType +from dstack._internal.core.models.configurations import ( + NodeGroup, + PortMapping, + RunConfigurationType, +) from dstack._internal.core.models.profiles import SpotPolicy from dstack._internal.core.models.runs import JobSpec -from dstack._internal.server.services.jobs.configurators.base import JobConfigurator +from dstack._internal.server.services.jobs.configurators.base import ( + JobConfigurator, + NodeGroupJobContext, +) class TaskJobConfigurator(JobConfigurator): @@ -11,18 +18,31 @@ class TaskJobConfigurator(JobConfigurator): async def get_job_specs(self, replica_num: int) -> List[JobSpec]: assert self.run_spec.configuration.type == "task" + groups = self.run_spec.configuration.node_groups + total = sum(group.nodes for group in groups) + job_specs = [] - for job_num in range(self.run_spec.configuration.nodes): - job_spec = await self._get_job_spec( - replica_num=replica_num, - job_num=job_num, - jobs_per_replica=self.run_spec.configuration.nodes, - ) - job_specs.append(job_spec) + job_num = 0 + for group_index, group in enumerate(groups): + for local_index in range(group.nodes): + job_spec = await self._get_job_spec( + replica_num=replica_num, + job_num=job_num, + jobs_per_replica=total, + node_group_context=NodeGroupJobContext( + group=group, + group_index=group_index, + job_index=local_index, + ), + ) + job_specs.append(job_spec) + job_num += 1 return job_specs - def _shell_commands(self) -> List[str]: + def _shell_commands(self, node_group: Optional[NodeGroup] = None) -> List[str]: assert self.run_spec.configuration.type == "task" + if node_group is not None and node_group.commands: + return node_group.commands return self.run_spec.configuration.commands def _default_single_branch(self) -> bool: @@ -34,6 +54,8 @@ def _default_max_duration(self) -> Optional[int]: def _spot_policy(self) -> SpotPolicy: return self.run_spec.merged_profile.spot_policy or SpotPolicy.ONDEMAND - def _ports(self) -> List[PortMapping]: + def _ports(self, node_group: Optional[NodeGroup] = None) -> List[PortMapping]: assert self.run_spec.configuration.type == "task" + if node_group is not None and node_group.ports: + return node_group.ports return self.run_spec.configuration.ports diff --git a/src/dstack/_internal/server/services/runs/__init__.py b/src/dstack/_internal/server/services/runs/__init__.py index 02b72c981f..3448392ce9 100644 --- a/src/dstack/_internal/server/services/runs/__init__.py +++ b/src/dstack/_internal/server/services/runs/__init__.py @@ -1182,7 +1182,7 @@ async def _validate_run_volumes( # that won't be created immediately (e.g. range of replicas or nodes). nodes = 1 if run_spec.configuration.type == "task": - nodes = run_spec.configuration.nodes + nodes = run_spec.configuration.nodes_num for job_num in range(nodes): volumes = await get_job_configured_volumes( session=session, project=project, run_spec=run_spec, job_num=job_num diff --git a/src/dstack/_internal/server/services/runs/spec.py b/src/dstack/_internal/server/services/runs/spec.py index 364f81769c..508b644ccf 100644 --- a/src/dstack/_internal/server/services/runs/spec.py +++ b/src/dstack/_internal/server/services/runs/spec.py @@ -256,7 +256,7 @@ def can_update_run_spec(current_run_spec: RunSpec, new_run_spec: RunSpec) -> boo def get_nodes_required_num(run_spec: RunSpec) -> int: nodes_required_num = 1 if run_spec.configuration.type == "task": - nodes_required_num = run_spec.configuration.nodes + 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 diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py index 46b51a189e..079ec6ee15 100644 --- a/src/dstack/_internal/server/testing/common.py +++ b/src/dstack/_internal/server/testing/common.py @@ -447,10 +447,16 @@ async def create_job( if deployment_num is None: deployment_num = run.deployment_num run_spec = validate_json_extra_ignore(RunSpec, run.run_spec) - job_spec = ( - await get_job_specs_from_run_spec(run_spec=run_spec, secrets={}, replica_num=replica_num) - )[0] - job_spec.job_num = job_num + job_specs = await get_job_specs_from_run_spec( + run_spec=run_spec, secrets={}, replica_num=replica_num + ) + if 0 <= job_num < len(job_specs): + job_spec = job_specs[job_num] + else: + job_spec = job_specs[0].model_copy(deep=True) + job_spec.job_num = job_num + job_spec.job_name = f"{run_spec.run_name}-{job_num}-{replica_num}" + job = JobModel( project_id=run.project_id, fleet=fleet, diff --git a/src/dstack/_internal/utils/interpolator.py b/src/dstack/_internal/utils/interpolator.py index 9a4e44659b..641b71dafb 100644 --- a/src/dstack/_internal/utils/interpolator.py +++ b/src/dstack/_internal/utils/interpolator.py @@ -63,10 +63,15 @@ def interpolate( raise InterpolatorError(f"No pattern closing: {s[opening:]}") name = s[opening + len(Pattern.opening) : closing].strip() - if not self.validate_name(name): - raise InterpolatorError(f"Illegal reference name: {name}") - if name.split(".")[0] in self.skip: + # Skip before validate_name so non-standard refs (e.g. groups[0].nodes[0].IP_ADDRESS) + # can be left for later interpolators. Invalid skipped names without brackets + # (e.g. secrets.pass-word) still raise. + root = name.split(".")[0] + skip_ns = root.split("[")[0] + if skip_ns in self.skip and ("[" in root or self.validate_name(name)): tokens.append(s[opening : closing + len(Pattern.closing)]) + elif not self.validate_name(name): + raise InterpolatorError(f"Illegal reference name: {name}") elif name in self.variables: tokens.append(self.variables[name]) else: diff --git a/src/dstack/_internal/utils/nodes_interpolator.py b/src/dstack/_internal/utils/nodes_interpolator.py new file mode 100644 index 0000000000..84512aa00a --- /dev/null +++ b/src/dstack/_internal/utils/nodes_interpolator.py @@ -0,0 +1,24 @@ +import re + +from dstack._internal.utils.interpolator import InterpolatorError + +_GROUPS_IP_REF = re.compile(r"\$\{\{\s*groups\[(\d+)\]\.nodes\[(\d+)\]\.IP_ADDRESS\s*\}\}") + + +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 interpolate_groups_ip_address(s: str, nodes: list[list[str]]) -> str: + def repl(m: re.Match) -> str: + gi, ni = int(m.group(1)), int(m.group(2)) + if gi >= len(nodes) or ni >= len(nodes[gi]): + raise InterpolatorError( + f"Invalid reference groups[{gi}].nodes[{ni}].IP_ADDRESS: out of range" + ) + ip = nodes[gi][ni] + if not ip: + raise InterpolatorError(f"IP not available for groups[{gi}].nodes[{ni}].IP_ADDRESS") + return ip + + return _GROUPS_IP_REF.sub(repl, s) diff --git a/src/tests/_internal/core/models/test_configurations.py b/src/tests/_internal/core/models/test_configurations.py index 98f6ac2b5b..3eb92a39a9 100644 --- a/src/tests/_internal/core/models/test_configurations.py +++ b/src/tests/_internal/core/models/test_configurations.py @@ -916,3 +916,120 @@ def test_ide_optional(self): def test_version_requires_ide(self): with pytest.raises(ValueError, match="`version` requires `ide` to be set"): DevEnvironmentConfigurationParams(version="1.80.0") + + +class TestNodeGroups: + def test_parses_int_nodes(self): + parsed = parse_run_configuration({"type": "task", "nodes": 2, "commands": ["true"]}) + assert parsed.type == "task" + assert parsed.nodes == 2 + assert parsed.groups is None + assert parsed.nodes_num == 2 + assert len(parsed.node_groups) == 1 + assert parsed.node_groups[0].nodes == 2 + assert parsed.node_groups[0].name == "0" + + def test_parses_groups_and_defaults_names(self): + parsed = parse_run_configuration( + { + "type": "task", + "image": "debian", + "groups": [ + {"nodes": 2, "commands": ["echo head"]}, + {"name": "workers", "nodes": 1, "commands": ["echo worker"]}, + ], + } + ) + assert parsed.type == "task" + assert parsed.groups is not None + assert parsed.nodes_num == 3 + assert parsed.node_groups[0].name == "0" + assert parsed.node_groups[1].name == "workers" + assert parsed.node_groups[0].commands == ["echo head"] + assert parsed.node_groups[1].commands == ["echo worker"] + + def test_accepts_default_nodes_with_groups(self): + # Serialized TaskConfiguration always includes nodes=1 (the field default). + # The xor validator must allow that so model round-trips succeed. + parsed = parse_run_configuration( + { + "type": "task", + "image": "debian", + "nodes": 1, + "groups": [ + {"nodes": 2, "commands": ["echo head"]}, + {"name": "workers", "nodes": 1, "commands": ["echo worker"]}, + ], + } + ) + assert parsed.type == "task" + assert parsed.nodes == 1 + assert parsed.groups is not None + assert parsed.nodes_num == 3 + + def test_groups_round_trip_via_model_dump(self): + parsed = parse_run_configuration( + { + "type": "task", + "image": "debian", + "groups": [ + {"nodes": 2, "commands": ["echo head"], "ports": [8000]}, + {"name": "workers", "nodes": 1, "commands": ["echo worker"]}, + ], + } + ) + assert parsed.type == "task" + dumped = parsed.model_dump(mode="json") + assert dumped["nodes"] == 1 + assert dumped["groups"] is not None + + reparsed = parse_run_configuration(dumped) + assert reparsed.type == "task" + assert reparsed.nodes == 1 + assert reparsed.nodes_num == 3 + assert [g.name for g in reparsed.node_groups] == ["0", "workers"] + assert reparsed.node_groups[0].commands == ["echo head"] + assert reparsed.node_groups[0].ports[0].container_port == 8000 + assert reparsed.node_groups[1].commands == ["echo worker"] + + def test_rejects_auto_name_collision_with_explicit_name(self): + # Unnamed group at index 1 becomes "1", colliding with an explicit name "1". + with pytest.raises(ConfigurationError, match="Duplicate node group names"): + parse_run_configuration( + { + "type": "task", + "image": "debian", + "groups": [ + {"name": "1", "nodes": 1, "commands": ["true"]}, + {"nodes": 1, "commands": ["true"]}, + ], + } + ) + + def test_rejects_duplicate_group_names(self): + with pytest.raises(ConfigurationError, match="Duplicate node group names"): + parse_run_configuration( + { + "type": "task", + "image": "debian", + "groups": [ + {"name": "head", "nodes": 1, "commands": ["true"]}, + {"name": "head", "nodes": 1, "commands": ["true"]}, + ], + } + ) + + def test_rejects_empty_groups(self): + with pytest.raises(ConfigurationError, match="cannot be an empty list"): + parse_run_configuration({"type": "task", "image": "debian", "groups": []}) + + def test_rejects_nodes_and_groups_together(self): + with pytest.raises(ConfigurationError, match="mutually exclusive"): + parse_run_configuration( + { + "type": "task", + "image": "debian", + "nodes": 2, + "groups": [{"nodes": 1, "commands": ["true"]}], + } + ) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_node_groups.py b/src/tests/_internal/server/background/pipeline_tasks/test_node_groups.py new file mode 100644 index 0000000000..6366a274bc --- /dev/null +++ b/src/tests/_internal/server/background/pipeline_tasks/test_node_groups.py @@ -0,0 +1,115 @@ +from datetime import datetime, timezone +from uuid import uuid4 + +import pytest + +from dstack._internal.core.models.runs import Job, JobSpec, JobSubmission +from dstack._internal.server.background.pipeline_tasks.jobs_running import ( + _build_nodes_ip_view, + _get_cluster_info, + _referenced_ips_ready, +) +from dstack._internal.server.testing.common import get_job_provisioning_data +from dstack._internal.utils.interpolator import InterpolatorError + + +def _job( + *, + job_num: int, + node_group_index: int, + node_group_job_index: int, + internal_ip: str, + gpu_count: int, +) -> Job: + return Job.model_construct( + job_spec=JobSpec.model_construct( + replica_num=0, + job_num=job_num, + node_group_index=node_group_index, + node_group_job_index=node_group_job_index, + commands=[], + ), + job_submissions=[ + JobSubmission.model_construct( + id=uuid4(), + submitted_at=datetime.now(timezone.utc), + job_provisioning_data=get_job_provisioning_data( + internal_ip=internal_ip, + gpu_count=gpu_count, + ), + job_runtime_data=None, + ) + ], + ) + + +class TestGetClusterInfo: + def test_fills_gpus_per_node(self): + jobs = [ + _job( + job_num=0, + node_group_index=0, + node_group_job_index=0, + internal_ip="10.0.0.1", + gpu_count=8, + ), + _job( + job_num=1, + node_group_index=1, + node_group_job_index=0, + internal_ip="10.0.0.2", + gpu_count=4, + ), + ] + this_jpd = get_job_provisioning_data(internal_ip="10.0.0.1", gpu_count=8) + info = _get_cluster_info( + jobs=jobs, + replica_num=0, + job_provisioning_data=this_jpd, + job_runtime_data=None, + ) + assert info.job_ips == ["10.0.0.1", "10.0.0.2"] + assert info.master_job_ip == "10.0.0.1" + assert info.gpus_per_job == 8 + assert info.gpus_per_node == [8, 4] + + +class TestNodesIpView: + def test_builds_group_view(self): + jobs = [ + _job( + job_num=0, + node_group_index=0, + node_group_job_index=0, + internal_ip="10.0.0.1", + gpu_count=1, + ), + _job( + job_num=1, + node_group_index=0, + node_group_job_index=1, + internal_ip="10.0.0.2", + gpu_count=1, + ), + _job( + job_num=2, + node_group_index=1, + node_group_job_index=0, + internal_ip="10.0.0.3", + gpu_count=1, + ), + ] + assert _build_nodes_ip_view(jobs, replica_num=0) == [ + ["10.0.0.1", "10.0.0.2"], + ["10.0.0.3"], + ] + + def test_referenced_ips_ready(self): + nodes_view = [["10.0.0.1"], [""]] + assert _referenced_ips_ready(["echo ${{ groups[0].nodes[0].IP_ADDRESS }}"], nodes_view) + assert not _referenced_ips_ready(["echo ${{ groups[1].nodes[0].IP_ADDRESS }}"], nodes_view) + + 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) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py index ee34ec5ad5..491868327f 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py @@ -17,7 +17,11 @@ RegistryAuth, validate_json_extra_ignore, ) -from dstack._internal.core.models.configurations import ServiceConfiguration, TaskConfiguration +from dstack._internal.core.models.configurations import ( + NodeGroup, + ServiceConfiguration, + TaskConfiguration, +) from dstack._internal.core.models.envs import Env from dstack._internal.core.models.fleets import FleetNodesSpec, InstanceGroupPlacement from dstack._internal.core.models.instances import InstanceStatus @@ -2692,13 +2696,15 @@ async def test_single_node_master_loads_only_current_job(self, test_db, session: assert not context.multinode assert context.jobs_to_provision == [context.job] - async def test_non_master_loads_master_and_current_job(self, test_db, session: AsyncSession): - """Non-master: run_model.jobs should contain master job + current job (latest submissions).""" + async def test_non_master_multinode_loads_master_and_current_job( + self, test_db, session: AsyncSession + ): + """Homogeneous multinode workers: run_model.jobs should contain job 0 + current.""" project = await create_project(session=session) user = await create_user(session=session) repo = await create_repo(session=session, project_id=project.id) fleet = await create_fleet(session=session, project=project) - configuration = TaskConfiguration(image="debian", nodes=2) + configuration = TaskConfiguration(image="debian", nodes=3) run_spec = get_run_spec(run_name="run", repo_id=repo.name, configuration=configuration) run = await create_run( session=session, @@ -2718,19 +2724,25 @@ async def test_non_master_loads_master_and_current_job(self, test_db, session: A job_provisioning_data=get_job_provisioning_data(), waiting_master_job=False, ) - worker_job = await create_job( + worker_job_1 = await create_job( session=session, run=run, job_num=1, status=JobStatus.SUBMITTED, waiting_master_job=False, ) + await create_job( + session=session, + run=run, + job_num=2, + status=JobStatus.SUBMITTED, + waiting_master_job=False, + ) await session.commit() - context = await _load_submitted_job_context(session=session, job_model=worker_job) - # Only master (job_num=0) and current job (job_num=1) should be loaded. + context = await _load_submitted_job_context(session=session, job_model=worker_job_1) loaded_job_ids = {jm.id for jm in context.run_model.jobs} - assert loaded_job_ids == {master_job.id, worker_job.id} + assert loaded_job_ids == {master_job.id, worker_job_1.id} assert context.jobs_to_provision == [context.job] async def test_multinode_master_loads_all_replica_jobs(self, test_db, session: AsyncSession): @@ -2774,6 +2786,87 @@ async def test_multinode_master_loads_all_replica_jobs(self, test_db, session: A assert len(context.jobs_to_provision) == 2 assert len(context.replica_job_model_ids) == 2 + async def test_node_group_master_provisions_only_its_group( + self, test_db, session: AsyncSession + ): + """Heterogeneous node groups: each group master batches only its own group.""" + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + fleet = await create_fleet(session=session, project=project) + configuration = TaskConfiguration( + image="debian", + groups=[ + NodeGroup( + name="small", + nodes=1, + resources=ResourcesSpec( + cpu=Range[int](min=2), memory=Range[Memory](min=Memory(4)) + ), + commands=["echo small"], + ), + NodeGroup( + name="large", + nodes=2, + resources=ResourcesSpec( + cpu=Range[int](min=4), memory=Range[Memory](min=Memory(8)) + ), + commands=["echo large"], + ), + ], + ) + run_spec = get_run_spec(run_name="run", repo_id=repo.name, configuration=configuration) + run = await create_run( + session=session, + run_name="run", + project=project, + repo=repo, + user=user, + run_spec=run_spec, + fleet=fleet, + ) + small_job = await create_job( + session=session, + run=run, + job_num=0, + status=JobStatus.SUBMITTED, + waiting_master_job=False, + ) + large_job_0 = await create_job( + session=session, + run=run, + job_num=1, + status=JobStatus.SUBMITTED, + waiting_master_job=False, + ) + large_job_1 = await create_job( + session=session, + run=run, + job_num=2, + status=JobStatus.SUBMITTED, + waiting_master_job=False, + ) + await session.commit() + + small_context = await _load_submitted_job_context(session=session, job_model=small_job) + # Multinode master: load all replica jobs. + assert {jm.job_num for jm in small_context.run_model.jobs} == {0, 1, 2} + assert [j.job_spec.job_num for j in small_context.jobs_to_provision] == [0] + assert small_context.jobs_to_provision[0].job_spec.node_group_name == "small" + + large_context = await _load_submitted_job_context(session=session, job_model=large_job_0) + # Heterogeneous group master: job 0 + its node group. + assert {jm.job_num for jm in large_context.run_model.jobs} == {0, 1, 2} + assert sorted(j.job_spec.job_num for j in large_context.jobs_to_provision) == [1, 2] + assert {j.job_spec.node_group_name for j in large_context.jobs_to_provision} == {"large"} + + large_worker_context = await _load_submitted_job_context( + session=session, job_model=large_job_1 + ) + # Non-master job in the group: job 0 + current. + assert {jm.job_num for jm in large_worker_context.run_model.jobs} == {0, 2} + assert [j.job_spec.job_num for j in large_worker_context.jobs_to_provision] == [2] + async def test_loads_only_latest_submission(self, test_db, session: AsyncSession): """Only the latest submission per (replica_num, job_num) should be loaded, not historical ones.""" project = await create_project(session=session) diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py index 46d530fff1..9319baba9f 100644 --- a/src/tests/_internal/server/routers/test_runs.py +++ b/src/tests/_internal/server/routers/test_runs.py @@ -321,6 +321,9 @@ def get_dev_env_run_plan_dict( "file_archives": [], "service_port": None, "probes": [], + "node_group_index": 0, + "node_group_name": "0", + "node_group_job_index": 0, }, "offers": [json.loads(o.model_dump_json()) for o in offers], "total_offers": total_offers, @@ -568,6 +571,9 @@ def get_dev_env_run_dict( "file_archives": [], "service_port": None, "probes": [], + "node_group_index": 0, + "node_group_name": "0", + "node_group_job_index": 0, }, "job_submissions": [ { diff --git a/src/tests/_internal/server/services/jobs/configurators/test_task.py b/src/tests/_internal/server/services/jobs/configurators/test_task.py index 54b8dd666d..383cd411e2 100644 --- a/src/tests/_internal/server/services/jobs/configurators/test_task.py +++ b/src/tests/_internal/server/services/jobs/configurators/test_task.py @@ -3,7 +3,8 @@ import pytest -from dstack._internal.core.models.configurations import TaskConfiguration +from dstack._internal.core.models.configurations import NodeGroup, TaskConfiguration +from dstack._internal.core.models.resources import GPUSpec, ResourcesSpec from dstack._internal.core.models.runs import JobSSHKey from dstack._internal.server.services.docker import ImageConfig from dstack._internal.server.services.jobs.configurators.task import TaskJobConfigurator @@ -37,6 +38,65 @@ async def test_multi_node(self): assert job_specs[1].ssh_key == JobSSHKey(private="private1", public="public1") +@pytest.mark.asyncio +@pytest.mark.usefixtures("image_config_mock") +class TestNodeGroups: + async def test_assigns_contiguous_ranks_and_metadata(self): + configuration = TaskConfiguration( + image="debian", + groups=[ + NodeGroup(name="head", nodes=2, commands=["echo head"]), + NodeGroup(name="workers", nodes=2, commands=["echo worker"]), + ], + ) + run_spec = get_run_spec(run_name="run", repo_id="id", configuration=configuration) + configurator = TaskJobConfigurator(run_spec) + + job_specs = await configurator.get_job_specs(replica_num=0) + + assert len(job_specs) == 4 + assert [j.job_num for j in job_specs] == [0, 1, 2, 3] + assert [j.jobs_per_replica for j in job_specs] == [4, 4, 4, 4] + assert [j.node_group_name for j in job_specs] == [ + "head", + "head", + "workers", + "workers", + ] + assert [j.node_group_index for j in job_specs] == [0, 0, 1, 1] + assert [j.node_group_job_index for j in job_specs] == [0, 1, 0, 1] + + async def test_uses_per_group_commands_and_resources(self): + configuration = TaskConfiguration( + image="debian", + groups=[ + NodeGroup( + name="head", + nodes=1, + commands=["echo head"], + resources=ResourcesSpec(gpu=GPUSpec(name=["H100"], count=1)), + ), + NodeGroup( + name="workers", + nodes=1, + commands=["echo worker"], + resources=ResourcesSpec(gpu=GPUSpec(name=["A100"], count=2)), + ), + ], + ) + run_spec = get_run_spec(run_name="run", repo_id="id", configuration=configuration) + configurator = TaskJobConfigurator(run_spec) + + job_specs = await configurator.get_job_specs(replica_num=0) + + assert "echo head" in job_specs[0].commands[-1] + assert "echo worker" in job_specs[1].commands[-1] + assert job_specs[0].requirements.resources.gpu.name == ["H100"] + assert job_specs[0].requirements.resources.gpu.count.min == 1 + assert job_specs[1].requirements.resources.gpu.name == ["A100"] + assert job_specs[1].requirements.resources.gpu.count.min == 2 + + @pytest.mark.asyncio @pytest.mark.usefixtures("image_config_mock") class TestServerAccess: diff --git a/src/tests/_internal/utils/test_interpolator.py b/src/tests/_internal/utils/test_interpolator.py index 50c3845832..2acc2eaefc 100644 --- a/src/tests/_internal/utils/test_interpolator.py +++ b/src/tests/_internal/utils/test_interpolator.py @@ -49,3 +49,8 @@ def test_illegal_name(self): get_interpolator().interpolate("${{ secrets.password.hash }}") with pytest.raises(InterpolatorError): 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"]) + assert interpolator.interpolate(s) == s diff --git a/src/tests/_internal/utils/test_nodes_interpolator.py b/src/tests/_internal/utils/test_nodes_interpolator.py new file mode 100644 index 0000000000..84ed9cdc85 --- /dev/null +++ b/src/tests/_internal/utils/test_nodes_interpolator.py @@ -0,0 +1,40 @@ +import pytest + +from dstack._internal.utils.interpolator import InterpolatorError +from dstack._internal.utils.nodes_interpolator import ( + find_groups_ip_refs, + interpolate_groups_ip_address, +) + + +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)] + + 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)] + + def test_no_refs(self): + assert find_groups_ip_refs("echo hello") == [] + + +class TestInterpolateGroupsIpAddress: + def test_replaces_ip(self): + s = "ray start --address=${{ groups[0].nodes[0].IP_ADDRESS }}:6379" + result = interpolate_groups_ip_address(s, [["10.0.0.1", "10.0.0.2"], ["10.0.0.3"]]) + assert result == "ray start --address=10.0.0.1:6379" + + def test_replaces_nested_node(self): + s = "${{ groups[1].nodes[0].IP_ADDRESS }}" + result = interpolate_groups_ip_address(s, [["10.0.0.1"], ["10.0.0.2"]]) + assert result == "10.0.0.2" + + 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 }}", [[""]]) + + def test_raises_when_out_of_range(self): + with pytest.raises(InterpolatorError, match="out of range"): + interpolate_groups_ip_address("${{ groups[1].nodes[0].IP_ADDRESS }}", [["10.0.0.1"]]) From 90cdbc0eac8fd6878d530813090ca6ec23afe27b Mon Sep 17 00:00:00 2001 From: Bihan Rana Date: Mon, 10 Aug 2026 10:10:29 +0545 Subject: [PATCH 2/4] Review Comments Resolved Co-authored-by: Cursor --- runner/internal/runner/executor/executor.go | 24 +- .../internal/runner/executor/executor_test.go | 58 ++++ .../cli/services/configurators/run.py | 6 +- .../_internal/core/backends/slurm/compute.py | 11 +- .../_internal/core/compatibility/runs.py | 2 + .../_internal/core/models/configurations.py | 83 +++++- .../background/pipeline_tasks/jobs_running.py | 28 +- .../pipeline_tasks/jobs_submitted.py | 134 +++++---- .../services/jobs/configurators/base.py | 4 +- .../services/jobs/configurators/task.py | 12 +- .../_internal/server/services/runs/plan.py | 68 +++-- .../_internal/server/services/runs/spec.py | 42 +++ src/dstack/_internal/utils/interpolator.py | 35 ++- .../_internal/utils/nodes_interpolator.py | 51 +++- .../core/models/test_configurations.py | 112 ++++++-- src/tests/_internal/core/models/test_runs.py | 11 + .../pipeline_tasks/test_node_groups.py | 115 -------- .../pipeline_tasks/test_running_jobs.py | 264 +++++++++++++++++- .../pipeline_tasks/test_submitted_jobs.py | 261 ++++++++++++++++- .../services/jobs/configurators/test_task.py | 16 ++ .../server/services/runs/test_plan.py | 157 ++++++++++- .../server/services/runs/test_spec.py | 102 ++++++- .../_internal/utils/test_interpolator.py | 18 +- .../utils/test_nodes_interpolator.py | 54 ++++ 24 files changed, 1399 insertions(+), 269 deletions(-) delete mode 100644 src/tests/_internal/server/background/pipeline_tasks/test_node_groups.py diff --git a/runner/internal/runner/executor/executor.go b/runner/internal/runner/executor/executor.go index bb86c991bd..c1533f4730 100644 --- a/runner/internal/runner/executor/executor.go +++ b/runner/internal/runner/executor/executor.go @@ -552,7 +552,15 @@ func (ex *RunExecutor) execJob(ctx context.Context, jobLogFile io.Writer) error log.Warning(ctx, "failed to include dstack_profile", "path", profilePath, "err", err) } - if err := writeMpiHostfile(ctx, ex.clusterInfo.JobIPs, ex.clusterInfo.GPUSPerNode, gpusPerNodeNum, mpiHostfilePath); err != nil { + slots := ex.clusterInfo.GPUSPerNode + if len(slots) == 0 { + // Old servers omit gpus_per_node; fall back to homogeneous per-node GPU count. + slots = make([]int, len(ex.clusterInfo.JobIPs)) + for i := range slots { + slots[i] = gpusPerNodeNum + } + } + if err := writeMpiHostfile(ctx, ex.clusterInfo.JobIPs, slots, mpiHostfilePath); err != nil { return fmt.Errorf("write MPI hostfile: %w", err) } @@ -767,7 +775,7 @@ func prepareUserSshDir(user *linuxuser.User) (string, error) { return sshDir, nil } -func writeMpiHostfile(ctx context.Context, ips []string, gpusPerNode []int, fallbackGpusPerJob int, path string) error { +func writeMpiHostfile(ctx context.Context, ips []string, slots []int, path string) error { if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { return fmt.Errorf("create MPI hostfile directory: %w", err) } @@ -783,23 +791,19 @@ func writeMpiHostfile(ctx context.Context, ips []string, gpusPerNode []int, fall } } if len(nonEmptyIps) == len(ips) { - if len(gpusPerNode) > 0 && len(gpusPerNode) != len(ips) { + if len(slots) != len(ips) { return fmt.Errorf( "gpus_per_node length %d != job_ips length %d", - len(gpusPerNode), len(ips), + len(slots), len(ips), ) } for i, ip := range nonEmptyIps { - n := fallbackGpusPerJob - if len(gpusPerNode) > 0 { - n = gpusPerNode[i] - } - if n == 0 { + if slots[i] == 0 { // CPU node: the number of slots defaults to the number of processor cores on that host // See: https://docs.open-mpi.org/en/main/launching-apps/scheduling.html#calculating-the-number-of-slots _, err = fmt.Fprintf(file, "%s\n", ip) } else { - _, err = fmt.Fprintf(file, "%s slots=%d\n", ip, n) + _, err = fmt.Fprintf(file, "%s slots=%d\n", ip, slots[i]) } if err != nil { return fmt.Errorf("write MPI hostfile line: %w", err) diff --git a/runner/internal/runner/executor/executor_test.go b/runner/internal/runner/executor/executor_test.go index 2330cd6f3c..d6878ea3ee 100644 --- a/runner/internal/runner/executor/executor_test.go +++ b/runner/internal/runner/executor/executor_test.go @@ -287,6 +287,64 @@ func TestWriteDstackProfile(t *testing.T) { } } +func TestWriteMpiHostfile(t *testing.T) { + tmp := t.TempDir() + + t.Run("heterogeneous_slots", func(t *testing.T) { + path := filepath.Join(tmp, "hostfile_hetero") + err := writeMpiHostfile( + t.Context(), + []string{"10.0.0.1", "10.0.0.2", "10.0.0.3"}, + []int{8, 4, 0}, + path, + ) + require.NoError(t, err) + content, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, "10.0.0.1 slots=8\n10.0.0.2 slots=4\n10.0.0.3\n", string(content)) + }) + + t.Run("homogeneous_slots", func(t *testing.T) { + path := filepath.Join(tmp, "hostfile_homo") + err := writeMpiHostfile( + t.Context(), + []string{"10.0.0.1", "10.0.0.2"}, + []int{4, 4}, + path, + ) + require.NoError(t, err) + content, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, "10.0.0.1 slots=4\n10.0.0.2 slots=4\n", string(content)) + }) + + t.Run("slots_length_mismatch", func(t *testing.T) { + path := filepath.Join(tmp, "hostfile_mismatch") + err := writeMpiHostfile( + t.Context(), + []string{"10.0.0.1", "10.0.0.2"}, + []int{8}, + path, + ) + require.Error(t, err) + assert.Contains(t, err.Error(), "gpus_per_node length 1 != job_ips length 2") + }) + + t.Run("empty_ip_writes_empty_hostfile", func(t *testing.T) { + path := filepath.Join(tmp, "hostfile_empty_ip") + err := writeMpiHostfile( + t.Context(), + []string{"10.0.0.1", ""}, + []int{8, 4}, + path, + ) + require.NoError(t, err) + content, err := os.ReadFile(path) + require.NoError(t, err) + assert.Equal(t, "", string(content)) + }) +} + func TestExecutor_Logs(t *testing.T) { var b bytes.Buffer ex := makeTestExecutor(t) diff --git a/src/dstack/_internal/cli/services/configurators/run.py b/src/dstack/_internal/cli/services/configurators/run.py index 9ace4f474d..d5cc48adf0 100644 --- a/src/dstack/_internal/cli/services/configurators/run.py +++ b/src/dstack/_internal/cli/services/configurators/run.py @@ -69,6 +69,7 @@ from dstack._internal.utils.interpolator import InterpolatorError, VariablesInterpolator from dstack._internal.utils.logging import get_logger from dstack._internal.utils.nested_list import NestedList, NestedListItem +from dstack._internal.utils.nodes_interpolator import is_valid_groups_ip_ref from dstack._internal.utils.path import is_absolute_posix_path from dstack.api._public.runs import Run @@ -696,7 +697,10 @@ def _interpolate_commands(self, commands: list[str], args: argparse.Namespace) - run_args = shlex.join(args.run_args) interpolator = VariablesInterpolator( {"run": {"args": run_args}}, - skip=["secrets", "groups"], + skip={ + "secrets": VariablesInterpolator.validate_name, + "groups": is_valid_groups_ip_ref, + }, ) try: for i, command in enumerate(commands): diff --git a/src/dstack/_internal/core/backends/slurm/compute.py b/src/dstack/_internal/core/backends/slurm/compute.py index 423de1f05d..0849b20ebf 100644 --- a/src/dstack/_internal/core/backends/slurm/compute.py +++ b/src/dstack/_internal/core/backends/slurm/compute.py @@ -138,12 +138,15 @@ def run_job( placement_group: Optional[PlacementGroup], requirements: Requirements, ) -> JobProvisioningData: + # run_job provisions a single dstack job → one Slurm node. Do not fall + # back to jobs_per_replica (total across hetero groups). compute_provisioning_data = self._run_slurm_job( run=run, job=job, instance_offer=instance_offer, project_ssh_public_key=project_ssh_public_key, requirements=requirements, + node_count=1, ) return compute_provisioning_data.job_provisioning_datas[0] @@ -211,11 +214,11 @@ def _run_slurm_job( assert run.run_spec.ssh_key_pub is not None authorized_keys = [project_ssh_public_key.strip(), run.run_spec.ssh_key_pub.strip()] - # Heterogeneous groups provision one shape at a time; Slurm allocation - # size must match that batch. Fall back to jobs_per_replica for - # run_job / homogeneous single-call paths. + # Allocation size must match the provision batch (run_jobs passes + # len(job_configurations)). Default to 1 for any caller that omits it + # — never jobs_per_replica (that is the replica total, not batch size). if node_count is None: - node_count = job.job_spec.jobs_per_replica + node_count = 1 resources_spec = requirements.resources requested_resources = get_requested_resources_from_resources_spec(resources_spec) diff --git a/src/dstack/_internal/core/compatibility/runs.py b/src/dstack/_internal/core/compatibility/runs.py index fa95dcd689..ef4c157e6c 100644 --- a/src/dstack/_internal/core/compatibility/runs.py +++ b/src/dstack/_internal/core/compatibility/runs.py @@ -102,6 +102,8 @@ def get_run_spec_excludes(run_spec: RunSpec) -> IncludeExcludeDictType: if isinstance(run_spec.configuration, TaskConfiguration): if run_spec.configuration.groups is None: configuration_excludes["groups"] = True + if run_spec.configuration.nodes is None: + configuration_excludes["nodes"] = True if isinstance(run_spec.configuration, ServiceConfiguration): if run_spec.configuration.probes: diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py index 04ea6a5385..81fd672edd 100644 --- a/src/dstack/_internal/core/models/configurations.py +++ b/src/dstack/_internal/core/models/configurations.py @@ -814,7 +814,12 @@ class NodeGroup(CoreModel): nodes: Annotated[int, Field(description="The number of nodes in this group", ge=1)] = 1 resources: Annotated[ ResourcesSpec, - Field(description="The resources requirements for nodes in this group"), + Field( + description=( + "The resources requirements for nodes in this group. " + "Does not inherit top-level `resources` (same as replica groups)" + ) + ), ] = ResourcesSpec() commands: Annotated[ CommandsList, @@ -830,21 +835,37 @@ class NodeGroup(CoreModel): def validate_name(cls, v: Optional[str]) -> Optional[str]: if v is not None: if not is_valid_replica_group_name(v): - raise ValueError("Resource name should match regex '^[a-z0-9][a-z0-9-]{0,39}$'") + raise ValueError("Node group name should match regex '^[a-z0-9][a-z0-9-]{0,39}$'") return v + @property + def required_name(self) -> str: + """Name after normalization. + + Omitted names are filled by TaskConfiguration.validate_groups; directly + constructed groups must set `name` explicitly. + """ + if self.name is None: + raise ValueError("NodeGroup.name must be set before use") + return self.name + class TaskConfigurationParams(CoreModel): nodes: Annotated[ - int, + Optional[int], Field(description="The number of nodes for homogeneous multi-node tasks", ge=1), - ] = 1 + ] = None groups: Annotated[ Optional[List[NodeGroup]], Field( description=( "A list of node groups for heterogeneous multi-node tasks. " - "Mutually exclusive with `nodes`." + "Mutually exclusive with `nodes`. " + "When `groups` is set, top-level `commands`, `ports`, and `entrypoint` are " + "not allowed; specify `commands` and `ports` in each node group instead. " + "Top-level `resources` is not rejected (same as replica groups; see server " + "defaults), but each group's `resources` is used for provisioning — omit " + "means default empty resources, not inheritance from the top level." ), ), ] = None @@ -854,12 +875,8 @@ class TaskConfigurationParams(CoreModel): def validate_nodes_xor_groups(cls, data): if not isinstance(data, dict): return data - # Allow groups with default nodes: 1 (serialized configs always include it). - # Reject nodes: N (N != 1) together with groups. - if data.get("groups") is not None and "nodes" in data: - nodes = data.get("nodes") - if nodes is not None and nodes != 1: - raise ValueError("`nodes` and `groups` are mutually exclusive") + if data.get("groups") is not None and data.get("nodes") is not None: + raise ValueError("`nodes` and `groups` are mutually exclusive") return data @field_validator("groups") @@ -891,6 +908,48 @@ class TaskConfiguration( ): type: Literal["task"] = "task" + @model_validator(mode="after") + def validate_top_level_properties_with_node_groups(self) -> Self: + """When groups is set, forbid top-level commands, ports, and entrypoint. + + Top-level `resources` is not rejected: the server may mutate it later + (defaults/plugins), and strict parse-time checks would break round-trips. + Provisioning still uses each group's `resources` (default ResourcesSpec()), + not top-level — same as replica groups. + """ + if self.groups is None: + return self + if self.commands: + raise ValueError( + "Top-level `commands` is not allowed when `groups` is set. " + "Specify `commands` in each node group instead." + ) + if self.ports: + raise ValueError( + "Top-level `ports` is not allowed when `groups` is set. " + "Specify `ports` in each node group instead." + ) + if self.entrypoint is not None: + raise ValueError( + "Top-level `entrypoint` is not allowed when `groups` is set. " + "Specify `commands` in each node group instead." + ) + return self + + @model_validator(mode="after") + def validate_node_groups_have_commands_or_image(self) -> Self: + """When groups is set, each group needs commands or a task-level image.""" + if self.groups is None: + return self + task_has_image = self.image is not None + for group in self.groups: + if not group.commands and not task_has_image: + raise ValueError( + f"Node group '{group.name}': either `commands` must be set in the group, " + "or `image` at the task level." + ) + return self + @property def node_groups(self) -> List[NodeGroup]: if self.groups is not None: @@ -898,7 +957,7 @@ def node_groups(self) -> List[NodeGroup]: return [ NodeGroup( name=DEFAULT_REPLICA_GROUP_NAME, - nodes=self.nodes, + nodes=self.nodes if self.nodes is not None else 1, commands=self.commands, resources=self.resources, ports=self.ports, 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 21d78da495..e3682807c8 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py @@ -124,6 +124,8 @@ from dstack._internal.utils.nodes_interpolator import ( find_groups_ip_refs, interpolate_groups_ip_address, + validate_groups_ref_bounds, + validate_groups_refs, ) logger = get_logger(__name__) @@ -519,6 +521,9 @@ async def _prepare_startup_context( other_job.job_spec.replica_num == context.job.job_spec.replica_num and other_job.job_submissions[-1].status == JobStatus.SUBMITTED ): + # Wait until all jobs in the replica leave SUBMITTED before starting. + # No hard timeout: TERMINATED_BY_SERVER is not retryable and would + # regress multinode retry-on-no-capacity. Follow-up: bound by retry. logger.debug( "%s: waiting for all jobs in the replica to be provisioned", fmt(context.job_model), @@ -616,10 +621,24 @@ async def _prepare_startup_context( return None commands = context.job.job_spec.commands + try: + for c in commands: + validate_groups_refs(c) + 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 + 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): + # Wait for referenced internal_ips. No hard timeout for now + # (same rationale as the replica SUBMITTED wait above). logger.debug( "%s: waiting for referenced node group IPs", fmt(context.job_model), @@ -1809,13 +1828,10 @@ def _build_nodes_ip_view(jobs: list[Job], replica_num: int) -> list[list[str]]: def _referenced_ips_ready(commands: list[str], nodes_view: list[list[str]]) -> 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): - if group_index >= len(nodes_view) or node_index >= len(nodes_view[group_index]): - raise InterpolatorError( - f"Invalid reference groups[{group_index}].nodes[{node_index}].IP_ADDRESS: " - "out of range" - ) # Wait until every referenced slot has a non-empty internal IP. if not nodes_view[group_index][node_index]: return False @@ -1848,7 +1864,7 @@ def _get_cluster_info( gpus_per_job = len(job_runtime_data.offer.instance.resources.gpus) return ClusterInfo( job_ips=job_ips, - master_job_ip=job_ips[0] if job_ips else "", + master_job_ip=job_ips[0], gpus_per_job=gpus_per_job, gpus_per_node=gpus_per_node, ) diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py index e7cf969c2a..990c3638ea 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py @@ -45,11 +45,11 @@ Job, JobProvisioningData, JobRuntimeData, + JobSpec, JobStatus, JobTerminationReason, Requirements, Run, - RunSpec, ) from dstack._internal.core.models.volumes import Volume from dstack._internal.core.services.profiles import get_termination @@ -782,9 +782,9 @@ async def _fetch_run_model_for_submitted_job( siblings for coordination, but still provisions only its own node group. Only a small subset is needed depending on the job type: - * Multinode master (job_num=0): all same-replica jobs. - * First job in a node group (not job 0): job 0 + jobs in its group - (same shape batch). + * Multinode global master (job_num=0) or node-group master: all + same-replica jobs (group masters need siblings to chain-unlock the + next waiting group master). * Other multinode jobs: job 0 + current job. * Single-node master: current job only. @@ -792,28 +792,21 @@ async def _fetch_run_model_for_submitted_job( submissions are never accessed in submitted job processing. """ job_spec = get_job_spec(job_model) - is_master = job_model.job_num == 0 is_multinode = job_spec.jobs_per_replica > 1 job_num_filters: list = [] if not is_multinode: - if is_master: + if job_model.job_num == 0: # Single-node master: only current job needed. job_num_filters.append(JobModel.job_num == 0) else: # Non-master single-node should not happen; keep master + current for safety. job_num_filters.append(JobModel.job_num.in_([0, job_model.job_num])) - elif is_master: - # Multinode master: load all jobs (fleet setup, release waiting_master_job). - # Provisioning still batches only this job's node group (same shape). + elif _is_node_group_master(job_spec): + # Global master or later node-group master: load all jobs (fleet setup, + # chain waiting_master_job release). Provisioning still batches only + # this job's node group (same shape). pass - elif job_spec.node_group_job_index == 0: - # First job in a node group (not job 0): job 0 + this group's jobs. - run_spec = await _get_run_spec(session, job_model.run_id) - group_job_nums = _job_nums_for_node_group( - run_spec.configuration, job_spec.node_group_index - ) - job_num_filters.append(JobModel.job_num.in_(sorted({0, *group_job_nums}))) else: # Other multinode jobs: job 0 + current job. job_num_filters.append(JobModel.job_num.in_([0, job_model.job_num])) @@ -860,21 +853,6 @@ async def _fetch_run_model_for_submitted_job( return res.unique().scalar_one() -async def _get_run_spec(session: AsyncSession, run_id: uuid.UUID) -> RunSpec: - res = await session.execute(select(RunModel.run_spec).where(RunModel.id == run_id)) - return RunSpec.model_validate_json(res.scalar_one()) - - -def _job_nums_for_node_group(configuration, group_index: int) -> list[int]: - assert configuration.type == "task" - job_num = 0 - for index, group in enumerate(configuration.node_groups): - if index == group_index: - return list(range(job_num, job_num + group.nodes)) - job_num += group.nodes - raise ValueError(f"node_group_index {group_index} out of range") - - def _get_job_models_for_jobs( job_models: list[JobModel], jobs: list[Job], @@ -1408,11 +1386,11 @@ async def _apply_existing_instance_provisioning( context.job_model.skip_min_processing_interval = True _release_replica_jobs_from_master_wait( job_model=context.job_model, + job=context.job, replica_job_models=_get_job_models_by_ids( job_models=context.run_model.jobs, job_model_ids=context.replica_job_model_ids, ), - jobs_to_provision=context.jobs_to_provision, ) await _unlock_related_volumes( session=session, @@ -1545,11 +1523,11 @@ async def _apply_new_capacity_provisioning( ) _release_replica_jobs_from_master_wait( job_model=fresh_context.job_model, + job=fresh_context.job, replica_job_models=_get_job_models_by_ids( job_models=fresh_context.run_model.jobs, job_model_ids=fresh_context.replica_job_model_ids, ), - jobs_to_provision=fresh_context.jobs_to_provision, ) await _unlock_related_volumes( session=session, @@ -2178,16 +2156,17 @@ def _hint_pipelines_fetch( pipeline_hinter.hint_fetch(FleetModel.__name__) -def _is_node_group_master(job: Job, replica_jobs: list[Job]) -> bool: - """True if `job` has the lowest job_num among loaded jobs in its node group. +def _is_node_group_master(job_spec: JobSpec) -> bool: + """True if this job_spec is a node-group master for unlock/batching/fetch. - `job` must be in `replica_jobs`. + Prefer `node_group_job_index == 0`, but that field defaults to 0 on pre-branch + job specs, so every legacy job would look like a master. Treat job_num 0 as + master always; only treat a non-zero group index + local index 0 as a later + group master (new heterogeneous specs). """ - group_index = job.job_spec.node_group_index - group_job_nums = [ - j.job_spec.job_num for j in replica_jobs if j.job_spec.node_group_index == group_index - ] - return job.job_spec.job_num == min(group_job_nums) + return job_spec.job_num == 0 or ( + job_spec.node_group_index > 0 and job_spec.node_group_job_index == 0 + ) def _job_needs_provisioning(job: Job) -> bool: @@ -2205,15 +2184,16 @@ def _select_jobs_to_provision(job: Job, replica_jobs: list[Job], job_model: JobM Heterogeneous node groups batch only jobs that share `node_group_index`, so ComputeGroup backends (`run_jobs`) receive a single-shape offer set. - Global `waiting_master_job` is unchanged: non-masters stay blocked until the - global master (job_num=0) finishes its provision attempt. + Global `waiting_master_job` blocks non-masters until the global master + (job_num=0) finishes; chain unlock then lets each group master run before + its workers (see `_release_replica_jobs_from_master_wait`). """ if not is_multinode_job(job): return [job] # Legacy rows without the master-wait protocol: provision one-by-one only. if job_model.waiting_master_job is None: return [job] - if not _is_node_group_master(job, replica_jobs): + if not _is_node_group_master(job.job_spec): return [job] group_index = job.job_spec.node_group_index @@ -2232,20 +2212,64 @@ def _get_required_targeted_instance_offers(context: _SubmittedJobContext) -> int return 1 +def _next_waiting_group_master_job_num( + job: Job, + replica_job_models: list[JobModel], +) -> Optional[int]: + """Lowest job_num of a still-waiting group master in a later node group.""" + my_group = job.job_spec.node_group_index + candidates: list[int] = [] + for m in replica_job_models: + if not m.waiting_master_job: + continue + spec = get_job_spec(m) + if spec.node_group_job_index == 0 and spec.node_group_index > my_group: + candidates.append(m.job_num) + return min(candidates) if candidates else None + + def _release_replica_jobs_from_master_wait( job_model: JobModel, + job: Job, replica_job_models: list[JobModel], - jobs_to_provision: list[Job], ) -> None: - # Global master may only provision its own node group (len == 1). Still release - # waiting workers so other groups can provision on later ticks. - if job_model.job_num != 0: - return + """Chain unlock of `waiting_master_job` after a provision attempt. + + Group masters unlock: + * remaining jobs in their own node group (non-ComputeGroup one-by-one), and + * the lowest still-waiting group master in a later node group (chain). + + Unlocking one later master at a time is deliberate: concurrent group-master + provisioning races on instance_num assignment (no fleet lock). N groups + therefore take N sequential provision rounds after job 0 finishes fleet setup. + + Only called on successful provision paths; a terminating group master leaves + later masters locked until run-level FAILED cleanup. + + Group masters load all replica jobs for this (same cost as the global master). + Non-masters are a no-op. + """ if not any(m.waiting_master_job for m in replica_job_models): return - logger.debug("%s: allow replica jobs to be provisioned one-by-one", fmt(job_model)) + if not _is_node_group_master(job.job_spec): + return + + my_group = job.job_spec.node_group_index + next_master = _next_waiting_group_master_job_num(job, replica_job_models) + logger.debug( + "%s: chain unlock (same-group workers + next waiting group master=%s)", + fmt(job_model), + next_master, + ) for replica_job_model in replica_job_models: - if replica_job_model.waiting_master_job: + if not replica_job_model.waiting_master_job: + continue + if replica_job_model.job_num == job_model.job_num: + continue + spec = get_job_spec(replica_job_model) + if spec.node_group_index == my_group or ( + next_master is not None and replica_job_model.job_num == next_master + ): replica_job_model.waiting_master_job = False @@ -2366,7 +2390,13 @@ async def _provision_new_capacity( known_placement_group_ids.add(placement_group_model.id) offers_tried += 1 try: - if len(jobs) > 1 and offer.backend in BACKENDS_WITH_GROUP_PROVISIONING_SUPPORT: + # Use run_jobs only for a real batch (2+ jobs). A 1-job batch must + # stay on run_job: Slurm gets node_count=1 there, and RunPod's + # Instant Cluster API rejects pod_count=1 (only sizes like 2/4/8). + use_group_provisioning = ( + offer.backend in BACKENDS_WITH_GROUP_PROVISIONING_SUPPORT and len(jobs) > 1 + ) + if use_group_provisioning: assert isinstance(compute, ComputeWithGroupProvisioningSupport) compute_group_provisioning_data = await run_async( compute.run_jobs, diff --git a/src/dstack/_internal/server/services/jobs/configurators/base.py b/src/dstack/_internal/server/services/jobs/configurators/base.py index 9ae8b38b3d..2b5b0c6d26 100644 --- a/src/dstack/_internal/server/services/jobs/configurators/base.py +++ b/src/dstack/_internal/server/services/jobs/configurators/base.py @@ -211,9 +211,7 @@ async def _get_job_spec( node_group_context.group_index if node_group_context is not None else 0 ), node_group_name=( - node_group.name - if node_group is not None and node_group.name is not None - else DEFAULT_REPLICA_GROUP_NAME + node_group.required_name if node_group is not None else DEFAULT_REPLICA_GROUP_NAME ), node_group_job_index=( node_group_context.job_index if node_group_context is not None else 0 diff --git a/src/dstack/_internal/server/services/jobs/configurators/task.py b/src/dstack/_internal/server/services/jobs/configurators/task.py index 03bdf9e8cb..be5eff4d5b 100644 --- a/src/dstack/_internal/server/services/jobs/configurators/task.py +++ b/src/dstack/_internal/server/services/jobs/configurators/task.py @@ -40,10 +40,8 @@ async def get_job_specs(self, replica_num: int) -> List[JobSpec]: return job_specs def _shell_commands(self, node_group: Optional[NodeGroup] = None) -> List[str]: - assert self.run_spec.configuration.type == "task" - if node_group is not None and node_group.commands: - return node_group.commands - return self.run_spec.configuration.commands + assert node_group is not None + return node_group.commands def _default_single_branch(self) -> bool: return True @@ -55,7 +53,5 @@ def _spot_policy(self) -> SpotPolicy: return self.run_spec.merged_profile.spot_policy or SpotPolicy.ONDEMAND def _ports(self, node_group: Optional[NodeGroup] = None) -> List[PortMapping]: - assert self.run_spec.configuration.type == "task" - if node_group is not None and node_group.ports: - return node_group.ports - return self.run_spec.configuration.ports + assert node_group is not None + return node_group.ports diff --git a/src/dstack/_internal/server/services/runs/plan.py b/src/dstack/_internal/server/services/runs/plan.py index 5da7a7fe77..ecfc1a1f59 100644 --- a/src/dstack/_internal/server/services/runs/plan.py +++ b/src/dstack/_internal/server/services/runs/plan.py @@ -1,3 +1,4 @@ +import itertools import math import uuid from collections.abc import Hashable, Mapping @@ -100,8 +101,9 @@ async def get_job_plans( best-fleet-candidate selection and collects offers directly: global offers when no fleets are specified, or offers from the selected fleets when `--fleet` is used. - Services are planned per replica group. Other run types are planned once and then expanded - into per-job `JobPlan` results. + Services are planned per replica group. Tasks are planned per node group so each + group's requirements get their own offers (heterogeneous `groups:`). Other run types + are planned once and then expanded into per-job `JobPlan` results. """ run_name = run_spec.run_name if run_spec.run_name is None: @@ -134,18 +136,13 @@ async def get_job_plans( or run_spec.merged_profile.instances is not None ) - if run_spec.configuration.type == "service": - replica_group_names = [g.name for g in run_spec.configuration.replica_groups] - else: - replica_group_names = [None] + job_batches = await _get_job_batches_for_planning( + run_spec=run_spec, + secrets=secrets, + ) - for replica_group_name in replica_group_names: - jobs = await get_jobs_from_run_spec( - run_spec=run_spec, - secrets=secrets, - replica_num=0, - replica_group_name=replica_group_name, - ) + for jobs in job_batches: + plan_job = jobs[0] if candidate_fleet_models is not None: # Regular job planning fleet_model, instance_offers, backend_offers = await find_optimal_fleet_with_offers( @@ -153,7 +150,7 @@ async def get_job_plans( fleet_models=candidate_fleet_models, run_model=None, run_spec=run_spec, - job=jobs[0], + job=plan_job, master_job_provisioning_data=None, volumes=volumes, exclude_not_available=False, @@ -167,7 +164,7 @@ async def get_job_plans( session=session, project=project, run_spec=run_spec, - job=jobs[0], + job=plan_job, volumes=volumes, ) backend_offers = [] @@ -177,7 +174,7 @@ async def get_job_plans( session=session, project=project, run_spec=run_spec, - job=jobs[0], + job=plan_job, volumes=volumes, skip_backend_offers=skip_backend_offers, full_offers=full_offers, @@ -189,7 +186,7 @@ async def get_job_plans( session=session, project=project, run_spec=run_spec, - job=jobs[0], + job=plan_job, volumes=volumes, skip_backend_offers=skip_backend_offers, full_offers=full_offers, @@ -209,6 +206,43 @@ async def get_job_plans( return job_plans +async def _get_job_batches_for_planning( + run_spec: RunSpec, + secrets: dict[str, str], +) -> list[list[Job]]: + """Split jobs into batches that share the same offer/fleet planning pass. + + Each batch is planned from its first job (group master / only job). Services + use one batch per replica group; tasks use one batch per node group. + """ + if run_spec.configuration.type == "service": + batches: list[list[Job]] = [] + for replica_group_name in [g.name for g in run_spec.configuration.replica_groups]: + jobs = await get_jobs_from_run_spec( + run_spec=run_spec, + secrets=secrets, + replica_num=0, + replica_group_name=replica_group_name, + ) + if jobs: + batches.append(jobs) + return batches + + jobs = await get_jobs_from_run_spec( + run_spec=run_spec, + secrets=secrets, + replica_num=0, + ) + if run_spec.configuration.type != "task" or not jobs: + return [jobs] if jobs else [] + + # Jobs are emitted in node-group order; group consecutive same index. + return [ + list(group_jobs) + for _, group_jobs in itertools.groupby(jobs, key=lambda j: j.job_spec.node_group_index) + ] + + async def get_run_candidate_fleet_models_filters( session: AsyncSession, project: ProjectModel, diff --git a/src/dstack/_internal/server/services/runs/spec.py b/src/dstack/_internal/server/services/runs/spec.py index 508b644ccf..e4ac4a9b3f 100644 --- a/src/dstack/_internal/server/services/runs/spec.py +++ b/src/dstack/_internal/server/services/runs/spec.py @@ -3,6 +3,7 @@ RUN_PRIORITY_DEFAULT, SERVICE_HTTPS_DEFAULT, ServiceConfiguration, + TaskConfiguration, ) from dstack._internal.core.models.profiles import ProfileRetry from dstack._internal.core.models.repos.virtual import DEFAULT_VIRTUAL_REPO_ID, VirtualRunRepoData @@ -18,7 +19,13 @@ set_gpu_vendor_default, set_resources_defaults, ) +from dstack._internal.utils.interpolator import InterpolatorError from dstack._internal.utils.logging import get_logger +from dstack._internal.utils.nodes_interpolator import ( + contains_groups_ref, + validate_groups_ref_bounds, + validate_groups_refs, +) logger = get_logger(__name__) @@ -82,6 +89,7 @@ def validate_run_spec_and_set_defaults( if run_spec.run_name is not None: validate_dstack_resource_name(run_spec.run_name) _validate_retry_duration(run_spec) + _validate_groups_ip_refs(run_spec) for mount_point in run_spec.configuration.volumes: if not is_valid_docker_volume_target(mount_point.path): raise ServerClientError(f"Invalid volume mount path: {mount_point.path}") @@ -125,6 +133,7 @@ def validate_run_spec_and_set_defaults( run_spec.configuration.priority = RUN_PRIORITY_DEFAULT # We do not reject top-level `resources` when `replicas` is a list. 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_resources_defaults(run_spec.configuration.resources) set_gpu_vendor_default( run_spec.configuration.resources, @@ -146,6 +155,39 @@ def _validate_retry_duration(run_spec: RunSpec) -> None: raise ServerClientError("retry.duration cannot be negative") +def _validate_groups_ip_refs(run_spec: RunSpec) -> None: + """Validate groups IP refs at submit time (CLI and API). + + Refs are only supported in commands. Typo'd and out-of-range refs are rejected. + """ + for value in run_spec.configuration.env.values(): + if isinstance(value, str) and contains_groups_ref(value): + raise ServerClientError( + "groups IP references are only supported in commands, not in `env`" + ) + try: + for command in _iter_configuration_commands(run_spec.configuration): + validate_groups_refs(command) + if isinstance(run_spec.configuration, TaskConfiguration): + 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) + except InterpolatorError as e: + raise ServerClientError(e.args[0]) from e + + +def _iter_configuration_commands(configuration: AnyRunConfiguration): + if isinstance(configuration, TaskConfiguration): + for group in configuration.node_groups: + yield from group.commands + return + yield from getattr(configuration, "commands", None) or [] + yield from getattr(configuration, "init", None) or [] + if isinstance(configuration, ServiceConfiguration): + for group in configuration.replica_groups: + yield from group.commands + + def _check_dynamo_in_place_update_compatibility( current_run_spec: RunSpec, new_run_spec: RunSpec ) -> None: diff --git a/src/dstack/_internal/utils/interpolator.py b/src/dstack/_internal/utils/interpolator.py index 641b71dafb..462c85674d 100644 --- a/src/dstack/_internal/utils/interpolator.py +++ b/src/dstack/_internal/utils/interpolator.py @@ -1,6 +1,8 @@ import string from collections.abc import Mapping -from typing import Iterable, List, Literal, Optional, Tuple, Union, overload +from typing import Callable, Iterable, List, Literal, Optional, Tuple, Union, overload + +NameValidator = Callable[[str], bool] class Pattern: @@ -24,11 +26,26 @@ class InterpolatorError(ValueError): pass +def namespace_root(name: str) -> str: + """Return the namespace of a ref name, e.g. groups[0].nodes[0].IP_ADDRESS -> groups.""" + return name.split(".")[0].split("[")[0] + + class VariablesInterpolator: def __init__( - self, namespaces: Mapping[str, Mapping[str, str]], *, skip: Optional[Iterable[str]] = None + self, + namespaces: Mapping[str, Mapping[str, str]], + *, + skip: Optional[Union[Iterable[str], Mapping[str, NameValidator]]] = None, ): - self.skip = set(skip) if skip is not None else set() + # Iterable[str] keeps old callers working and uses validate_name. + # Mapping[str, validator] lets callers plug feature-specific rules. + if skip is None: + self.skip_validators: dict[str, NameValidator] = {} + elif isinstance(skip, Mapping): + self.skip_validators = dict(skip) + else: + self.skip_validators = {ns: self.validate_name for ns in skip} self.variables = {f"{ns}.{k}": v for ns in namespaces for k, v in namespaces[ns].items()} @overload @@ -63,12 +80,12 @@ def interpolate( raise InterpolatorError(f"No pattern closing: {s[opening:]}") name = s[opening + len(Pattern.opening) : closing].strip() - # Skip before validate_name so non-standard refs (e.g. groups[0].nodes[0].IP_ADDRESS) - # can be left for later interpolators. Invalid skipped names without brackets - # (e.g. secrets.pass-word) still raise. - root = name.split(".")[0] - skip_ns = root.split("[")[0] - if skip_ns in self.skip and ("[" in root or self.validate_name(name)): + # Skip before validate_name so deferred refs can be left for later + # interpolators. Invalid skipped names still raise. + skip_ns = namespace_root(name) + if skip_ns in self.skip_validators: + if not self.skip_validators[skip_ns](name): + raise InterpolatorError(f"Illegal reference name: {name}") tokens.append(s[opening : closing + len(Pattern.closing)]) elif not self.validate_name(name): raise InterpolatorError(f"Illegal reference name: {name}") diff --git a/src/dstack/_internal/utils/nodes_interpolator.py b/src/dstack/_internal/utils/nodes_interpolator.py index 84512aa00a..86ec1ca40d 100644 --- a/src/dstack/_internal/utils/nodes_interpolator.py +++ b/src/dstack/_internal/utils/nodes_interpolator.py @@ -1,21 +1,60 @@ import re -from dstack._internal.utils.interpolator import InterpolatorError +from dstack._internal.utils.interpolator import InterpolatorError, namespace_root -_GROUPS_IP_REF = re.compile(r"\$\{\{\s*groups\[(\d+)\]\.nodes\[(\d+)\]\.IP_ADDRESS\s*\}\}") +# Shared grammar for groups[i].nodes[j].IP_ADDRESS refs. +_GROUPS_IP_INNER = r"groups\[(\d+)\]\.nodes\[(\d+)\]\.IP_ADDRESS" +_GROUPS_IP_REF_NAME = re.compile(rf"^{_GROUPS_IP_INNER}$") +# (? bool: + return _GROUPS_IP_REF_NAME.fullmatch(name.strip()) is not None + + +def is_groups_namespace(name: str) -> bool: + return namespace_root(name) == "groups" + + +def validate_groups_refs(s: str) -> None: + """Reject typo'd / unknown groups refs so they never reach the container.""" + for m in _ANY_REF.finditer(s): + name = m.group(1).strip() + if is_groups_namespace(name): + if not is_valid_groups_ip_ref(name): + raise InterpolatorError(f"Illegal reference name: {name}") + + +def contains_groups_ref(s: str) -> bool: + for m in _ANY_REF.finditer(s): + name = m.group(1).strip() + if is_groups_namespace(name): + return True + 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 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]: + raise InterpolatorError( + f"Invalid reference groups[{group_index}].nodes[{node_index}].IP_ADDRESS: " + "out of range" + ) + + def interpolate_groups_ip_address(s: str, nodes: list[list[str]]) -> str: + validate_groups_refs(s) + validate_groups_ref_bounds(s, [len(g) for g in nodes]) + def repl(m: re.Match) -> str: gi, ni = int(m.group(1)), int(m.group(2)) - if gi >= len(nodes) or ni >= len(nodes[gi]): - raise InterpolatorError( - f"Invalid reference groups[{gi}].nodes[{ni}].IP_ADDRESS: out of range" - ) ip = nodes[gi][ni] if not ip: raise InterpolatorError(f"IP not available for groups[{gi}].nodes[{ni}].IP_ADDRESS") diff --git a/src/tests/_internal/core/models/test_configurations.py b/src/tests/_internal/core/models/test_configurations.py index 3eb92a39a9..ee483ecea2 100644 --- a/src/tests/_internal/core/models/test_configurations.py +++ b/src/tests/_internal/core/models/test_configurations.py @@ -11,7 +11,7 @@ ServiceConfiguration, parse_run_configuration, ) -from dstack._internal.core.models.resources import Range +from dstack._internal.core.models.resources import Range, ResourcesSpec from dstack._internal.core.models.routers import ReplicaGroupRouterConfig @@ -948,24 +948,27 @@ def test_parses_groups_and_defaults_names(self): assert parsed.node_groups[0].commands == ["echo head"] assert parsed.node_groups[1].commands == ["echo worker"] - def test_accepts_default_nodes_with_groups(self): - # Serialized TaskConfiguration always includes nodes=1 (the field default). - # The xor validator must allow that so model round-trips succeed. - parsed = parse_run_configuration( - { - "type": "task", - "image": "debian", - "nodes": 1, - "groups": [ - {"nodes": 2, "commands": ["echo head"]}, - {"name": "workers", "nodes": 1, "commands": ["echo worker"]}, - ], - } - ) - assert parsed.type == "task" - assert parsed.nodes == 1 - assert parsed.groups is not None - assert parsed.nodes_num == 3 + def test_rejects_nodes_with_groups(self): + with pytest.raises( + ConfigurationError, match="`nodes` and `groups` are mutually exclusive" + ): + parse_run_configuration( + { + "type": "task", + "image": "debian", + "nodes": 1, + "groups": [ + {"nodes": 2, "commands": ["echo head"]}, + {"name": "workers", "nodes": 1, "commands": ["echo worker"]}, + ], + } + ) + + def test_omitted_nodes_defaults_to_one_node_group(self): + parsed = parse_run_configuration({"type": "task", "commands": ["true"]}) + assert parsed.nodes is None + assert parsed.nodes_num == 1 + assert parsed.node_groups[0].nodes == 1 def test_groups_round_trip_via_model_dump(self): parsed = parse_run_configuration( @@ -980,12 +983,12 @@ def test_groups_round_trip_via_model_dump(self): ) assert parsed.type == "task" dumped = parsed.model_dump(mode="json") - assert dumped["nodes"] == 1 + assert dumped["nodes"] is None assert dumped["groups"] is not None reparsed = parse_run_configuration(dumped) assert reparsed.type == "task" - assert reparsed.nodes == 1 + assert reparsed.nodes is None assert reparsed.nodes_num == 3 assert [g.name for g in reparsed.node_groups] == ["0", "workers"] assert reparsed.node_groups[0].commands == ["echo head"] @@ -1019,6 +1022,16 @@ def test_rejects_duplicate_group_names(self): } ) + def test_rejects_invalid_group_name(self): + with pytest.raises(ConfigurationError, match="Node group name should match regex"): + parse_run_configuration( + { + "type": "task", + "image": "debian", + "groups": [{"name": "Bad_Name", "nodes": 1, "commands": ["true"]}], + } + ) + def test_rejects_empty_groups(self): with pytest.raises(ConfigurationError, match="cannot be an empty list"): parse_run_configuration({"type": "task", "image": "debian", "groups": []}) @@ -1033,3 +1046,60 @@ def test_rejects_nodes_and_groups_together(self): "groups": [{"nodes": 1, "commands": ["true"]}], } ) + + def test_rejects_top_level_commands_with_groups(self): + with pytest.raises(ConfigurationError, match="Top-level `commands` is not allowed"): + parse_run_configuration( + { + "type": "task", + "image": "debian", + "commands": ["echo top"], + "groups": [{"nodes": 1, "commands": ["echo group"]}], + } + ) + + def test_rejects_top_level_ports_with_groups(self): + with pytest.raises(ConfigurationError, match="Top-level `ports` is not allowed"): + parse_run_configuration( + { + "type": "task", + "image": "debian", + "ports": [8000], + "groups": [{"nodes": 1, "commands": ["true"]}], + } + ) + + def test_rejects_top_level_entrypoint_with_groups(self): + with pytest.raises(ConfigurationError, match="Top-level `entrypoint` is not allowed"): + parse_run_configuration( + { + "type": "task", + "image": "debian", + "entrypoint": "python", + "groups": [{"nodes": 1, "commands": ["echo ok"]}], + } + ) + + def test_rejects_group_without_commands_or_image(self): + with pytest.raises(ConfigurationError, match="either `commands` must be set"): + parse_run_configuration( + { + "type": "task", + "groups": [{"name": "head", "nodes": 1}], + } + ) + + def test_accepts_top_level_resources_with_groups(self): + """Top-level resources is allowed (not rejected) but not used for group jobs.""" + parsed = parse_run_configuration( + { + "type": "task", + "image": "debian", + "resources": {"gpu": "H100"}, + "groups": [{"nodes": 1, "commands": ["true"]}], + } + ) + assert parsed.groups is not None + assert parsed.groups[0].resources == ResourcesSpec() + assert parsed.resources.gpu is not None + assert parsed.resources.gpu.name == ["H100"] diff --git a/src/tests/_internal/core/models/test_runs.py b/src/tests/_internal/core/models/test_runs.py index e0bb9fbcee..02e7fff0a6 100644 --- a/src/tests/_internal/core/models/test_runs.py +++ b/src/tests/_internal/core/models/test_runs.py @@ -46,6 +46,17 @@ def test_server_access_run_spec_compatibility(configuration_type: str, dstack: b assert ("dstack" in configuration_excludes) is not dstack +def test_unset_task_nodes_are_excluded_for_compatibility(): + configuration = TaskConfiguration(commands=["true"]) + + configuration_excludes = get_run_spec_excludes(RunSpec(configuration=configuration)).get( + "configuration" + ) + + assert isinstance(configuration_excludes, dict) + assert configuration_excludes["nodes"] is True + + def test_job_termination_reason_to_status_works_with_all_enum_variants(): for job_termination_reason in JobTerminationReason: job_status = job_termination_reason.to_status() diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_node_groups.py b/src/tests/_internal/server/background/pipeline_tasks/test_node_groups.py deleted file mode 100644 index 6366a274bc..0000000000 --- a/src/tests/_internal/server/background/pipeline_tasks/test_node_groups.py +++ /dev/null @@ -1,115 +0,0 @@ -from datetime import datetime, timezone -from uuid import uuid4 - -import pytest - -from dstack._internal.core.models.runs import Job, JobSpec, JobSubmission -from dstack._internal.server.background.pipeline_tasks.jobs_running import ( - _build_nodes_ip_view, - _get_cluster_info, - _referenced_ips_ready, -) -from dstack._internal.server.testing.common import get_job_provisioning_data -from dstack._internal.utils.interpolator import InterpolatorError - - -def _job( - *, - job_num: int, - node_group_index: int, - node_group_job_index: int, - internal_ip: str, - gpu_count: int, -) -> Job: - return Job.model_construct( - job_spec=JobSpec.model_construct( - replica_num=0, - job_num=job_num, - node_group_index=node_group_index, - node_group_job_index=node_group_job_index, - commands=[], - ), - job_submissions=[ - JobSubmission.model_construct( - id=uuid4(), - submitted_at=datetime.now(timezone.utc), - job_provisioning_data=get_job_provisioning_data( - internal_ip=internal_ip, - gpu_count=gpu_count, - ), - job_runtime_data=None, - ) - ], - ) - - -class TestGetClusterInfo: - def test_fills_gpus_per_node(self): - jobs = [ - _job( - job_num=0, - node_group_index=0, - node_group_job_index=0, - internal_ip="10.0.0.1", - gpu_count=8, - ), - _job( - job_num=1, - node_group_index=1, - node_group_job_index=0, - internal_ip="10.0.0.2", - gpu_count=4, - ), - ] - this_jpd = get_job_provisioning_data(internal_ip="10.0.0.1", gpu_count=8) - info = _get_cluster_info( - jobs=jobs, - replica_num=0, - job_provisioning_data=this_jpd, - job_runtime_data=None, - ) - assert info.job_ips == ["10.0.0.1", "10.0.0.2"] - assert info.master_job_ip == "10.0.0.1" - assert info.gpus_per_job == 8 - assert info.gpus_per_node == [8, 4] - - -class TestNodesIpView: - def test_builds_group_view(self): - jobs = [ - _job( - job_num=0, - node_group_index=0, - node_group_job_index=0, - internal_ip="10.0.0.1", - gpu_count=1, - ), - _job( - job_num=1, - node_group_index=0, - node_group_job_index=1, - internal_ip="10.0.0.2", - gpu_count=1, - ), - _job( - job_num=2, - node_group_index=1, - node_group_job_index=0, - internal_ip="10.0.0.3", - gpu_count=1, - ), - ] - assert _build_nodes_ip_view(jobs, replica_num=0) == [ - ["10.0.0.1", "10.0.0.2"], - ["10.0.0.3"], - ] - - def test_referenced_ips_ready(self): - nodes_view = [["10.0.0.1"], [""]] - assert _referenced_ips_ready(["echo ${{ groups[0].nodes[0].IP_ADDRESS }}"], nodes_view) - assert not _referenced_ips_ready(["echo ${{ groups[1].nodes[0].IP_ADDRESS }}"], nodes_view) - - 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) 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 ffb7d18028..491c5a1131 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 @@ -1,6 +1,6 @@ import asyncio import uuid -from contextlib import asynccontextmanager +from contextlib import ExitStack, asynccontextmanager from dataclasses import dataclass from datetime import datetime, timedelta, timezone from pathlib import Path @@ -30,8 +30,11 @@ from dstack._internal.core.models.runs import ( ClusterInfo, ImagePullProgress, + Job, JobRuntimeData, + JobSpec, JobStatus, + JobSubmission, JobTerminationReason, RunSpec, RunStatus, @@ -46,10 +49,13 @@ JobRunningPipeline, JobRunningPipelineItem, JobRunningWorker, + _build_nodes_ip_view, _fetch_run_model, + _get_cluster_info, _prepare_startup_context, _ProcessContext, _ProcessResult, + _referenced_ips_ready, _RunnerAvailability, _SubmitJobToRunnerResult, ) @@ -91,6 +97,7 @@ list_events, ) from dstack._internal.utils.common import get_current_datetime, get_or_error +from dstack._internal.utils.interpolator import InterpolatorError pytestmark = pytest.mark.usefixtures("image_config_mock", "test_log_storage") @@ -2605,6 +2612,125 @@ async def _fake_session_ctx(): assert out.router_env == router_env +@pytest.mark.asyncio +class TestPrepareStartupContextClusterWait: + def _make_context(self) -> _ProcessContext: + job_model = MagicMock() + job_model.submitted_at = datetime(2023, 1, 1, 10, 0, 0, tzinfo=timezone.utc) + sibling = MagicMock() + sibling.job_spec.replica_num = 0 + sibling.job_submissions = [MagicMock(status=JobStatus.SUBMITTED)] + job = MagicMock() + job.job_spec.replica_num = 0 + job.job_spec.commands = ["echo ok"] + run = MagicMock() + run.jobs = [sibling] + run.run_spec = MagicMock() + 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, + ) + + @freeze_time("2023-01-01 12:00:00Z") + async def test_submitted_sibling_defers(self): + context = self._make_context() + result = _ProcessResult() + out = await _prepare_startup_context(context=context, result=result) + assert out is None + assert result.job_update_map == {} + + +@pytest.mark.asyncio +class TestPrepareStartupContextGroupsIpWait: + def _make_context(self) -> _ProcessContext: + job_model = MagicMock() + job_model.submitted_at = datetime(2023, 1, 1, 10, 0, 0, tzinfo=timezone.utc) + peer = MagicMock() + peer.job_spec.replica_num = 0 + peer.job_spec.node_group_index = 0 + peer.job_spec.node_group_job_index = 0 + peer.job_submissions = [ + MagicMock( + status=JobStatus.PROVISIONING, job_provisioning_data=MagicMock(internal_ip="") + ) + ] + job = MagicMock() + job.job_spec.replica_num = 0 + job.job_spec.commands = ["echo ${{ groups[0].nodes[0].IP_ADDRESS }}"] + run = MagicMock() + run.jobs = [peer] + run.run_spec = MagicMock() + 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_groups_ip_not_ready_defers(self): + context = self._make_context() + 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 == {} + + @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) class TestFetchRunModelDynamoBranch: @@ -2680,3 +2806,139 @@ async def test_non_dynamo_loads_only_own_replica(self, test_db, session: AsyncSe run_spec=parsed, ) assert {j.replica_num for j in run_model.jobs} == {0} + + +def _node_group_job( + *, + job_num: int, + node_group_index: int, + node_group_job_index: int, + internal_ip: str, + gpu_count: int, +) -> Job: + return Job.model_construct( + job_spec=JobSpec.model_construct( + replica_num=0, + job_num=job_num, + node_group_index=node_group_index, + node_group_job_index=node_group_job_index, + commands=[], + ), + job_submissions=[ + JobSubmission.model_construct( + id=uuid.uuid4(), + submitted_at=datetime.now(timezone.utc), + job_provisioning_data=get_job_provisioning_data( + internal_ip=internal_ip, + gpu_count=gpu_count, + ), + job_runtime_data=None, + ) + ], + ) + + +class TestGetClusterInfo: + def test_fills_gpus_per_node(self): + jobs = [ + _node_group_job( + job_num=0, + node_group_index=0, + node_group_job_index=0, + internal_ip="10.0.0.1", + gpu_count=8, + ), + _node_group_job( + job_num=1, + node_group_index=1, + node_group_job_index=0, + internal_ip="10.0.0.2", + gpu_count=4, + ), + ] + this_jpd = get_job_provisioning_data(internal_ip="10.0.0.1", gpu_count=8) + info = _get_cluster_info( + jobs=jobs, + replica_num=0, + job_provisioning_data=this_jpd, + job_runtime_data=None, + ) + assert info.job_ips == ["10.0.0.1", "10.0.0.2"] + assert info.master_job_ip == "10.0.0.1" + assert info.gpus_per_job == 8 + assert info.gpus_per_node == [8, 4] + + def test_raises_when_sibling_missing_provisioning_data(self): + provisioned = _node_group_job( + job_num=0, + node_group_index=0, + node_group_job_index=0, + internal_ip="10.0.0.1", + gpu_count=1, + ) + unprovisioned = Job.model_construct( + job_spec=JobSpec.model_construct( + replica_num=0, + job_num=1, + node_group_index=1, + node_group_job_index=0, + commands=[], + ), + job_submissions=[ + JobSubmission.model_construct( + id=uuid.uuid4(), + submitted_at=datetime.now(timezone.utc), + job_provisioning_data=None, + job_runtime_data=None, + ) + ], + ) + this_jpd = get_job_provisioning_data(internal_ip="10.0.0.1", gpu_count=1) + with pytest.raises(ValueError, match="Optional value is None"): + _get_cluster_info( + jobs=[provisioned, unprovisioned], + replica_num=0, + job_provisioning_data=this_jpd, + job_runtime_data=None, + ) + + +class TestNodesIpView: + def test_builds_group_view(self): + jobs = [ + _node_group_job( + job_num=0, + node_group_index=0, + node_group_job_index=0, + internal_ip="10.0.0.1", + gpu_count=1, + ), + _node_group_job( + job_num=1, + node_group_index=0, + node_group_job_index=1, + internal_ip="10.0.0.2", + gpu_count=1, + ), + _node_group_job( + job_num=2, + node_group_index=1, + node_group_job_index=0, + internal_ip="10.0.0.3", + gpu_count=1, + ), + ] + assert _build_nodes_ip_view(jobs, replica_num=0) == [ + ["10.0.0.1", "10.0.0.2"], + ["10.0.0.3"], + ] + + def test_referenced_ips_ready(self): + nodes_view = [["10.0.0.1"], [""]] + assert _referenced_ips_ready(["echo ${{ groups[0].nodes[0].IP_ADDRESS }}"], nodes_view) + assert not _referenced_ips_ready(["echo ${{ groups[1].nodes[0].IP_ADDRESS }}"], nodes_view) + + 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) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py index 491868327f..09dce3cf46 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py @@ -49,6 +49,7 @@ JobSubmittedPipelineItem, JobSubmittedWorker, _load_submitted_job_context, + _release_replica_jobs_from_master_wait, ) from dstack._internal.server.models import ( ComputeGroupModel, @@ -2028,6 +2029,68 @@ async def test_provisions_compute_group( res = await session.execute(select(ComputeGroupModel)) assert res.scalar_one_or_none() is not None + async def test_provisions_one_job_node_group_via_run_job( + self, test_db, session: AsyncSession, worker: JobSubmittedWorker + ): + """1-node hetero groups use run_job (not run_jobs). + + RunPod Instant Clusters reject pod_count=1; Slurm gets node_count=1 from run_job. + """ + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + fleet = await create_fleet(session=session, project=project) + configuration = TaskConfiguration( + image="debian", + groups=[ + NodeGroup(name="prefill", nodes=1, commands=["echo prefill"]), + NodeGroup(name="decode", nodes=1, commands=["echo decode"]), + ], + ) + run_spec = get_run_spec(repo_id=repo.name, configuration=configuration) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + fleet=fleet, + run_spec=run_spec, + ) + job0 = await create_job( + session=session, + run=run, + instance_assigned=True, + job_num=0, + waiting_master_job=False, + ) + await create_job( + session=session, + run=run, + instance_assigned=False, + job_num=1, + waiting_master_job=True, + ) + + offer = get_instance_offer_with_availability(backend=BackendType.RUNPOD) + with patch("dstack._internal.server.services.backends.get_project_backends") as m: + backend_mock = Mock() + compute_mock = Mock(spec=ComputeMockSpec) + backend_mock.compute.return_value = compute_mock + m.return_value = [backend_mock] + backend_mock.TYPE = BackendType.RUNPOD + compute_mock.get_offers.return_value = [offer] + compute_mock.run_job.return_value = get_job_provisioning_data( + dockerized=True, backend=BackendType.RUNPOD + ) + + await _process_job(session=session, worker=worker, job_model=job0) + + compute_mock.run_job.assert_called_once() + compute_mock.run_jobs.assert_not_called() + assert compute_mock.run_job.call_args[0][1].job_spec.job_num == 0 + job0 = await _get_job(session, job0.id) + assert job0.status == JobStatus.PROVISIONING + async def test_defers_job_while_waiting_for_master_provisioning( self, test_db, session: AsyncSession, worker: JobSubmittedWorker ): @@ -2855,7 +2918,7 @@ async def test_node_group_master_provisions_only_its_group( assert small_context.jobs_to_provision[0].job_spec.node_group_name == "small" large_context = await _load_submitted_job_context(session=session, job_model=large_job_0) - # Heterogeneous group master: job 0 + its node group. + # Node-group masters load all replica jobs (for chain unlock). assert {jm.job_num for jm in large_context.run_model.jobs} == {0, 1, 2} assert sorted(j.job_spec.job_num for j in large_context.jobs_to_provision) == [1, 2] assert {j.job_spec.node_group_name for j in large_context.jobs_to_provision} == {"large"} @@ -2867,6 +2930,202 @@ async def test_node_group_master_provisions_only_its_group( assert {jm.job_num for jm in large_worker_context.run_model.jobs} == {0, 2} assert [j.job_spec.job_num for j in large_worker_context.jobs_to_provision] == [2] + async def test_global_master_unlocks_only_next_waiting_group_master( + self, test_db, session: AsyncSession + ): + """After job 0, unlock only the next waiting group master — not later masters or their workers.""" + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + fleet = await create_fleet(session=session, project=project) + configuration = TaskConfiguration( + image="debian", + groups=[ + NodeGroup(name="small", nodes=1, commands=["echo small"]), + NodeGroup(name="large", nodes=2, commands=["echo large"]), + NodeGroup(name="other", nodes=1, commands=["echo other"]), + ], + ) + run_spec = get_run_spec(run_name="run", repo_id=repo.name, configuration=configuration) + run = await create_run( + session=session, + run_name="run", + project=project, + repo=repo, + user=user, + run_spec=run_spec, + fleet=fleet, + ) + small_job = await create_job( + session=session, + run=run, + job_num=0, + status=JobStatus.SUBMITTED, + waiting_master_job=False, + ) + large_master = await create_job( + session=session, + run=run, + job_num=1, + status=JobStatus.SUBMITTED, + waiting_master_job=True, + ) + large_worker = await create_job( + session=session, + run=run, + job_num=2, + status=JobStatus.SUBMITTED, + waiting_master_job=True, + ) + other_master = await create_job( + session=session, + run=run, + job_num=3, + status=JobStatus.SUBMITTED, + waiting_master_job=True, + ) + await session.commit() + + context = await _load_submitted_job_context(session=session, job_model=small_job) + _release_replica_jobs_from_master_wait( + job_model=context.job_model, + job=context.job, + replica_job_models=list(context.run_model.jobs), + ) + await session.commit() + await session.refresh(large_master) + await session.refresh(large_worker) + await session.refresh(other_master) + + assert large_master.waiting_master_job is False + assert large_worker.waiting_master_job is True + assert other_master.waiting_master_job is True + + async def test_homogeneous_master_unlocks_same_group_workers( + self, test_db, session: AsyncSession + ): + """Homogeneous nodes:N — job 0 unlocks the other ranks in the single group.""" + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + fleet = await create_fleet(session=session, project=project) + configuration = TaskConfiguration(image="debian", nodes=3, commands=["true"]) + run_spec = get_run_spec(run_name="run", repo_id=repo.name, configuration=configuration) + run = await create_run( + session=session, + run_name="run", + project=project, + repo=repo, + user=user, + run_spec=run_spec, + fleet=fleet, + ) + master = await create_job( + session=session, + run=run, + job_num=0, + status=JobStatus.SUBMITTED, + waiting_master_job=False, + ) + worker_1 = await create_job( + session=session, + run=run, + job_num=1, + status=JobStatus.SUBMITTED, + waiting_master_job=True, + ) + worker_2 = await create_job( + session=session, + run=run, + job_num=2, + status=JobStatus.SUBMITTED, + waiting_master_job=True, + ) + await session.commit() + + context = await _load_submitted_job_context(session=session, job_model=master) + _release_replica_jobs_from_master_wait( + job_model=context.job_model, + job=context.job, + replica_job_models=list(context.run_model.jobs), + ) + await session.commit() + await session.refresh(worker_1) + await session.refresh(worker_2) + + assert worker_1.waiting_master_job is False + assert worker_2.waiting_master_job is False + + async def test_group_master_unlocks_same_group_workers_and_next_master( + self, test_db, session: AsyncSession + ): + """After a non-zero group master: unlock its workers and the next waiting group master.""" + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + fleet = await create_fleet(session=session, project=project) + configuration = TaskConfiguration( + image="debian", + groups=[ + NodeGroup(name="small", nodes=1, commands=["echo small"]), + NodeGroup(name="large", nodes=2, commands=["echo large"]), + NodeGroup(name="other", nodes=1, commands=["echo other"]), + ], + ) + run_spec = get_run_spec(run_name="run", repo_id=repo.name, configuration=configuration) + run = await create_run( + session=session, + run_name="run", + project=project, + repo=repo, + user=user, + run_spec=run_spec, + fleet=fleet, + ) + await create_job( + session=session, + run=run, + job_num=0, + status=JobStatus.SUBMITTED, + waiting_master_job=False, + ) + large_master = await create_job( + session=session, + run=run, + job_num=1, + status=JobStatus.SUBMITTED, + waiting_master_job=False, + ) + large_worker = await create_job( + session=session, + run=run, + job_num=2, + status=JobStatus.SUBMITTED, + waiting_master_job=True, + ) + other_master = await create_job( + session=session, + run=run, + job_num=3, + status=JobStatus.SUBMITTED, + waiting_master_job=True, + ) + await session.commit() + + context = await _load_submitted_job_context(session=session, job_model=large_master) + assert {jm.job_num for jm in context.run_model.jobs} == {0, 1, 2, 3} + _release_replica_jobs_from_master_wait( + job_model=context.job_model, + job=context.job, + replica_job_models=list(context.run_model.jobs), + ) + await session.commit() + await session.refresh(large_worker) + await session.refresh(other_master) + + assert large_worker.waiting_master_job is False + assert other_master.waiting_master_job is False + async def test_loads_only_latest_submission(self, test_db, session: AsyncSession): """Only the latest submission per (replica_num, job_num) should be loaded, not historical ones.""" project = await create_project(session=session) diff --git a/src/tests/_internal/server/services/jobs/configurators/test_task.py b/src/tests/_internal/server/services/jobs/configurators/test_task.py index 383cd411e2..6abb17e6da 100644 --- a/src/tests/_internal/server/services/jobs/configurators/test_task.py +++ b/src/tests/_internal/server/services/jobs/configurators/test_task.py @@ -96,6 +96,22 @@ async def test_uses_per_group_commands_and_resources(self): assert job_specs[1].requirements.resources.gpu.name == ["A100"] assert job_specs[1].requirements.resources.gpu.count.min == 2 + async def test_group_without_resources_does_not_inherit_top_level(self): + """Same as replica groups: omitted group resources → ResourcesSpec(), not top-level.""" + configuration = TaskConfiguration( + image="debian", + resources=ResourcesSpec(gpu=GPUSpec(name=["H100"], count=1)), + groups=[ + NodeGroup(name="head", nodes=1, commands=["echo head"]), + ], + ) + run_spec = get_run_spec(run_name="run", repo_id="id", configuration=configuration) + configurator = TaskJobConfigurator(run_spec) + + job_specs = await configurator.get_job_specs(replica_num=0) + + assert job_specs[0].requirements.resources.gpu.name is None + @pytest.mark.asyncio @pytest.mark.usefixtures("image_config_mock") diff --git a/src/tests/_internal/server/services/runs/test_plan.py b/src/tests/_internal/server/services/runs/test_plan.py index df457d33f9..c22cf43912 100644 --- a/src/tests/_internal/server/services/runs/test_plan.py +++ b/src/tests/_internal/server/services/runs/test_plan.py @@ -8,6 +8,9 @@ from dstack._internal.core.models.common import EntityReference from dstack._internal.core.models.configurations import ( DevEnvironmentConfiguration, + NodeGroup, + ReplicaGroup, + ServiceConfiguration, TaskConfiguration, ) from dstack._internal.core.models.fleets import FleetNodesSpec, InstanceGroupPlacement @@ -19,7 +22,7 @@ InstanceNameSelector, Profile, ) -from dstack._internal.core.models.resources import CPUSpec, Memory, Range, ResourcesSpec +from dstack._internal.core.models.resources import CPUSpec, GPUSpec, Memory, Range, ResourcesSpec from dstack._internal.server.services.jobs import get_jobs_from_run_spec from dstack._internal.server.services.projects import get_project_model_by_name from dstack._internal.server.services.runs import get_plan @@ -184,6 +187,158 @@ async def test_excludes_backend_offers_when_instances_specified( assert job_plans[0].offers == [instance_offer] +class TestGetJobPlansNodeGroups: + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_plans_each_node_group_with_its_own_requirements( + self, + test_db, + session: AsyncSession, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + repo = await create_repo(session=session, project_id=project.id) + run_spec = get_run_spec( + repo_id=repo.name, + configuration=TaskConfiguration( + image="debian", + groups=[ + NodeGroup( + name="router", + nodes=1, + commands=["echo router"], + resources=ResourcesSpec(cpu=CPUSpec(count=Range[int](min=4, max=4))), + ), + NodeGroup( + name="prefill", + nodes=2, + commands=["echo prefill"], + resources=ResourcesSpec( + gpu=GPUSpec(name=["L40S"], count=1), + ), + ), + ], + ), + ) + cpu_offer = get_instance_offer_with_availability(price=1.0) + gpu_offer = get_instance_offer_with_availability(price=2.0, gpu_count=1, gpu_name="L40S") + + async def find_optimal_fleet_with_offers_side_effect(*, job, **kwargs): + gpu = job.job_spec.requirements.resources.gpu + if gpu is not None and gpu.name: + return Mock(), [], [(Mock(), gpu_offer)] + return Mock(), [], [(Mock(), cpu_offer)] + + monkeypatch.setattr( + "dstack._internal.server.services.runs.plan._select_candidate_fleet_models", + AsyncMock(return_value=[Mock()]), + ) + find_optimal_mock = AsyncMock(side_effect=find_optimal_fleet_with_offers_side_effect) + monkeypatch.setattr( + "dstack._internal.server.services.runs.plan.find_optimal_fleet_with_offers", + find_optimal_mock, + ) + + job_plans = await get_job_plans( + session=session, + project=project, + run_spec=run_spec, + max_offers=None, + full_offers=False, + unallocated_resources=False, + ) + + assert find_optimal_mock.await_count == 2 + planned_jobs = [call.kwargs["job"] for call in find_optimal_mock.await_args_list] + assert [j.job_spec.node_group_name for j in planned_jobs] == ["router", "prefill"] + assert planned_jobs[0].job_spec.requirements.resources.gpu.name is None + assert planned_jobs[1].job_spec.requirements.resources.gpu.name == ["L40S"] + + assert len(job_plans) == 3 + assert [p.job_spec.node_group_name for p in job_plans] == [ + "router", + "prefill", + "prefill", + ] + assert job_plans[0].offers == [cpu_offer] + assert job_plans[1].offers == [gpu_offer] + assert job_plans[2].offers == [gpu_offer] + + +class TestGetJobPlansReplicaGroups: + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_plans_each_replica_group_with_its_own_requirements( + self, + test_db, + session: AsyncSession, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + repo = await create_repo(session=session, project_id=project.id) + run_spec = get_run_spec( + repo_id=repo.name, + configuration=ServiceConfiguration( + port=8080, + gateway=False, + replicas=[ + ReplicaGroup( + name="gpu-group", + count=Range[int](min=1, max=1), + resources=ResourcesSpec(gpu=GPUSpec(name=["L40S"], count=1)), + commands=["python server.py"], + ), + ReplicaGroup( + name="cpu-group", + count=Range[int](min=1, max=1), + resources=ResourcesSpec(cpu=CPUSpec(count=Range[int](min=4, max=4))), + commands=["python router.py"], + ), + ], + ), + ) + gpu_offer = get_instance_offer_with_availability(price=2.0, gpu_count=1, gpu_name="L40S") + cpu_offer = get_instance_offer_with_availability(price=1.0) + + async def find_optimal_fleet_with_offers_side_effect(*, job, **kwargs): + gpu = job.job_spec.requirements.resources.gpu + if gpu is not None and gpu.name: + return Mock(), [], [(Mock(), gpu_offer)] + return Mock(), [], [(Mock(), cpu_offer)] + + monkeypatch.setattr( + "dstack._internal.server.services.runs.plan._select_candidate_fleet_models", + AsyncMock(return_value=[Mock()]), + ) + find_optimal_mock = AsyncMock(side_effect=find_optimal_fleet_with_offers_side_effect) + monkeypatch.setattr( + "dstack._internal.server.services.runs.plan.find_optimal_fleet_with_offers", + find_optimal_mock, + ) + + job_plans = await get_job_plans( + session=session, + project=project, + run_spec=run_spec, + max_offers=None, + full_offers=False, + unallocated_resources=False, + ) + + assert find_optimal_mock.await_count == 2 + planned_jobs = [call.kwargs["job"] for call in find_optimal_mock.await_args_list] + assert [j.job_spec.replica_group for j in planned_jobs] == ["gpu-group", "cpu-group"] + assert planned_jobs[0].job_spec.requirements.resources.gpu.name == ["L40S"] + assert planned_jobs[1].job_spec.requirements.resources.gpu.name is None + + assert len(job_plans) == 2 + assert [p.job_spec.replica_group for p in job_plans] == ["gpu-group", "cpu-group"] + assert job_plans[0].offers == [gpu_offer] + assert job_plans[1].offers == [cpu_offer] + + class TestGetPlan: @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) diff --git a/src/tests/_internal/server/services/runs/test_spec.py b/src/tests/_internal/server/services/runs/test_spec.py index 0c62ad7219..3bdc1ee76d 100644 --- a/src/tests/_internal/server/services/runs/test_spec.py +++ b/src/tests/_internal/server/services/runs/test_spec.py @@ -5,7 +5,11 @@ import pytest from dstack._internal.core.errors import ServerClientError -from dstack._internal.core.models.configurations import ServiceConfiguration +from dstack._internal.core.models.configurations import ( + NodeGroup, + ServiceConfiguration, + TaskConfiguration, +) from dstack._internal.core.models.files import FileArchiveMapping from dstack._internal.core.models.profiles import Profile, ProfileRetry from dstack._internal.core.models.repos.local import LocalRunRepoData @@ -98,6 +102,102 @@ def test_rejects_negative_retry_duration_for_new_run_specs(self): ) +class TestValidateRunSpecGroupsIpRefs: + def test_rejects_typo_groups_ref_in_node_group_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].nodes[0].IP }}"], + ), + ], + ), + ) + + 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_valid_groups_ref(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 }}"], + ), + ], + ), + ) + + validate_run_spec_and_set_defaults( + SimpleNamespace(ssh_public_key="ssh-rsa test"), run_spec + ) + + def test_rejects_groups_ref_in_env(self): + run_spec = get_run_spec( + repo_id="test-repo", + configuration=TaskConfiguration( + image="debian", + commands=["echo ok"], + env={"PREFILL_URL": "http://${{ groups[1].nodes[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 + ) + + def test_rejects_out_of_range_group_index(self): + run_spec = get_run_spec( + repo_id="test-repo", + configuration=TaskConfiguration( + image="debian", + groups=[ + NodeGroup( + name="head", + nodes=1, + commands=["echo ${{ groups[1].nodes[0].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_out_of_range_node_index(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[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 + ) + + class TestCheckCanUpdateConfigurationRouterType: def test_sglang_to_dynamo_router_type_change_is_rejected(self): current = _run_spec(_service_configuration(router_type="sglang")) diff --git a/src/tests/_internal/utils/test_interpolator.py b/src/tests/_internal/utils/test_interpolator.py index 2acc2eaefc..3b545ecd08 100644 --- a/src/tests/_internal/utils/test_interpolator.py +++ b/src/tests/_internal/utils/test_interpolator.py @@ -1,6 +1,7 @@ import pytest from dstack._internal.utils.interpolator import InterpolatorError, VariablesInterpolator +from dstack._internal.utils.nodes_interpolator import is_valid_groups_ip_ref def get_interpolator(): @@ -52,5 +53,20 @@ def test_illegal_name(self): def test_skips_groups_refs(self): s = "ray start --address=${{ groups[0].nodes[0].IP_ADDRESS }}:6379" - interpolator = VariablesInterpolator({"run": {"args": "x"}}, skip=["groups"]) + interpolator = VariablesInterpolator( + {"run": {"args": "x"}}, + skip={"groups": is_valid_groups_ip_ref}, + ) assert interpolator.interpolate(s) == s + + def test_rejects_invalid_groups_refs(self): + interpolator = VariablesInterpolator( + {"run": {"args": "x"}}, + skip={"groups": is_valid_groups_ip_ref}, + ) + with pytest.raises(InterpolatorError, match="Illegal reference name"): + interpolator.interpolate("${{ groups[0].nodes[0].IP }}") + with pytest.raises(InterpolatorError, match="Illegal reference name"): + interpolator.interpolate("${{ groups[0].node[0].IP_ADDRESS }}") + with pytest.raises(InterpolatorError, match="Illegal reference name"): + interpolator.interpolate("${{ groups.prefill.nodes[0].IP_ADDRESS }}") diff --git a/src/tests/_internal/utils/test_nodes_interpolator.py b/src/tests/_internal/utils/test_nodes_interpolator.py index 84ed9cdc85..e4adfa2ef3 100644 --- a/src/tests/_internal/utils/test_nodes_interpolator.py +++ b/src/tests/_internal/utils/test_nodes_interpolator.py @@ -2,8 +2,11 @@ from dstack._internal.utils.interpolator import InterpolatorError from dstack._internal.utils.nodes_interpolator import ( + contains_groups_ref, find_groups_ip_refs, interpolate_groups_ip_address, + validate_groups_ref_bounds, + validate_groups_refs, ) @@ -19,6 +22,53 @@ def test_finds_multiple_refs(self): 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 }}") == [] + + +class TestValidateGroupsRefBounds: + def test_accepts_in_range(self): + validate_groups_ref_bounds("${{ groups[0].nodes[0].IP_ADDRESS }}", [1, 2]) + + def test_rejects_group_out_of_range(self): + with pytest.raises(InterpolatorError, match="out of range"): + validate_groups_ref_bounds("${{ groups[2].nodes[0].IP_ADDRESS }}", [1, 2]) + + 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]) + + +class TestValidateGroupsRefs: + def test_accepts_valid_ref(self): + validate_groups_refs("ray start --address=${{ groups[0].nodes[0].IP_ADDRESS }}:6379") + + def test_rejects_typo_field(self): + with pytest.raises(InterpolatorError, match="Illegal reference name"): + validate_groups_refs("${{ groups[0].nodes[0].IP }}") + + def test_rejects_typo_path(self): + with pytest.raises(InterpolatorError, match="Illegal reference name"): + validate_groups_refs("${{ groups[0].node[0].IP_ADDRESS }}") + + def test_rejects_named_group_ref(self): + with pytest.raises(InterpolatorError, match="Illegal reference name"): + validate_groups_refs("${{ groups.prefill.nodes[0].IP_ADDRESS }}") + + def test_ignores_escaped_refs(self): + validate_groups_refs("echo $${{ groups[0].bad }}") + + +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].nodes[0].IP }}") + assert not contains_groups_ref("${{ secrets.token }}") + assert not contains_groups_ref("${{ groups_config.x }}") + + def test_ignores_escaped_refs(self): + assert not contains_groups_ref("echo $${{ groups[0].nodes[0].IP_ADDRESS }}") + class TestInterpolateGroupsIpAddress: def test_replaces_ip(self): @@ -38,3 +88,7 @@ def test_raises_when_ip_missing(self): def test_raises_when_out_of_range(self): with pytest.raises(InterpolatorError, match="out of range"): interpolate_groups_ip_address("${{ groups[1].nodes[0].IP_ADDRESS }}", [["10.0.0.1"]]) + + 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"]]) From 2541c36e9f72975f16f948bd0d501f16fdeb3504 Mon Sep 17 00:00:00 2001 From: Bihan Rana Date: Mon, 10 Aug 2026 12:12:54 +0545 Subject: [PATCH 3/4] Fix ResourcesSpec cpu type in hetero node group test --- .../server/background/pipeline_tasks/test_submitted_jobs.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py index 3313ac25b1..312d379ffd 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py @@ -2900,7 +2900,8 @@ async def test_node_group_master_provisions_only_its_group( name="small", nodes=1, resources=ResourcesSpec( - cpu=Range[int](min=2), memory=Range[Memory](min=Memory(4)) + cpu=CPUSpec(count=Range[int](min=2)), + memory=Range[Memory](min=Memory(4)), ), commands=["echo small"], ), @@ -2908,7 +2909,8 @@ async def test_node_group_master_provisions_only_its_group( name="large", nodes=2, resources=ResourcesSpec( - cpu=Range[int](min=4), memory=Range[Memory](min=Memory(8)) + cpu=CPUSpec(count=Range[int](min=4)), + memory=Range[Memory](min=Memory(8)), ), commands=["echo large"], ), From 83d164852e6a4f98f967d095f1fb2a3ec415ed14 Mon Sep 17 00:00:00 2001 From: Bihan Rana Date: Mon, 10 Aug 2026 12:29:27 +0545 Subject: [PATCH 4/4] Pass for_offers_only in get_job_plans tests --- src/tests/_internal/server/services/runs/test_plan.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/tests/_internal/server/services/runs/test_plan.py b/src/tests/_internal/server/services/runs/test_plan.py index 6aacf3af9b..e55ff6de67 100644 --- a/src/tests/_internal/server/services/runs/test_plan.py +++ b/src/tests/_internal/server/services/runs/test_plan.py @@ -249,6 +249,7 @@ async def find_optimal_fleet_with_offers_side_effect(*, job, **kwargs): max_offers=None, full_offers=False, unallocated_resources=False, + for_offers_only=False, ) assert find_optimal_mock.await_count == 2 @@ -327,6 +328,7 @@ async def find_optimal_fleet_with_offers_side_effect(*, job, **kwargs): max_offers=None, full_offers=False, unallocated_resources=False, + for_offers_only=False, ) assert find_optimal_mock.await_count == 2