Support Heterogenous Node Groups

[Bihan]

• Add heterogeneous node groups for multinode tasks (groups:), mutually exclusive with homogeneous nodes:, so one run can mix roles/resources (e.g. CPU router + GPU prefill/decode for PD-disaggregation).

• Expand groups into a single global rank space with per-job node_group_* metadata; provision each group as its own same-shape batch; wire cross-group addressing via ${{ groups[i].nodes[j].IP_ADDRESS }} after IPs are ready (workers-first supported).

• Runner: accept ClusterInfo.gpus_per_node, sum it for $DSTACK_GPUS_NUM, and write the MPI hostfile with per-IP slots= (CPU nodes omit slots); fall back to homogeneous gpus_per_job when the list is empty (older servers).

Design doc: Heterogeneous Node Groups

Sample config (PD-disaggregation)

type: task
name: pd-disagg
image: lmsysorg/sglang:latest
env:
  - HF_TOKEN
  - MODEL_ID=meta-llama/Llama-3.2-3B-Instruct

startup_order: workers-first
groups:
# ---- Router (CPU) — wires prefill + decode workers by name. This is master ----
  - name: router
    nodes: 1

    commands:
      - pip install smg
      - |
        echo "prefill=${{ groups[1].nodes[0].IP_ADDRESS }}" 
        echo "decode=${{ groups[2].nodes[0].IP_ADDRESS }}" 
        smg launch\
          --pd-disaggregation \
          --prefill http://${{ groups[1].nodes[0].IP_ADDRESS }}:8000 8998 \
          --decode  http://${{ groups[2].nodes[0].IP_ADDRESS }}:8000 \
          --prefill-policy cache_aware \
          --host 0.0.0.0 --port 8000
    ports:
      - 8000
    resources:
      cpu: 4

  # ---- Prefill workers ----
  - name: prefill
    nodes: 1
    commands:
      - |
        python -m sglang.launch_server \
          --model-path $MODEL_ID \
          --disaggregation-mode prefill \
          --disaggregation-transfer-backend nixl \
          --host 0.0.0.0 --port 8000 \
          --disaggregation-bootstrap-port 8998
    resources:
      gpu: L40S
 
  # ---- Decode workers ----
  - name: decode
    nodes: 1
    commands:
      - |
        python -m sglang.launch_server \
          --model-path $MODEL_ID \
          --disaggregation-mode decode \
          --disaggregation-transfer-backend nixl \
          --host 0.0.0.0 --port 8000
    resources:
      gpu: L40S

SSH-Fleet

type: fleet
name: pd-disagg

placement: cluster

ssh_config:
  user: dstack
  identity_file: ~/.ssh/id_rsa
  hosts:
    - 89.169.102.124 #cpu
    - 89.169.127.168 # L40S
    - 89.169.122.52  # L40S

[r4victor]

So group masters can now provision the jobs within its group? This seems to conflict with _release_replica_jobs_from_master_wait that still releases all the jobs when the global master is provisioned – both group masters and group non-master jobs provision concurrently.

[Bihan]

@r4victor

Proposed staged unlock / provision flow

Layout

• Group 0: jobs 0, 1, 2 — job 0 = global master + group-0 master
• Group 1: jobs 3, 4 — job 3 = group-1 master
• Group 2: jobs 5, 6 — job 5 = group-2 master
  Invariant: later groups stay blocked (waiting_master_job) until the global master (job 0) has finished its provision attempt. After that, unlock is staged (not “release everyone”).

---

1. Job 0 runs

jobs_to_provision = [0, 1, 2]

Backend	Result
ComputeGroup	One run_jobs → VMs for 0, 1, 2
Non-ComputeGroup	This call only creates a VM for job 0

2. Job 0 unlocks (after its provision attempt)

Unlock	Why
Jobs 1, 2 (rest of same group)	Non-ComputeGroup: they still need to provision one-by-one
Job 3 (group-1 master only)	Lets group 1 start its batch
Job 5 (group-2 master only)	Lets group 2 start when ready
Not 4 or 6	Avoids racing the group master’s batch

3. Job 3 runs (group 1)

jobs_to_provision = [3, 4]

Backend	Result
ComputeGroup	One run_jobs → VMs for 3, 4; job 4 already done
Non-ComputeGroup	VM for job 3 only → job 3 then unlocks job 4 → 4 provisions alone

4. Job 5 runs (group 2)

Same pattern as step 3 with jobs_to_provision = [5, 6].

[r4victor]

Please update the docs on node groups in the same PR or submit another PR for an easier review.

[r4victor]

group resources/image never get server-side defaults or validation

set_run_spec_resources_defaults, _validate_gpu_vendor_and_image, and _validate_cpu_arch_and_image in server/services/runs/spec.py all special-case service replica_groups but were not extended to task node_groups.

• top-level resources: {gpu: MI300X}, no image -> rejected ("image must be set ... ['amd']")
• group resources: {gpu: MI300X}, no image -> ACCEPTED; group gpu.vendor=None, cpu.arch=None
• group resources: {gpu: GB200} -> group cpu.arch=None (should be ARM), gpu.vendor=None
• top-level resources (unused for groups) -> gpu.vendor=nvidia, cpu.arch=x86

[r4victor]

What if some other nodes never get an internal ip?

[Bihan]

If a peer nodes never gets an IP, that peer should still fail through normal job/run logic:

1. PROVISIONING / PULLING — if the runner never becomes ready, the peer hits WAITING_RUNNER_LIMIT_EXCEEDED → peer FAILED → run torn down.

2. No capacity — peer fails with FAILED_TO_START_DUE_TO_NO_CAPACITY, retries until retry.duration then RETRY_LIMIT_EXCEEDED → run torn down.