dstack 0.19.1__py3-none-any.whl → 0.19.3__py3-none-any.whl

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

@@ -1,8 +1,12 @@
 import json
 
+from datacrunch import DataCrunchClient
+from datacrunch.exceptions import APIException
+
 from dstack._internal.core.backends.base.configurator import (
     BackendRecord,
     Configurator,
+    raise_invalid_credentials_error,
 )
 from dstack._internal.core.backends.datacrunch.backend import DataCrunchBackend
 from dstack._internal.core.backends.datacrunch.models import (
@@ -17,13 +21,6 @@ from dstack._internal.core.models.backends.base import (
     BackendType,
 )
 
-REGIONS = [
-    "FIN-01",
-    "ICE-01",
-]
-
-DEFAULT_REGION = "FIN-01"
-
 
 class DataCrunchConfigurator(Configurator):
     TYPE = BackendType.DATACRUNCH
@@ -32,14 +29,11 @@ class DataCrunchConfigurator(Configurator):
     def validate_config(
         self, config: DataCrunchBackendConfigWithCreds, default_creds_enabled: bool
     ):
-        # FIXME: validate datacrunch creds
-        return
+        self._validate_creds(config.creds)
 
     def create_backend(
         self, project_name: str, config: DataCrunchBackendConfigWithCreds
     ) -> BackendRecord:
-        if config.regions is None:
-            config.regions = REGIONS
         return BackendRecord(
             config=DataCrunchStoredConfig(
                 **DataCrunchBackendConfig.__response__.parse_obj(config).dict()
@@ -64,3 +58,14 @@ class DataCrunchConfigurator(Configurator):
             **json.loads(record.config),
             creds=DataCrunchCreds.parse_raw(record.auth),
         )
+
+    def _validate_creds(self, creds: DataCrunchCreds):
+        try:
+            DataCrunchClient(
+                client_id=creds.client_id,
+                client_secret=creds.client_secret,
+            )
+        except APIException as e:
+            if e.code == "unauthorized_request":
+                raise_invalid_credentials_error(fields=[["creds", "api_key"]])
+            raise

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

@@ -1,7 +1,7 @@
 import concurrent.futures
 import json
 from collections import defaultdict
-from typing import Callable, Dict, List, Literal, Optional
+from typing import Callable, Dict, List, Literal, Optional, Tuple
 
 import google.api_core.exceptions
 import google.cloud.compute_v1 as compute_v1
@@ -10,11 +10,13 @@ from gpuhunt import KNOWN_TPUS
 
 import dstack._internal.core.backends.gcp.auth as auth
 import dstack._internal.core.backends.gcp.resources as gcp_resources
+from dstack import version
 from dstack._internal.core.backends.base.compute import (
     Compute,
     ComputeWithCreateInstanceSupport,
     ComputeWithGatewaySupport,
     ComputeWithMultinodeSupport,
+    ComputeWithPlacementGroupSupport,
     ComputeWithVolumeSupport,
     generate_unique_gateway_instance_name,
     generate_unique_instance_name,
@@ -25,11 +27,13 @@ from dstack._internal.core.backends.base.compute import (
     merge_tags,
 )
 from dstack._internal.core.backends.base.offers import get_catalog_offers
+from dstack._internal.core.backends.gcp.features import tcpx as tcpx_features
 from dstack._internal.core.backends.gcp.models import GCPConfig
 from dstack._internal.core.errors import (
     ComputeError,
     ComputeResourceNotFoundError,
     NoCapacityError,
+    PlacementGroupInUseError,
     ProvisioningError,
 )
 from dstack._internal.core.models.backends.base import BackendType
@@ -46,6 +50,7 @@ from dstack._internal.core.models.instances import (
     InstanceType,
     Resources,
 )
+from dstack._internal.core.models.placement import PlacementGroup, PlacementGroupProvisioningData
 from dstack._internal.core.models.resources import Memory, Range
 from dstack._internal.core.models.runs import JobProvisioningData, Requirements
 from dstack._internal.core.models.volumes import (
@@ -74,6 +79,7 @@ class GCPVolumeDiskBackendData(CoreModel):
 class GCPCompute(
     ComputeWithCreateInstanceSupport,
     ComputeWithMultinodeSupport,
+    ComputeWithPlacementGroupSupport,
     ComputeWithGatewaySupport,
     ComputeWithVolumeSupport,
     Compute,
@@ -89,6 +95,9 @@ class GCPCompute(
         self.routers_client = compute_v1.RoutersClient(credentials=self.credentials)
         self.tpu_client = tpu_v2.TpuClient(credentials=self.credentials)
         self.disk_client = compute_v1.DisksClient(credentials=self.credentials)
+        self.resource_policies_client = compute_v1.ResourcePoliciesClient(
+            credentials=self.credentials
+        )
 
     def get_offers(
         self, requirements: Optional[Requirements] = None
@@ -183,6 +192,19 @@ class GCPCompute(
             config=self.config,
             region=instance_offer.region,
         )
+        extra_subnets = _get_extra_subnets(
+            subnetworks_client=self.subnetworks_client,
+            config=self.config,
+            region=instance_offer.region,
+            instance_type_name=instance_offer.instance.name,
+        )
+        placement_policy = None
+        if instance_config.placement_group_name is not None:
+            placement_policy = gcp_resources.get_placement_policy_resource_name(
+                project_id=self.config.project_id,
+                region=instance_offer.region,
+                placement_policy=instance_config.placement_group_name,
+            )
         labels = {
             "owner": "dstack",
             "dstack_project": instance_config.project_name.lower(),
@@ -259,8 +281,9 @@ class GCPCompute(
             request.project = self.config.project_id
             request.instance_resource = gcp_resources.create_instance_struct(
                 disk_size=disk_size,
-                image_id=gcp_resources.get_image_id(
-                    len(instance_offer.instance.resources.gpus) > 0,
+                image_id=_get_image_id(
+                    instance_type_name=instance_offer.instance.name,
+                    cuda=len(instance_offer.instance.resources.gpus) > 0,
                 ),
                 machine_type=instance_offer.instance.name,
                 accelerators=gcp_resources.get_accelerators(
@@ -269,7 +292,12 @@ class GCPCompute(
                     gpus=instance_offer.instance.resources.gpus,
                 ),
                 spot=instance_offer.instance.resources.spot,
-                user_data=get_user_data(authorized_keys),
+                user_data=get_user_data(
+                    authorized_keys,
+                    backend_specific_commands=_get_backend_specific_commands(
+                        instance_offer.instance.name
+                    ),
+                ),
                 authorized_keys=authorized_keys,
                 labels=labels,
                 tags=[gcp_resources.DSTACK_INSTANCE_TAG],
@@ -278,7 +306,9 @@ class GCPCompute(
                 service_account=self.config.vm_service_account,
                 network=self.config.vpc_resource_name,
                 subnetwork=subnetwork,
+                extra_subnetworks=extra_subnets,
                 allocate_public_ip=allocate_public_ip,
+                placement_policy=placement_policy,
             )
             try:
                 # GCP needs some time to return an error in case of no capacity (< 30s).
@@ -371,6 +401,43 @@ class GCPCompute(
             f"Failed to get instance IP address. Instance status: {instance.status}"
         )
 
+    def create_placement_group(
+        self,
+        placement_group: PlacementGroup,
+    ) -> PlacementGroupProvisioningData:
+        policy = compute_v1.ResourcePolicy(
+            name=placement_group.name,
+            region=placement_group.configuration.region,
+            group_placement_policy=compute_v1.ResourcePolicyGroupPlacementPolicy(
+                availability_domain_count=1,
+                collocation="COLLOCATED",
+            ),
+        )
+        self.resource_policies_client.insert(
+            project=self.config.project_id,
+            region=placement_group.configuration.region,
+            resource_policy_resource=policy,
+        )
+        return PlacementGroupProvisioningData(backend=BackendType.GCP)
+
+    def delete_placement_group(
+        self,
+        placement_group: PlacementGroup,
+    ):
+        try:
+            operation = self.resource_policies_client.delete(
+                project=self.config.project_id,
+                region=placement_group.configuration.region,
+                resource_policy=placement_group.name,
+            )
+            operation.result()  # Wait for operation to complete
+        except google.api_core.exceptions.NotFound:
+            logger.debug("Placement group %s not found", placement_group.name)
+        except google.api_core.exceptions.BadRequest as e:
+            if "is already being used by" in e.message:
+                raise PlacementGroupInUseError()
+            raise
+
     def create_gateway(
         self,
         configuration: GatewayComputeConfiguration,
@@ -412,7 +479,7 @@ class GCPCompute(
         request.project = self.config.project_id
         request.instance_resource = gcp_resources.create_instance_struct(
             disk_size=10,
-            image_id=gcp_resources.get_gateway_image_id(),
+            image_id=_get_gateway_image_id(),
             machine_type="e2-small",
             accelerators=[],
             spot=False,
@@ -681,21 +748,6 @@ class GCPCompute(
         )
 
 
-def _get_vpc_subnet(
-    subnetworks_client: compute_v1.SubnetworksClient,
-    config: GCPConfig,
-    region: str,
-) -> Optional[str]:
-    if config.vpc_name is None:
-        return None
-    return gcp_resources.get_vpc_subnet_or_error(
-        subnetworks_client=subnetworks_client,
-        vpc_project_id=config.vpc_project_id or config.project_id,
-        vpc_name=config.vpc_name,
-        region=region,
-    )
-
-
 def _supported_instances_and_zones(
     regions: List[str],
 ) -> Optional[Callable[[InstanceOffer], bool]]:
@@ -754,6 +806,74 @@ def _unique_instance_name(instance: InstanceType) -> str:
     return f"{name}-{gpu.name}-{gpu.memory_mib}"
 
 
+def _get_vpc_subnet(
+    subnetworks_client: compute_v1.SubnetworksClient,
+    config: GCPConfig,
+    region: str,
+) -> Optional[str]:
+    if config.vpc_name is None:
+        return None
+    return gcp_resources.get_vpc_subnet_or_error(
+        subnetworks_client=subnetworks_client,
+        vpc_project_id=config.vpc_project_id or config.project_id,
+        vpc_name=config.vpc_name,
+        region=region,
+    )
+
+
+def _get_extra_subnets(
+    subnetworks_client: compute_v1.SubnetworksClient,
+    config: GCPConfig,
+    region: str,
+    instance_type_name: str,
+) -> List[Tuple[str, str]]:
+    if config.extra_vpcs is None:
+        return []
+    if instance_type_name != "a3-megagpu-8g":
+        return []
+    extra_subnets = []
+    for vpc_name in config.extra_vpcs:
+        subnet = gcp_resources.get_vpc_subnet_or_error(
+            subnetworks_client=subnetworks_client,
+            vpc_project_id=config.vpc_project_id or config.project_id,
+            vpc_name=vpc_name,
+            region=region,
+        )
+        vpc_resource_name = gcp_resources.vpc_name_to_vpc_resource_name(
+            project_id=config.vpc_project_id or config.project_id,
+            vpc_name=vpc_name,
+        )
+        extra_subnets.append((vpc_resource_name, subnet))
+    return extra_subnets[:8]
+
+
+def _get_image_id(instance_type_name: str, cuda: bool) -> str:
+    if instance_type_name == "a3-megagpu-8g":
+        image_name = "dstack-a3mega-5"
+    elif cuda:
+        image_name = f"dstack-cuda-{version.base_image}"
+    else:
+        image_name = f"dstack-{version.base_image}"
+    image_name = image_name.replace(".", "-")
+    return f"projects/dstack/global/images/{image_name}"
+
+
+def _get_gateway_image_id() -> str:
+    return "projects/ubuntu-os-cloud/global/images/ubuntu-2204-jammy-v20230714"
+
+
+def _get_backend_specific_commands(instance_type_name: str) -> List[str]:
+    if instance_type_name == "a3-megagpu-8g":
+        return tcpx_features.get_backend_specific_commands_tcpxo()
+    return []
+
+
+def _get_volume_price(size: int) -> float:
+    # https://cloud.google.com/compute/disks-image-pricing#persistentdisk
+    # The price is different in different regions. Take max across supported regions.
+    return size * 0.12
+
+
 def _get_tpu_startup_script(authorized_keys: List[str]) -> str:
     commands = get_shim_commands(
         authorized_keys=authorized_keys, is_privileged=True, pjrt_device="TPU"
@@ -805,12 +925,6 @@ def _is_single_host_tpu(instance_name: str) -> bool:
         return False
 
 
-def _get_volume_price(size: int) -> float:
-    # https://cloud.google.com/compute/disks-image-pricing#persistentdisk
-    # The price is different in different regions. Take max across supported regions.
-    return size * 0.12
-
-
 def _get_tpu_data_disks(
     project_id: str, volumes: Optional[List[Volume]]
 ) -> List[tpu_v2.AttachedDisk]:

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

@@ -199,3 +199,5 @@ class GCPConfigurator(Configurator):
             )
         except BackendError as e:
             raise ServerClientError(e.args[0])
+        # Not checking config.extra_vpc so that users are not required to configure subnets for all regions
+        # but only for regions they intend to use. Validation will be done on provisioning.

dstack/_internal/core/backends/gcp/features/tcpx.py

@@ -0,0 +1,34 @@
+from typing import List
+
+
+def get_backend_specific_commands_tcpxo() -> List[str]:
+    return [
+        "modprobe import-helper",
+        "gcloud -q auth configure-docker us-docker.pkg.dev",
+        # Install the nccl, nccl-net lib into /var/lib/tcpxo/lib64/.
+        (
+            "docker run --rm "
+            "--name nccl-installer "
+            "--pull=never "
+            "--network=host "
+            "--volume /var/lib:/var/lib "
+            "us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/nccl-plugin-gpudirecttcpx-dev:v1.0.8-1 "
+            "install --install-nccl"
+        ),
+        # Start FasTrak receive-datapath-manager
+        (
+            "docker run "
+            "--name receive-datapath-manager "
+            "--detach "
+            "--pull=never "
+            "--cap-add=NET_ADMIN "
+            "--network=host "
+            "--privileged "
+            "--gpus all "
+            "--volume /usr/lib32:/usr/local/nvidia/lib64 "
+            "--volume /dev/dmabuf_import_helper:/dev/dmabuf_import_helper "
+            "--env LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu "
+            "us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/tcpgpudmarxd-dev:v1.0.14 "
+            "--num_hops=2 --num_nics=8 --uid= --alsologtostderr"
+        ),
+    ]

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

@@ -33,7 +33,19 @@ class GCPBackendConfig(CoreModel):
     regions: Annotated[
         Optional[List[str]], Field(description="The list of GCP regions. Omit to use all regions")
     ] = None
-    vpc_name: Annotated[Optional[str], Field(description="The name of a custom VPC")] = None
+    vpc_name: Annotated[
+        Optional[str],
+        Field(description="The name of a custom VPC. If not specified, the default VPC is used"),
+    ] = None
+    extra_vpcs: Annotated[
+        Optional[List[str]],
+        Field(
+            description=(
+                "The names of additional VPCs used for GPUDirect. Specify eight VPCs to maximize bandwidth."
+                " Each VPC must have a subnet and a firewall rule allowing internal traffic across all subnets"
+            )
+        ),
+    ] = None
     vpc_project_id: Annotated[
         Optional[str],
         Field(description="The shared VPC hosted project ID. Required for shared VPC only"),

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

@@ -1,6 +1,6 @@
 import concurrent.futures
 import re
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 
 import google.api_core.exceptions
 import google.cloud.compute_v1 as compute_v1
@@ -8,7 +8,6 @@ from google.api_core.extended_operation import ExtendedOperation
 from google.api_core.operation import Operation
 from google.cloud import tpu_v2
 
-import dstack.version as version
 from dstack._internal.core.errors import BackendError, ComputeError
 from dstack._internal.core.models.instances import Gpu
 from dstack._internal.utils.common import remove_prefix
@@ -54,12 +53,16 @@ def check_vpc(
     if shared_vpc_project_id:
         vpc_project_id = shared_vpc_project_id
     try:
+        usable_subnets = list_project_usable_subnets(
+            subnetworks_client=subnetworks_client, project_id=vpc_project_id
+        )
         for region in regions:
             get_vpc_subnet_or_error(
                 subnetworks_client=subnetworks_client,
                 vpc_project_id=vpc_project_id,
                 vpc_name=vpc_name,
                 region=region,
+                usable_subnets=usable_subnets,
             )
     except google.api_core.exceptions.NotFound:
         raise ComputeError(f"Failed to find VPC project {vpc_project_id}")
@@ -117,26 +120,19 @@ def create_instance_struct(
     service_account: Optional[str] = None,
     network: str = "global/networks/default",
     subnetwork: Optional[str] = None,
+    extra_subnetworks: Optional[List[Tuple[str, str]]] = None,
     allocate_public_ip: bool = True,
+    placement_policy: Optional[str] = None,
 ) -> compute_v1.Instance:
-    network_interface = compute_v1.NetworkInterface()
-    network_interface.network = network
-    if subnetwork is not None:
-        network_interface.subnetwork = subnetwork
-
-    if allocate_public_ip:
-        access = compute_v1.AccessConfig()
-        access.type_ = compute_v1.AccessConfig.Type.ONE_TO_ONE_NAT.name
-        access.name = "External NAT"
-        access.network_tier = access.NetworkTier.PREMIUM.name
-        network_interface.access_configs = [access]
-    else:
-        network_interface.access_configs = []
-
     instance = compute_v1.Instance()
-    instance.network_interfaces = [network_interface]
     instance.name = instance_name
     instance.machine_type = f"zones/{zone}/machineTypes/{machine_type}"
+    instance.network_interfaces = _get_network_interfaces(
+        network=network,
+        subnetwork=subnetwork,
+        allocate_public_ip=allocate_public_ip,
+        extra_subnetworks=extra_subnetworks,
+    )
 
     disk = compute_v1.AttachedDisk()
     disk.auto_delete = True
@@ -160,6 +156,9 @@ def create_instance_struct(
         # Attachable GPUs, H100, A100, and L4
         instance.scheduling.on_host_maintenance = "TERMINATE"
 
+    if placement_policy is not None:
+        instance.resource_policies = [placement_policy]
+
     if spot:
         instance.scheduling = compute_v1.Scheduling()
         instance.scheduling.provisioning_model = compute_v1.Scheduling.ProvisioningModel.SPOT.name
@@ -187,18 +186,42 @@ def create_instance_struct(
     return instance
 
 
-def get_image_id(cuda: bool) -> str:
-    if not cuda:
-        image_name = f"dstack-{version.base_image}"
+def _get_network_interfaces(
+    network: str,
+    subnetwork: Optional[str],
+    allocate_public_ip: bool,
+    extra_subnetworks: Optional[List[Tuple[str, str]]],
+) -> List[compute_v1.NetworkInterface]:
+    network_interface = compute_v1.NetworkInterface()
+    network_interface.network = network
+    if subnetwork is not None:
+        network_interface.subnetwork = subnetwork
+    if allocate_public_ip:
+        access = compute_v1.AccessConfig()
+        access.type_ = compute_v1.AccessConfig.Type.ONE_TO_ONE_NAT.name
+        access.name = "External NAT"
+        access.network_tier = access.NetworkTier.PREMIUM.name
+        network_interface.access_configs = [access]
     else:
-        image_name = f"dstack-cuda-{version.base_image}"
-    image_name = image_name.replace(".", "-")
+        network_interface.access_configs = []
 
-    return f"projects/dstack/global/images/{image_name}"
+    network_interfaces = [network_interface]
+    for network, subnetwork in extra_subnetworks or []:
+        network_interfaces.append(
+            compute_v1.NetworkInterface(
+                network=network,
+                subnetwork=subnetwork,
+            )
+        )
+    return network_interfaces
 
 
-def get_gateway_image_id() -> str:
-    return "projects/ubuntu-os-cloud/global/images/ubuntu-2204-jammy-v20230714"
+def list_project_usable_subnets(
+    subnetworks_client: compute_v1.SubnetworksClient,
+    project_id: str,
+) -> List[compute_v1.UsableSubnetwork]:
+    request = compute_v1.ListUsableSubnetworksRequest(project=project_id)
+    return [s for s in subnetworks_client.list_usable(request=request)]
 
 
 def get_vpc_subnet_or_error(
@@ -206,13 +229,15 @@ def get_vpc_subnet_or_error(
     vpc_project_id: str,
     vpc_name: str,
     region: str,
+    usable_subnets: Optional[List[compute_v1.UsableSubnetwork]] = None,
 ) -> str:
     """
     Returns resource name of any usable subnet in a given VPC
     (e.g. "projects/example-project/regions/europe-west4/subnetworks/example-subnet")
     """
-    request = compute_v1.ListUsableSubnetworksRequest(project=vpc_project_id)
-    for subnet in subnetworks_client.list_usable(request=request):
+    if usable_subnets is None:
+        usable_subnets = list_project_usable_subnets(subnetworks_client, vpc_project_id)
+    for subnet in usable_subnets:
         network_name = subnet.network.split("/")[-1]
         subnet_url = subnet.subnetwork
         subnet_resource_name = remove_prefix(subnet_url, "https://www.googleapis.com/compute/v1/")
@@ -410,3 +435,15 @@ def wait_for_operation(operation: Operation, verbose_name: str = "operation", ti
 
 def full_resource_name_to_name(full_resource_name: str) -> str:
     return full_resource_name.split("/")[-1]
+
+
+def vpc_name_to_vpc_resource_name(project_id: str, vpc_name: str) -> str:
+    return f"projects/{project_id}/global/networks/{vpc_name}"
+
+
+def get_placement_policy_resource_name(
+    project_id: str,
+    region: str,
+    placement_policy: str,
+) -> str:
+    return f"projects/{project_id}/regions/{region}/resourcePolicies/{placement_policy}"

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

@@ -39,7 +39,7 @@ class LambdaCompute(
     ) -> List[InstanceOfferWithAvailability]:
         offers = get_catalog_offers(
             backend=BackendType.LAMBDA,
-            locations=self.config.regions,
+            locations=self.config.regions or None,
             requirements=requirements,
         )
         offers_with_availability = self._get_offers_with_availability(offers)
@@ -90,7 +90,7 @@ class LambdaCompute(
         if instance_info is not None and instance_info["status"] != "booting":
             provisioning_data.hostname = instance_info["ip"]
             commands = get_shim_commands(authorized_keys=[project_ssh_public_key])
-            # shim is asssumed to be run under root
+            # shim is assumed to be run under root
             launch_command = "sudo sh -c '" + "&& ".join(commands) + "'"
             thread = Thread(
                 target=_start_runner,
@@ -119,8 +119,6 @@ class LambdaCompute(
         }
         availability_offers = []
         for offer in offers:
-            if offer.region not in self.config.regions:
-                continue
             availability = InstanceAvailability.NOT_AVAILABLE
             if offer.region in instance_availability.get(offer.instance.name, []):
                 availability = InstanceAvailability.AVAILABLE

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

@@ -19,25 +19,6 @@ from dstack._internal.core.models.backends.base import (
     BackendType,
 )
 
-REGIONS = [
-    "us-south-1",
-    "us-south-2",
-    "us-south-3",
-    "us-west-2",
-    "us-west-1",
-    "us-midwest-1",
-    "us-west-3",
-    "us-east-1",
-    "us-east-2",
-    "europe-central-1",
-    "asia-south-1",
-    "me-west-1",
-    "asia-northeast-1",
-    "asia-northeast-2",
-]
-
-DEFAULT_REGION = "us-east-1"
-
 
 class LambdaConfigurator(Configurator):
     TYPE = BackendType.LAMBDA
@@ -49,8 +30,6 @@ class LambdaConfigurator(Configurator):
     def create_backend(
         self, project_name: str, config: LambdaBackendConfigWithCreds
     ) -> BackendRecord:
-        if config.regions is None:
-            config.regions = REGIONS
         return BackendRecord(
             config=LambdaStoredConfig(
                 **LambdaBackendConfig.__response__.parse_obj(config).dict()

dstack/_internal/core/backends/models.py

@@ -34,6 +34,11 @@ from dstack._internal.core.backends.lambdalabs.models import (
     LambdaBackendConfig,
     LambdaBackendConfigWithCreds,
 )
+from dstack._internal.core.backends.nebius.models import (
+    NebiusBackendConfig,
+    NebiusBackendConfigWithCreds,
+    NebiusBackendFileConfigWithCreds,
+)
 from dstack._internal.core.backends.oci.models import (
     OCIBackendConfig,
     OCIBackendConfigWithCreds,
@@ -65,6 +70,7 @@ AnyBackendConfigWithoutCreds = Union[
     GCPBackendConfig,
     KubernetesBackendConfig,
     LambdaBackendConfig,
+    NebiusBackendConfig,
     OCIBackendConfig,
     RunpodBackendConfig,
     TensorDockBackendConfig,
@@ -86,6 +92,7 @@ AnyBackendConfigWithCreds = Union[
     KubernetesBackendConfigWithCreds,
     LambdaBackendConfigWithCreds,
     OCIBackendConfigWithCreds,
+    NebiusBackendConfigWithCreds,
     RunpodBackendConfigWithCreds,
     TensorDockBackendConfigWithCreds,
     VastAIBackendConfigWithCreds,
@@ -105,6 +112,7 @@ AnyBackendFileConfigWithCreds = Union[
     KubernetesBackendFileConfigWithCreds,
     LambdaBackendConfigWithCreds,
     OCIBackendConfigWithCreds,
+    NebiusBackendFileConfigWithCreds,
     RunpodBackendConfigWithCreds,
     TensorDockBackendConfigWithCreds,
     VastAIBackendConfigWithCreds,

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

@@ -0,0 +1,16 @@
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.nebius.compute import NebiusCompute
+from dstack._internal.core.backends.nebius.models import NebiusConfig
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class NebiusBackend(Backend):
+    TYPE = BackendType.NEBIUS
+    COMPUTE_CLASS = NebiusCompute
+
+    def __init__(self, config: NebiusConfig):
+        self.config = config
+        self._compute = NebiusCompute(self.config)
+
+    def compute(self) -> NebiusCompute:
+        return self._compute