dstack 0.19.33__py3-none-any.whl → 0.19.35__py3-none-any.whl

dstack/_internal/server/testing/common.py

@@ -13,6 +13,7 @@ from dstack._internal.core.backends.base.compute import (
     Compute,
     ComputeWithCreateInstanceSupport,
     ComputeWithGatewaySupport,
+    ComputeWithGroupProvisioningSupport,
     ComputeWithMultinodeSupport,
     ComputeWithPlacementGroupSupport,
     ComputeWithPrivateGatewaySupport,
@@ -22,6 +23,10 @@ from dstack._internal.core.backends.base.compute import (
 )
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import NetworkMode
+from dstack._internal.core.models.compute_groups import (
+    ComputeGroupProvisioningData,
+    ComputeGroupStatus,
+)
 from dstack._internal.core.models.configurations import (
     AnyRunConfiguration,
     DevEnvironmentConfiguration,
@@ -83,6 +88,7 @@ from dstack._internal.core.models.volumes import (
 )
 from dstack._internal.server.models import (
     BackendModel,
+    ComputeGroupModel,
     DecryptedString,
     FileArchiveModel,
     FleetModel,
@@ -126,6 +132,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 +145,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)
@@ -349,6 +359,7 @@ async def create_job(
     instance_assigned: bool = False,
     disconnected_at: Optional[datetime] = None,
     registered: bool = False,
+    waiting_master_job: Optional[bool] = None,
 ) -> JobModel:
     if deployment_num is None:
         deployment_num = run.deployment_num
@@ -380,6 +391,7 @@ async def create_job(
         disconnected_at=disconnected_at,
         probes=[],
         registered=registered,
+        waiting_master_job=waiting_master_job,
     )
     session.add(job)
     await session.commit()
@@ -451,6 +463,48 @@ def get_job_runtime_data(
     )
 
 
+def get_compute_group_provisioning_data(
+    compute_group_id: str = "test_compute_group",
+    compute_group_name: str = "test_compute_group",
+    backend: BackendType = BackendType.RUNPOD,
+    region: str = "US",
+    job_provisioning_datas: Optional[list[JobProvisioningData]] = None,
+    backend_data: Optional[str] = None,
+) -> ComputeGroupProvisioningData:
+    if job_provisioning_datas is None:
+        job_provisioning_datas = []
+    return ComputeGroupProvisioningData(
+        compute_group_id=compute_group_id,
+        compute_group_name=compute_group_name,
+        backend=backend,
+        region=region,
+        job_provisioning_datas=job_provisioning_datas,
+        backend_data=backend_data,
+    )
+
+
+async def create_compute_group(
+    session: AsyncSession,
+    project: ProjectModel,
+    fleet: FleetModel,
+    status: ComputeGroupStatus = ComputeGroupStatus.RUNNING,
+    provisioning_data: Optional[ComputeGroupProvisioningData] = None,
+    last_processed_at: datetime = datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc),
+):
+    if provisioning_data is None:
+        provisioning_data = get_compute_group_provisioning_data()
+    compute_group = ComputeGroupModel(
+        project=project,
+        fleet=fleet,
+        status=status,
+        provisioning_data=provisioning_data.json(),
+        last_processed_at=last_processed_at,
+    )
+    session.add(compute_group)
+    await session.commit()
+    return compute_group
+
+
 async def create_probe(
     session: AsyncSession,
     job: JobModel,
@@ -1132,6 +1186,7 @@ class AsyncContextManager:
 class ComputeMockSpec(
     Compute,
     ComputeWithCreateInstanceSupport,
+    ComputeWithGroupProvisioningSupport,
     ComputeWithPrivilegedSupport,
     ComputeWithMultinodeSupport,
     ComputeWithReservationSupport,

dstack/_internal/server/utils/routers.py

@@ -1,12 +1,13 @@
 from typing import Any, Dict, List, Optional
 
 import orjson
+import packaging.version
 from fastapi import HTTPException, Request, Response, status
-from packaging import version
 
 from dstack._internal.core.errors import ServerClientError, ServerClientErrorCode
 from dstack._internal.core.models.common import CoreModel
 from dstack._internal.utils.json_utils import get_orjson_default_options, orjson_default
+from dstack._internal.utils.version import parse_version
 
 
 class CustomORJSONResponse(Response):
@@ -122,8 +123,15 @@ def get_request_size(request: Request) -> int:
     return int(request.headers["content-length"])
 
 
+def get_client_version(request: Request) -> Optional[packaging.version.Version]:
+    version = request.headers.get("x-api-version")
+    if version is None:
+        return None
+    return parse_version(version)
+
+
 def check_client_server_compatibility(
-    client_version: Optional[str],
+    client_version: Optional[packaging.version.Version],
     server_version: Optional[str],
 ) -> Optional[CustomORJSONResponse]:
     """
@@ -132,28 +140,18 @@ def check_client_server_compatibility(
     """
     if client_version is None or server_version is None:
         return None
-    parsed_server_version = version.parse(server_version)
-    # latest allows client to bypass compatibility check (e.g. frontend)
-    if client_version == "latest":
+    parsed_server_version = parse_version(server_version)
+    if parsed_server_version is None:
         return None
-    try:
-        parsed_client_version = version.parse(client_version)
-    except version.InvalidVersion:
-        return CustomORJSONResponse(
-            status_code=status.HTTP_400_BAD_REQUEST,
-            content={
-                "detail": get_server_client_error_details(
-                    ServerClientError("Bad API version specified")
-                )
-            },
-        )
     # We preserve full client backward compatibility across patch releases.
     # Server is always partially backward-compatible (so no check).
-    if parsed_client_version > parsed_server_version and (
-        parsed_client_version.major > parsed_server_version.major
-        or parsed_client_version.minor > parsed_server_version.minor
+    if client_version > parsed_server_version and (
+        client_version.major > parsed_server_version.major
+        or client_version.minor > parsed_server_version.minor
     ):
-        return error_incompatible_versions(client_version, server_version, ask_cli_update=False)
+        return error_incompatible_versions(
+            str(client_version), server_version, ask_cli_update=False
+        )
     return None
 
 

dstack/_internal/settings.py

@@ -1,9 +1,10 @@
 import os
 
 from dstack import version
+from dstack._internal.utils.version import parse_version
 
 DSTACK_VERSION = os.getenv("DSTACK_VERSION", version.__version__)
-if DSTACK_VERSION == "0.0.0":
+if parse_version(DSTACK_VERSION) is None:
     # The build backend (hatching) requires not None for versions,
     # but the code currently treats None as dev version.
     # TODO: update the code to treat 0.0.0 as dev version.
@@ -33,3 +34,5 @@ class FeatureFlags:
     large features. This class may be empty if there are no such features in
     development. Feature flags are environment variables of the form DSTACK_FF_*
     """
+
+    pass

dstack/_internal/utils/version.py

@@ -0,0 +1,22 @@
+from typing import Optional
+
+import packaging.version
+
+
+def parse_version(version_string: str) -> Optional[packaging.version.Version]:
+    """
+    Returns a `packaging.version.Version` instance or `None` if the version is dev/latest.
+
+    Values parsed as the dev/latest version:
+    * the "latest" literal
+    * any "0.0.0" release, e.g., "0.0.0", "0.0.0a1", "0.0.0.dev0"
+    """
+    if version_string == "latest":
+        return None
+    try:
+        version = packaging.version.parse(version_string)
+    except packaging.version.InvalidVersion as e:
+        raise ValueError(f"Invalid version: {version_string}") from e
+    if version.release == (0, 0, 0):
+        return None
+    return version

dstack/api/_public/__init__.py

@@ -6,7 +6,7 @@ from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.path import PathLike as PathLike
 from dstack.api._public.backends import BackendCollection
 from dstack.api._public.repos import RepoCollection
-from dstack.api._public.runs import RunCollection, warn
+from dstack.api._public.runs import RunCollection
 from dstack.api.server import APIClient
 
 logger = get_logger(__name__)
@@ -42,7 +42,7 @@ class Client:
         self._backends = BackendCollection(api_client, project_name)
         self._runs = RunCollection(api_client, project_name, self)
         if ssh_identity_file is not None:
-            warn(
+            logger.warning(
                 "[code]ssh_identity_file[/code] in [code]Client[/code] is deprecated and ignored; will be removed"
                 " since 0.19.40"
             )

dstack/api/_public/runs.py

@@ -1,6 +1,4 @@
 import base64
-import hashlib
-import os
 import queue
 import tempfile
 import threading
@@ -17,7 +15,6 @@ from urllib.parse import urlparse
 from websocket import WebSocketApp
 
 import dstack.api as api
-from dstack._internal.cli.utils.common import warn
 from dstack._internal.core.consts import DSTACK_RUNNER_HTTP_PORT, DSTACK_RUNNER_SSH_PORT
 from dstack._internal.core.errors import ClientError, ConfigurationError, ResourceNotExistsError
 from dstack._internal.core.models.backends.base import BackendType
@@ -48,10 +45,10 @@ from dstack._internal.core.models.runs import (
     get_service_port,
 )
 from dstack._internal.core.models.runs import Run as RunModel
-from dstack._internal.core.models.users import UserWithCreds
 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
@@ -88,7 +85,7 @@ class Run(ABC):
         self._ports_lock: Optional[PortsLock] = ports_lock
         self._ssh_attach: Optional[SSHAttach] = None
         if ssh_identity_file is not None:
-            warn(
+            logger.warning(
                 "[code]ssh_identity_file[/code] in [code]Run[/code] is deprecated and ignored; will be removed"
                 " since 0.19.40"
             )
@@ -281,31 +278,20 @@ class Run(ABC):
             dstack.api.PortUsedError: If ports are in use or the run is attached by another process.
         """
         if not ssh_identity_file:
-            user = self._api_client.users.get_my_user()
-            run_ssh_key_pub = self._run.run_spec.ssh_key_pub
             config_manager = ConfigManager()
-            if isinstance(user, UserWithCreds) and user.ssh_public_key == run_ssh_key_pub:
-                token_hash = hashlib.sha1(user.creds.token.encode()).hexdigest()[:8]
-                config_manager.dstack_ssh_dir.mkdir(parents=True, exist_ok=True)
-                ssh_identity_file = config_manager.dstack_ssh_dir / token_hash
-
-                def key_opener(path, flags):
-                    return os.open(path, flags, 0o600)
-
-                with open(ssh_identity_file, "wb", opener=key_opener) as f:
-                    assert user.ssh_private_key
-                    f.write(user.ssh_private_key.encode())
+            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
-                    warn(
-                        f"Using legacy [code]{config_manager.dstack_key_path}[/code]."
-                        " Future versions will use the user SSH key from the server.",
-                    )
+                    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 doen't match; default SSH key ({config_manager.dstack_key_path}) doesn't exist"
+                        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)
 
@@ -504,15 +490,19 @@ class RunCollection:
             ssh_key_pub = Path(ssh_identity_file).with_suffix(".pub").read_text()
         else:
             config_manager = ConfigManager()
-            if not config_manager.dstack_key_path.exists():
-                generate_rsa_key_pair(private_key_path=config_manager.dstack_key_path)
-            warn(
-                f"Using legacy [code]{config_manager.dstack_key_path.with_suffix('.pub')}[/code]."
-                " Future versions will use the user SSH key from the server.",
-            )
-            ssh_key_pub = config_manager.dstack_key_path.with_suffix(".pub").read_text()
-            # TODO: Uncomment after 0.19.40
-            # ssh_key_pub = None
+            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,
@@ -760,12 +750,19 @@ class RunCollection:
             idle_duration=idle_duration,  # type: ignore[assignment]
         )
         config_manager = ConfigManager()
-        if not config_manager.dstack_key_path.exists():
-            generate_rsa_key_pair(private_key_path=config_manager.dstack_key_path)
-        warn(
-            f"Using legacy [code]{config_manager.dstack_key_path.with_suffix('.pub')}[/code]."
-            " Future versions will use the user SSH key from the server.",
-        )
+        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,
@@ -775,7 +772,7 @@ class RunCollection:
             configuration_path=configuration_path,
             configuration=configuration,
             profile=profile,
-            ssh_key_pub=config_manager.dstack_key_path.with_suffix(".pub").read_text(),
+            ssh_key_pub=ssh_key_pub,
         )
         logger.debug("Getting run plan")
         run_plan = self._api_client.runs.get_plan(self._project, run_spec)

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/version.py

@@ -1,4 +1,4 @@
-__version__ = "0.19.33"
+__version__ = "0.19.35"
 __is_release__ = True
 base_image = "0.11"
 base_image_ubuntu_version = "22.04"

{dstack-0.19.33.dist-info → dstack-0.19.35.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dstack
-Version: 0.19.33
+Version: 0.19.35
 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.10
+Requires-Dist: gpuhunt==0.1.11
 Requires-Dist: ignore-python>=0.2.0
 Requires-Dist: jsonschema
 Requires-Dist: orjson
@@ -73,7 +73,7 @@ Requires-Dist: grpcio>=1.50; extra == 'all'
 Requires-Dist: httpx; extra == 'all'
 Requires-Dist: jinja2; extra == 'all'
 Requires-Dist: kubernetes; extra == 'all'
-Requires-Dist: nebius<=0.2.72,>=0.2.40; (python_version >= '3.10') and extra == 'all'
+Requires-Dist: nebius<0.4,>=0.3.4; (python_version >= '3.10') and extra == 'all'
 Requires-Dist: oci>=2.150.0; extra == 'all'
 Requires-Dist: prometheus-client; extra == 'all'
 Requires-Dist: pyopenssl>=23.2.0; extra == 'all'
@@ -259,7 +259,7 @@ Requires-Dist: fastapi; extra == 'nebius'
 Requires-Dist: grpcio>=1.50; extra == 'nebius'
 Requires-Dist: httpx; extra == 'nebius'
 Requires-Dist: jinja2; extra == 'nebius'
-Requires-Dist: nebius<=0.2.72,>=0.2.40; (python_version >= '3.10') and extra == 'nebius'
+Requires-Dist: nebius<0.4,>=0.3.4; (python_version >= '3.10') and extra == 'nebius'
 Requires-Dist: prometheus-client; extra == 'nebius'
 Requires-Dist: python-dxf==12.1.0; extra == 'nebius'
 Requires-Dist: python-json-logger>=3.1.0; extra == 'nebius'