dstack 0.19.23rc1__py3-none-any.whl → 0.19.25__py3-none-any.whl

dstack/_internal/core/backends/aws/compute.py

@@ -383,10 +383,7 @@ class AWSCompute(
     ) -> bool:
         if not _offer_supports_placement_group(instance_offer, placement_group):
             return False
-        return (
-            placement_group.configuration.backend == BackendType.AWS
-            and placement_group.configuration.region == instance_offer.region
-        )
+        return placement_group.configuration.region == instance_offer.region
 
     def create_gateway(
         self,

dstack/_internal/core/backends/base/compute.py

@@ -263,10 +263,6 @@ class ComputeWithPlacementGroupSupport(ABC):
         Checks if the instance offer can be provisioned in the placement group.
 
         Should return immediately, without performing API calls.
-
-        Can be called with an offer originating from a different backend, because some backends
-        (BackendType.DSTACK) produce offers on behalf of other backends. Should return `False`
-        in that case.
         """
         pass
 

dstack/_internal/core/backends/gcp/compute.py

@@ -448,10 +448,7 @@ class GCPCompute(
         placement_group: PlacementGroup,
         instance_offer: InstanceOffer,
     ) -> bool:
-        return (
-            placement_group.configuration.backend == BackendType.GCP
-            and placement_group.configuration.region == instance_offer.region
-        )
+        return placement_group.configuration.region == instance_offer.region
 
     def create_gateway(
         self,

dstack/_internal/core/backends/nebius/compute.py

@@ -298,10 +298,7 @@ class NebiusCompute(
         placement_group: PlacementGroup,
         instance_offer: InstanceOffer,
     ) -> bool:
-        if not (
-            placement_group.configuration.backend == BackendType.NEBIUS
-            and placement_group.configuration.region == instance_offer.region
-        ):
+        if placement_group.configuration.region != instance_offer.region:
             return False
         assert placement_group.provisioning_data is not None
         backend_data = NebiusPlacementGroupBackendData.load(

dstack/_internal/core/models/common.py

@@ -102,7 +102,7 @@ class RegistryAuth(CoreModel):
         password (str): The password or access token
     """
 
-    class Config:
+    class Config(CoreModel.Config):
         frozen = True
 
     username: Annotated[str, Field(description="The username")]

dstack/_internal/core/models/config.py

@@ -16,7 +16,9 @@ class RepoConfig(CoreModel):
     path: str
     repo_id: str
     repo_type: RepoType
-    ssh_key_path: str
+    # Deprecated since 0.19.25, not used. Can be removed when most users update their `config.yml`
+    # (it's updated each time a project or repo is added)
+    ssh_key_path: Annotated[Optional[str], Field(exclude=True)] = None
 
 
 class GlobalConfig(CoreModel):

dstack/_internal/core/models/configurations.py

@@ -20,6 +20,7 @@ from dstack._internal.core.models.services import AnyModel, OpenAIChatModel
 from dstack._internal.core.models.unix import UnixUser
 from dstack._internal.core.models.volumes import MountPoint, VolumeConfiguration, parse_mount_point
 from dstack._internal.utils.common import has_duplicates
+from dstack._internal.utils.json_schema import add_extra_schema_types
 from dstack._internal.utils.json_utils import (
     pydantic_orjson_dumps_with_indent,
 )
@@ -561,7 +562,7 @@ class ServiceConfigurationParams(CoreModel):
     )
     auth: Annotated[bool, Field(description="Enable the authorization")] = True
     replicas: Annotated[
-        Union[conint(ge=1), constr(regex=r"^[0-9]+..[1-9][0-9]*$"), Range[int]],
+        Range[int],
         Field(
             description="The number of replicas. Can be a number (e.g. `2`) or a range (`0..4` or `1..8`). "
             "If it's a range, the `scaling` property is required"
@@ -592,20 +593,13 @@ class ServiceConfigurationParams(CoreModel):
         return v
 
     @validator("replicas")
-    def convert_replicas(cls, v: Any) -> Range[int]:
-        if isinstance(v, str) and ".." in v:
-            min, max = v.replace(" ", "").split("..")
-            v = Range(min=min or 0, max=max or None)
-        elif isinstance(v, (int, float)):
-            v = Range(min=v, max=v)
+    def convert_replicas(cls, v: Range[int]) -> Range[int]:
         if v.max is None:
             raise ValueError("The maximum number of replicas is required")
+        if v.min is None:
+            v.min = 0
         if v.min < 0:
             raise ValueError("The minimum number of replicas must be greater than or equal to 0")
-        if v.max < v.min:
-            raise ValueError(
-                "The maximum number of replicas must be greater than or equal to the minimum number of replicas"
-            )
         return v
 
     @validator("gateway")
@@ -622,9 +616,9 @@ class ServiceConfigurationParams(CoreModel):
     def validate_scaling(cls, values):
         scaling = values.get("scaling")
         replicas = values.get("replicas")
-        if replicas.min != replicas.max and not scaling:
+        if replicas and replicas.min != replicas.max and not scaling:
             raise ValueError("When you set `replicas` to a range, ensure to specify `scaling`.")
-        if replicas.min == replicas.max and scaling:
+        if replicas and replicas.min == replicas.max and scaling:
             raise ValueError("To use `scaling`, `replicas` must be set to a range.")
         return values
 
@@ -655,6 +649,14 @@ class ServiceConfiguration(
 ):
     type: Literal["service"] = "service"
 
+    class Config(CoreModel.Config):
+        @staticmethod
+        def schema_extra(schema: Dict[str, Any]):
+            add_extra_schema_types(
+                schema["properties"]["replicas"],
+                extra_types=[{"type": "integer"}, {"type": "string"}],
+            )
+
 
 AnyRunConfiguration = Union[DevEnvironmentConfiguration, TaskConfiguration, ServiceConfiguration]
 
@@ -715,7 +717,7 @@ class DstackConfiguration(CoreModel):
         Field(discriminator="type"),
     ]
 
-    class Config:
+    class Config(CoreModel.Config):
         json_loads = orjson.loads
         json_dumps = pydantic_orjson_dumps_with_indent
 

dstack/_internal/core/models/fleets.py

@@ -234,7 +234,7 @@ class InstanceGroupParams(CoreModel):
     termination_policy: Annotated[Optional[TerminationPolicy], Field(exclude=True)] = None
     termination_idle_time: Annotated[Optional[Union[str, int]], Field(exclude=True)] = None
 
-    class Config:
+    class Config(CoreModel.Config):
         @staticmethod
         def schema_extra(schema: Dict[str, Any], model: Type):
             del schema["properties"]["termination_policy"]
@@ -279,7 +279,7 @@ class FleetSpec(CoreModel):
     # TODO: make merged_profile a computed field after migrating to pydanticV2
     merged_profile: Annotated[Profile, Field(exclude=True)] = None
 
-    class Config:
+    class Config(CoreModel.Config):
         @staticmethod
         def schema_extra(schema: Dict[str, Any], model: Type) -> None:
             prop = schema.get("properties", {})

dstack/_internal/core/models/instances.py

@@ -122,7 +122,7 @@ class SSHConnectionParams(CoreModel):
     username: str
     port: int
 
-    class Config:
+    class Config(CoreModel.Config):
         frozen = True
 
 
@@ -165,6 +165,9 @@ class InstanceAvailability(Enum):
     AVAILABLE = "available"
     NOT_AVAILABLE = "not_available"
     NO_QUOTA = "no_quota"
+    NO_BALANCE = (
+        "no_balance"  # Introduced in 0.19.24, may be used after a short compatibility period
+    )
     IDLE = "idle"
     BUSY = "busy"
 

dstack/_internal/core/models/profiles.py

@@ -339,7 +339,7 @@ class ProfileParams(CoreModel):
     termination_policy: Annotated[Optional[TerminationPolicy], Field(exclude=True)] = None
     termination_idle_time: Annotated[Optional[Union[str, int]], Field(exclude=True)] = None
 
-    class Config:
+    class Config(CoreModel.Config):
         @staticmethod
         def schema_extra(schema: Dict[str, Any]) -> None:
             del schema["properties"]["pool_name"]
@@ -379,7 +379,7 @@ class Profile(ProfileProps, ProfileParams):
 class ProfilesConfig(CoreModel):
     profiles: List[Profile]
 
-    class Config:
+    class Config(CoreModel.Config):
         json_loads = orjson.loads
         json_dumps = pydantic_orjson_dumps_with_indent
 

dstack/_internal/core/models/repos/remote.py

@@ -32,7 +32,7 @@ class RemoteRepoCreds(CoreModel):
     # TODO: remove in 0.20. Left for compatibility with CLI <=0.18.44
     protocol: Annotated[Optional[str], Field(exclude=True)] = None
 
-    class Config:
+    class Config(CoreModel.Config):
         @staticmethod
         def schema_extra(schema: Dict[str, Any]) -> None:
             del schema["properties"]["protocol"]
@@ -47,7 +47,7 @@ class RemoteRepoInfo(BaseRepoInfo):
     repo_port: Annotated[Optional[int], Field(exclude=True)] = None
     repo_user_name: Annotated[Optional[str], Field(exclude=True)] = None
 
-    class Config:
+    class Config(BaseRepoInfo.Config):
         @staticmethod
         def schema_extra(schema: Dict[str, Any]) -> None:
             del schema["properties"]["repo_host_name"]

dstack/_internal/core/models/resources.py

@@ -130,7 +130,7 @@ DEFAULT_GPU_COUNT = Range[int](min=1)
 
 
 class CPUSpec(CoreModel):
-    class Config:
+    class Config(CoreModel.Config):
         @staticmethod
         def schema_extra(schema: Dict[str, Any]):
             add_extra_schema_types(
@@ -191,7 +191,7 @@ class CPUSpec(CoreModel):
 
 
 class GPUSpec(CoreModel):
-    class Config:
+    class Config(CoreModel.Config):
         @staticmethod
         def schema_extra(schema: Dict[str, Any]):
             add_extra_schema_types(
@@ -314,7 +314,7 @@ class GPUSpec(CoreModel):
 
 
 class DiskSpec(CoreModel):
-    class Config:
+    class Config(CoreModel.Config):
         @staticmethod
         def schema_extra(schema: Dict[str, Any]):
             add_extra_schema_types(
@@ -340,7 +340,7 @@ DEFAULT_DISK = DiskSpec(size=Range[Memory](min=Memory.parse("100GB"), max=None))
 
 
 class ResourcesSpec(CoreModel):
-    class Config:
+    class Config(CoreModel.Config):
         @staticmethod
         def schema_extra(schema: Dict[str, Any]):
             add_extra_schema_types(

dstack/_internal/core/models/runs.py

@@ -350,15 +350,17 @@ class JobSubmission(CoreModel):
     deployment_num: int = 0  # default for compatibility with pre-0.19.14 servers
     submitted_at: datetime
     last_processed_at: datetime
-    finished_at: Optional[datetime]
-    inactivity_secs: Optional[int]
+    finished_at: Optional[datetime] = None
+    inactivity_secs: Optional[int] = None
     status: JobStatus
     status_message: str = ""  # default for backward compatibility
-    termination_reason: Optional[JobTerminationReason]
-    termination_reason_message: Optional[str]
-    exit_status: Optional[int]
-    job_provisioning_data: Optional[JobProvisioningData]
-    job_runtime_data: Optional[JobRuntimeData]
+    # termination_reason stores JobTerminationReason.
+    # str allows adding new enum members without breaking compatibility with old clients.
+    termination_reason: Optional[str] = None
+    termination_reason_message: Optional[str] = None
+    exit_status: Optional[int] = None
+    job_provisioning_data: Optional[JobProvisioningData] = None
+    job_runtime_data: Optional[JobRuntimeData] = None
     error: Optional[str] = None
     probes: list[Probe] = []
 
@@ -442,7 +444,7 @@ class RunSpec(CoreModel):
     # TODO: make merged_profile a computed field after migrating to pydanticV2
     merged_profile: Annotated[Profile, Field(exclude=True)] = None
 
-    class Config:
+    class Config(CoreModel.Config):
         @staticmethod
         def schema_extra(schema: Dict[str, Any], model: Type) -> None:
             prop = schema.get("properties", {})
@@ -508,7 +510,9 @@ class Run(CoreModel):
     last_processed_at: datetime
     status: RunStatus
     status_message: str = ""  # default for backward compatibility
-    termination_reason: Optional[RunTerminationReason] = None
+    # termination_reason stores RunTerminationReason.
+    # str allows adding new enum members without breaking compatibility with old clients.
+    termination_reason: Optional[str] = None
     run_spec: RunSpec
     jobs: List[Job]
     latest_job_submission: Optional[JobSubmission] = None

dstack/_internal/core/services/configs/__init__.py

@@ -38,7 +38,10 @@ class ConfigManager:
             with open(self.config_filepath, "r") as f:
                 config = yaml.safe_load(f)
             self.config = GlobalConfig.parse_obj(config)
-        except (FileNotFoundError, ValidationError):
+        except FileNotFoundError:
+            self.config = GlobalConfig()
+        except ValidationError:
+            logger.error(f"Error in `{self.config_filepath}`", exc_info=True)
             self.config = GlobalConfig()
 
     def get_project_config(self, name: Optional[str] = None) -> Optional[ProjectConfig]:
@@ -71,19 +74,15 @@ class ConfigManager:
     def delete_project(self, name: str):
         self.config.projects = [p for p in self.config.projects if p.name != name]
 
-    def save_repo_config(
-        self, repo_path: PathLike, repo_id: str, repo_type: RepoType, ssh_key_path: PathLike
-    ):
+    def save_repo_config(self, repo_path: PathLike, repo_id: str, repo_type: RepoType):
         self.config_filepath.parent.mkdir(parents=True, exist_ok=True)
         with filelock.FileLock(str(self.config_filepath) + ".lock"):
             self.load()
             repo_path = os.path.abspath(repo_path)
-            ssh_key_path = os.path.abspath(ssh_key_path)
             for repo in self.config.repos:
                 if repo.path == repo_path:
                     repo.repo_id = repo_id
                     repo.repo_type = repo_type
-                    repo.ssh_key_path = ssh_key_path
                     break
             else:
                 self.config.repos.append(
@@ -91,7 +90,6 @@ class ConfigManager:
                         path=repo_path,
                         repo_id=repo_id,
                         repo_type=repo_type,
-                        ssh_key_path=ssh_key_path,
                     )
                 )
             self.save()
@@ -110,6 +108,9 @@ class ConfigManager:
             return repo_config
         raise DstackError("No repo config found")
 
+    def delete_repo_config(self, repo_id: str):
+        self.config.repos = [p for p in self.config.repos if p.repo_id != repo_id]
+
     @property
     def dstack_ssh_dir(self) -> Path:
         return self.dstack_dir / "ssh"

dstack/_internal/proxy/gateway/services/registry.py

@@ -152,6 +152,8 @@ async def register_replica(
             )
 
         if old_service.find_replica(replica_id) is not None:
+            # NOTE: as of 0.19.25, the dstack server relies on the exact text of this error.
+            # See dstack._internal.server.services.services.register_replica
             raise ProxyError(f"Replica {replica_id} already exists in service {old_service.fmt()}")
 
         service = old_service.with_replicas(old_service.replicas + (replica,))

dstack/_internal/server/app.py

@@ -29,6 +29,7 @@ from dstack._internal.server.routers import (
     files,
     fleets,
     gateways,
+    gpus,
     instances,
     logs,
     metrics,
@@ -204,6 +205,7 @@ def register_routes(app: FastAPI, ui: bool = True):
     app.include_router(repos.router)
     app.include_router(runs.root_router)
     app.include_router(runs.project_router)
+    app.include_router(gpus.project_router)
     app.include_router(metrics.router)
     app.include_router(logs.router)
     app.include_router(secrets.router)

dstack/_internal/server/background/tasks/process_fleets.py

@@ -15,6 +15,7 @@ from dstack._internal.server.models import (
     RunModel,
 )
 from dstack._internal.server.services.fleets import (
+    get_fleet_spec,
     is_fleet_empty,
     is_fleet_in_use,
 )
@@ -92,11 +93,18 @@ async def _process_fleets(session: AsyncSession, fleet_models: List[FleetModel])
 
 
 def _autodelete_fleet(fleet_model: FleetModel) -> bool:
-    # Currently all empty fleets are autodeleted.
-    # TODO: If fleets with `nodes: 0..` are supported, their deletion should be skipped.
     if is_fleet_in_use(fleet_model) or not is_fleet_empty(fleet_model):
         return False
 
+    fleet_spec = get_fleet_spec(fleet_model)
+    if (
+        fleet_model.status != FleetStatus.TERMINATING
+        and fleet_spec.configuration.nodes is not None
+        and (fleet_spec.configuration.nodes.min is None or fleet_spec.configuration.nodes.min == 0)
+    ):
+        # Empty fleets that allow 0 nodes should not be auto-deleted
+        return False
+
     logger.info("Automatic cleanup of an empty fleet %s", fleet_model.name)
     fleet_model.status = FleetStatus.TERMINATED
     fleet_model.deleted = True

dstack/_internal/server/background/tasks/process_running_jobs.py

@@ -32,6 +32,7 @@ from dstack._internal.core.models.runs import (
     JobSpec,
     JobStatus,
     JobTerminationReason,
+    ProbeSpec,
     Run,
     RunSpec,
     RunStatus,
@@ -70,6 +71,7 @@ from dstack._internal.server.services.repos import (
 from dstack._internal.server.services.runner import client
 from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel
 from dstack._internal.server.services.runs import (
+    is_job_ready,
     run_model_to_run,
 )
 from dstack._internal.server.services.secrets import get_project_secrets_mapping
@@ -140,6 +142,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
         select(JobModel)
         .where(JobModel.id == job_model.id)
         .options(joinedload(JobModel.instance).joinedload(InstanceModel.project))
+        .options(joinedload(JobModel.probes).load_only(ProbeModel.success_streak))
         .execution_options(populate_existing=True)
     )
     job_model = res.unique().scalar_one()
@@ -382,52 +385,21 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                         job_submission.age,
                     )
 
-    if (
-        initial_status != job_model.status
-        and job_model.status == JobStatus.RUNNING
-        and job_model.job_num == 0  # gateway connects only to the first node
-        and run.run_spec.configuration.type == "service"
-    ):
-        ssh_head_proxy: Optional[SSHConnectionParams] = None
-        ssh_head_proxy_private_key: Optional[str] = None
-        instance = common_utils.get_or_error(job_model.instance)
-        if instance.remote_connection_info is not None:
-            rci = RemoteConnectionInfo.__response__.parse_raw(instance.remote_connection_info)
-            if rci.ssh_proxy is not None:
-                ssh_head_proxy = rci.ssh_proxy
-                ssh_head_proxy_keys = common_utils.get_or_error(rci.ssh_proxy_keys)
-                ssh_head_proxy_private_key = ssh_head_proxy_keys[0].private
-        try:
-            await services.register_replica(
-                session,
-                run_model.gateway_id,
-                run,
-                job_model,
-                ssh_head_proxy,
-                ssh_head_proxy_private_key,
-            )
-        except GatewayError as e:
-            logger.warning(
-                "%s: failed to register service replica: %s, age=%s",
-                fmt(job_model),
-                e,
-                job_submission.age,
-            )
-            job_model.status = JobStatus.TERMINATING
-            job_model.termination_reason = JobTerminationReason.GATEWAY_ERROR
-        else:
-            for probe_num in range(len(job.job_spec.probes)):
-                session.add(
-                    ProbeModel(
-                        name=f"{job_model.job_name}-{probe_num}",
-                        job=job_model,
-                        probe_num=probe_num,
-                        due=common_utils.get_current_datetime(),
-                        success_streak=0,
-                        active=True,
-                    )
+    if initial_status != job_model.status and job_model.status == JobStatus.RUNNING:
+        job_model.probes = []
+        for probe_num in range(len(job.job_spec.probes)):
+            job_model.probes.append(
+                ProbeModel(
+                    name=f"{job_model.job_name}-{probe_num}",
+                    probe_num=probe_num,
+                    due=common_utils.get_current_datetime(),
+                    success_streak=0,
+                    active=True,
                 )
+            )
 
+    if job_model.status == JobStatus.RUNNING:
+        await _maybe_register_replica(session, run_model, run, job_model, job.job_spec.probes)
     if job_model.status == JobStatus.RUNNING:
         await _check_gpu_utilization(session, job_model, job)
 
@@ -455,8 +427,7 @@ async def _wait_for_instance_provisioning_data(job_model: JobModel):
 
     if job_model.instance.status == InstanceStatus.TERMINATED:
         job_model.status = JobStatus.TERMINATING
-        # TODO use WAITING_INSTANCE_LIMIT_EXCEEDED after 0.19.x
-        job_model.termination_reason = JobTerminationReason.FAILED_TO_START_DUE_TO_NO_CAPACITY
+        job_model.termination_reason = JobTerminationReason.WAITING_INSTANCE_LIMIT_EXCEEDED
         return
 
     job_model.job_provisioning_data = job_model.instance.job_provisioning_data
@@ -823,6 +794,55 @@ def _should_terminate_job_due_to_disconnect(job_model: JobModel) -> bool:
     )
 
 
+async def _maybe_register_replica(
+    session: AsyncSession,
+    run_model: RunModel,
+    run: Run,
+    job_model: JobModel,
+    probe_specs: Iterable[ProbeSpec],
+) -> None:
+    """
+    Register the replica represented by this job to receive service requests if it is ready.
+    """
+
+    if (
+        run.run_spec.configuration.type != "service"
+        or job_model.registered
+        or job_model.job_num != 0  # only the first job in the replica receives service requests
+        or not is_job_ready(job_model.probes, probe_specs)
+    ):
+        return
+
+    ssh_head_proxy: Optional[SSHConnectionParams] = None
+    ssh_head_proxy_private_key: Optional[str] = None
+    instance = common_utils.get_or_error(job_model.instance)
+    if instance.remote_connection_info is not None:
+        rci: RemoteConnectionInfo = RemoteConnectionInfo.__response__.parse_raw(
+            instance.remote_connection_info
+        )
+        if rci.ssh_proxy is not None:
+            ssh_head_proxy = rci.ssh_proxy
+            ssh_head_proxy_keys = common_utils.get_or_error(rci.ssh_proxy_keys)
+            ssh_head_proxy_private_key = ssh_head_proxy_keys[0].private
+    try:
+        await services.register_replica(
+            session,
+            run_model.gateway_id,
+            run,
+            job_model,
+            ssh_head_proxy,
+            ssh_head_proxy_private_key,
+        )
+    except GatewayError as e:
+        logger.warning(
+            "%s: failed to register service replica: %s",
+            fmt(job_model),
+            e,
+        )
+        job_model.status = JobStatus.TERMINATING
+        job_model.termination_reason = JobTerminationReason.GATEWAY_ERROR
+
+
 async def _check_gpu_utilization(session: AsyncSession, job_model: JobModel, job: Job) -> None:
     policy = job.job_spec.utilization_policy
     if policy is None:

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

@@ -23,7 +23,6 @@ from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import (
     InstanceModel,
     JobModel,
-    ProbeModel,
     ProjectModel,
     RunModel,
     UserModel,
@@ -37,7 +36,7 @@ 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 (
     fmt,
-    is_replica_ready,
+    is_replica_registered,
     process_terminating_run,
     retry_run_replica_jobs,
     run_model_to_run,
@@ -151,11 +150,6 @@ async def _process_run(session: AsyncSession, run_model: RunModel):
             .joinedload(JobModel.instance)
             .load_only(InstanceModel.fleet_id)
         )
-        .options(
-            selectinload(RunModel.jobs)
-            .joinedload(JobModel.probes)
-            .load_only(ProbeModel.success_streak)
-        )
         .execution_options(populate_existing=True)
     )
     run_model = res.unique().scalar_one()
@@ -402,6 +396,8 @@ async def _process_active_run(session: AsyncSession, run_model: RunModel):
 
         if new_status == RunStatus.PENDING:
             run_metrics.increment_pending_runs(run_model.project.name, run_spec.configuration.type)
+            # Unassign run from fleet so that the new fleet can be chosen when retrying
+            run_model.fleet = None
 
         run_model.status = new_status
         run_model.termination_reason = termination_reason
@@ -465,6 +461,9 @@ async def _handle_run_replicas(
         run_spec=run_spec,
     )
     if _has_out_of_date_replicas(run_model):
+        assert run_spec.configuration.type == "service", (
+            "Rolling deployment is only supported for services"
+        )
         non_terminated_replica_count = len(
             {j.replica_num for j in run_model.jobs if not j.status.is_finished()}
         )
@@ -479,22 +478,24 @@ async def _handle_run_replicas(
             )
 
         replicas_to_stop_count = 0
-        # stop any out-of-date replicas that are not ready
+        # stop any out-of-date replicas that are not registered
         replicas_to_stop_count += sum(
             any(j.deployment_num < run_model.deployment_num for j in jobs)
             and any(
                 j.status not in [JobStatus.TERMINATING] + JobStatus.finished_statuses()
                 for j in jobs
             )
-            and not is_replica_ready(jobs)
+            and not is_replica_registered(jobs)
+            for _, jobs in group_jobs_by_replica_latest(run_model.jobs)
+        )
+        # stop excessive registered out-of-date replicas, except those that are already `terminating`
+        non_terminating_registered_replicas_count = sum(
+            is_replica_registered(jobs) and all(j.status != JobStatus.TERMINATING for j in jobs)
             for _, jobs in group_jobs_by_replica_latest(run_model.jobs)
         )
-        ready_replica_count = sum(
-            is_replica_ready(jobs) for _, jobs in group_jobs_by_replica_latest(run_model.jobs)
+        replicas_to_stop_count += max(
+            0, non_terminating_registered_replicas_count - run_model.desired_replica_count
         )
-        if ready_replica_count > run_model.desired_replica_count:
-            # stop excessive ready out-of-date replicas
-            replicas_to_stop_count += ready_replica_count - run_model.desired_replica_count
         if replicas_to_stop_count:
             await scale_run_replicas(
                 session,
@@ -574,7 +575,7 @@ def _should_retry_job(run: Run, job: Job, job_model: JobModel) -> Optional[datet
 
     if (
         last_provisioned_submission.termination_reason is not None
-        and last_provisioned_submission.termination_reason.to_retry_event()
+        and JobTerminationReason(last_provisioned_submission.termination_reason).to_retry_event()
         in job.job_spec.retry.on_events
     ):
         return common.get_current_datetime() - last_provisioned_submission.last_processed_at