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

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,
 )
@@ -71,7 +80,6 @@ from dstack._internal.server.models import (
     JobModel,
     JobPrometheusMetrics,
     PlacementGroupModel,
-    PoolModel,
     ProjectModel,
     RepoCredsModel,
     RepoModel,
@@ -181,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(
@@ -206,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",
@@ -325,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(
@@ -344,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,
@@ -439,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,
@@ -463,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,
@@ -507,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,
@@ -522,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,
@@ -531,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()
@@ -564,7 +564,6 @@ async def create_instance(
         id=instance_id,
         name=name,
         instance_num=instance_num,
-        pool=pool,
         fleet=fleet,
         project=project,
         status=status,
@@ -574,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,
@@ -611,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,
@@ -618,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),
@@ -634,7 +643,7 @@ def get_instance_offer_with_availability(
             ),
         ),
         region=region,
-        price=1,
+        price=price,
         availability=InstanceAvailability.AVAILABLE,
         availability_zones=availability_zones,
         blocks=blocks,
@@ -947,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"""