Remove create_group_run_spec and use Job Configurator instead

Author: Bihan Rana

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

@@ -611,7 +611,7 @@ def convert_ports(cls, v) -> PortMapping:
 class ConfigurationWithCommandsParams(CoreModel):
     commands: Annotated[CommandsList, Field(description="The shell commands to run")] = []
 
-    @root_validator(pre=True)
+    @root_validator
     def check_image_or_commands_present(cls, values):
         # If replicas is list, skip validation - commands come from replica groups
         replicas = values.get("replicas")
@@ -957,7 +957,6 @@ def validate_top_level_properties_with_replica_groups(cls, values):
             )
 
         resources = values.get("resources")
-        from dstack._internal.core.models.resources import ResourcesSpec
 
         default_resources = ResourcesSpec()
         if resources and resources.dict() != default_resources.dict():

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

@@ -40,7 +40,6 @@
 from dstack._internal.server.services.locking import get_locker
 from dstack._internal.server.services.prometheus.client_metrics import run_metrics
 from dstack._internal.server.services.runs import (
-    create_group_run_spec,
     fmt,
     process_terminating_run,
     run_model_to_run,
@@ -482,7 +481,7 @@ async def _handle_run_replicas(
                     session=session,
                     run_model=run_model,
                     group=group,
-                    base_run_spec=run_spec,
+                    run_spec=run_spec,
                     desired_replica_counts=counts,
                 )
         return
@@ -563,7 +562,6 @@ async def _update_jobs_to_new_deployment_in_place(
         session=session,
         project=run_model.project,
     )
-    base_run_spec = run_spec
 
     for replica_num, job_models in group_jobs_by_replica_latest(run_model.jobs):
         if all(j.status.is_finished() for j in job_models):
@@ -573,22 +571,17 @@ async def _update_jobs_to_new_deployment_in_place(
 
         # Determine which group this replica belongs to
         replica_group_name = None
-        group_run_spec = base_run_spec
 
         if replicas:
             job_spec = JobSpec.__response__.parse_raw(job_models[0].job_spec_data)
             replica_group_name = job_spec.replica_group
 
-            for group in replicas:
-                if group.name == replica_group_name:
-                    group_run_spec = create_group_run_spec(base_run_spec, group)
-                    break
-
         # FIXME: Handle getting image configuration errors or skip it.
         new_job_specs = await get_job_specs_from_run_spec(
-            run_spec=group_run_spec,
+            run_spec=run_spec,
             secrets=secrets,
             replica_num=replica_num,
+            replica_group_name=replica_group_name,
         )
         assert len(new_job_specs) == len(job_models), (
             "Changing the number of jobs within a replica is not yet supported"
@@ -681,7 +674,7 @@ async def _handle_rolling_deployment_for_group(
     session: AsyncSession,
     run_model: RunModel,
     group: ReplicaGroup,
-    base_run_spec: RunSpec,
+    run_spec: RunSpec,
     desired_replica_counts: dict,
 ) -> None:
     """
@@ -726,7 +719,7 @@ async def _handle_rolling_deployment_for_group(
             run_model=run_model,
             group=group,
             replicas_diff=group_max_replica_count - non_terminated_replica_count,
-            base_run_spec=base_run_spec,
+            run_spec=run_spec,
             active_replicas=active_replicas,
             inactive_replicas=inactive_replicas,
         )
@@ -774,7 +767,7 @@ async def _handle_rolling_deployment_for_group(
             run_model=run_model,
             group=group,
             replicas_diff=-replicas_to_stop_count,
-            base_run_spec=base_run_spec,
+            run_spec=run_spec,
             active_replicas=active_replicas,
             inactive_replicas=inactive_replicas,
         )

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

@@ -98,22 +98,31 @@ def switch_job_status(
 
 
 async def get_jobs_from_run_spec(
-    run_spec: RunSpec, secrets: Dict[str, str], replica_num: int
+    run_spec: RunSpec,
+    secrets: Dict[str, str],
+    replica_num: int,
+    replica_group_name: Optional[str] = None,
 ) -> List[Job]:
     return [
         Job(job_spec=s, job_submissions=[])
         for s in await get_job_specs_from_run_spec(
             run_spec=run_spec,
             secrets=secrets,
             replica_num=replica_num,
+            replica_group_name=replica_group_name,
         )
     ]
 
 
 async def get_job_specs_from_run_spec(
-    run_spec: RunSpec, secrets: Dict[str, str], replica_num: int
+    run_spec: RunSpec,
+    secrets: Dict[str, str],
+    replica_num: int,
+    replica_group_name: Optional[str] = None,
 ) -> List[JobSpec]:
-    job_configurator = _get_job_configurator(run_spec=run_spec, secrets=secrets)
+    job_configurator = _get_job_configurator(
+        run_spec=run_spec, secrets=secrets, replica_group_name=replica_group_name
+    )
     job_specs = await job_configurator.get_job_specs(replica_num=replica_num)
     return job_specs
 
@@ -241,10 +250,14 @@ def is_master_job(job: Job) -> bool:
     return job.job_spec.job_num == 0
 
 
-def _get_job_configurator(run_spec: RunSpec, secrets: Dict[str, str]) -> JobConfigurator:
+def _get_job_configurator(
+    run_spec: RunSpec, secrets: Dict[str, str], replica_group_name: Optional[str] = None
+) -> JobConfigurator:
     configuration_type = RunConfigurationType(run_spec.configuration.type)
     configurator_class = _configuration_type_to_configurator_class_map[configuration_type]
-    return configurator_class(run_spec=run_spec, secrets=secrets)
+    return configurator_class(
+        run_spec=run_spec, secrets=secrets, replica_group_name=replica_group_name
+    )
 
 
 _job_configurator_classes = [

src/dstack/_internal/server/services/jobs/configurators/base.py

@@ -90,9 +90,11 @@ def __init__(
         self,
         run_spec: RunSpec,
         secrets: Optional[Dict[str, str]] = None,
+        replica_group_name: Optional[str] = None,
     ):
         self.run_spec = run_spec
         self.secrets = secrets or {}
+        self.replica_group_name = replica_group_name
 
     async def get_job_specs(self, replica_num: int) -> List[JobSpec]:
         job_spec = await self._get_job_spec(replica_num=replica_num, job_num=0, jobs_per_replica=1)
@@ -150,6 +152,7 @@ async def _get_job_spec(
             job_num=job_num,
             job_name=f"{self.run_spec.run_name}-{job_num}-{replica_num}",
             jobs_per_replica=jobs_per_replica,
+            replica_group=self.replica_group_name or "default",
             app_specs=self._app_specs(),
             commands=await self._commands(),
             env=self._env(),
@@ -298,9 +301,15 @@ def _registry_auth(self) -> Optional[RegistryAuth]:
         return self.run_spec.configuration.registry_auth
 
     def _requirements(self, jobs_per_replica: int) -> Requirements:
+        resources = self.run_spec.configuration.resources
+        if self.run_spec.configuration.type == "service":
+            for group in self.run_spec.configuration.replica_groups:
+                if group.name == self.replica_group_name:
+                    resources = group.resources
+                    break
         spot_policy = self._spot_policy()
         return Requirements(
-            resources=self.run_spec.configuration.resources,
+            resources=resources,
             max_price=self.run_spec.merged_profile.max_price,
             spot=None if spot_policy == SpotPolicy.AUTO else (spot_policy == SpotPolicy.SPOT),
             reservation=self.run_spec.merged_profile.reservation,

src/dstack/_internal/server/services/jobs/configurators/service.py

@@ -10,6 +10,9 @@ class ServiceJobConfigurator(JobConfigurator):
 
     def _shell_commands(self) -> List[str]:
         assert self.run_spec.configuration.type == "service"
+        for group in self.run_spec.configuration.replica_groups:
+            if group.name == self.replica_group_name:
+                return group.commands
         return self.run_spec.configuration.commands
 
     def _default_single_branch(self) -> bool:

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

@@ -18,7 +18,6 @@
     ServerClientError,
 )
 from dstack._internal.core.models.common import ApplyAction
-from dstack._internal.core.models.configurations import ReplicaGroup
 from dstack._internal.core.models.profiles import (
     RetryEvent,
 )
@@ -520,29 +519,22 @@ async def submit_run(
 
             global_replica_num = 0  # Global counter across all groups for unique replica_num
 
-            for replica_group in service_config.replica_groups or []:
+            for replica_group in service_config.replica_groups:
                 if run_spec.merged_profile.schedule is not None:
                     group_initial_replicas = 0
                 else:
                     group_initial_replicas = replica_group.count.min or 0
 
                 # Each replica in this group gets the same group-specific configuration
                 for group_replica_num in range(group_initial_replicas):
-                    group_run_spec = create_group_run_spec(
-                        base_run_spec=run_spec,
-                        replica_group=replica_group,
-                    )
                     jobs = await get_jobs_from_run_spec(
-                        run_spec=group_run_spec,
+                        run_spec=run_spec,
                         secrets=secrets,
                         replica_num=global_replica_num,
+                        replica_group_name=replica_group.name,
                     )
 
                     for job in jobs:
-                        assert replica_group.name is not None, (
-                            "ReplicaGroup name should be set by validator"
-                        )
-                        job.job_spec.replica_group = replica_group.name
                         job_model = create_job_model_for_new_submission(
                             run_model=run_model,
                             job=job,
@@ -590,36 +582,6 @@ async def submit_run(
         return common_utils.get_or_error(run)
 
 
-def create_group_run_spec(
-    base_run_spec: RunSpec,
-    replica_group: ReplicaGroup,
-) -> RunSpec:
-    # Create a copy of the configuration as a dict
-    config_dict = base_run_spec.configuration.dict()
-
-    # Remove replicas and scaling fields since we're creating a single-replica config
-    # This prevents validation errors when commands/resources are added
-    config_dict.pop("replicas", None)
-    config_dict.pop("scaling", None)
-
-    # Override with group-specific values (only if provided)
-    if replica_group.commands:
-        config_dict["commands"] = replica_group.commands
-
-    if replica_group.resources:
-        config_dict["resources"] = replica_group.resources
-
-    # Create new configuration object with merged values
-    # Use the same class as base (ServiceConfiguration)
-    new_config = base_run_spec.configuration.__class__.parse_obj(config_dict)
-
-    # Create new RunSpec with modified configuration
-    # Preserve all other RunSpec properties (repo_data, file_archives, etc.)
-    run_spec_dict = base_run_spec.dict()
-    run_spec_dict["configuration"] = new_config
-    return RunSpec.parse_obj(run_spec_dict)
-
-
 def create_job_model_for_new_submission(
     run_model: RunModel,
     job: Job,