dstack 0.18.40rc1__py3-none-any.whl → 0.18.41__py3-none-any.whl

dstack/_internal/server/services/pools.py

@@ -2,7 +2,7 @@ import ipaddress
 import uuid
 from collections.abc import Container, Iterable
 from datetime import datetime, timezone
-from typing import List, Optional
+from typing import List, Literal, Optional, Union
 
 import gpuhunt
 from sqlalchemy import and_, or_, select
@@ -30,6 +30,7 @@ from dstack._internal.core.models.instances import (
     InstanceType,
     RemoteConnectionInfo,
     Resources,
+    SSHConnectionParams,
     SSHKey,
 )
 from dstack._internal.core.models.pools import Instance, Pool, PoolInstances
@@ -51,6 +52,7 @@ from dstack._internal.server.models import (
     UserModel,
 )
 from dstack._internal.server.services.locking import get_locker
+from dstack._internal.server.services.offers import generate_shared_offer
 from dstack._internal.server.services.projects import list_project_models, list_user_project_models
 from dstack._internal.utils import common as common_utils
 from dstack._internal.utils import random_names
@@ -206,7 +208,7 @@ async def remove_instance(
     terminated = False
     for instance in pool.instances:
         if instance.name == instance_name:
-            if force or instance.job_id is None:
+            if force or not instance.jobs:
                 instance.status = InstanceStatus.TERMINATING
                 terminated = True
     await session.commit()
@@ -249,6 +251,8 @@ def instance_model_to_instance(instance_model: InstanceModel) -> Instance:
         unreachable=instance_model.unreachable,
         termination_reason=instance_model.termination_reason,
         created=instance_model.created_at.replace(tzinfo=timezone.utc),
+        total_blocks=instance_model.total_blocks,
+        busy_blocks=instance_model.busy_blocks,
     )
 
     offer = get_instance_offer(instance_model)
@@ -261,9 +265,7 @@ def instance_model_to_instance(instance_model: InstanceModel) -> Instance:
     if jpd is not None:
         instance.instance_type = jpd.instance_type
         instance.hostname = jpd.hostname
-
-    if instance_model.job is not None:
-        instance.job_name = instance_model.job.job_name
+        instance.availability_zone = jpd.availability_zone
 
     return instance
 
@@ -292,6 +294,27 @@ def get_instance_requirements(instance_model: InstanceModel) -> Requirements:
     return Requirements.__response__.parse_raw(instance_model.requirements)
 
 
+def get_instance_ssh_private_keys(instance_model: InstanceModel) -> tuple[str, Optional[str]]:
+    """
+    Returns a pair of SSH private keys: host key and optional proxy jump key.
+    """
+    host_private_key = instance_model.project.ssh_private_key
+    if instance_model.remote_connection_info is None:
+        # Cloud instance
+        return host_private_key, None
+    # SSH instance
+    rci = RemoteConnectionInfo.__response__.parse_raw(instance_model.remote_connection_info)
+    if rci.ssh_proxy is None:
+        return host_private_key, None
+    if rci.ssh_proxy_keys is None:
+        # Inconsistent RemoteConnectionInfo structure - proxy without keys
+        raise ValueError("Missing instance SSH proxy private keys")
+    proxy_private_keys = [key.private for key in rci.ssh_proxy_keys if key.private is not None]
+    if not proxy_private_keys:
+        raise ValueError("No instance SSH proxy private key found")
+    return host_private_key, proxy_private_keys[0]
+
+
 async def generate_instance_name(
     session: AsyncSession,
     project: ProjectModel,
@@ -409,20 +432,21 @@ async def add_remote(
 def filter_pool_instances(
     pool_instances: List[InstanceModel],
     profile: Profile,
-    requirements: Requirements,
     *,
+    requirements: Optional[Requirements] = None,
     status: Optional[InstanceStatus] = None,
     fleet_model: Optional[FleetModel] = None,
     multinode: bool = False,
     master_job_provisioning_data: Optional[JobProvisioningData] = None,
     volumes: Optional[List[List[Volume]]] = None,
+    shared: bool = False,
 ) -> List[InstanceModel]:
     instances: List[InstanceModel] = []
     candidates: List[InstanceModel] = []
 
     backend_types = profile.backends
     regions = profile.regions
-    zones = None
+    zones = profile.availability_zones
 
     if volumes:
         mount_point_volumes = volumes[0]
@@ -433,8 +457,9 @@ def filter_pool_instances(
             for v in mount_point_volumes
             if v.provisioning_data is not None
         ]
-        if volume_zones:
+        if zones is None:
             zones = volume_zones
+        zones = [z for z in zones if z in volume_zones]
 
     if multinode:
         if not backend_types:
@@ -480,9 +505,17 @@ def filter_pool_instances(
             and jpd.availability_zone not in zones
         ):
             continue
+        if instance.total_blocks is None:
+            # Still provisioning, we don't know yet if it shared or not
+            continue
+        if (instance.total_blocks > 1) != shared:
+            continue
 
         candidates.append(instance)
 
+    if requirements is None:
+        return candidates
+
     query_filter = requirements_to_query_filter(requirements)
     for instance in candidates:
         if instance.offer is None:
@@ -494,6 +527,47 @@ def filter_pool_instances(
     return instances
 
 
+def get_shared_pool_instances_with_offers(
+    pool_instances: List[InstanceModel],
+    profile: Profile,
+    requirements: Requirements,
+    *,
+    idle_only: bool = False,
+    fleet_model: Optional[FleetModel] = None,
+    volumes: Optional[List[List[Volume]]] = None,
+) -> list[tuple[InstanceModel, InstanceOfferWithAvailability]]:
+    instances_with_offers: list[tuple[InstanceModel, InstanceOfferWithAvailability]] = []
+    query_filter = requirements_to_query_filter(requirements)
+    filtered_instances = filter_pool_instances(
+        pool_instances=pool_instances,
+        profile=profile,
+        fleet_model=fleet_model,
+        multinode=False,
+        volumes=volumes,
+        shared=True,
+    )
+    for instance in filtered_instances:
+        if idle_only and instance.status not in [InstanceStatus.IDLE, InstanceStatus.BUSY]:
+            continue
+        offer = get_instance_offer(instance)
+        if offer is None:
+            continue
+        total_blocks = common_utils.get_or_error(instance.total_blocks)
+        idle_blocks = total_blocks - instance.busy_blocks
+        for blocks in range(1, total_blocks + 1):
+            shared_offer = generate_shared_offer(offer, blocks, total_blocks)
+            catalog_item = offer_to_catalog_item(shared_offer)
+            if gpuhunt.matches(catalog_item, query_filter):
+                if blocks <= idle_blocks:
+                    shared_offer.availability = InstanceAvailability.IDLE
+                else:
+                    shared_offer.availability = InstanceAvailability.BUSY
+                if shared_offer.availability == InstanceAvailability.IDLE or not idle_only:
+                    instances_with_offers.append((instance, shared_offer))
+                break
+    return instances_with_offers
+
+
 async def list_pools_instance_models(
     session: AsyncSession,
     projects: List[ProjectModel],
@@ -557,7 +631,7 @@ async def list_pools_instance_models(
         .limit(limit)
         .options(joinedload(InstanceModel.pool), joinedload(InstanceModel.fleet))
     )
-    instance_models = list(res.scalars().all())
+    instance_models = list(res.unique().scalars().all())
     return instance_models
 
 
@@ -618,7 +692,7 @@ async def list_active_remote_instances(
     res = await session.execute(
         select(InstanceModel).where(*filters).order_by(InstanceModel.created_at.asc())
     )
-    instance_models = list(res.scalars().all())
+    instance_models = list(res.unique().scalars().all())
     return instance_models
 
 
@@ -633,6 +707,7 @@ async def create_instance_model(
     instance_num: int,
     placement_group_name: Optional[str],
     reservation: Optional[str],
+    blocks: Union[Literal["auto"], int],
 ) -> InstanceModel:
     termination_policy, termination_idle_time = get_termination(
         profile, DEFAULT_POOL_TERMINATION_IDLE_TIME
@@ -646,8 +721,8 @@ async def create_instance_model(
         project_name=project.name,
         instance_name=instance_name,
         user=user.name,
-        instance_id=str(instance_id),
         ssh_keys=[project_ssh_key],
+        instance_id=str(instance_id),
         placement_group_name=placement_group_name,
         reservation=reservation,
     )
@@ -665,6 +740,8 @@ async def create_instance_model(
         instance_configuration=instance_config.json(),
         termination_policy=termination_policy,
         termination_idle_time=termination_idle_time,
+        total_blocks=None if blocks == "auto" else blocks,
+        busy_blocks=0,
     )
     session.add(instance)
     return instance
@@ -682,7 +759,10 @@ async def create_ssh_instance_model(
     port: int,
     ssh_user: str,
     ssh_keys: List[SSHKey],
+    ssh_proxy: Optional[SSHConnectionParams],
+    ssh_proxy_keys: Optional[list[SSHKey]],
     env: Env,
+    blocks: Union[Literal["auto"], int],
 ) -> InstanceModel:
     # TODO: doc - will overwrite after remote connected
     instance_resource = Resources(cpus=2, memory_mib=8, gpus=[], spot=False)
@@ -717,6 +797,8 @@ async def create_ssh_instance_model(
         port=port,
         ssh_user=ssh_user,
         ssh_keys=ssh_keys,
+        ssh_proxy=ssh_proxy,
+        ssh_proxy_keys=ssh_proxy_keys,
         env=env,
     )
     im = InstanceModel(
@@ -737,5 +819,7 @@ async def create_ssh_instance_model(
         price=offer.price,
         termination_policy=TerminationPolicy.DONT_DESTROY,
         termination_idle_time=0,
+        total_blocks=None if blocks == "auto" else blocks,
+        busy_blocks=0,
     )
     return im

dstack/_internal/server/services/proxy/repo.py

@@ -9,7 +9,7 @@ import dstack._internal.server.services.jobs as jobs_services
 from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT
 from dstack._internal.core.models.common import is_core_model_instance
 from dstack._internal.core.models.configurations import ServiceConfiguration
-from dstack._internal.core.models.instances import SSHConnectionParams
+from dstack._internal.core.models.instances import RemoteConnectionInfo, SSHConnectionParams
 from dstack._internal.core.models.runs import (
     JobProvisioningData,
     JobStatus,
@@ -30,6 +30,7 @@ from dstack._internal.proxy.lib.models import (
 from dstack._internal.proxy.lib.repo import BaseProxyRepo
 from dstack._internal.server.models import JobModel, ProjectModel, RunModel
 from dstack._internal.server.settings import DEFAULT_SERVICE_CLIENT_MAX_BODY_SIZE
+from dstack._internal.utils.common import get_or_error
 
 
 class ServerProxyRepo(BaseProxyRepo):
@@ -53,9 +54,12 @@ class ServerProxyRepo(BaseProxyRepo):
                 JobModel.status == JobStatus.RUNNING,
                 JobModel.job_num == 0,
             )
-            .options(joinedload(JobModel.run))
+            .options(
+                joinedload(JobModel.run),
+                joinedload(JobModel.instance),
+            )
         )
-        jobs = res.scalars().all()
+        jobs = res.unique().scalars().all()
         if not len(jobs):
             return None
         run = jobs[0].run
@@ -83,12 +87,22 @@ class ServerProxyRepo(BaseProxyRepo):
                     username=jpd.username,
                     port=jpd.ssh_port,
                 )
+            ssh_head_proxy: Optional[SSHConnectionParams] = None
+            ssh_head_proxy_private_key: Optional[str] = None
+            instance = get_or_error(job.instance)
+            if instance.remote_connection_info is not None:
+                rci = RemoteConnectionInfo.__response__.parse_raw(instance.remote_connection_info)
+                if rci.ssh_proxy is not None:
+                    ssh_head_proxy = rci.ssh_proxy
+                    ssh_head_proxy_private_key = get_or_error(rci.ssh_proxy_keys)[0].private
             replica = Replica(
                 id=job.id.hex,
                 app_port=run_spec.configuration.port.container_port,
                 ssh_destination=ssh_destination,
                 ssh_port=ssh_port,
                 ssh_proxy=ssh_proxy,
+                ssh_head_proxy=ssh_head_proxy,
+                ssh_head_proxy_private_key=ssh_head_proxy_private_key,
             )
             replicas.append(replica)
         return Service(

dstack/_internal/server/services/runner/client.py

@@ -30,7 +30,7 @@ from dstack._internal.server.schemas.runner import (
 from dstack._internal.utils.common import get_or_error
 from dstack._internal.utils.logging import get_logger
 
-REQUEST_TIMEOUT = 15
+REQUEST_TIMEOUT = 9
 
 logger = get_logger(__name__)
 

dstack/_internal/server/services/runner/ssh.py

@@ -17,22 +17,36 @@ from dstack._internal.utils.path import FileContent
 logger = get_logger(__name__)
 P = ParamSpec("P")
 R = TypeVar("R")
+# A host private key or pair of (host private key, optional proxy jump private key)
+PrivateKeyOrPair = Union[str, tuple[str, Optional[str]]]
 
 
 def runner_ssh_tunnel(
     ports: List[int], retries: int = 3, retry_interval: float = 1
 ) -> Callable[
     [Callable[Concatenate[Dict[int, int], P], R]],
-    Callable[Concatenate[str, JobProvisioningData, Optional[JobRuntimeData], P], Union[bool, R]],
+    Callable[
+        Concatenate[PrivateKeyOrPair, JobProvisioningData, Optional[JobRuntimeData], P],
+        Union[bool, R],
+    ],
 ]:
+    """
+    A decorator that opens an SSH tunnel to the runner.
+
+    NOTE: connections from dstack-server to running jobs are expected to be short.
+    The runner uses a heuristic to differentiate dstack-server connections from
+    client connections based on their duration. See `ConnectionTracker` for details.
+    """
+
     def decorator(
         func: Callable[Concatenate[Dict[int, int], P], R],
     ) -> Callable[
-        Concatenate[str, JobProvisioningData, Optional[JobRuntimeData], P], Union[bool, R]
+        Concatenate[PrivateKeyOrPair, JobProvisioningData, Optional[JobRuntimeData], P],
+        Union[bool, R],
     ]:
         @functools.wraps(func)
         def wrapper(
-            ssh_private_key: str,
+            ssh_private_key: PrivateKeyOrPair,
             job_provisioning_data: JobProvisioningData,
             job_runtime_data: Optional[JobRuntimeData],
             *args: P.args,
@@ -51,6 +65,20 @@ def runner_ssh_tunnel(
                 # without SSH
                 return func(container_ports_map, *args, **kwargs)
 
+            if isinstance(ssh_private_key, str):
+                ssh_proxy_private_key = None
+            else:
+                ssh_private_key, ssh_proxy_private_key = ssh_private_key
+            identity = FileContent(ssh_private_key)
+            if ssh_proxy_private_key is not None:
+                proxy_identity = FileContent(ssh_proxy_private_key)
+            else:
+                proxy_identity = None
+
+            ssh_proxies = []
+            if job_provisioning_data.ssh_proxy is not None:
+                ssh_proxies.append((job_provisioning_data.ssh_proxy, proxy_identity))
+
             for attempt in range(retries):
                 last = attempt == retries - 1
                 # remote_host:local mapping
@@ -66,8 +94,8 @@ def runner_ssh_tunnel(
                         ),
                         port=job_provisioning_data.ssh_port,
                         forwarded_sockets=ports_to_forwarded_sockets(tunnel_ports_map),
-                        identity=FileContent(ssh_private_key),
-                        ssh_proxy=job_provisioning_data.ssh_proxy,
+                        identity=identity,
+                        ssh_proxies=ssh_proxies,
                     ):
                         return func(runner_ports_map, *args, **kwargs)
                 except SSHError: