Support Heterogenous Node Groups

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

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 {

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

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)

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)
 

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}

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."""

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