Support elastic fleets

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

@@ -383,10 +383,7 @@ def is_suitable_placement_group(
     ) -> 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,

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

@@ -263,10 +263,6 @@ def is_suitable_placement_group(
         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
 

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

@@ -448,10 +448,7 @@ def is_suitable_placement_group(
         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,

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

@@ -298,10 +298,7 @@ def is_suitable_placement_group(
         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(

src/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")]

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

@@ -20,6 +20,7 @@
 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 @@ def convert_model(cls, v: Optional[Union[AnyModel, str]]) -> Optional[AnyModel]:
         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 @@ def validate_gateway(
     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
 

src/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", {})

src/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
 
 

src/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
 

src/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"]

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

@@ -130,7 +130,7 @@ def __str__(self):
 
 
 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 @@ def _validate_arch(cls, v: Any) -> Any:
 
 
 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 @@ def _vendor_from_string(cls, v: str) -> gpuhunt.AcceleratorVendor:
 
 
 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 @@ def _parse(cls, v: Any) -> Any:
 
 
 class ResourcesSpec(CoreModel):
-    class Config:
+    class Config(CoreModel.Config):
         @staticmethod
         def schema_extra(schema: Dict[str, Any]):
             add_extra_schema_types(

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

@@ -444,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", {})

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

@@ -15,6 +15,7 @@
     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

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

@@ -396,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