dstack 0.19.15rc1__py3-none-any.whl → 0.19.17__py3-none-any.whl

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

@@ -0,0 +1,66 @@
+import json
+
+from dstack._internal.core.backends.base.configurator import (
+    BackendRecord,
+    Configurator,
+    raise_invalid_credentials_error,
+)
+from dstack._internal.core.backends.cloudrift.api_client import RiftClient
+from dstack._internal.core.backends.cloudrift.backend import CloudRiftBackend
+from dstack._internal.core.backends.cloudrift.models import (
+    AnyCloudRiftBackendConfig,
+    AnyCloudRiftCreds,
+    CloudRiftBackendConfig,
+    CloudRiftBackendConfigWithCreds,
+    CloudRiftConfig,
+    CloudRiftCreds,
+    CloudRiftStoredConfig,
+)
+from dstack._internal.core.models.backends.base import (
+    BackendType,
+)
+
+
+class CloudRiftConfigurator(Configurator):
+    TYPE = BackendType.CLOUDRIFT
+    BACKEND_CLASS = CloudRiftBackend
+
+    def validate_config(
+        self, config: CloudRiftBackendConfigWithCreds, default_creds_enabled: bool
+    ):
+        self._validate_creds(config.creds)
+
+    def create_backend(
+        self, project_name: str, config: CloudRiftBackendConfigWithCreds
+    ) -> BackendRecord:
+        return BackendRecord(
+            config=CloudRiftStoredConfig(
+                **CloudRiftBackendConfig.__response__.parse_obj(config).dict()
+            ).json(),
+            auth=CloudRiftCreds.parse_obj(config.creds).json(),
+        )
+
+    def get_backend_config(
+        self, record: BackendRecord, include_creds: bool
+    ) -> AnyCloudRiftBackendConfig:
+        config = self._get_config(record)
+        if include_creds:
+            return CloudRiftBackendConfigWithCreds.__response__.parse_obj(config)
+        return CloudRiftBackendConfig.__response__.parse_obj(config)
+
+    def get_backend(self, record: BackendRecord) -> CloudRiftBackend:
+        config = self._get_config(record)
+        return CloudRiftBackend(config=config)
+
+    def _get_config(self, record: BackendRecord) -> CloudRiftConfig:
+        return CloudRiftConfig.__response__(
+            **json.loads(record.config),
+            creds=CloudRiftCreds.parse_raw(record.auth),
+        )
+
+    def _validate_creds(self, creds: AnyCloudRiftCreds):
+        if not isinstance(creds, CloudRiftCreds):
+            raise_invalid_credentials_error(fields=[["creds"]])
+        client = RiftClient(creds.api_key)
+        if not client.validate_api_key():
+            raise_invalid_credentials_error(fields=[["creds", "api_key"]])

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

@@ -0,0 +1,40 @@
+from typing import Annotated, List, Literal, Optional, Union
+
+from pydantic import Field
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class CloudRiftAPIKeyCreds(CoreModel):
+    type: Annotated[Literal["api_key"], Field(description="The type of credentials")] = "api_key"
+    api_key: Annotated[str, Field(description="The API key")]
+
+
+AnyCloudRiftCreds = CloudRiftAPIKeyCreds
+CloudRiftCreds = AnyCloudRiftCreds
+
+
+class CloudRiftBackendConfig(CoreModel):
+    type: Annotated[
+        Literal["cloudrift"],
+        Field(description="The type of backend"),
+    ] = "cloudrift"
+    regions: Annotated[
+        Optional[List[str]],
+        Field(description="The list of CloudRift regions. Omit to use all regions"),
+    ] = None
+
+
+class CloudRiftBackendConfigWithCreds(CloudRiftBackendConfig):
+    creds: Annotated[AnyCloudRiftCreds, Field(description="The credentials")]
+
+
+AnyCloudRiftBackendConfig = Union[CloudRiftBackendConfig, CloudRiftBackendConfigWithCreds]
+
+
+class CloudRiftStoredConfig(CloudRiftBackendConfig):
+    pass
+
+
+class CloudRiftConfig(CloudRiftStoredConfig):
+    creds: AnyCloudRiftCreds

dstack/_internal/core/backends/configurators.py

@@ -20,6 +20,15 @@ try:
 except ImportError:
     pass
 
+try:
+    from dstack._internal.core.backends.cloudrift.configurator import (
+        CloudRiftConfigurator,
+    )
+
+    _CONFIGURATOR_CLASSES.append(CloudRiftConfigurator)
+except ImportError:
+    pass
+
 try:
     from dstack._internal.core.backends.cudo.configurator import (
         CudoConfigurator,

dstack/_internal/core/backends/models.py

@@ -8,6 +8,10 @@ from dstack._internal.core.backends.azure.models import (
     AzureBackendConfig,
     AzureBackendConfigWithCreds,
 )
+from dstack._internal.core.backends.cloudrift.models import (
+    CloudRiftBackendConfig,
+    CloudRiftBackendConfigWithCreds,
+)
 from dstack._internal.core.backends.cudo.models import (
     CudoBackendConfig,
     CudoBackendConfigWithCreds,
@@ -65,6 +69,7 @@ from dstack._internal.core.models.common import CoreModel
 AnyBackendConfigWithoutCreds = Union[
     AWSBackendConfig,
     AzureBackendConfig,
+    CloudRiftBackendConfig,
     CudoBackendConfig,
     DataCrunchBackendConfig,
     GCPBackendConfig,
@@ -86,6 +91,7 @@ AnyBackendConfigWithoutCreds = Union[
 AnyBackendConfigWithCreds = Union[
     AWSBackendConfigWithCreds,
     AzureBackendConfigWithCreds,
+    CloudRiftBackendConfigWithCreds,
     CudoBackendConfigWithCreds,
     DataCrunchBackendConfigWithCreds,
     GCPBackendConfigWithCreds,
@@ -106,6 +112,7 @@ AnyBackendConfigWithCreds = Union[
 AnyBackendFileConfigWithCreds = Union[
     AWSBackendConfigWithCreds,
     AzureBackendConfigWithCreds,
+    CloudRiftBackendConfigWithCreds,
     CudoBackendConfigWithCreds,
     DataCrunchBackendConfigWithCreds,
     GCPBackendFileConfigWithCreds,

dstack/_internal/core/compatibility/logs.py

@@ -0,0 +1,15 @@
+from typing import Dict, Optional
+
+from dstack._internal.server.schemas.logs import PollLogsRequest
+
+
+def get_poll_logs_excludes(request: PollLogsRequest) -> Optional[Dict]:
+    """
+    Returns exclude mapping to exclude certain fields from the request.
+    Use this method to exclude new fields when they are not set to keep
+    clients backward-compatibility with older servers.
+    """
+    excludes = {}
+    if request.next_token is None:
+        excludes["next_token"] = True
+    return excludes if excludes else None

dstack/_internal/core/compatibility/runs.py

@@ -1,7 +1,7 @@
 from typing import Any, Dict, Optional
 
 from dstack._internal.core.models.configurations import ServiceConfiguration
-from dstack._internal.core.models.runs import ApplyRunPlanInput, JobSubmission, RunSpec
+from dstack._internal.core.models.runs import ApplyRunPlanInput, JobSpec, JobSubmission, RunSpec
 from dstack._internal.server.schemas.runs import GetRunPlanRequest
 
 
@@ -25,7 +25,10 @@ def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[Dict]:
         current_resource_excludes["run_spec"] = get_run_spec_excludes(current_resource.run_spec)
         job_submissions_excludes = {}
         current_resource_excludes["jobs"] = {
-            "__all__": {"job_submissions": {"__all__": job_submissions_excludes}}
+            "__all__": {
+                "job_spec": get_job_spec_excludes([job.job_spec for job in current_resource.jobs]),
+                "job_submissions": {"__all__": job_submissions_excludes},
+            }
         }
         job_submissions = [js for j in current_resource.jobs for js in j.job_submissions]
         if all(map(_should_exclude_job_submission_jpd_cpu_arch, job_submissions)):
@@ -97,6 +100,8 @@ def get_run_spec_excludes(run_spec: RunSpec) -> Optional[Dict]:
         configuration_excludes["rate_limits"] = True
     if configuration.shell is None:
         configuration_excludes["shell"] = True
+    if configuration.docker is None:
+        configuration_excludes["docker"] = True
     if configuration.priority is None:
         configuration_excludes["priority"] = True
     if configuration.startup_order is None:
@@ -107,6 +112,10 @@ def get_run_spec_excludes(run_spec: RunSpec) -> Optional[Dict]:
         configuration_excludes["stop_criteria"] = True
     if profile is not None and profile.stop_criteria is None:
         profile_excludes.add("stop_criteria")
+    if not configuration.files:
+        configuration_excludes["files"] = True
+    if not run_spec.file_archives:
+        spec_excludes["file_archives"] = True
 
     if configuration_excludes:
         spec_excludes["configuration"] = configuration_excludes
@@ -117,6 +126,26 @@ def get_run_spec_excludes(run_spec: RunSpec) -> Optional[Dict]:
     return None
 
 
+def get_job_spec_excludes(job_specs: list[JobSpec]) -> Optional[dict]:
+    """
+    Returns `job_spec` exclude mapping to exclude certain fields from the request.
+    Use this method to exclude new fields when they are not set to keep
+    clients backward-compatibility with older servers.
+    """
+    spec_excludes: dict[str, Any] = {}
+
+    if all(s.repo_code_hash is None for s in job_specs):
+        spec_excludes["repo_code_hash"] = True
+    if all(s.repo_data is None for s in job_specs):
+        spec_excludes["repo_data"] = True
+    if all(not s.file_archives for s in job_specs):
+        spec_excludes["file_archives"] = True
+
+    if spec_excludes:
+        return spec_excludes
+    return None
+
+
 def _should_exclude_job_submission_jpd_cpu_arch(job_submission: JobSubmission) -> bool:
     try:
         return job_submission.job_provisioning_data.instance_type.resources.cpu_arch is None

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

@@ -6,6 +6,7 @@ class BackendType(str, enum.Enum):
     Attributes:
         AWS (BackendType): Amazon Web Services
         AZURE (BackendType): Microsoft Azure
+        CLOUDRIFT (BackendType): CloudRift
         CUDO (BackendType): Cudo
         DSTACK (BackendType): dstack Sky
         GCP (BackendType): Google Cloud Platform
@@ -22,6 +23,7 @@ class BackendType(str, enum.Enum):
 
     AWS = "aws"
     AZURE = "azure"
+    CLOUDRIFT = "cloudrift"
     CUDO = "cudo"
     DATACRUNCH = "datacrunch"
     DSTACK = "dstack"

dstack/_internal/core/models/configurations.py

@@ -10,6 +10,7 @@ from typing_extensions import Annotated, Literal
 from dstack._internal.core.errors import ConfigurationError
 from dstack._internal.core.models.common import CoreModel, Duration, RegistryAuth
 from dstack._internal.core.models.envs import Env
+from dstack._internal.core.models.files import FilePathMapping
 from dstack._internal.core.models.fleets import FleetConfiguration
 from dstack._internal.core.models.gateways import GatewayConfiguration
 from dstack._internal.core.models.profiles import ProfileParams, parse_off_duration
@@ -194,12 +195,14 @@ class BaseRunConfiguration(CoreModel):
     ] = None
     python: Annotated[
         Optional[PythonVersion],
-        Field(description="The major version of Python. Mutually exclusive with `image`"),
+        Field(
+            description="The major version of Python. Mutually exclusive with `image` and `docker`"
+        ),
     ] = None
     nvcc: Annotated[
         Optional[bool],
         Field(
-            description="Use image with NVIDIA CUDA Compiler (NVCC) included. Mutually exclusive with `image`"
+            description="Use image with NVIDIA CUDA Compiler (NVCC) included. Mutually exclusive with `image` and `docker`"
         ),
     ] = None
     single_branch: Annotated[
@@ -244,6 +247,16 @@ class BaseRunConfiguration(CoreModel):
     volumes: Annotated[
         List[Union[MountPoint, str]], Field(description="The volumes mount points")
     ] = []
+    docker: Annotated[
+        Optional[bool],
+        Field(
+            description="Use Docker inside the container. Mutually exclusive with `image`, `python`, and `nvcc`. Overrides `privileged`"
+        ),
+    ] = None
+    files: Annotated[
+        list[Union[FilePathMapping, str]],
+        Field(description="The local to container file path mappings"),
+    ] = []
     # deprecated since 0.18.31; task, service -- no effect; dev-environment -- executed right before `init`
     setup: CommandsList = []
 
@@ -259,12 +272,30 @@ class BaseRunConfiguration(CoreModel):
             return PythonVersion(v)
         return v
 
+    @validator("docker", pre=True, always=True)
+    def _docker(cls, v, values) -> Optional[bool]:
+        if v is True and values.get("image"):
+            raise KeyError("`image` and `docker` are mutually exclusive fields")
+        if v is True and values.get("python"):
+            raise KeyError("`python` and `docker` are mutually exclusive fields")
+        if v is True and values.get("nvcc"):
+            raise KeyError("`nvcc` and `docker` are mutually exclusive fields")
+        # Ideally, we'd like to also prohibit privileged=False when docker=True,
+        #   but it's not possible to do so without breaking backwards compatibility.
+        return v
+
     @validator("volumes", each_item=True)
     def convert_volumes(cls, v) -> MountPoint:
         if isinstance(v, str):
             return parse_mount_point(v)
         return v
 
+    @validator("files", each_item=True)
+    def convert_files(cls, v) -> FilePathMapping:
+        if isinstance(v, str):
+            return FilePathMapping.parse(v)
+        return v
+
     @validator("user")
     def validate_user(cls, v) -> Optional[str]:
         if v is None:

dstack/_internal/core/models/files.py

@@ -0,0 +1,67 @@
+import pathlib
+import string
+from uuid import UUID
+
+from pydantic import Field, validator
+from typing_extensions import Annotated, Self
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class FileArchive(CoreModel):
+    id: UUID
+    hash: str
+
+
+class FilePathMapping(CoreModel):
+    local_path: Annotated[
+        str,
+        Field(
+            description=(
+                "The path on the user's machine. Relative paths are resolved relative to"
+                " the parent directory of the the configuration file"
+            )
+        ),
+    ]
+    path: Annotated[
+        str,
+        Field(
+            description=(
+                "The path in the container. Relative paths are resolved relative to"
+                " the repo directory (`/workflow`)"
+            )
+        ),
+    ]
+
+    @classmethod
+    def parse(cls, v: str) -> Self:
+        local_path: str
+        path: str
+        parts = v.split(":")
+        # A special case for Windows paths, e.g., `C:\path\to`, 'c:/path/to'
+        if (
+            len(parts) > 1
+            and len(parts[0]) == 1
+            and parts[0] in string.ascii_letters
+            and parts[1][:1] in ["\\", "/"]
+        ):
+            parts = [f"{parts[0]}:{parts[1]}", *parts[2:]]
+        if len(parts) == 1:
+            local_path = path = parts[0]
+        elif len(parts) == 2:
+            local_path, path = parts
+        else:
+            raise ValueError(f"invalid file path mapping: {v}")
+        return cls(local_path=local_path, path=path)
+
+    @validator("path")
+    def validate_path(cls, v) -> str:
+        # True for `C:/.*`, False otherwise, including `/abs/unix/path`, `rel\windows\path`, etc.
+        if pathlib.PureWindowsPath(v).is_absolute():
+            raise ValueError(f"path must be a Unix file path: {v}")
+        return v
+
+
+class FileArchiveMapping(CoreModel):
+    id: Annotated[UUID, Field(description="The File archive ID")]
+    path: Annotated[str, Field(description="The path in the container")]

dstack/_internal/core/models/logs.py

@@ -1,6 +1,6 @@
 from datetime import datetime
 from enum import Enum
-from typing import List
+from typing import List, Optional
 
 from dstack._internal.core.models.common import CoreModel
 
@@ -23,3 +23,4 @@ class LogEvent(CoreModel):
 
 class JobSubmissionLogs(CoreModel):
     logs: List[LogEvent]
+    next_token: Optional[str]

dstack/_internal/core/models/runs.py

@@ -12,6 +12,7 @@ from dstack._internal.core.models.configurations import (
     AnyRunConfiguration,
     RunConfiguration,
 )
+from dstack._internal.core.models.files import FileArchiveMapping
 from dstack._internal.core.models.instances import (
     InstanceOfferWithAvailability,
     InstanceType,
@@ -217,6 +218,15 @@ class JobSpec(CoreModel):
     volumes: Optional[List[MountPoint]] = None
     ssh_key: Optional[JobSSHKey] = None
     working_dir: Optional[str]
+    # `repo_data` is optional for client compatibility with pre-0.19.17 servers and for compatibility
+    # with jobs submitted before 0.19.17. All new jobs are expected to have non-None `repo_data`.
+    # For --no-repo runs, `repo_data` is `VirtualRunRepoData()`.
+    repo_data: Annotated[Optional[AnyRunRepoData], Field(discriminator="repo_type")] = None
+    # `repo_code_hash` can be None because it is not used for the repo or because the job was
+    # submitted before 0.19.17. See `_get_repo_code_hash` on how to get the correct `repo_code_hash`
+    # TODO: drop this comment when supporting jobs submitted before 0.19.17 is no longer relevant.
+    repo_code_hash: Optional[str] = None
+    file_archives: list[FileArchiveMapping] = []
 
 
 class JobProvisioningData(CoreModel):
@@ -301,7 +311,7 @@ class JobSubmission(CoreModel):
     job_provisioning_data: Optional[JobProvisioningData]
     job_runtime_data: Optional[JobRuntimeData]
     # TODO: make status_message and error a computed field after migrating to pydanticV2
-    status_message: Optional[str]
+    status_message: Optional[str] = None
     error: Optional[str] = None
 
     @property
@@ -413,6 +423,10 @@ class RunSpec(CoreModel):
         Optional[str],
         Field(description="The hash of the repo diff. Can be omitted if there is no repo diff."),
     ] = None
+    file_archives: Annotated[
+        list[FileArchiveMapping],
+        Field(description="The list of file archive ID to container path mappings"),
+    ] = []
     working_dir: Annotated[
         Optional[str],
         Field(
@@ -548,11 +562,17 @@ class Run(CoreModel):
             retry_on_events = (
                 jobs[0].job_spec.retry.on_events if jobs and jobs[0].job_spec.retry else []
             )
+            job_status = (
+                jobs[0].job_submissions[-1].status
+                if len(jobs) == 1 and jobs[0].job_submissions
+                else None
+            )
             termination_reason = Run.get_last_termination_reason(jobs[0]) if jobs else None
         except KeyError:
             return values
         values["status_message"] = Run._get_status_message(
             status=status,
+            job_status=job_status,
             retry_on_events=retry_on_events,
             termination_reason=termination_reason,
         )
@@ -568,9 +588,12 @@ class Run(CoreModel):
     @staticmethod
     def _get_status_message(
         status: RunStatus,
+        job_status: Optional[JobStatus],
         retry_on_events: List[RetryEvent],
         termination_reason: Optional[JobTerminationReason],
     ) -> str:
+        if job_status == JobStatus.PULLING:
+            return "pulling"
         # Currently, `retrying` is shown only for `no-capacity` events
         if (
             status in [RunStatus.SUBMITTED, RunStatus.PENDING]

dstack/_internal/core/models/secrets.py

@@ -1,9 +1,16 @@
+from typing import Optional
+from uuid import UUID
+
 from dstack._internal.core.models.common import CoreModel
 
 
 class Secret(CoreModel):
+    id: UUID
     name: str
-    value: str
+    value: Optional[str] = None
 
     def __str__(self) -> str:
-        return f'Secret(name="{self.name}", value={"*" * len(self.value)})'
+        displayed_value = "*"
+        if self.value is not None:
+            displayed_value = "*" * len(self.value)
+        return f'Secret(name="{self.name}", value={displayed_value})'

dstack/_internal/server/app.py

@@ -23,6 +23,7 @@ from dstack._internal.server.background import start_background_tasks
 from dstack._internal.server.db import get_db, get_session_ctx, migrate
 from dstack._internal.server.routers import (
     backends,
+    files,
     fleets,
     gateways,
     instances,
@@ -197,6 +198,7 @@ def register_routes(app: FastAPI, ui: bool = True):
     app.include_router(service_proxy.router, prefix="/proxy/services", tags=["service-proxy"])
     app.include_router(model_proxy.router, prefix="/proxy/models", tags=["model-proxy"])
     app.include_router(prometheus.router)
+    app.include_router(files.router)
 
     @app.exception_handler(ForbiddenError)
     async def forbidden_error_handler(request: Request, exc: ForbiddenError):

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

@@ -29,7 +29,7 @@ async def process_fleets():
                 )
                 .order_by(FleetModel.last_processed_at.asc())
                 .limit(1)
-                .with_for_update(skip_locked=True)
+                .with_for_update(skip_locked=True, key_share=True)
             )
             fleet_model = res.scalar()
             if fleet_model is None:

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

@@ -40,7 +40,7 @@ async def process_submitted_gateways():
                 .options(lazyload(GatewayModel.gateway_compute))
                 .order_by(GatewayModel.last_processed_at.asc())
                 .limit(1)
-                .with_for_update(skip_locked=True)
+                .with_for_update(skip_locked=True, key_share=True)
             )
             gateway_model = res.scalar()
             if gateway_model is None:

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

@@ -149,7 +149,7 @@ async def _process_next_instance():
                 .options(lazyload(InstanceModel.jobs))
                 .order_by(InstanceModel.last_processed_at.asc())
                 .limit(1)
-                .with_for_update(skip_locked=True)
+                .with_for_update(skip_locked=True, key_share=True)
             )
             instance = res.scalar()
             if instance is None:

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

@@ -30,7 +30,7 @@ async def process_placement_groups():
                     PlacementGroupModel.id.not_in(lockset),
                 )
                 .order_by(PlacementGroupModel.id)  # take locks in order
-                .with_for_update(skip_locked=True)
+                .with_for_update(skip_locked=True, key_share=True)
             )
             placement_group_models = res.scalars().all()
             if len(placement_group_models) == 0: