Add replica groups in dstack-service #3408

Bihan · 2025-12-20T03:14:52Z

Steps To Test

Step1: Create replica-groups-service.yml

# replica-groups-service.yml
type: service
name: replica-groups-test
python: 3.12

replicas:
  - name: replica-1 # replica name is optional
    count: 0..2
    scaling:
      metric: rps
      target: 2
    commands:
      - echo "Group 1 - Version 0" > /tmp/version.txt
      - python3 -m http.server 8000
    resources:
      cpu: 2

  - name: replica-2
    count: 0..3
    scaling:
      metric: rps
      target: 2
    commands:
      - echo "Group 2 - Version 0" > /tmp/version.txt
      - python3 -m http.server 8000
    resources:
      cpu: 2

port: 8000

Step2: dstack apply -f replica-groups-service.yml

Step3: Run load_test_replica_groups.py by subsituting your URL and TOKEN

import asyncio
import aiohttp
import time

# ==== Configuration ====
URL = "<URL>"
TOKEN = "<TOKEN>"
RPS = 8          # Requests per second
DURATION = 1800       # Duration in seconds
METHOD = "GET"     # or "POST"
# =======================

HEADERS = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {TOKEN}"
}


async def send_request(session, idx):
    """Send a request and print response"""
    try:
        async with session.request(METHOD, URL, headers=HEADERS) as resp:
            text = await resp.text()
            print(f"\n[{idx}] Status: {resp.status}")
            # print small part of response (HTML preview)
            print(text[:200].strip(), "...\n")
    except Exception as e:
        print(f"[{idx}] Error: {e}")


async def run_load_test():
    total_requests = RPS * DURATION
    interval = 1.0 / RPS

    async with aiohttp.ClientSession() as session:
        start_time = time.perf_counter()
        tasks = []

        for i in range(total_requests):
            task = asyncio.create_task(send_request(session, i + 1))
            tasks.append(task)
            await asyncio.sleep(interval)

        await asyncio.gather(*tasks)
        elapsed = time.perf_counter() - start_time
        print(f"\n✅ Sent {total_requests} requests in {elapsed:.2f}s "
              f"(~{total_requests/elapsed:.2f} RPS)")


if __name__ == "__main__":
    asyncio.run(run_load_test())

Expected Output
Each group gets one replica

Submit the run replica-groups-test? [y/n]: y
 NAME                  BACKEND          GPU  PRICE    STATUS   SUBMITTED 
 replica-groups-test                    -    -        running  07:31     
    group=0 replica=0  aws (us-east-2)  -    $0.0832  running  07:32     
    group=1 replica=1  aws (us-east-2)  -    $0.0832  running  07:32

Later, both groups scale respecting group configs.
group0 scales to 2 replicas,
and group1 scales to 3.

Below is the expected output

NAME                  BACKEND          GPU  PRICE    STATUS   SUBMITTED  
 replica-groups-test                    -    -        running  9 mins ago 
    group=0 replica=0  aws (us-east-2)  -    $0.0832  running  8 mins ago 
            replica=2  aws (us-east-2)  -    $0.0832  running  3 mins ago 
    group=1 replica=1  aws (us-east-2)  -    $0.0832  running  8 mins ago 
            replica=3  aws (us-east-2)  -    $0.0832  running  3 mins ago 
            replica=4  aws (us-east-2)  -    $0.0832  running  3 mins ago

Step4: Check whether replica specific commands were executed.
Attach to the desired replica
Eg:
dstack attach -replica 2 replica-groups-test
ssh replica-groups-test-0-2 'cat /tmp/version.txt'
output: Group 1 - Version 0

Step5: Check rolling deployment.
Important:
Rolling deployments are currently affected by a race condition that also impacts the non–replica group implementation and must be addressed separately (issue). However, when each replica group is configured with a single replica, this race condition does not affect rolling deployments.

Testing instructions:

Scale down each replica group to 1 replica.

Restart the load-testing script with RPS = 2.

After all groups have scaled down to a single replica, re-apply the configuration:

Re-apply
dstack apply -f replica-groups-service.yml

Active run replica-groups-test already exists. Detected changes that can be updated in-place:
- Configuration properties:
  - replica_groups

Update the run? [y/n]: y
 NAME                  BACKEND          GPU  PRICE    STATUS      SUBMITTED 
 replica-groups-test                    -    -        running     07:51     
    group=0 replica=0  aws (us-east-2)  -    $0.0832  terminated  07:51     
            replica=2  aws (us-east-2)  -    $0.0832  running     07:53     
    group=1 replica=1  aws (us-east-2)  -    $0.0832  terminated  07:51     
            replica=3  aws (us-east-2)  -    $0.0832  running     07:53

Bihan · 2025-12-20T03:15:47Z

Will be solving merge conflicts as review continues.

Bihan · 2025-12-20T03:19:20Z

Related PRs

#3205 from @DragonStuff

peterschmidt85 · 2025-12-20T09:52:11Z

@Bihan Do we really need replica group names?

peterschmidt85 · 2025-12-20T09:52:41Z

@Bihan Also please check the conflicts with master.

peterschmidt85 · 2025-12-20T09:54:28Z

Cosmetics only: I would rename replica_groups to replicas and also rename replicas under replica_groups to count.

Bihan · 2025-12-22T07:42:15Z

Cosmetics only: I would rename replica_groups to replicas and also rename replicas under replica_groups to count.

Yes. will rename it.

Bihan · 2025-12-22T09:00:46Z

@Bihan Do we really need replica group names?

Yes.

Without replica names, we would rely on indices, which are position-dependent. If groups are reordered by users during manual scaling, indices shift, but existing jobs and persisted state (like desired_replica_counts) still reference the old positions. This mismatch prevents reliable identification of which group a job belongs to, leading to incorrect scaling decisions. Replica names are not affected by reordering in the YAML file.

Initial

replica_groups:
  - replicas: 1        # Index 0 is prefill
    commands: ["--disaggregation-mode prefill"]
  - replicas: 2        # Index 1 is decode
    commands: ["--disaggregation-mode decode"]

Manual Scaling

replica_groups:
  - replicas: 3        # Index 0 now becomes decode
    commands: ["--disaggregation-mode decode"]
  - replicas: 2        # Index 1 now becomes prefill
    commands: ["--disaggregation-mode prefill"]

Instead of relying on replica group's position in the config, another possibility is matching job specs to identify replicas; but this approach fails during rolling deployments because old and new jobs from the same group have different specs.

peterschmidt85 · 2025-12-22T09:08:30Z

As a user I find it unnecessary to give names. I would prefer not to ask names if this is possible technically.
Why not simply treat indices as names. If the user for some reason changes the order, we re-deploy?

Bihan · 2025-12-22T10:01:57Z

As a user I find it unnecessary to give names. I would prefer not to ask names if this is possible technically. Why not simply treat indices as names. If the user for some reason changes the order, we re-deploy?

If a user changes commands for group 0 and reorders groups at the same time, they expect a rolling deployment for group 0 only. However, the system detects the order change and triggers a full redeployment for all groups. Users may find this implicit behavior annoying because it provisions extra instances for each groups.

peterschmidt85 · 2025-12-22T10:15:35Z

Perhaps we could make these names optional?

Bihan · 2025-12-22T10:34:52Z

Perhaps we could make these names optional?

Yes, we can make it optional.

add_replica_groups_model Replica Groups AutoScaling Rolling deployment and UI Replica Groups implementation clean up

jvstme · 2025-12-23T21:47:00Z

src/dstack/_internal/core/models/configurations.py

        )


+class ReplicaGroup(ConfigurationWithCommandsParams, CoreModel):


(nit) I think we could allow to set many more properties per replica group. If the user can set commands, they may also want to set entrypoint, working_dir, image, volumes, repos, etc. And if the user can set resources, they may also want to set instance_types, spot_policy, reservation, etc.

Although it may be a good idea to leave this to a future iteration, because some properties may be non-trivial to support correctly

src/dstack/_internal/core/models/configurations.py

src/dstack/_internal/server/services/runs/replicas.py

src/dstack/_internal/core/models/configurations.py

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

src/dstack/_internal/server/services/runs/replicas.py

src/dstack/_internal/server/background/tasks/process_runs.py

src/dstack/_internal/core/models/configurations.py

src/dstack/_internal/server/services/runs/spec.py

src/dstack/_internal/server/services/services/__init__.py

src/dstack/_internal/core/models/configurations.py

src/dstack/_internal/server/background/tasks/process_runs.py

src/dstack/_internal/core/models/configurations.py

src/dstack/_internal/server/background/tasks/process_runs.py

src/dstack/_internal/server/services/runs/plan.py

src/dstack/_internal/server/services/runs/__init__.py

peterschmidt85 · 2026-01-01T16:51:30Z

@Bihan the test seem to be broken. Have you seen it?

r4victor · 2026-01-21T07:22:45Z

src/dstack/_internal/server/services/runs/plan.py

+            volumes = await get_job_configured_volumes(
+                session=session,
+                project=project,
+                run_spec=run_spec,
+                job_num=0,
+            )
+            candidate_fleet_models = await _select_candidate_fleet_models(
+                session=session,
+                project=project,
+                run_model=None,
+                run_spec=run_spec,
+            )


Why doing these calls for every replica_group?

@r4victor Thanks for pointing it. Both the calls don't depend on replica_group and can be called once for all replica groups. I will update it.

r4victor · 2026-01-21T07:45:00Z

src/dstack/_internal/core/models/configurations.py

+    name: Annotated[
+        Optional[str],
+        Field(
+            description="The name of the replica group. If not provided, defaults to 'replica0', 'replica1', etc. based on position."


IIUC, there can be multiple different replicas per replica group, right? So then you can have replica group 'replica0' with replica=0, replica=1. The naming is very confusing. Maybe name replica groups replica_group_0, replica_group_1 or group_0,group_1?

@r4victor Yes. You are right. I will update the replica groups name to replica_group_0, replica_group_1. This will make it clear.

I think we prefer kebab-case for generated names, so that'd be replica-group-X.

Alternatively, I can suggest to use just 0, 1, etc (but still as str). The fact that it refers to the replica group should be clear from the context where it is displayed (e.g., group=0 in dstack ps).

src/dstack/_internal/core/models/configurations.py

src/dstack/_internal/server/services/runs/spec.py

src/dstack/_internal/server/background/tasks/process_runs.py

jvstme · 2026-01-21T13:27:31Z

src/dstack/_internal/server/background/tasks/process_runs.py

+                    run_spec=run_spec,
+                    desired_replica_counts=counts,
+                )
+        return


(nit) Since this returns for any service, most of the code below is redundant, because it is (or was) only applicable to services.

If I'm not mistaken, only the _update_jobs_to_new_deployment_in_place call below is relevant for non-services, and the rest can be removed.

I will check in detail whether the code below is redundant and is only used by services.

jvstme · 2026-01-21T23:20:10Z

src/dstack/_internal/cli/utils/run.py

+            if group_index != last_shown_group_index:
+                # First job in group: use 3 spaces indent
+                prefix = "   "
+                name_parts.append(f"group={group_index} replica={job.job_spec.replica_num}")


(nit) Maybe show the group name instead of the index? For example, group=prefill could be more informative than group=0

jvstme · 2026-01-21T23:24:58Z

src/dstack/_internal/core/models/configurations.py

 DEFAULT_PROBE_READY_AFTER = 1
 DEFAULT_PROBE_METHOD = "get"
 MAX_PROBE_URL_LEN = 2048
+DEFAULT_REPLICA_GROUP_NAME = "default"


(nit) Maybe use replica0 (or other name of the first replica group if you decide to change the naming scheme) as the default? That would allow switching from replicas: int to replicas: list in-place, without replica redeployment

jvstme · 2026-01-21T23:32:45Z

src/dstack/_internal/server/background/tasks/process_runs.py

+    from dstack._internal.server.services.runs.replicas import (
+        _build_replica_lists,
+        scale_run_replicas_for_group,
+    )


(nit)

Not sure if there's a reason to import here rather than at the top of the module.

By convention, functions that start with an underscore are private functions for internal use within one module, they are not supposed be imported to other modules. But you can rename _build_replica_lists -> build_replica_lists.

The imports are still in the function body though, as of 8aee1dd

src/dstack/_internal/server/services/runs/replicas.py

jvstme · 2026-01-22T00:11:07Z

src/dstack/_internal/core/models/configurations.py

+    name: Annotated[
+        Optional[str],
+        Field(
+            description="The name of the replica group. If not provided, defaults to 'replica0', 'replica1', etc. based on position."


I think we prefer kebab-case for generated names, so that'd be replica-group-X.

Alternatively, I can suggest to use just 0, 1, etc (but still as str). The fact that it refers to the replica group should be clear from the context where it is displayed (e.g., group=0 in dstack ps).

jvstme · 2026-01-23T13:14:56Z

src/dstack/_internal/cli/utils/run.py

+        if run.run_spec.configuration.type == "service" and hasattr(
+            run.run_spec.configuration, "replica_groups"
+        ):


(nit) ServiceConfiguration always has the replica_groups attribute, so hasattr is redundant

jvstme · 2026-01-23T13:19:05Z

src/dstack/_internal/server/background/tasks/process_runs.py

+            try:
+                job_spec = JobSpec.__response__.parse_raw(job.job_spec_data)
+                existing_group_names.add(job_spec.replica_group)
+            except Exception:
+                continue


(nit) Same comment as here

jvstme · 2026-01-23T15:29:14Z

src/dstack/_internal/server/services/runs/replicas.py

    # Determine replica group from existing job
    run_spec = RunSpec.__response__.parse_raw(run_model.run_spec)
-    job_spec = JobSpec.parse_raw(latest_jobs[0].job_spec_data)
+    job_spec = JobSpec.__response__.parse_raw(latest_jobs[0].job_spec_data)


👍, important fix

Bihan requested review from jvstme and peterschmidt85 December 20, 2025 03:16

Bihan Rana added 2 commits December 23, 2025 06:59

Add replica groups in dstack-service

22c1410

add_replica_groups_model Replica Groups AutoScaling Rolling deployment and UI Replica Groups implementation clean up

Resolve Merge Conflict & Rename replica_groups to replicas

5abbcad

Bihan force-pushed the add_replica_groups_in_dstack_service branch from 86139c5 to 5abbcad Compare December 23, 2025 09:29

Resolve pyright type check

abba7da

jvstme reviewed Dec 23, 2025

View reviewed changes

Bihan Rana and others added 4 commits December 24, 2025 15:04

Rename replicas to count and make replica names optional

d974292

Merge branch 'master' into add_replica_groups_in_dstack_service

caa4283

Resolve review comments on probes and rate limits

1ec1d6d

Resolve tests

7b4bc52

jvstme reviewed Dec 24, 2025

View reviewed changes

src/dstack/_internal/core/models/configurations.py Outdated Show resolved Hide resolved

jvstme reviewed Dec 24, 2025

View reviewed changes

Bihan Rana added 3 commits December 25, 2025 17:30

Transform to ReplicaGroup in the replica_groups property

8c5589d

Resolve review comments

0a54e07

Resolve test_runs

f4c9fdf

jvstme reviewed Dec 25, 2025

View reviewed changes

Resolved major comments

a0e13f6

jvstme reviewed Dec 30, 2025

View reviewed changes

Remove create_group_run_spec and use Job Configurator instead

24d976e

Bihan Rana and others added 4 commits January 4, 2026 07:51

Resolve Minor Issues

263e312

Resolve Minor Issues - Additional fixes

5c71f76

Resolve conflict with master branch

940cded

Merge branch 'master' into add_replica_groups_in_dstack_service

c1364f2

r4victor reviewed Jan 21, 2026

View reviewed changes

jvstme reviewed Jan 22, 2026

View reviewed changes

Resolve Major Comments

8aee1dd

jvstme approved these changes Jan 23, 2026

View reviewed changes

Bihan Rana added 2 commits January 23, 2026 20:35

Resolve some minor comments

0b034ff

Resolve some minor comments

5bdb90c

jvstme approved these changes Jan 23, 2026

View reviewed changes

Bihan merged commit bd2d485 into dstackai:master Jan 23, 2026
28 checks passed

		)


		class ReplicaGroup(ConfigurationWithCommandsParams, CoreModel):

Add replica groups in dstack-service #3408

Add replica groups in dstack-service #3408

Uh oh!

Conversation

Bihan commented Dec 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Bihan commented Dec 20, 2025

Uh oh!

Bihan commented Dec 20, 2025

Related PRs

Uh oh!

peterschmidt85 commented Dec 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

peterschmidt85 commented Dec 20, 2025

Uh oh!

peterschmidt85 commented Dec 20, 2025

Uh oh!

Bihan commented Dec 22, 2025

Uh oh!

Bihan commented Dec 22, 2025

Uh oh!

peterschmidt85 commented Dec 22, 2025

Uh oh!

Bihan commented Dec 22, 2025

Uh oh!

peterschmidt85 commented Dec 22, 2025

Uh oh!

Bihan commented Dec 22, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

peterschmidt85 commented Jan 1, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Bihan commented Dec 20, 2025 •

edited

Loading

peterschmidt85 commented Dec 20, 2025 •

edited

Loading