dstack 0.19.0rc1__py3-none-any.whl → 0.19.2__py3-none-any.whl

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

@@ -0,0 +1,74 @@
+import json
+
+from nebius.aio.service_error import RequestError
+
+from dstack._internal.core.backends.base.configurator import (
+    BackendRecord,
+    Configurator,
+    raise_invalid_credentials_error,
+)
+from dstack._internal.core.backends.nebius import resources
+from dstack._internal.core.backends.nebius.backend import NebiusBackend
+from dstack._internal.core.backends.nebius.models import (
+    AnyNebiusBackendConfig,
+    NebiusBackendConfig,
+    NebiusBackendConfigWithCreds,
+    NebiusConfig,
+    NebiusCreds,
+    NebiusServiceAccountCreds,
+    NebiusStoredConfig,
+)
+from dstack._internal.core.models.backends.base import BackendType
+
+
+class NebiusConfigurator(Configurator):
+    TYPE = BackendType.NEBIUS
+    BACKEND_CLASS = NebiusBackend
+
+    def validate_config(self, config: NebiusBackendConfigWithCreds, default_creds_enabled: bool):
+        assert isinstance(config.creds, NebiusServiceAccountCreds)
+        try:
+            sdk = resources.make_sdk(config.creds)
+            available_regions = set(resources.get_region_to_project_id_map(sdk))
+        except (ValueError, RequestError) as e:
+            raise_invalid_credentials_error(
+                fields=[["creds"]],
+                details=str(e),
+            )
+        if invalid_regions := set(config.regions or []) - available_regions:
+            raise_invalid_credentials_error(
+                fields=[["regions"]],
+                details=(
+                    f"Configured regions {invalid_regions} do not exist in this Nebius tenancy."
+                    " Omit `regions` to use all regions or select some of the available regions:"
+                    f" {available_regions}"
+                ),
+            )
+
+    def create_backend(
+        self, project_name: str, config: NebiusBackendConfigWithCreds
+    ) -> BackendRecord:
+        return BackendRecord(
+            config=NebiusStoredConfig(
+                **NebiusBackendConfig.__response__.parse_obj(config).dict()
+            ).json(),
+            auth=NebiusCreds.parse_obj(config.creds).json(),
+        )
+
+    def get_backend_config(
+        self, record: BackendRecord, include_creds: bool
+    ) -> AnyNebiusBackendConfig:
+        config = self._get_config(record)
+        if include_creds:
+            return NebiusBackendConfigWithCreds.__response__.parse_obj(config)
+        return NebiusBackendConfig.__response__.parse_obj(config)
+
+    def get_backend(self, record: BackendRecord) -> NebiusBackend:
+        config = self._get_config(record)
+        return NebiusBackend(config=config)
+
+    def _get_config(self, record: BackendRecord) -> NebiusConfig:
+        return NebiusConfig.__response__(
+            **json.loads(record.config),
+            creds=NebiusCreds.parse_raw(record.auth),
+        )

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

@@ -0,0 +1,108 @@
+from typing import Annotated, Literal, Optional, Union
+
+from pydantic import Field, root_validator
+
+from dstack._internal.core.backends.base.models import fill_data
+from dstack._internal.core.models.common import CoreModel
+
+
+class NebiusServiceAccountCreds(CoreModel):
+    type: Annotated[Literal["service_account"], Field(description="The type of credentials")] = (
+        "service_account"
+    )
+    service_account_id: Annotated[str, Field(description="Service account ID")]
+    public_key_id: Annotated[str, Field(description="ID of the service account public key")]
+    private_key_file: Annotated[
+        Optional[str], Field(description=("Path to the service account private key"))
+    ] = None
+    private_key_content: Annotated[
+        str,
+        Field(
+            description=(
+                "Content of the service account private key. When configuring via"
+                " `server/config.yml`, it's automatically filled from `private_key_file`."
+                " When configuring via UI, it has to be specified explicitly."
+            )
+        ),
+    ]
+
+
+class NebiusServiceAccountFileCreds(CoreModel):
+    type: Annotated[Literal["service_account"], Field(description="The type of credentials")] = (
+        "service_account"
+    )
+    service_account_id: Annotated[str, Field(description="Service account ID")]
+    public_key_id: Annotated[str, Field(description="ID of the service account public key")]
+    private_key_file: Annotated[
+        Optional[str], Field(description=("Path to the service account private key"))
+    ] = None
+    private_key_content: Annotated[
+        Optional[str],
+        Field(
+            description=(
+                "Content of the service account private key. When configuring via"
+                " `server/config.yml`, it's automatically filled from `private_key_file`."
+                " When configuring via UI, it has to be specified explicitly."
+            )
+        ),
+    ] = None
+
+    @root_validator
+    def fill_data(cls, values):
+        return fill_data(
+            values, filename_field="private_key_file", data_field="private_key_content"
+        )
+
+
+AnyNebiusCreds = NebiusServiceAccountCreds
+NebiusCreds = AnyNebiusCreds
+AnyNebiusFileCreds = NebiusServiceAccountFileCreds
+
+
+class NebiusBackendConfig(CoreModel):
+    """
+    The backend config used in the API, server/config.yml, `NebiusConfigurator`.
+    It also serves as a base class for other backend config models.
+    Should not include creds.
+    """
+
+    type: Annotated[
+        Literal["nebius"],
+        Field(description="The type of backend"),
+    ] = "nebius"
+    regions: Annotated[
+        Optional[list[str]],
+        Field(description="The list of Nebius regions. Omit to use all regions"),
+    ] = None
+
+
+class NebiusBackendConfigWithCreds(NebiusBackendConfig):
+    """
+    Same as `NebiusBackendConfig` but also includes creds.
+    """
+
+    creds: Annotated[AnyNebiusCreds, Field(description="The credentials")]
+
+
+class NebiusBackendFileConfigWithCreds(NebiusBackendConfig):
+    creds: AnyNebiusFileCreds = Field(description="The credentials")
+
+
+AnyNebiusBackendConfig = Union[NebiusBackendConfig, NebiusBackendConfigWithCreds]
+
+
+class NebiusStoredConfig(NebiusBackendConfig):
+    """
+    The backend config used for config parameters in the DB.
+    Can extend `NebiusBackendConfig` with additional parameters.
+    """
+
+    pass
+
+
+class NebiusConfig(NebiusStoredConfig):
+    """
+    The backend config used by `NebiusBackend` and `NebiusCompute`.
+    """
+
+    creds: AnyNebiusCreds

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

@@ -0,0 +1,222 @@
+import logging
+import time
+from collections.abc import Container as ContainerT
+from collections.abc import Generator
+from contextlib import contextmanager
+from tempfile import NamedTemporaryFile
+
+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
+from nebius.api.nebius.common.v1 import Operation, ResourceMetadata
+from nebius.api.nebius.compute.v1 import (
+    AttachedDiskSpec,
+    CreateDiskRequest,
+    CreateInstanceRequest,
+    DeleteDiskRequest,
+    DeleteInstanceRequest,
+    DiskServiceClient,
+    DiskSpec,
+    ExistingDisk,
+    GetInstanceRequest,
+    Instance,
+    InstanceServiceClient,
+    InstanceSpec,
+    IPAddress,
+    NetworkInterfaceSpec,
+    PublicIPAddress,
+    ResourcesSpec,
+    SourceImageFamily,
+)
+from nebius.api.nebius.iam.v1 import (
+    ListProjectsRequest,
+    ListTenantsRequest,
+    ProjectServiceClient,
+    TenantServiceClient,
+)
+from nebius.api.nebius.vpc.v1 import ListSubnetsRequest, Subnet, SubnetServiceClient
+from nebius.sdk import SDK
+
+from dstack._internal.core.backends.nebius.models import NebiusServiceAccountCreds
+from dstack._internal.core.errors import BackendError, NoCapacityError
+from dstack._internal.utils.event_loop import DaemonEventLoop
+
+#
+# Guidelines on using the Nebius SDK:
+#
+# Do not use Request.wait() or other sync SDK methods, they suffer from deadlocks.
+# Instead, use async methods and await them with LOOP.await_()
+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",
+    }
+)
+
+# disables log messages about errors such as invalid creds or expired timeouts
+logging.getLogger("nebius").setLevel(logging.CRITICAL)
+
+
+@contextmanager
+def wrap_capacity_errors() -> Generator[None, None, None]:
+    try:
+        yield
+    except RequestError as e:
+        if e.status.code == StatusCode.RESOURCE_EXHAUSTED or "Quota limit exceeded" in str(e):
+            raise NoCapacityError(e)
+        raise
+
+
+@contextmanager
+def ignore_errors(status_codes: ContainerT[StatusCode]) -> Generator[None, None, None]:
+    try:
+        yield
+    except RequestError as e:
+        if e.status.code not in status_codes:
+            raise
+
+
+def make_sdk(creds: NebiusServiceAccountCreds) -> SDK:
+    with NamedTemporaryFile("w") as f:
+        f.write(creds.private_key_content)
+        f.flush()
+        return SDK(
+            service_account_private_key_file_name=f.name,
+            service_account_public_key_id=creds.public_key_id,
+            service_account_id=creds.service_account_id,
+        )
+
+
+def wait_for_operation(
+    op: SDKOperation[Operation],
+    timeout: float,
+    interval: float = 1,
+) -> None:
+    # Re-implementation of SDKOperation.wait() to avoid https://github.com/nebius/pysdk/issues/74
+    deadline = time.monotonic() + timeout
+    while not op.done():
+        if time.monotonic() + interval > deadline:
+            raise TimeoutError(f"Operation {op.id} wait timeout")
+        time.sleep(interval)
+        LOOP.await_(op.update(timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD))
+
+
+def get_region_to_project_id_map(sdk: SDK) -> dict[str, str]:
+    tenants = LOOP.await_(
+        TenantServiceClient(sdk).list(
+            ListTenantsRequest(), timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD
+        )
+    )
+    if len(tenants.items) != 1:
+        raise ValueError(f"Expected to find 1 tenant, found {(len(tenants.items))}")
+    projects = LOOP.await_(
+        ProjectServiceClient(sdk).list(
+            ListProjectsRequest(parent_id=tenants.items[0].metadata.id, page_size=999),
+            timeout=REQUEST_TIMEOUT,
+            metadata=REQUEST_MD,
+        )
+    )
+    result = {}
+    for project in projects.items:
+        if project.metadata.name == f"default-project-{project.status.region}":
+            result[project.status.region] = project.metadata.id
+    return result
+
+
+def get_default_subnet(sdk: SDK, project_id: str) -> Subnet:
+    subnets = LOOP.await_(
+        SubnetServiceClient(sdk).list(
+            ListSubnetsRequest(parent_id=project_id, page_size=999),
+            timeout=REQUEST_TIMEOUT,
+            metadata=REQUEST_MD,
+        )
+    )
+    for subnet in subnets.items:
+        if subnet.metadata.name.startswith("default-subnet"):
+            return subnet
+    raise BackendError(f"Could not find default subnet in project {project_id}")
+
+
+def create_disk(
+    sdk: SDK, name: str, project_id: str, size_mib: int, image_family: str
+) -> SDKOperation[Operation]:
+    client = DiskServiceClient(sdk)
+    request = CreateDiskRequest(
+        metadata=ResourceMetadata(
+            name=name,
+            parent_id=project_id,
+        ),
+        spec=DiskSpec(
+            size_mebibytes=size_mib,
+            type=DiskSpec.DiskType.NETWORK_SSD,
+            source_image_family=SourceImageFamily(image_family=image_family),
+        ),
+    )
+    with wrap_capacity_errors():
+        return LOOP.await_(client.create(request, timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD))
+
+
+def delete_disk(sdk: SDK, disk_id: str) -> None:
+    LOOP.await_(
+        DiskServiceClient(sdk).delete(
+            DeleteDiskRequest(id=disk_id), timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD
+        )
+    )
+
+
+def create_instance(
+    sdk: SDK,
+    name: str,
+    project_id: str,
+    user_data: str,
+    platform: str,
+    preset: str,
+    disk_id: str,
+    subnet_id: str,
+) -> SDKOperation[Operation]:
+    client = InstanceServiceClient(sdk)
+    request = CreateInstanceRequest(
+        metadata=ResourceMetadata(
+            name=name,
+            parent_id=project_id,
+        ),
+        spec=InstanceSpec(
+            cloud_init_user_data=user_data,
+            resources=ResourcesSpec(platform=platform, preset=preset),
+            boot_disk=AttachedDiskSpec(
+                attach_mode=AttachedDiskSpec.AttachMode.READ_WRITE,
+                existing_disk=ExistingDisk(id=disk_id),
+            ),
+            network_interfaces=[
+                NetworkInterfaceSpec(
+                    name="dstack-default-interface",
+                    subnet_id=subnet_id,
+                    ip_address=IPAddress(),
+                    public_ip_address=PublicIPAddress(static=True),
+                )
+            ],
+        ),
+    )
+    with wrap_capacity_errors():
+        return LOOP.await_(client.create(request, timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD))
+
+
+def get_instance(sdk: SDK, instance_id: str) -> Instance:
+    return LOOP.await_(
+        InstanceServiceClient(sdk).get(
+            GetInstanceRequest(id=instance_id), timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD
+        )
+    )
+
+
+def delete_instance(sdk: SDK, instance_id: str) -> SDKOperation[Operation]:
+    return LOOP.await_(
+        InstanceServiceClient(sdk).delete(
+            DeleteInstanceRequest(id=instance_id), timeout=REQUEST_TIMEOUT, metadata=REQUEST_MD
+        )
+    )

dstack/_internal/core/errors.py

@@ -102,6 +102,20 @@ class PlacementGroupInUseError(ComputeError):
     pass
 
 
+class NotYetTerminated(ComputeError):
+    """
+    Used by Compute.terminate_instance to signal that instance termination is not complete
+    and the method should be called again after some time to continue termination.
+    """
+
+    def __init__(self, details: str) -> None:
+        """
+        Args:
+            details: some details about the termination status
+        """
+        return super().__init__(details)
+
+
 class CLIError(DstackError):
     pass
 

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

@@ -12,6 +12,7 @@ class BackendType(str, enum.Enum):
         DATACRUNCH (BackendType): DataCrunch
         KUBERNETES (BackendType): Kubernetes
         LAMBDA (BackendType): Lambda Cloud
+        NEBIUS (BackendType): Nebius AI Cloud
         OCI (BackendType): Oracle Cloud Infrastructure
         RUNPOD (BackendType): Runpod Cloud
         TENSORDOCK (BackendType): TensorDock Marketplace
@@ -29,6 +30,7 @@ class BackendType(str, enum.Enum):
     LAMBDA = "lambda"
     LOCAL = "local"
     REMOTE = "remote"  # TODO: replace for LOCAL
+    NEBIUS = "nebius"
     OCI = "oci"
     RUNPOD = "runpod"
     TENSORDOCK = "tensordock"

dstack/_internal/proxy/lib/schemas/model_proxy.py

@@ -57,11 +57,11 @@ class ChatCompletionsResponse(CoreModel):
 
 
 class ChatCompletionsChunk(CoreModel):
-    id: str
+    id: Optional[str] = None
     choices: List[ChatCompletionsChunkChoice]
-    created: int
+    created: Optional[int] = None
     model: str
-    system_fingerprint: str = ""
+    system_fingerprint: Optional[str] = ""
     object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
 
 

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

@@ -39,7 +39,7 @@ from dstack._internal.core.backends.remote.provisioning import (
 from dstack._internal.core.consts import DSTACK_SHIM_HTTP_PORT
 
 # FIXME: ProvisioningError is a subclass of ComputeError and should not be used outside of Compute
-from dstack._internal.core.errors import BackendError, ProvisioningError
+from dstack._internal.core.errors import BackendError, NotYetTerminated, ProvisioningError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.fleets import InstanceGroupPlacement
 from dstack._internal.core.models.instances import (
@@ -64,6 +64,7 @@ from dstack._internal.core.models.runs import (
     Retry,
 )
 from dstack._internal.core.services.profiles import get_retry
+from dstack._internal.server import settings as server_settings
 from dstack._internal.server.background.tasks.common import get_provisioning_timeout
 from dstack._internal.server.db import get_session_ctx
 from dstack._internal.server.models import (
@@ -529,7 +530,9 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
             session=session, fleet_id=instance.fleet_id
         )
 
-    for backend, instance_offer in offers:
+    # Limit number of offers tried to prevent long-running processing
+    # in case all offers fail.
+    for backend, instance_offer in offers[: server_settings.MAX_OFFERS_TRIED]:
         if instance_offer.backend not in BACKENDS_WITH_CREATE_INSTANCE_SUPPORT:
             continue
         compute = backend.compute()
@@ -578,8 +581,13 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
                 extra={"instance_name": instance.name},
             )
             continue
-        except NotImplementedError:
-            # skip a backend without create_instance support, continue with next backend and offer
+        except Exception:
+            logger.exception(
+                "Got exception when launching %s in %s/%s",
+                instance_offer.instance.name,
+                instance_offer.backend.value,
+                instance_offer.region,
+            )
             continue
 
         instance.status = InstanceStatus.PROVISIONING
@@ -607,10 +615,11 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
 
     if not should_retry:
         instance.status = InstanceStatus.TERMINATED
-        instance.termination_reason = "No offers found"
+        instance.termination_reason = "All offers failed" if offers else "No offers found"
         logger.info(
-            "No offers found. Terminated instance %s",
+            "Terminated instance %s: %s",
             instance.name,
+            instance.termination_reason,
             extra={
                 "instance_name": instance.name,
                 "instance_status": InstanceStatus.TERMINATED.value,
@@ -837,12 +846,17 @@ async def _terminate(instance: InstanceModel) -> None:
                         instance.first_termination_retry_at = get_current_datetime()
                     instance.last_termination_retry_at = get_current_datetime()
                     if _next_termination_retry_at(instance) < _get_termination_deadline(instance):
-                        logger.warning(
-                            "Failed to terminate instance %s. Will retry. Error: %r",
-                            instance.name,
-                            e,
-                            exc_info=not isinstance(e, BackendError),
-                        )
+                        if isinstance(e, NotYetTerminated):
+                            logger.debug(
+                                "Instance %s termination in progress: %s", instance.name, e
+                            )
+                        else:
+                            logger.warning(
+                                "Failed to terminate instance %s. Will retry. Error: %r",
+                                instance.name,
+                                e,
+                                exc_info=not isinstance(e, BackendError),
+                            )
                         return
                     logger.error(
                         "Failed all attempts to terminate instance %s."

dstack/_internal/server/routers/prometheus.py

@@ -1,3 +1,4 @@
+import os
 from typing import Annotated
 
 from fastapi import APIRouter, Depends
@@ -6,14 +7,16 @@ from sqlalchemy.ext.asyncio import AsyncSession
 
 from dstack._internal.server import settings
 from dstack._internal.server.db import get_session
-from dstack._internal.server.deps import Project
-from dstack._internal.server.models import ProjectModel
+from dstack._internal.server.security.permissions import OptionalServiceAccount
 from dstack._internal.server.services import prometheus
 from dstack._internal.server.utils.routers import error_not_found
 
+_auth = OptionalServiceAccount(os.getenv("DSTACK_PROMETHEUS_AUTH_TOKEN"))
+
 router = APIRouter(
     tags=["prometheus"],
     default_response_class=PlainTextResponse,
+    dependencies=[Depends(_auth)],
 )
 
 
@@ -24,13 +27,3 @@ async def get_prometheus_metrics(
     if not settings.ENABLE_PROMETHEUS_METRICS:
         raise error_not_found()
     return await prometheus.get_metrics(session=session)
-
-
-@router.get("/metrics/project/{project_name}", deprecated=True)
-async def get_project_prometheus_metrics(
-    session: Annotated[AsyncSession, Depends(get_session)],
-    project: Annotated[ProjectModel, Depends(Project())],
-) -> str:
-    if not settings.ENABLE_PROMETHEUS_METRICS:
-        raise error_not_found()
-    return await prometheus.get_project_metrics(session=session, project=project)

dstack/_internal/server/security/permissions.py

@@ -1,4 +1,4 @@
-from typing import Tuple
+from typing import Annotated, Optional, Tuple
 
 from fastapi import Depends, HTTPException, Security
 from fastapi.security import HTTPBearer
@@ -99,6 +99,24 @@ class ProjectMember:
         return await get_project_member(session, project_name, token.credentials)
 
 
+class OptionalServiceAccount:
+    def __init__(self, token: Optional[str]) -> None:
+        self._token = token
+
+    async def __call__(
+        self,
+        token: Annotated[
+            Optional[HTTPAuthorizationCredentials], Security(HTTPBearer(auto_error=False))
+        ],
+    ) -> None:
+        if self._token is None:
+            return
+        if token is None:
+            raise error_forbidden()
+        if token.credentials != self._token:
+            raise error_invalid_token()
+
+
 async def get_project_member(
     session: AsyncSession, project_name: str, token: str
 ) -> Tuple[UserModel, ProjectModel]:

dstack/_internal/server/services/jobs/configurators/extensions/cursor.py

@@ -35,7 +35,7 @@ class CursorDesktop:
 
     def get_print_readme_commands(self) -> List[str]:
         return [
-            "echo To open in VS Code Desktop, use link below:",
+            "echo To open in Cursor, use link below:",
             "echo ''",
             f"echo '  cursor://vscode-remote/ssh-remote+{self.run_name}/workflow'",  # TODO use $REPO_DIR
             "echo ''",