dstack 0.19.34__py3-none-any.whl → 0.19.35__py3-none-any.whl

dstack/_internal/cli/services/configurators/run.py

@@ -941,7 +941,7 @@ def _warn_fleet_autocreated(api: APIClient, run: Run):
     if not fleet.spec.autocreated:
         return
     warn(
-        f"\nNo existing fleet matched, so the run created a new fleet [code]{fleet.name}[/code].\n"
+        f"\nThe run is using automatically created fleet [code]{fleet.name}[/code].\n"
         "Future dstack versions won't create fleets automatically.\n"
         "Create a fleet explicitly: https://dstack.ai/docs/concepts/fleets/"
     )

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

@@ -17,6 +17,7 @@ from cachetools import TTLCache, cachedmethod
 from gpuhunt import CPUArchitecture
 
 from dstack._internal import settings
+from dstack._internal.core.backends.base.models import JobConfiguration
 from dstack._internal.core.backends.base.offers import OfferModifier, filter_offers_by_requirements
 from dstack._internal.core.consts import (
     DSTACK_RUNNER_HTTP_PORT,
@@ -24,6 +25,7 @@ from dstack._internal.core.consts import (
     DSTACK_SHIM_HTTP_PORT,
 )
 from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.compute_groups import ComputeGroup, ComputeGroupProvisioningData
 from dstack._internal.core.models.configurations import LEGACY_REPO_DIR
 from dstack._internal.core.models.gateways import (
     GatewayComputeConfiguration,
@@ -324,6 +326,23 @@ class ComputeWithCreateInstanceSupport(ABC):
             ]
 
 
+class ComputeWithGroupProvisioningSupport(ABC):
+    @abstractmethod
+    def run_jobs(
+        self,
+        run: Run,
+        job_configurations: List[JobConfiguration],
+        instance_offer: InstanceOfferWithAvailability,
+        project_ssh_public_key: str,
+        project_ssh_private_key: str,
+    ) -> ComputeGroupProvisioningData:
+        pass
+
+    @abstractmethod
+    def terminate_compute_group(self, compute_group: ComputeGroup):
+        pass
+
+
 class ComputeWithPrivilegedSupport:
     """
     Must be subclassed to support runs with `privileged: true`.
@@ -680,7 +699,7 @@ def get_shim_env(
     backend_shim_env: Optional[Dict[str, str]] = None,
     arch: Optional[str] = None,
 ) -> Dict[str, str]:
-    log_level = "6"  # Trace
+    log_level = "5"  # Debug
     envs = {
         "DSTACK_SHIM_HOME": get_dstack_working_dir(base_path),
         "DSTACK_SHIM_HTTP_PORT": str(DSTACK_SHIM_HTTP_PORT),

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

@@ -1,4 +1,14 @@
 from pathlib import Path
+from typing import List
+
+from dstack._internal.core.models.common import CoreModel
+from dstack._internal.core.models.runs import Job
+from dstack._internal.core.models.volumes import Volume
+
+
+class JobConfiguration(CoreModel):
+    job: Job
+    volumes: List[Volume]
 
 
 def fill_data(values: dict, filename_field: str = "filename", data_field: str = "data") -> dict:

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

@@ -25,6 +25,7 @@ SUPPORTED_GPUHUNT_FLAGS = [
     "gcp-a4",
     "gcp-g4",
     "gcp-dws-calendar-mode",
+    "runpod-cluster",
 ]
 
 

dstack/_internal/core/backends/features.py

@@ -1,6 +1,7 @@
 from dstack._internal.core.backends.base.compute import (
     ComputeWithCreateInstanceSupport,
     ComputeWithGatewaySupport,
+    ComputeWithGroupProvisioningSupport,
     ComputeWithMultinodeSupport,
     ComputeWithPlacementGroupSupport,
     ComputeWithPrivateGatewaySupport,
@@ -39,6 +40,10 @@ BACKENDS_WITH_CREATE_INSTANCE_SUPPORT = _get_backends_with_compute_feature(
     configurator_classes=_configurator_classes,
     compute_feature_class=ComputeWithCreateInstanceSupport,
 )
+BACKENDS_WITH_GROUP_PROVISIONING_SUPPORT = _get_backends_with_compute_feature(
+    configurator_classes=_configurator_classes,
+    compute_feature_class=ComputeWithGroupProvisioningSupport,
+)
 BACKENDS_WITH_PRIVILEGED_SUPPORT = _get_backends_with_compute_feature(
     configurator_classes=_configurator_classes,
     compute_feature_class=ComputeWithPrivilegedSupport,

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

@@ -28,8 +28,11 @@ from dstack._internal.core.backends.base.offers import (
     get_offers_disk_modifier,
 )
 from dstack._internal.core.backends.nebius import resources
-from dstack._internal.core.backends.nebius.fabrics import get_suitable_infiniband_fabrics
-from dstack._internal.core.backends.nebius.models import NebiusConfig, NebiusServiceAccountCreds
+from dstack._internal.core.backends.nebius.models import (
+    NebiusConfig,
+    NebiusOfferBackendData,
+    NebiusServiceAccountCreds,
+)
 from dstack._internal.core.errors import (
     BackendError,
     NotYetTerminated,
@@ -281,12 +284,16 @@ class NebiusCompute(
         master_instance_offer: InstanceOffer,
     ) -> PlacementGroupProvisioningData:
         assert placement_group.configuration.placement_strategy == PlacementStrategy.CLUSTER
-        backend_data = NebiusPlacementGroupBackendData(cluster=None)
+        master_instance_offer_backend_data: NebiusOfferBackendData = (
+            NebiusOfferBackendData.__response__.parse_obj(master_instance_offer.backend_data)
+        )
+        fabrics = list(master_instance_offer_backend_data.fabrics)
+        if self.config.fabrics is not None:
+            fabrics = [f for f in fabrics if f in self.config.fabrics]
+        placement_group_backend_data = NebiusPlacementGroupBackendData(cluster=None)
         # Only create a Nebius cluster if the instance supports it.
         # For other instances, return dummy PlacementGroupProvisioningData.
-        if fabrics := get_suitable_infiniband_fabrics(
-            master_instance_offer, allowed_fabrics=self.config.fabrics
-        ):
+        if fabrics:
             fabric = random.choice(fabrics)
             op = resources.create_cluster(
                 self._sdk,
@@ -294,10 +301,13 @@ class NebiusCompute(
                 project_id=self._region_to_project_id[placement_group.configuration.region],
                 fabric=fabric,
             )
-            backend_data.cluster = NebiusClusterBackendData(id=op.resource_id, fabric=fabric)
+            placement_group_backend_data.cluster = NebiusClusterBackendData(
+                id=op.resource_id,
+                fabric=fabric,
+            )
         return PlacementGroupProvisioningData(
             backend=BackendType.NEBIUS,
-            backend_data=backend_data.json(),
+            backend_data=placement_group_backend_data.json(),
         )
 
     def delete_placement_group(self, placement_group: PlacementGroup) -> None:
@@ -317,16 +327,15 @@ class NebiusCompute(
         if placement_group.configuration.region != instance_offer.region:
             return False
         assert placement_group.provisioning_data is not None
-        backend_data = NebiusPlacementGroupBackendData.load(
+        placement_group_backend_data = NebiusPlacementGroupBackendData.load(
             placement_group.provisioning_data.backend_data
         )
+        instance_offer_backend_data: NebiusOfferBackendData = (
+            NebiusOfferBackendData.__response__.parse_obj(instance_offer.backend_data)
+        )
         return (
-            backend_data.cluster is None
-            or backend_data.cluster.fabric
-            in get_suitable_infiniband_fabrics(
-                instance_offer,
-                allowed_fabrics=None,  # enforced at cluster creation time, no need to enforce here
-            )
+            placement_group_backend_data.cluster is None
+            or placement_group_backend_data.cluster.fabric in instance_offer_backend_data.fabrics
         )
 
 
@@ -380,7 +389,10 @@ def _wait_for_instance(sdk: SDK, op: SDKOperation[Operation]) -> None:
         )
         time.sleep(WAIT_FOR_INSTANCE_UPDATE_INTERVAL)
         resources.LOOP.await_(
-            op.update(per_retry_timeout=resources.REQUEST_TIMEOUT, metadata=resources.REQUEST_MD)
+            op.update(
+                per_retry_timeout=resources.REQUEST_TIMEOUT,
+                auth_options=resources.REQUEST_AUTH_OPTIONS,
+            )
         )
 
 

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

@@ -10,7 +10,6 @@ from dstack._internal.core.backends.base.configurator import (
 )
 from dstack._internal.core.backends.nebius import resources
 from dstack._internal.core.backends.nebius.backend import NebiusBackend
-from dstack._internal.core.backends.nebius.fabrics import get_all_infiniband_fabrics
 from dstack._internal.core.backends.nebius.models import (
     NebiusBackendConfig,
     NebiusBackendConfigWithCreds,
@@ -19,6 +18,7 @@ from dstack._internal.core.backends.nebius.models import (
     NebiusServiceAccountCreds,
     NebiusStoredConfig,
 )
+from dstack._internal.core.backends.nebius.resources import get_all_infiniband_fabrics
 from dstack._internal.core.errors import BackendError, ServerClientError
 from dstack._internal.core.models.backends.base import BackendType
 

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

@@ -179,3 +179,7 @@ class NebiusConfig(NebiusStoredConfig):
     """
 
     creds: AnyNebiusCreds
+
+
+class NebiusOfferBackendData(CoreModel):
+    fabrics: set[str] = set()

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

@@ -8,7 +8,6 @@ from contextlib import contextmanager
 from tempfile import NamedTemporaryFile
 from typing import Dict, Optional
 
-from nebius.aio.authorization.options import options_to_metadata
 from nebius.aio.operation import Operation as SDKOperation
 from nebius.aio.service_error import RequestError, StatusCode
 from nebius.aio.token.renewable import OPTION_RENEW_REQUEST_TIMEOUT, OPTION_RENEW_SYNCHRONOUS
@@ -50,11 +49,14 @@ from nebius.api.nebius.vpc.v1 import ListSubnetsRequest, Subnet, SubnetServiceCl
 from nebius.sdk import SDK
 
 from dstack._internal.core.backends.base.configurator import raise_invalid_credentials_error
+from dstack._internal.core.backends.base.offers import get_catalog_offers
 from dstack._internal.core.backends.nebius.models import (
     DEFAULT_PROJECT_NAME_PREFIX,
+    NebiusOfferBackendData,
     NebiusServiceAccountCreds,
 )
 from dstack._internal.core.errors import BackendError, NoCapacityError
+from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.utils.event_loop import DaemonEventLoop
 from dstack._internal.utils.logging import get_logger
 
@@ -66,13 +68,11 @@ from dstack._internal.utils.logging import get_logger
 LOOP = DaemonEventLoop()
 # Pass a timeout to all methods to avoid infinite waiting
 REQUEST_TIMEOUT = 10
-# Pass REQUEST_MD to all methods to avoid infinite retries in case of invalid credentials
-REQUEST_MD = options_to_metadata(
-    {
-        OPTION_RENEW_SYNCHRONOUS: "true",
-        OPTION_RENEW_REQUEST_TIMEOUT: "5",
-    }
-)
+# Pass REQUEST_AUTH_OPTIONS to all methods to avoid infinite retries in case of invalid credentials
+REQUEST_AUTH_OPTIONS = {
+    OPTION_RENEW_SYNCHRONOUS: "true",
+    OPTION_RENEW_REQUEST_TIMEOUT: "5",
+}
 
 # disables log messages about errors such as invalid creds or expired timeouts
 logging.getLogger("nebius").setLevel(logging.CRITICAL)
@@ -120,7 +120,9 @@ def wait_for_operation(
         if time.monotonic() + interval > deadline:
             raise TimeoutError(f"Operation {op.id} wait timeout")
         time.sleep(interval)
-        LOOP.await_(op.update(per_retry_timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD))
+        LOOP.await_(
+            op.update(per_retry_timeout=REQUEST_TIMEOUT, auth_options=REQUEST_AUTH_OPTIONS)
+        )
 
 
 def get_region_to_project_id_map(
@@ -156,7 +158,9 @@ def validate_regions(configured: set[str], available: set[str]) -> None:
 def list_tenant_projects(sdk: SDK) -> Sequence[Container]:
     tenants = LOOP.await_(
         TenantServiceClient(sdk).list(
-            ListTenantsRequest(), per_retry_timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD
+            ListTenantsRequest(),
+            per_retry_timeout=REQUEST_TIMEOUT,
+            auth_options=REQUEST_AUTH_OPTIONS,
         )
     )
     if len(tenants.items) != 1:
@@ -166,7 +170,7 @@ def list_tenant_projects(sdk: SDK) -> Sequence[Container]:
         ProjectServiceClient(sdk).list(
             ListProjectsRequest(parent_id=tenant_id, page_size=999),
             per_retry_timeout=REQUEST_TIMEOUT,
-            metadata=REQUEST_MD,
+            auth_options=REQUEST_AUTH_OPTIONS,
         )
     )
     return projects.items
@@ -240,7 +244,7 @@ def get_default_subnet(sdk: SDK, project_id: str) -> Subnet:
         SubnetServiceClient(sdk).list(
             ListSubnetsRequest(parent_id=project_id, page_size=999),
             per_retry_timeout=REQUEST_TIMEOUT,
-            metadata=REQUEST_MD,
+            auth_options=REQUEST_AUTH_OPTIONS,
         )
     )
     for subnet in subnets.items:
@@ -249,6 +253,17 @@ def get_default_subnet(sdk: SDK, project_id: str) -> Subnet:
     raise BackendError(f"Could not find default subnet in project {project_id}")
 
 
+def get_all_infiniband_fabrics() -> set[str]:
+    offers = get_catalog_offers(backend=BackendType.NEBIUS)
+    result = set()
+    for offer in offers:
+        backend_data: NebiusOfferBackendData = NebiusOfferBackendData.__response__.parse_obj(
+            offer.backend_data
+        )
+        result |= backend_data.fabrics
+    return result
+
+
 def create_disk(
     sdk: SDK, name: str, project_id: str, size_mib: int, image_family: str, labels: Dict[str, str]
 ) -> SDKOperation[Operation]:
@@ -267,14 +282,18 @@ def create_disk(
     )
     with wrap_capacity_errors():
         return LOOP.await_(
-            client.create(request, per_retry_timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD)
+            client.create(
+                request, per_retry_timeout=REQUEST_TIMEOUT, auth_options=REQUEST_AUTH_OPTIONS
+            )
         )
 
 
 def delete_disk(sdk: SDK, disk_id: str) -> None:
     LOOP.await_(
         DiskServiceClient(sdk).delete(
-            DeleteDiskRequest(id=disk_id), per_retry_timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD
+            DeleteDiskRequest(id=disk_id),
+            per_retry_timeout=REQUEST_TIMEOUT,
+            auth_options=REQUEST_AUTH_OPTIONS,
         )
     )
 
@@ -325,7 +344,9 @@ def create_instance(
     )
     with wrap_capacity_errors():
         return LOOP.await_(
-            client.create(request, per_retry_timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD)
+            client.create(
+                request, per_retry_timeout=REQUEST_TIMEOUT, auth_options=REQUEST_AUTH_OPTIONS
+            )
         )
 
 
@@ -334,7 +355,7 @@ def get_instance(sdk: SDK, instance_id: str) -> Instance:
         InstanceServiceClient(sdk).get(
             GetInstanceRequest(id=instance_id),
             per_retry_timeout=REQUEST_TIMEOUT,
-            metadata=REQUEST_MD,
+            auth_options=REQUEST_AUTH_OPTIONS,
         )
     )
 
@@ -344,7 +365,7 @@ def delete_instance(sdk: SDK, instance_id: str) -> SDKOperation[Operation]:
         InstanceServiceClient(sdk).delete(
             DeleteInstanceRequest(id=instance_id),
             per_retry_timeout=REQUEST_TIMEOUT,
-            metadata=REQUEST_MD,
+            auth_options=REQUEST_AUTH_OPTIONS,
         )
     )
 
@@ -358,17 +379,17 @@ def create_cluster(sdk: SDK, name: str, project_id: str, fabric: str) -> SDKOper
                     spec=GpuClusterSpec(infiniband_fabric=fabric),
                 ),
                 per_retry_timeout=REQUEST_TIMEOUT,
-                metadata=REQUEST_MD,
+                auth_options=REQUEST_AUTH_OPTIONS,
             )
         )
 
 
 def delete_cluster(sdk: SDK, cluster_id: str) -> None:
-    return LOOP.await_(
+    LOOP.await_(
         GpuClusterServiceClient(sdk).delete(
             DeleteGpuClusterRequest(id=cluster_id),
             per_retry_timeout=REQUEST_TIMEOUT,
-            metadata=REQUEST_MD,
+            auth_options=REQUEST_AUTH_OPTIONS,
         )
     )