feat: enhance replica group handling and migration for legacy jobs

- Implemented migration for legacy jobs that lack a replica_group_name, ensuring they are correctly assigned to the appropriate replica groups.
- Updated CLI output to display group-specific properties such as spot policy, regions, and backends for better clarity.
- Enhanced tests to validate the migration process and ensure that jobs are correctly assigned to their respective groups.
- Improved handling of pool offers to accommodate multiple jobs in replica groups, ensuring all GPU types are considered.

This update improves the robustness of the service configuration and enhances user experience by providing clearer information in the CLI.

Author: DragonStuff

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

@@ -122,31 +122,52 @@ def th(s: str) -> str:
 
     from dstack._internal.core.models.configurations import ServiceConfiguration
 
-    if (
+    has_replica_groups = (
         include_run_properties
         and isinstance(run_spec.configuration, ServiceConfiguration)
         and run_spec.configuration.replica_groups
-    ):
+    )
+
+    if has_replica_groups:
         groups_info = []
         for group in run_spec.configuration.replica_groups:
             group_parts = [f"[cyan]{group.name}[/cyan]"]
 
+            # Replica count
             if group.replicas.min == group.replicas.max:
                 group_parts.append(f"×{group.replicas.max}")
             else:
                 group_parts.append(f"×{group.replicas.min}..{group.replicas.max}")
                 group_parts.append("[dim](autoscalable)[/dim]")
 
+            # Resources
             group_parts.append(f"[dim]({group.resources.pretty_format()})[/dim]")
 
+            # Group-specific overrides
+            overrides = []
+            if group.spot_policy is not None:
+                overrides.append(f"spot={group.spot_policy.value}")
+            if group.regions:
+                regions_str = ",".join(group.regions[:2])  # Show first 2
+                if len(group.regions) > 2:
+                    regions_str += f",+{len(group.regions) - 2}"
+                overrides.append(f"regions={regions_str}")
+            if group.backends:
+                backends_str = ",".join([b.value for b in group.backends[:2]])
+                if len(group.backends) > 2:
+                    backends_str += f",+{len(group.backends) - 2}"
+                overrides.append(f"backends={backends_str}")
+
+            if overrides:
+                group_parts.append(f"[dim]({'; '.join(overrides)})[/dim]")
+
             groups_info.append(" ".join(group_parts))
 
         props.add_row(th("Replica groups"), "\n".join(groups_info))
     else:
         props.add_row(th("Resources"), pretty_req)
-
-    props.add_row(th("Spot policy"), spot_policy)
-    props.add_row(th("Max price"), max_price)
+        props.add_row(th("Spot policy"), spot_policy)
+        props.add_row(th("Max price"), max_price)
     if include_run_properties:
         props.add_row(th("Retry policy"), retry)
         props.add_row(th("Creation policy"), creation_policy)

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

@@ -156,6 +156,10 @@ async def _process_run(session: AsyncSession, run_model: RunModel):
     )
     run_model = res.unique().scalar_one()
     logger.debug("%s: processing run", fmt(run_model))
+
+    # Migrate legacy jobs without replica_group_name (one-time fix)
+    await _migrate_legacy_job_replica_groups(session, run_model)
+
     try:
         if run_model.status == RunStatus.PENDING:
             await _process_pending_run(session, run_model)
@@ -176,6 +180,70 @@ async def _process_run(session: AsyncSession, run_model: RunModel):
     await session.commit()
 
 
+async def _migrate_legacy_job_replica_groups(session: AsyncSession, run_model: RunModel):
+    """
+    Migrate jobs from old runs that don't have replica_group_name set.
+    This fixes jobs created before the replica_groups feature was added.
+    """
+    run_spec = RunSpec.__response__.parse_raw(run_model.run_spec)
+
+    # Only migrate service runs with replica_groups
+    if run_spec.configuration.type != "service":
+        return
+
+    # Check if run uses replica_groups
+    if not getattr(run_spec.configuration, "replica_groups", None):
+        return
+
+    from dstack._internal.core.models.runs import get_normalized_replica_groups
+
+    normalized_groups = get_normalized_replica_groups(run_spec.configuration)
+
+    # Check if any jobs need migration
+    needs_migration = any(job.replica_group_name is None for job in run_model.jobs)
+
+    if not needs_migration:
+        return
+
+    logger.info(
+        "%s: Migrating legacy jobs to assign replica_group_name",
+        fmt(run_model),
+    )
+
+    # Build a map of replica_num -> group_name based on how jobs were originally created
+    replica_num_to_group = {}
+    current_replica_num = 0
+
+    for group in normalized_groups:
+        group_min = group.replicas.min or 0
+        for _ in range(group_min):
+            replica_num_to_group[current_replica_num] = group.name
+            current_replica_num += 1
+
+    # Update jobs
+    migrated_count = 0
+    for job in run_model.jobs:
+        if job.replica_group_name is None:
+            expected_group = replica_num_to_group.get(job.replica_num)
+            if expected_group:
+                job.replica_group_name = expected_group
+                migrated_count += 1
+                logger.info(
+                    "%s: Migrated job replica_num=%d to group '%s'",
+                    fmt(run_model),
+                    job.replica_num,
+                    expected_group,
+                )
+
+    if migrated_count > 0:
+        await session.commit()
+        logger.info(
+            "%s: Migrated %d job(s) to replica groups",
+            fmt(run_model),
+            migrated_count,
+        )
+
+
 async def _process_pending_run(session: AsyncSession, run_model: RunModel):
     """Jobs are not created yet"""
     run = run_model_to_run(run_model)

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

@@ -740,11 +740,23 @@ def _get_profile_for_job(run_spec: RunSpec, job: Job) -> Profile:
     base_profile = run_spec.merged_profile
 
     group_name = job.job_spec.replica_group_name
+    logger.info(
+        "Getting profile for job %s: replica_group_name=%s, config_type=%s",
+        job.job_spec.job_name,
+        group_name,
+        run_spec.configuration.type,
+    )
+
     if not group_name or run_spec.configuration.type != "service":
+        logger.info("Using base profile (no group_name or not a service)")
         return base_profile
 
     # Find the group
     normalized_groups = get_normalized_replica_groups(run_spec.configuration)
+    logger.info(
+        "Normalized groups: %s",
+        [f"{g.name} (regions={g.regions})" for g in normalized_groups],
+    )
     group = next((g for g in normalized_groups if g.name == group_name), None)
 
     if not group:
@@ -762,7 +774,7 @@ def _get_profile_for_job(run_spec: RunSpec, job: Job) -> Profile:
         if group_value is not None:
             setattr(merged, field_name, group_value)
 
-    logger.debug(
+    logger.info(
         "Profile for group '%s': regions=%s, backends=%s, spot_policy=%s (base had: regions=%s, backends=%s, spot_policy=%s)",
         group_name,
         merged.regions,
@@ -832,6 +844,21 @@ async def _run_job_on_new_instance(
     multinode = job.job_spec.jobs_per_replica > 1 or (
         fleet is not None and fleet.spec.configuration.placement == InstanceGroupPlacement.CLUSTER
     )
+
+    # Log the requirements and profile being used
+    gpu_requirement = (
+        requirements.resources.gpu.name
+        if requirements.resources and requirements.resources.gpu
+        else None
+    )
+    logger.info(
+        "%s: Fetching offers with GPU=%s, regions=%s, backends=%s",
+        fmt(job_model),
+        gpu_requirement,
+        profile.regions,
+        [b.value for b in profile.backends] if profile.backends else None,
+    )
+
     offers = await get_offers_by_requirements(
         project=project,
         profile=profile,
@@ -845,14 +872,12 @@ async def _run_job_on_new_instance(
     )
 
     # Debug logging for offers
-    if replica_group_name and len(offers) > 0:
-        logger.debug(
-            "%s: Got %d offers for group '%s'. First 3: %s",
-            fmt(job_model),
-            len(offers),
-            replica_group_name,
-            [f"{o.instance.name} ({o.backend.value}/{o.region})" for _, o in offers[:3]],
-        )
+    logger.info(
+        "%s: Got %d offers. First 3: %s",
+        fmt(job_model),
+        len(offers),
+        [f"{o.instance.name} ({o.backend.value}/{o.region})" for _, o in offers[:3]],
+    )
     # Limit number of offers tried to prevent long-running processing
     # in case all offers fail.
     for backend, offer in offers[: settings.MAX_OFFERS_TRIED]:

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

@@ -417,13 +417,36 @@ async def get_plan(
         job_num=0,
     )
 
-    pool_offers = await _get_pool_offers(
-        session=session,
-        project=project,
-        run_spec=effective_run_spec,
-        job=jobs[0],
-        volumes=volumes,
-    )
+    # For replica groups, we need pool offers for all GPU types, not just the first job's type
+    # So we fetch pool offers for each job separately and aggregate them
+    all_pool_offers = []
+    if len(jobs) > 1:
+        # Multiple jobs (likely replica groups) - get pool offers per job to include all GPU types
+        for job in jobs:
+            job_pool_offers = await _get_pool_offers(
+                session=session,
+                project=project,
+                run_spec=effective_run_spec,
+                job=job,
+                volumes=volumes,
+            )
+            all_pool_offers.extend(job_pool_offers)
+        # Deduplicate by (backend, instance_name, region) tuple
+        seen_offers = set()
+        pool_offers = []
+        for offer in all_pool_offers:
+            offer_key = (offer.backend, offer.instance.name, offer.region)
+            if offer_key not in seen_offers:
+                seen_offers.add(offer_key)
+                pool_offers.append(offer)
+    else:
+        pool_offers = await _get_pool_offers(
+            session=session,
+            project=project,
+            run_spec=effective_run_spec,
+            job=jobs[0],
+            volumes=volumes,
+        )
     effective_run_spec.run_name = "dry-run"  # will regenerate jobs on submission
 
     # Check if all jobs have identical requirements (optimization for single-type jobs)
@@ -1561,10 +1584,29 @@ async def retry_run_replica_jobs(
         session=session,
         project=run_model.project,
     )
+
+    # Determine which replica group this job belongs to
+    run_spec = RunSpec.__response__.parse_raw(run_model.run_spec)
+    replica_group = None
+    if run_spec.configuration.type == "service" and latest_jobs:
+        from dstack._internal.core.models.runs import get_normalized_replica_groups
+
+        group_name = latest_jobs[0].replica_group_name
+        if group_name:
+            normalized_groups = get_normalized_replica_groups(run_spec.configuration)
+            replica_group = next((g for g in normalized_groups if g.name == group_name), None)
+            if replica_group:
+                logger.info(
+                    "%s: retrying job from replica group '%s'",
+                    fmt(run_model),
+                    replica_group.name,
+                )
+
     new_jobs = await get_jobs_from_run_spec(
-        run_spec=RunSpec.__response__.parse_raw(run_model.run_spec),
+        run_spec=run_spec,
         secrets=secrets,
         replica_num=latest_jobs[0].replica_num,
+        replica_group=replica_group,
     )
     assert len(new_jobs) == len(latest_jobs), (
         "Changing the number of jobs within a replica is not yet supported"

src/dstack/_internal/server/testing/common.py

@@ -354,8 +354,19 @@ async def create_job(
     if deployment_num is None:
         deployment_num = run.deployment_num
     run_spec = RunSpec.parse_raw(run.run_spec)
+
+    # Look up replica group if specified
+    replica_group = None
+    if replica_group_name and run_spec.configuration.type == "service":
+        from dstack._internal.core.models.runs import get_normalized_replica_groups
+
+        normalized_groups = get_normalized_replica_groups(run_spec.configuration)
+        replica_group = next((g for g in normalized_groups if g.name == replica_group_name), None)
+
     job_spec = (
-        await get_job_specs_from_run_spec(run_spec=run_spec, secrets={}, replica_num=replica_num)
+        await get_job_specs_from_run_spec(
+            run_spec=run_spec, secrets={}, replica_num=replica_num, replica_group=replica_group
+        )
     )[0]
     job_spec.job_num = job_num
     job = JobModel(