dstack 0.18.43__py3-none-any.whl → 0.19.0__py3-none-any.whl

dstack/_internal/server/statics/static/media/entraID.d65d1f3e9486a8e56d24fc07b3230885.svg

@@ -0,0 +1,9 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg id="uuid-f8d4d392-7c12-4bd9-baff-66fbf7814b91" data-name="Layer 1" xmlns="http://www.w3.org/2000/svg" width="14" height="14" viewBox="0 0 18 18">
+  <path d="m3.802,14.032c.388.242,1.033.511,1.715.511.621,0,1.198-.18,1.676-.487,0,0,.001,0,.002-.001l1.805-1.128v4.073c-.286,0-.574-.078-.824-.234l-4.374-2.734Z" fill="#225086"/>
+  <path d="m7.853,1.507L.353,9.967c-.579.654-.428,1.642.323,2.111,0,0,2.776,1.735,3.126,1.954.388.242,1.033.511,1.715.511.621,0,1.198-.18,1.676-.487,0,0,.001,0,.002-.001l1.805-1.128-4.364-2.728,4.365-4.924V1s0,0,0,0c-.424,0-.847.169-1.147.507Z" fill="#6df"/>
+  <polygon points="4.636 10.199 4.688 10.231 9 12.927 9.001 12.927 9.001 12.927 9.001 5.276 9 5.275 4.636 10.199" fill="#cbf8ff"/>
+  <path d="m17.324,12.078c.751-.469.902-1.457.323-2.111l-4.921-5.551c-.397-.185-.842-.291-1.313-.291-.925,0-1.752.399-2.302,1.026l-.109.123h0s4.364,4.924,4.364,4.924h0s0,0,0,0l-4.365,2.728v4.073c.287,0,.573-.078.823-.234l7.5-4.688Z" fill="#074793"/>
+  <path d="m9.001,1v4.275s.109-.123.109-.123c.55-.627,1.377-1.026,2.302-1.026.472,0,.916.107,1.313.291l-2.579-2.909c-.299-.338-.723-.507-1.146-.507Z" fill="#0294e4"/>
+  <polygon points="13.365 10.199 13.365 10.199 13.365 10.199 9.001 5.276 9.001 12.926 13.365 10.199" fill="#96bcc2"/>
+</svg>

dstack/_internal/server/testing/common.py

@@ -8,6 +8,16 @@ from uuid import UUID
 import gpuhunt
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from dstack._internal.core.backends.base.compute import (
+    Compute,
+    ComputeWithCreateInstanceSupport,
+    ComputeWithGatewaySupport,
+    ComputeWithMultinodeSupport,
+    ComputeWithPlacementGroupSupport,
+    ComputeWithPrivateGatewaySupport,
+    ComputeWithReservationSupport,
+    ComputeWithVolumeSupport,
+)
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import NetworkMode
 from dstack._internal.core.models.configurations import (
@@ -35,8 +45,7 @@ from dstack._internal.core.models.placement import (
     PlacementStrategy,
 )
 from dstack._internal.core.models.profiles import (
-    DEFAULT_POOL_NAME,
-    DEFAULT_POOL_TERMINATION_IDLE_TIME,
+    DEFAULT_FLEET_TERMINATION_IDLE_TIME,
     Profile,
     TerminationPolicy,
 )
@@ -69,8 +78,8 @@ from dstack._internal.server.models import (
     InstanceModel,
     JobMetricsPoint,
     JobModel,
+    JobPrometheusMetrics,
     PlacementGroupModel,
-    PoolModel,
     ProjectModel,
     RepoCredsModel,
     RepoModel,
@@ -180,9 +189,6 @@ async def create_repo(
     if info is None:
         info = {
             "repo_type": "remote",
-            "repo_host_name": "",
-            "repo_port": None,
-            "repo_user_name": "",
             "repo_name": "dstack",
         }
     repo = RepoModel(
@@ -205,7 +211,6 @@ async def create_repo_creds(
 ) -> RepoCredsModel:
     if creds is None:
         creds = {
-            "protocol": "https",
             "clone_url": "https://github.com/dstackai/dstack.git",
             "private_key": None,
             "oauth_token": "test_token",
@@ -324,13 +329,22 @@ def get_job_provisioning_data(
     backend: BackendType = BackendType.AWS,
     region: str = "us-east-1",
     gpu_count: int = 0,
+    gpu_memory_gib: float = 16,
+    gpu_name: str = "T4",
     cpu_count: int = 1,
     memory_gib: float = 0.5,
     spot: bool = False,
     hostname: str = "127.0.0.4",
     internal_ip: Optional[str] = "127.0.0.4",
+    price: float = 10.5,
 ) -> JobProvisioningData:
-    gpus = [Gpu(name="T4", memory_mib=16384, vendor=gpuhunt.AcceleratorVendor.NVIDIA)] * gpu_count
+    gpus = [
+        Gpu(
+            name=gpu_name,
+            memory_mib=int(gpu_memory_gib * 1024),
+            vendor=gpuhunt.AcceleratorVendor.NVIDIA,
+        )
+    ] * gpu_count
     return JobProvisioningData(
         backend=backend,
         instance_type=InstanceType(
@@ -343,7 +357,7 @@ def get_job_provisioning_data(
         hostname=hostname,
         internal_ip=internal_ip,
         region=region,
-        price=10.5,
+        price=price,
         username="ubuntu",
         ssh_port=22,
         dockerized=dockerized,
@@ -438,22 +452,6 @@ def get_gateway_compute_configuration(
     )
 
 
-async def create_pool(
-    session: AsyncSession,
-    project: ProjectModel,
-    pool_name: Optional[str] = None,
-) -> PoolModel:
-    pool_name = pool_name if pool_name is not None else DEFAULT_POOL_NAME
-    pool = PoolModel(
-        name=pool_name,
-        project=project,
-        project_id=project.id,
-    )
-    session.add(pool)
-    await session.commit()
-    return pool
-
-
 async def create_fleet(
     session: AsyncSession,
     project: ProjectModel,
@@ -462,11 +460,14 @@ async def create_fleet(
     fleet_id: Optional[UUID] = None,
     status: FleetStatus = FleetStatus.ACTIVE,
     deleted: bool = False,
+    name: Optional[str] = None,
 ) -> FleetModel:
     if fleet_id is None:
         fleet_id = uuid.uuid4()
     if spec is None:
         spec = get_fleet_spec()
+    if name is not None:
+        spec.configuration.name = name
     fm = FleetModel(
         id=fleet_id,
         project=project,
@@ -506,7 +507,6 @@ def get_fleet_configuration(
 async def create_instance(
     session: AsyncSession,
     project: ProjectModel,
-    pool: PoolModel,
     fleet: Optional[FleetModel] = None,
     status: InstanceStatus = InstanceStatus.IDLE,
     unreachable: bool = False,
@@ -521,7 +521,7 @@ async def create_instance(
     instance_num: int = 0,
     backend: BackendType = BackendType.DATACRUNCH,
     termination_policy: Optional[TerminationPolicy] = None,
-    termination_idle_time: int = DEFAULT_POOL_TERMINATION_IDLE_TIME,
+    termination_idle_time: int = DEFAULT_FLEET_TERMINATION_IDLE_TIME,
     region: str = "eu-west",
     remote_connection_info: Optional[RemoteConnectionInfo] = None,
     offer: Optional[InstanceOfferWithAvailability] = None,
@@ -530,6 +530,7 @@ async def create_instance(
     busy_blocks: int = 0,
     name: str = "test_instance",
     volumes: Optional[List[VolumeModel]] = None,
+    price: float = 1.0,
 ) -> InstanceModel:
     if instance_id is None:
         instance_id = uuid.uuid4()
@@ -563,7 +564,6 @@ async def create_instance(
         id=instance_id,
         name=name,
         instance_num=instance_num,
-        pool=pool,
         fleet=fleet,
         project=project,
         status=status,
@@ -573,7 +573,7 @@ async def create_instance(
         finished_at=finished_at,
         job_provisioning_data=job_provisioning_data.json(),
         offer=offer.json(),
-        price=1,
+        price=price,
         region=region,
         backend=backend,
         termination_policy=termination_policy,
@@ -610,6 +610,8 @@ def get_instance_offer_with_availability(
     backend: BackendType = BackendType.AWS,
     region: str = "eu-west",
     gpu_count: int = 0,
+    gpu_name: str = "T4",
+    gpu_memory_gib: float = 16,
     cpu_count: int = 2,
     memory_gib: float = 12,
     disk_gib: float = 100.0,
@@ -617,12 +619,20 @@ def get_instance_offer_with_availability(
     blocks: int = 1,
     total_blocks: int = 1,
     availability_zones: Optional[List[str]] = None,
+    price: float = 1.0,
+    instance_type: str = "instance",
 ):
-    gpus = [Gpu(name="T4", memory_mib=16384, vendor=gpuhunt.AcceleratorVendor.NVIDIA)] * gpu_count
+    gpus = [
+        Gpu(
+            name=gpu_name,
+            memory_mib=int(gpu_memory_gib * 1024),
+            vendor=gpuhunt.AcceleratorVendor.NVIDIA,
+        )
+    ] * gpu_count
     return InstanceOfferWithAvailability(
         backend=backend,
         instance=InstanceType(
-            name="instance",
+            name=instance_type,
             resources=Resources(
                 cpus=cpu_count,
                 memory_mib=int(memory_gib * 1024),
@@ -633,7 +643,7 @@ def get_instance_offer_with_availability(
             ),
         ),
         region=region,
-        price=1,
+        price=price,
         availability=InstanceAvailability.AVAILABLE,
         availability_zones=availability_zones,
         blocks=blocks,
@@ -858,6 +868,22 @@ async def create_job_metrics_point(
     return jmp
 
 
+async def create_job_prometheus_metrics(
+    session: AsyncSession,
+    job: JobModel,
+    collected_at: datetime = datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc),
+    text: str = "# Prometheus metrics\n",
+):
+    metrics = JobPrometheusMetrics(
+        job_id=job.id,
+        collected_at=collected_at,
+        text=text,
+    )
+    session.add(metrics)
+    await session.commit()
+    return metrics
+
+
 def get_private_key_string() -> str:
     return """
 -----BEGIN RSA PRIVATE KEY-----
@@ -930,3 +956,20 @@ class AsyncContextManager:
 
     async def __aexit__(self, exc_type, exc, traceback):
         pass
+
+
+class ComputeMockSpec(
+    Compute,
+    ComputeWithCreateInstanceSupport,
+    ComputeWithMultinodeSupport,
+    ComputeWithReservationSupport,
+    ComputeWithPlacementGroupSupport,
+    ComputeWithGatewaySupport,
+    ComputeWithPrivateGatewaySupport,
+    ComputeWithVolumeSupport,
+):
+    """
+    Can be used to create Compute mocks that pass all isinstance asserts.
+    """
+
+    pass

dstack/_internal/utils/json_schema.py

@@ -0,0 +1,6 @@
+def add_extra_schema_types(schema_property: dict, extra_types: list[dict]):
+    if "allOf" in schema_property:
+        ref = schema_property.pop("allOf")[0]
+    else:
+        ref = {"type": schema_property.pop("type")}
+    schema_property["anyOf"] = [ref, *extra_types]

dstack/_internal/utils/ssh.py

@@ -128,7 +128,8 @@ def include_ssh_config(path: PathLike, ssh_config_path: PathLike = default_ssh_c
     except PermissionError:
         logger.warning(
             f"Couldn't update `{ssh_config_path}` due to a permissions problem.\n"
-            f"The `vscode://vscode-remote/ssh-remote+<run name>/workflow` link and "
+            f"The `vscode://vscode-remote/ssh-remote+<run name>/workflow` and "
+            f"`cursor://vscode-remote/ssh-remote+<run name>/workflow` links and "
             f"the `ssh <run name>` command won't work.\n"
             f"To fix this, make sure `{ssh_config_path}` is writable, or add "
             f"`Include {path}` to the top of `{ssh_config_path}` manually.",

dstack/api/__init__.py

@@ -2,6 +2,9 @@
 from dstack._internal.core.errors import ClientError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import RegistryAuth
+from dstack._internal.core.models.configurations import (
+    DevEnvironmentConfiguration as _DevEnvironmentConfiguration,
+)
 from dstack._internal.core.models.configurations import ScalingSpec as Scaling
 from dstack._internal.core.models.configurations import (
     ServiceConfiguration as _ServiceConfiguration,
@@ -22,3 +25,4 @@ from dstack.api._public.runs import Run, RunStatus
 
 Service = _ServiceConfiguration
 Task = _TaskConfiguration
+DevEnvironment = _DevEnvironmentConfiguration

dstack/api/_public/__init__.py

@@ -6,7 +6,6 @@ from dstack._internal.core.services.configs import ConfigManager
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.path import PathLike
 from dstack.api._public.backends import BackendCollection
-from dstack.api._public.pools import PoolCollection
 from dstack.api._public.repos import RepoCollection, get_ssh_keypair
 from dstack.api._public.runs import RunCollection
 from dstack.api.server import APIClient
@@ -16,12 +15,14 @@ logger = get_logger(__name__)
 
 class Client:
     """
-    High-level API client for interacting with dstack server
+    High-level API client for interacting with the `dstack` server
 
     Attributes:
+        project: The project name.
         runs: Operations with runs.
         repos: Operations with repositories.
         backends: Operations with backends.
+        client: Low-level API client that supports all API endpoints.
     """
 
     def __init__(
@@ -41,7 +42,6 @@ class Client:
         self._repos = RepoCollection(api_client, project_name)
         self._backends = BackendCollection(api_client, project_name)
         self._runs = RunCollection(api_client, project_name, self)
-        self._pool = PoolCollection(api_client, project_name)
         if ssh_identity_file:
             self.ssh_identity_file = str(ssh_identity_file)
         else:
@@ -58,13 +58,13 @@ class Client:
         Creates a Client using the default configuration from `~/.dstack/config.yml` if it exists.
 
         Args:
-            project_name: The name of the project, required if `server_url` and `user_token` are specified
-            server_url: The dstack server URL (e.g. `http://localhost:3000/` or `https://sky.dstack.ai`)
-            user_token: The dstack user token
-            ssh_identity_file: The private SSH key path for SSH tunneling
+            project_name: The name of the project. required if `server_url` and `user_token` are specified.
+            server_url: The dstack server URL (e.g. `http://localhost:3000/` or `https://sky.dstack.ai`).
+            user_token: The dstack user token.
+            ssh_identity_file: The private SSH key path for SSH tunneling.
 
         Returns:
-            A client instance
+            A client instance.
         """
         if server_url is not None and user_token is not None:
             if project_name is None:
@@ -79,25 +79,21 @@ class Client:
         )
 
     @property
-    def repos(self) -> RepoCollection:
-        return self._repos
-
-    @property
-    def backends(self) -> BackendCollection:
-        return self._backends
+    def project(self) -> str:
+        return self._project
 
     @property
     def runs(self) -> RunCollection:
         return self._runs
 
     @property
-    def client(self) -> APIClient:
-        return self._client
+    def repos(self) -> RepoCollection:
+        return self._repos
 
     @property
-    def project(self) -> str:
-        return self._project
+    def backends(self) -> BackendCollection:
+        return self._backends
 
     @property
-    def pool(self) -> PoolCollection:
-        return self._pool
+    def client(self) -> APIClient:
+        return self._client

dstack/api/_public/backends.py

@@ -1,6 +1,6 @@
 from typing import List
 
-from dstack._internal.core.models.backends import BackendInfo
+from dstack._internal.core.backends.models import BackendInfo
 from dstack.api.server import APIClient
 
 

dstack/api/_public/repos.py

@@ -55,7 +55,7 @@ class RepoCollection:
         Once the repo is initialized, you can pass the repo object to the run:
 
         ```python
-        run = client.runs.submit(
+        run = client.runs.apply_configuration(
             configuration=...,
             repo=repo,
         )
@@ -78,25 +78,6 @@ class RepoCollection:
                 raise ConfigurationError(*e.args)
         self._api_client.repos.init(self._project, repo.repo_id, repo.get_repo_info(), creds)
 
-    def is_initialized(
-        self,
-        repo: Repo,
-    ) -> bool:
-        # """
-        # Checks if the remote repo is initialized in the project
-        #
-        # Args:
-        #     repo: repo to check
-        #
-        # Returns:
-        #     repo is initialized
-        # """
-        try:
-            self._api_client.repos.get(self._project, repo.repo_id, include_creds=False)
-            return True
-        except ResourceNotExistsError:
-            return False
-
     def load(
         self,
         repo_dir: PathLike,
@@ -105,22 +86,22 @@ class RepoCollection:
         git_identity_file: Optional[PathLike] = None,
         oauth_token: Optional[str] = None,
     ) -> Union[RemoteRepo, LocalRepo]:
-        # """
-        # Loads the repo from the local directory using global config
-        #
-        # Args:
-        #     repo_dir: repo root directory
-        #     local: do not try to load `RemoteRepo` first
-        #     init: initialize the repo if it's not initialized
-        #     git_identity_file: path to an SSH private key to access the remote repo
-        #     oauth_token: GitHub OAuth token to access the remote repo
-        #
-        # Raises:
-        #     ConfigurationError: if the repo is not initialized and `init` is `False`
-        #
-        # Returns:
-        #     repo: initialized repo
-        # """
+        """
+        Loads the repo from the local directory using global config
+
+        Args:
+            repo_dir: Repo root directory.
+            local: Do not try to load `RemoteRepo` first.
+            init: Initialize the repo if it's not initialized.
+            git_identity_file: Path to an SSH private key to access the remote repo.
+            oauth_token: GitHub OAuth token to access the remote repo.
+
+        Raises:
+            ConfigurationError: If the repo is not initialized and `init` is `False`.
+
+        Returns:
+            repo: Initialized repo.
+        """
         config = ConfigManager()
         if not init:
             logger.debug("Loading repo config")
@@ -155,6 +136,25 @@ class RepoCollection:
             )
         return repo
 
+    def is_initialized(
+        self,
+        repo: Repo,
+    ) -> bool:
+        """
+        Checks if the remote repo is initialized in the project
+
+        Args:
+            repo: The repo to check.
+
+        Returns:
+            Whether the repo is initialized or not.
+        """
+        try:
+            self._api_client.repos.get(self._project, repo.repo_id, include_creds=False)
+            return True
+        except ResourceNotExistsError:
+            return False
+
 
 def get_ssh_keypair(key_path: Optional[PathLike], dstack_key_path: Path) -> str:
     """Returns a path to the private key"""