Support Heterogenous Node Groups (#4094)

* Support Heterogenous Node Groups

* Review Comments Resolved

Co-authored-by: Cursor <cursoragent@cursor.com>

* Fix ResourcesSpec cpu type in hetero node group test

* Pass for_offers_only in get_job_plans tests

* Minor Updates

* Apply resource defaults and image validation to node groups

* Add docs for heterogenous node groups

* Fix hetero tests for CPU arch defaults and replica provisioning gate

* Remove node groups docs from PR

---------

Co-authored-by: Bihan  Rana

Author: Bihan

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,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, 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)
 	}
 
@@ -759,7 +775,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, slots []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 +791,21 @@ 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(slots) != len(ips) {
+			return fmt.Errorf(
+				"gpus_per_node length %d != job_ips length %d",
+				len(slots), len(ips),
+			)
 		}
-		for _, ip := range nonEmptyIps {
-			if _, err = fmt.Fprintf(file, template, ip); err != nil {
+		for i, ip := range nonEmptyIps {
+			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, slots[i])
+			}
+			if err != nil {
 				return fmt.Errorf("write MPI hostfile line: %w", err)
 			}
 		}

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)

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

@@ -73,6 +73,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
 
@@ -585,20 +586,28 @@ 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": VariablesInterpolator.validate_name,
+                "groups": is_valid_groups_ip_ref,
+            },
+        )
         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
@@ -615,6 +624,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

@@ -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]
 
@@ -164,6 +167,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 +190,7 @@ def _run_slurm_job(
         instance_offer: InstanceOfferWithAvailability,
         project_ssh_public_key: str,
         requirements: Requirements,
+        node_count: int,
     ) -> ComputeGroupProvisioningData:
         if job.job_spec.registry_auth is not None:
             self._skip_offer_cache.add(run, job, instance_offer)
@@ -209,7 +214,7 @@ 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
+            # Slurm --nodes for this call (1 from run_job, len(batch) from run_jobs).
             resources_spec = requirements.resources
             requested_resources = get_requested_resources_from_resources_spec(resources_spec)
 

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

@@ -5,7 +5,9 @@
     IncludeExcludeDictType,
     IncludeExcludeSetType,
 )
+from dstack._internal.core.models.configurations import TaskConfiguration
 from dstack._internal.core.models.runs import (
+    DEFAULT_REPLICA_GROUP_NAME,
     ApplyRunPlanInput,
     JobSpec,
     JobSubmission,
@@ -80,6 +82,14 @@ def get_run_spec_excludes(run_spec: RunSpec) -> IncludeExcludeDictType:
     profile_excludes = get_profile_excludes(run_spec.profile)
     for field in get_profile_excludes(run_spec.configuration):
         configuration_excludes[field] = True
+
+    if isinstance(run_spec.configuration, TaskConfiguration):
+        if run_spec.configuration.groups is None:
+            configuration_excludes["groups"] = True
+        if run_spec.configuration.nodes is None:
+            # Omit nodes when unset so old servers never see null (pre-hetero nodes was int=1).
+            configuration_excludes["nodes"] = True
+
     if configuration_excludes:
         spec_excludes["configuration"] = configuration_excludes
     if profile_excludes:
@@ -94,6 +104,12 @@ def get_job_spec_excludes(job_specs: list[JobSpec]) -> IncludeExcludeDictType:
     clients backward-compatibility with older servers.
     """
     spec_excludes: IncludeExcludeDictType = {}
+    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
     return spec_excludes