dstack 0.19.32__py3-none-any.whl → 0.19.34__py3-none-any.whl

dstack/_internal/server/testing/common.py

@@ -126,6 +126,8 @@ async def create_user(
     global_role: GlobalRole = GlobalRole.ADMIN,
     token: Optional[str] = None,
     email: Optional[str] = None,
+    ssh_public_key: Optional[str] = None,
+    ssh_private_key: Optional[str] = None,
     active: bool = True,
 ) -> UserModel:
     if token is None:
@@ -137,6 +139,8 @@ async def create_user(
         token=DecryptedString(plaintext=token),
         token_hash=get_token_hash(token),
         email=email,
+        ssh_public_key=ssh_public_key,
+        ssh_private_key=ssh_private_key,
         active=active,
     )
     session.add(user)

dstack/api/_public/__init__.py

@@ -2,11 +2,10 @@ from typing import Optional
 
 import dstack._internal.core.services.api_client as api_client_service
 from dstack._internal.core.errors import ConfigurationError
-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._internal.utils.path import PathLike as PathLike
 from dstack.api._public.backends import BackendCollection
-from dstack.api._public.repos import RepoCollection, get_ssh_keypair
+from dstack.api._public.repos import RepoCollection
 from dstack.api._public.runs import RunCollection
 from dstack.api.server import APIClient
 
@@ -35,24 +34,24 @@ class Client:
         # Args:
         #     api_client: low-level server API client
         #     project_name: project name used for runs
-        #     ssh_identity_file: SSH keypair to access instances
+        #     ssh_identity_file: deprecated and will be removed in 0.19.40
         # """
         self._client = api_client
         self._project = project_name
         self._repos = RepoCollection(api_client, project_name)
         self._backends = BackendCollection(api_client, project_name)
         self._runs = RunCollection(api_client, project_name, self)
-        if ssh_identity_file:
-            self.ssh_identity_file = str(ssh_identity_file)
-        else:
-            self.ssh_identity_file = get_ssh_keypair(None, ConfigManager().dstack_key_path)
+        if ssh_identity_file is not None:
+            logger.warning(
+                "[code]ssh_identity_file[/code] in [code]Client[/code] is deprecated and ignored; will be removed"
+                " since 0.19.40"
+            )
 
     @staticmethod
     def from_config(
         project_name: Optional[str] = None,
         server_url: Optional[str] = None,
         user_token: Optional[str] = None,
-        ssh_identity_file: Optional[PathLike] = None,
     ) -> "Client":
         """
         Creates a Client using the default configuration from `~/.dstack/config.yml` if it exists.
@@ -61,7 +60,6 @@ class Client:
             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.
@@ -75,7 +73,6 @@ class Client:
         return Client(
             api_client=api_client,
             project_name=project_name,
-            ssh_identity_file=ssh_identity_file,
         )
 
     @property

dstack/api/_public/repos.py

@@ -1,4 +1,3 @@
-from pathlib import Path
 from typing import Literal, Optional, Union, overload
 
 from git import InvalidGitRepositoryError
@@ -18,7 +17,6 @@ from dstack._internal.core.services.repos import (
     get_repo_creds_and_default_branch,
     load_repo,
 )
-from dstack._internal.utils.crypto import generate_rsa_key_pair
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.path import PathLike
 from dstack.api.server import APIClient
@@ -209,22 +207,3 @@ class RepoCollection:
             return method(self._project, repo_id)
         except ResourceNotExistsError:
             return None
-
-
-def get_ssh_keypair(key_path: Optional[PathLike], dstack_key_path: Path) -> str:
-    """Returns a path to the private key"""
-    if key_path is not None:
-        key_path = Path(key_path).expanduser().resolve()
-        pub_key = (
-            key_path
-            if key_path.suffix == ".pub"
-            else key_path.with_suffix(key_path.suffix + ".pub")
-        )
-        private_key = pub_key.with_suffix("")
-        if pub_key.exists() and private_key.exists():
-            return str(private_key)
-        raise ConfigurationError(f"Make sure valid keypair exists: {private_key}(.pub)")
-
-    if not dstack_key_path.exists():
-        generate_rsa_key_pair(private_key_path=dstack_key_path)
-    return str(dstack_key_path)

dstack/api/_public/runs.py

@@ -45,11 +45,14 @@ from dstack._internal.core.models.runs import (
     get_service_port,
 )
 from dstack._internal.core.models.runs import Run as RunModel
+from dstack._internal.core.services.configs import ConfigManager
 from dstack._internal.core.services.logs import URLReplacer
 from dstack._internal.core.services.ssh.attach import SSHAttach
+from dstack._internal.core.services.ssh.key_manager import UserSSHKeyManager
 from dstack._internal.core.services.ssh.ports import PortsLock
 from dstack._internal.server.schemas.logs import PollLogsRequest
 from dstack._internal.utils.common import get_or_error, make_proxy_url
+from dstack._internal.utils.crypto import generate_rsa_key_pair
 from dstack._internal.utils.files import create_file_archive
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.path import PathLike, path_in_dir
@@ -72,16 +75,20 @@ class Run(ABC):
         self,
         api_client: APIClient,
         project: str,
-        ssh_identity_file: Optional[PathLike],
         run: RunModel,
         ports_lock: Optional[PortsLock] = None,
+        ssh_identity_file: Optional[PathLike] = None,
     ):
         self._api_client = api_client
         self._project = project
-        self._ssh_identity_file = ssh_identity_file
         self._run = run
         self._ports_lock: Optional[PortsLock] = ports_lock
         self._ssh_attach: Optional[SSHAttach] = None
+        if ssh_identity_file is not None:
+            logger.warning(
+                "[code]ssh_identity_file[/code] in [code]Run[/code] is deprecated and ignored; will be removed"
+                " since 0.19.40"
+            )
 
     @property
     def name(self) -> str:
@@ -270,9 +277,22 @@ class Run(ABC):
         Raises:
             dstack.api.PortUsedError: If ports are in use or the run is attached by another process.
         """
-        ssh_identity_file = ssh_identity_file or self._ssh_identity_file
-        if ssh_identity_file is None:
-            raise ConfigurationError("SSH identity file is required to attach to the run")
+        if not ssh_identity_file:
+            config_manager = ConfigManager()
+            key_manager = UserSSHKeyManager(self._api_client, config_manager.dstack_ssh_dir)
+            if (
+                user_key := key_manager.get_user_key()
+            ) and user_key.public_key == self._run.run_spec.ssh_key_pub:
+                ssh_identity_file = user_key.private_key_path
+            else:
+                if config_manager.dstack_key_path.exists():
+                    # TODO: Remove since 0.19.40
+                    logger.debug(f"Using legacy [code]{config_manager.dstack_key_path}[/code].")
+                    ssh_identity_file = config_manager.dstack_key_path
+                else:
+                    raise ConfigurationError(
+                        f"User SSH key doesn't match; default SSH key ({config_manager.dstack_key_path}) doesn't exist"
+                    )
         ssh_identity_file = str(ssh_identity_file)
 
         job = self._find_job(replica_num=replica_num, job_num=job_num)
@@ -434,6 +454,7 @@ class RunCollection:
         profile: Optional[Profile] = None,
         configuration_path: Optional[str] = None,
         repo_dir: Optional[str] = None,
+        ssh_identity_file: Optional[PathLike] = None,
     ) -> RunPlan:
         """
         Get a run plan.
@@ -465,6 +486,23 @@ class RunCollection:
         if repo_dir is None and configuration.repos:
             repo_dir = configuration.repos[0].path
 
+        if ssh_identity_file:
+            ssh_key_pub = Path(ssh_identity_file).with_suffix(".pub").read_text()
+        else:
+            config_manager = ConfigManager()
+            key_manager = UserSSHKeyManager(self._api_client, config_manager.dstack_ssh_dir)
+            if key_manager.get_user_key():
+                ssh_key_pub = None  # using the server-managed user key
+            else:
+                if not config_manager.dstack_key_path.exists():
+                    generate_rsa_key_pair(private_key_path=config_manager.dstack_key_path)
+                logger.warning(
+                    f"Using legacy [code]{config_manager.dstack_key_path.with_suffix('.pub')}[/code]."
+                    " You will only be able to attach to the run from this client."
+                    " Update the [code]dstack[/] server to [code]0.19.34[/]+ to switch to user keys"
+                    " automatically replicated to all clients.",
+                )
+                ssh_key_pub = config_manager.dstack_key_path.with_suffix(".pub").read_text()
         run_spec = RunSpec(
             run_name=configuration.name,
             repo_id=repo.repo_id,
@@ -477,7 +515,7 @@ class RunCollection:
             configuration_path=configuration_path,
             configuration=configuration,
             profile=profile,
-            ssh_key_pub=Path(self._client.ssh_identity_file + ".pub").read_text().strip(),
+            ssh_key_pub=ssh_key_pub,
         )
         logger.debug("Getting run plan")
         run_plan = self._api_client.runs.get_plan(self._project, run_spec)
@@ -546,6 +584,7 @@ class RunCollection:
         profile: Optional[Profile] = None,
         configuration_path: Optional[str] = None,
         reserve_ports: bool = True,
+        ssh_identity_file: Optional[PathLike] = None,
     ) -> Run:
         """
         Apply the run configuration.
@@ -567,6 +606,7 @@ class RunCollection:
             repo=repo,
             profile=profile,
             configuration_path=configuration_path,
+            ssh_identity_file=ssh_identity_file,
         )
         run = self.apply_plan(
             run_plan=run_plan,
@@ -709,6 +749,20 @@ class RunCollection:
             creation_policy=creation_policy,
             idle_duration=idle_duration,  # type: ignore[assignment]
         )
+        config_manager = ConfigManager()
+        key_manager = UserSSHKeyManager(self._api_client, config_manager.dstack_ssh_dir)
+        if key_manager.get_user_key():
+            ssh_key_pub = None  # using the server-managed user key
+        else:
+            if not config_manager.dstack_key_path.exists():
+                generate_rsa_key_pair(private_key_path=config_manager.dstack_key_path)
+            logger.warning(
+                f"Using legacy [code]{config_manager.dstack_key_path.with_suffix('.pub')}[/code]."
+                " You will only be able to attach to the run from this client."
+                " Update the [code]dstack[/] server to [code]0.19.34[/]+ to switch to user keys"
+                " automatically replicated to all clients.",
+            )
+            ssh_key_pub = config_manager.dstack_key_path.with_suffix(".pub").read_text()
         run_spec = RunSpec(
             run_name=run_name,
             repo_id=repo.repo_id,
@@ -718,7 +772,7 @@ class RunCollection:
             configuration_path=configuration_path,
             configuration=configuration,
             profile=profile,
-            ssh_key_pub=Path(self._client.ssh_identity_file + ".pub").read_text().strip(),
+            ssh_key_pub=ssh_key_pub,
         )
         logger.debug("Getting run plan")
         run_plan = self._api_client.runs.get_plan(self._project, run_spec)
@@ -800,7 +854,6 @@ class RunCollection:
         return Run(
             self._api_client,
             self._project,
-            self._client.ssh_identity_file,
             run,
         )
 
@@ -808,7 +861,6 @@ class RunCollection:
         return Run(
             self._api_client,
             self._project,
-            self._client.ssh_identity_file,
             run,
             ports_lock,
         )

dstack/api/server/__init__.py

@@ -1,3 +1,4 @@
+import hashlib
 import os
 import pprint
 import time
@@ -121,6 +122,9 @@ class APIClient:
     def files(self) -> FilesAPIClient:
         return FilesAPIClient(self._request, self._logger)
 
+    def get_token_hash(self) -> str:
+        return hashlib.sha1(self._token.encode()).hexdigest()[:8]
+
     def _request(
         self,
         path: str,

dstack/api/server/_users.py

@@ -1,6 +1,6 @@
 from typing import List
 
-from pydantic import parse_obj_as
+from pydantic import ValidationError, parse_obj_as
 
 from dstack._internal.core.models.users import GlobalRole, User, UserWithCreds
 from dstack._internal.server.schemas.users import (
@@ -18,8 +18,23 @@ class UsersAPIClient(APIClientGroup):
         return parse_obj_as(List[User.__response__], resp.json())
 
     def get_my_user(self) -> User:
+        """
+        Returns `User` with pre-0.19.33 servers, or `UserWithCreds` with newer servers.
+        """
+
         resp = self._request("/api/users/get_my_user")
-        return parse_obj_as(User.__response__, resp.json())
+        try:
+            return parse_obj_as(UserWithCreds.__response__, resp.json())
+        except ValidationError as e:
+            # Compatibility with pre-0.19.33 server
+            if (
+                len(e.errors()) == 1
+                and e.errors()[0]["loc"] == ("__root__", "creds")
+                and e.errors()[0]["type"] == "value_error.missing"
+            ):
+                return parse_obj_as(User.__response__, resp.json())
+            else:
+                raise
 
     def get_user(self, username: str) -> User:
         body = GetUserRequest(username=username)

dstack/version.py

@@ -1,4 +1,4 @@
-__version__ = "0.19.32"
+__version__ = "0.19.34"
 __is_release__ = True
-base_image = "0.11rc2"
+base_image = "0.11"
 base_image_ubuntu_version = "22.04"

{dstack-0.19.32.dist-info → dstack-0.19.34.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dstack
-Version: 0.19.32
+Version: 0.19.34
 Summary: dstack is an open-source orchestration engine for running AI workloads on any cloud or on-premises.
 Project-URL: Homepage, https://dstack.ai
 Project-URL: Source, https://github.com/dstackai/dstack
@@ -22,7 +22,7 @@ Requires-Dist: cryptography
 Requires-Dist: cursor
 Requires-Dist: filelock
 Requires-Dist: gitpython
-Requires-Dist: gpuhunt==0.1.8
+Requires-Dist: gpuhunt==0.1.11
 Requires-Dist: ignore-python>=0.2.0
 Requires-Dist: jsonschema
 Requires-Dist: orjson