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

dstack/_internal/core/services/ssh/tunnel.py

@@ -69,13 +69,16 @@ class SSHTunnel:
         options: Dict[str, str] = SSH_DEFAULT_OPTIONS,
         ssh_config_path: Union[PathLike, Literal["none"]] = "none",
         port: Optional[int] = None,
-        ssh_proxy: Optional[SSHConnectionParams] = None,
+        ssh_proxies: Iterable[tuple[SSHConnectionParams, Optional[FilePathOrContent]]] = (),
     ):
         """
         :param forwarded_sockets: Connections to the specified local sockets will be
             forwarded to their corresponding remote sockets
         :param reverse_forwarded_sockets: Connections to the specified remote sockets
             will be forwarded to their corresponding local sockets
+        :param ssh_proxies: pairs of SSH connections params and optional identities,
+            in order from outer to inner. If an identity is `None`, the `identity` param
+            is used instead.
         """
         self.destination = destination
         self.forwarded_sockets = list(forwarded_sockets)
@@ -83,21 +86,21 @@ class SSHTunnel:
         self.options = options
         self.port = port
         self.ssh_config_path = normalize_path(ssh_config_path)
-        self.ssh_proxy = ssh_proxy
         temp_dir = tempfile.TemporaryDirectory()
         self.temp_dir = temp_dir
         if control_sock_path is None:
             control_sock_path = os.path.join(temp_dir.name, "control.sock")
         self.control_sock_path = normalize_path(control_sock_path)
-        if isinstance(identity, FilePath):
-            identity_path = identity.path
-        else:
-            identity_path = os.path.join(temp_dir.name, "identity")
-            with open(
-                identity_path, opener=lambda path, flags: os.open(path, flags, 0o600), mode="w"
-            ) as f:
-                f.write(identity.content)
-        self.identity_path = normalize_path(identity_path)
+        self.identity_path = normalize_path(self._get_identity_path(identity, "identity"))
+        self.ssh_proxies: list[tuple[SSHConnectionParams, PathLike]] = []
+        for proxy_index, (proxy_params, proxy_identity) in enumerate(ssh_proxies):
+            if proxy_identity is None:
+                proxy_identity_path = self.identity_path
+            else:
+                proxy_identity_path = self._get_identity_path(
+                    proxy_identity, f"proxy_identity_{proxy_index}"
+                )
+            self.ssh_proxies.append((proxy_params, proxy_identity_path))
         self.log_path = normalize_path(os.path.join(temp_dir.name, "tunnel.log"))
         self.ssh_client_info = get_ssh_client_info()
         self.ssh_exec_path = str(self.ssh_client_info.path)
@@ -142,8 +145,8 @@ class SSHTunnel:
             command += ["-p", str(self.port)]
         for k, v in self.options.items():
             command += ["-o", f"{k}={v}"]
-        if proxy_command := self.proxy_command():
-            command += ["-o", "ProxyCommand=" + shlex.join(proxy_command)]
+        if proxy_command := self._get_proxy_command():
+            command += ["-o", proxy_command]
         for socket_pair in self.forwarded_sockets:
             command += ["-L", f"{socket_pair.local.render()}:{socket_pair.remote.render()}"]
         for socket_pair in self.reverse_forwarded_sockets:
@@ -160,24 +163,6 @@ class SSHTunnel:
     def exec_command(self) -> List[str]:
         return [self.ssh_exec_path, "-S", self.control_sock_path, self.destination]
 
-    def proxy_command(self) -> Optional[List[str]]:
-        if self.ssh_proxy is None:
-            return None
-        return [
-            self.ssh_exec_path,
-            "-i",
-            self.identity_path,
-            "-W",
-            "%h:%p",
-            "-o",
-            "StrictHostKeyChecking=no",
-            "-o",
-            "UserKnownHostsFile=/dev/null",
-            "-p",
-            str(self.ssh_proxy.port),
-            f"{self.ssh_proxy.username}@{self.ssh_proxy.hostname}",
-        ]
-
     def open(self) -> None:
         # We cannot use `stderr=subprocess.PIPE` here since the forked process (daemon) does not
         # close standard streams if ProxyJump is used, therefore we will wait EOF from the pipe
@@ -251,6 +236,38 @@ class SSHTunnel:
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
+    def _get_proxy_command(self) -> Optional[str]:
+        proxy_command: Optional[str] = None
+        for params, identity_path in self.ssh_proxies:
+            proxy_command = self._build_proxy_command(params, identity_path, proxy_command)
+        return proxy_command
+
+    def _build_proxy_command(
+        self,
+        params: SSHConnectionParams,
+        identity_path: PathLike,
+        prev_proxy_command: Optional[str],
+    ) -> Optional[str]:
+        command = [
+            self.ssh_exec_path,
+            "-i",
+            identity_path,
+            "-W",
+            "%h:%p",
+            "-o",
+            "StrictHostKeyChecking=no",
+            "-o",
+            "UserKnownHostsFile=/dev/null",
+        ]
+        if prev_proxy_command is not None:
+            command += ["-o", prev_proxy_command.replace("%", "%%")]
+        command += [
+            "-p",
+            str(params.port),
+            f"{params.username}@{params.hostname}",
+        ]
+        return "ProxyCommand=" + shlex.join(command)
+
     def _read_log_file(self) -> bytes:
         with open(self.log_path, "rb") as f:
             return f.read()
@@ -263,6 +280,16 @@ class SSHTunnel:
         except OSError as e:
             logger.debug("Failed to remove SSH tunnel log file %s: %s", self.log_path, e)
 
+    def _get_identity_path(self, identity: FilePathOrContent, tmp_filename: str) -> PathLike:
+        if isinstance(identity, FilePath):
+            return identity.path
+        identity_path = os.path.join(self.temp_dir.name, tmp_filename)
+        with open(
+            identity_path, opener=lambda path, flags: os.open(path, flags, 0o600), mode="w"
+        ) as f:
+            f.write(identity.content)
+        return identity_path
+
 
 def ports_to_forwarded_sockets(
     ports: Dict[int, int], bind_local: str = "localhost"

dstack/_internal/proxy/gateway/routers/registry.py

@@ -76,6 +76,8 @@ async def register_replica(
         ssh_destination=body.ssh_host,
         ssh_port=body.ssh_port,
         ssh_proxy=body.ssh_proxy,
+        ssh_head_proxy=body.ssh_head_proxy,
+        ssh_head_proxy_private_key=body.ssh_head_proxy_private_key,
         repo=repo,
         nginx=nginx,
         service_conn_pool=service_conn_pool,

dstack/_internal/proxy/gateway/schemas/registry.py

@@ -50,6 +50,8 @@ class RegisterReplicaRequest(BaseModel):
     ssh_host: str
     ssh_port: int
     ssh_proxy: Optional[SSHConnectionParams]
+    ssh_head_proxy: Optional[SSHConnectionParams]
+    ssh_head_proxy_private_key: Optional[str]
 
 
 class RegisterEntrypointRequest(BaseModel):

dstack/_internal/proxy/gateway/services/registry.py

@@ -123,6 +123,8 @@ async def register_replica(
     ssh_destination: str,
     ssh_port: int,
     ssh_proxy: Optional[SSHConnectionParams],
+    ssh_head_proxy: Optional[SSHConnectionParams],
+    ssh_head_proxy_private_key: Optional[str],
     repo: GatewayProxyRepo,
     nginx: Nginx,
     service_conn_pool: ServiceConnectionPool,
@@ -133,6 +135,8 @@ async def register_replica(
         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,
     )
 
     async with lock:

dstack/_internal/proxy/lib/models.py

@@ -23,6 +23,9 @@ class Replica(ImmutableModel):
     ssh_destination: str
     ssh_port: int
     ssh_proxy: Optional[SSHConnectionParams]
+    # Optional outer proxy, a head node/bastion
+    ssh_head_proxy: Optional[SSHConnectionParams] = None
+    ssh_head_proxy_private_key: Optional[str] = None
 
 
 class Service(ImmutableModel):

dstack/_internal/proxy/lib/services/service_connection.py

@@ -18,6 +18,7 @@ from dstack._internal.core.services.ssh.tunnel import (
 from dstack._internal.proxy.lib.errors import UnexpectedProxyError
 from dstack._internal.proxy.lib.models import Project, Replica, Service
 from dstack._internal.proxy.lib.repo import BaseProxyRepo
+from dstack._internal.utils.common import get_or_error
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.path import FileContent
 
@@ -45,10 +46,16 @@ class ServiceConnection:
             os.chmod(self._temp_dir.name, 0o755)
             options["StreamLocalBindMask"] = "0111"
         self._app_socket_path = (Path(self._temp_dir.name) / "replica.sock").absolute()
+        ssh_proxies = []
+        if replica.ssh_head_proxy is not None:
+            ssh_head_proxy_private_key = get_or_error(replica.ssh_head_proxy_private_key)
+            ssh_proxies.append((replica.ssh_head_proxy, FileContent(ssh_head_proxy_private_key)))
+        if replica.ssh_proxy is not None:
+            ssh_proxies.append((replica.ssh_proxy, None))
         self._tunnel = SSHTunnel(
             destination=replica.ssh_destination,
             port=replica.ssh_port,
-            ssh_proxy=replica.ssh_proxy,
+            ssh_proxies=ssh_proxies,
             identity=FileContent(project.ssh_private_key),
             forwarded_sockets=[
                 SocketPair(

dstack/_internal/server/background/tasks/process_instances.py

@@ -42,12 +42,12 @@ from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.fleets import InstanceGroupPlacement
 from dstack._internal.core.models.instances import (
     InstanceAvailability,
-    InstanceConfiguration,
     InstanceOfferWithAvailability,
     InstanceRuntime,
     InstanceStatus,
     InstanceType,
     RemoteConnectionInfo,
+    SSHKey,
 )
 from dstack._internal.core.models.placement import (
     PlacementGroup,
@@ -77,6 +77,7 @@ from dstack._internal.server.services.fleets import (
     get_create_instance_offers,
 )
 from dstack._internal.server.services.locking import get_locker
+from dstack._internal.server.services.offers import is_divisible_into_blocks
 from dstack._internal.server.services.placement import (
     get_fleet_placement_groups,
     placement_group_model_to_placement_group,
@@ -86,6 +87,7 @@ from dstack._internal.server.services.pools import (
     get_instance_profile,
     get_instance_provisioning_data,
     get_instance_requirements,
+    get_instance_ssh_private_keys,
 )
 from dstack._internal.server.services.runner import client as runner_client
 from dstack._internal.server.services.runner.client import HealthStatus
@@ -133,7 +135,7 @@ async def _process_next_instance():
                     ),
                     InstanceModel.id.not_in(lockset),
                 )
-                .options(lazyload(InstanceModel.job))
+                .options(lazyload(InstanceModel.jobs))
                 .order_by(InstanceModel.last_processed_at.asc())
                 .limit(1)
                 .with_for_update(skip_locked=True)
@@ -156,7 +158,7 @@ async def _process_instance(session: AsyncSession, instance: InstanceModel):
         select(InstanceModel)
         .where(InstanceModel.id == instance.id)
         .options(joinedload(InstanceModel.project).joinedload(ProjectModel.backends))
-        .options(joinedload(InstanceModel.job))
+        .options(joinedload(InstanceModel.jobs))
         .options(joinedload(InstanceModel.fleet).joinedload(FleetModel.instances))
         .execution_options(populate_existing=True)
     )
@@ -164,7 +166,7 @@ async def _process_instance(session: AsyncSession, instance: InstanceModel):
     if (
         instance.status == InstanceStatus.IDLE
         and instance.termination_policy == TerminationPolicy.DESTROY_AFTER_IDLE
-        and instance.job_id is None
+        and not instance.jobs
     ):
         await _mark_terminating_if_idle_duration_expired(instance)
     if instance.status == InstanceStatus.PENDING:
@@ -232,11 +234,11 @@ async def _add_remote(instance: InstanceModel) -> None:
         remote_details = RemoteConnectionInfo.parse_raw(cast(str, instance.remote_connection_info))
         # Prepare connection key
         try:
-            pkeys = [
-                pkey_from_str(sk.private)
-                for sk in remote_details.ssh_keys
-                if sk.private is not None
-            ]
+            pkeys = _ssh_keys_to_pkeys(remote_details.ssh_keys)
+            if remote_details.ssh_proxy_keys is not None:
+                ssh_proxy_pkeys = _ssh_keys_to_pkeys(remote_details.ssh_proxy_keys)
+            else:
+                ssh_proxy_pkeys = None
         except (ValueError, PasswordRequiredException):
             instance.status = InstanceStatus.TERMINATED
             instance.termination_reason = "Unsupported private SSH key type"
@@ -254,7 +256,9 @@ async def _add_remote(instance: InstanceModel) -> None:
         authorized_keys.append(instance.project.ssh_public_key.strip())
 
         try:
-            future = run_async(_deploy_instance, remote_details, pkeys, authorized_keys)
+            future = run_async(
+                _deploy_instance, remote_details, pkeys, ssh_proxy_pkeys, authorized_keys
+            )
             deploy_timeout = 20 * 60  # 20 minutes
             result = await asyncio.wait_for(future, timeout=deploy_timeout)
             health, host_info = result
@@ -322,6 +326,26 @@ async def _add_remote(instance: InstanceModel) -> None:
             )
             return
 
+    divisible, blocks = is_divisible_into_blocks(
+        cpu_count=instance_type.resources.cpus,
+        gpu_count=len(instance_type.resources.gpus),
+        blocks="auto" if instance.total_blocks is None else instance.total_blocks,
+    )
+    if divisible:
+        instance.total_blocks = blocks
+    else:
+        instance.status = InstanceStatus.TERMINATED
+        instance.termination_reason = "Cannot split into blocks"
+        logger.warning(
+            "Failed to add instance %s: cannot split into blocks",
+            instance.name,
+            extra={
+                "instance_name": instance.name,
+                "instance_status": InstanceStatus.TERMINATED.value,
+            },
+        )
+        return
+
     region = instance.region
     jpd = JobProvisioningData(
         backend=BackendType.REMOTE,
@@ -336,7 +360,7 @@ async def _add_remote(instance: InstanceModel) -> None:
         ssh_port=remote_details.port,
         dockerized=True,
         backend_data=None,
-        ssh_proxy=None,
+        ssh_proxy=remote_details.ssh_proxy,
     )
 
     instance.status = InstanceStatus.IDLE if health else InstanceStatus.PROVISIONING
@@ -359,10 +383,16 @@ async def _add_remote(instance: InstanceModel) -> None:
 def _deploy_instance(
     remote_details: RemoteConnectionInfo,
     pkeys: List[PKey],
+    ssh_proxy_pkeys: Optional[list[PKey]],
     authorized_keys: List[str],
 ) -> Tuple[HealthStatus, Dict[str, Any]]:
     with get_paramiko_connection(
-        remote_details.ssh_user, remote_details.host, remote_details.port, pkeys
+        remote_details.ssh_user,
+        remote_details.host,
+        remote_details.port,
+        pkeys,
+        remote_details.ssh_proxy,
+        ssh_proxy_pkeys,
     ) as client:
         logger.info(f"Connected to {remote_details.ssh_user} {remote_details.host}")
 
@@ -479,6 +509,7 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
         requirements=requirements,
         exclude_not_available=True,
         fleet_model=instance.fleet,
+        blocks="auto" if instance.total_blocks is None else instance.total_blocks,
     )
 
     if not offers and should_retry:
@@ -496,11 +527,10 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
             session=session, fleet_id=instance.fleet_id
         )
 
-    instance_configuration = _patch_instance_configuration(instance)
-
     for backend, instance_offer in offers:
         if instance_offer.backend not in BACKENDS_WITH_CREATE_INSTANCE_SUPPORT:
             continue
+        instance_offer = _get_instance_offer_for_instance(instance_offer, instance)
         if (
             instance_offer.backend in BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT
             and instance.fleet
@@ -554,6 +584,7 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
         instance.instance_configuration = instance_configuration.json()
         instance.job_provisioning_data = job_provisioning_data.json()
         instance.offer = instance_offer.json()
+        instance.total_blocks = instance_offer.total_blocks
         instance.started_at = get_current_datetime()
         instance.last_retry_at = get_current_datetime()
 
@@ -585,8 +616,8 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
 async def _check_instance(instance: InstanceModel) -> None:
     if (
         instance.status == InstanceStatus.BUSY
-        and instance.job is not None
-        and instance.job.status.is_finished()
+        and instance.jobs
+        and all(job.status.is_finished() for job in instance.jobs)
     ):
         # A busy instance could have no active jobs due to this bug: https://github.com/dstackai/dstack/issues/2068
         instance.status = InstanceStatus.TERMINATING
@@ -617,18 +648,14 @@ async def _check_instance(instance: InstanceModel) -> None:
             instance.status = InstanceStatus.BUSY
         return
 
-    ssh_private_key = instance.project.ssh_private_key
-    # TODO: Drop this logic and always use project key once it's safe to assume that most on-prem
-    # fleets are (re)created after this change: https://github.com/dstackai/dstack/pull/1716
-    if instance.remote_connection_info is not None:
-        remote_conn_info: RemoteConnectionInfo = RemoteConnectionInfo.__response__.parse_raw(
-            instance.remote_connection_info
-        )
-        ssh_private_key = remote_conn_info.ssh_keys[0].private
+    ssh_private_keys = get_instance_ssh_private_keys(instance)
 
     # May return False if fails to establish ssh connection
     health_status_response = await run_async(
-        _instance_healthcheck, ssh_private_key, job_provisioning_data, None
+        _instance_healthcheck,
+        ssh_private_keys,
+        job_provisioning_data,
+        None,
     )
     if isinstance(health_status_response, bool) or health_status_response is None:
         health_status = HealthStatus(healthy=False, reason="SSH or tunnel error")
@@ -648,9 +675,7 @@ async def _check_instance(instance: InstanceModel) -> None:
         instance.unreachable = False
 
         if instance.status == InstanceStatus.PROVISIONING:
-            instance.status = (
-                InstanceStatus.IDLE if instance.job_id is None else InstanceStatus.BUSY
-            )
+            instance.status = InstanceStatus.IDLE if not instance.jobs else InstanceStatus.BUSY
             logger.info(
                 "Instance %s has switched to %s status",
                 instance.name,
@@ -869,21 +894,31 @@ def _need_to_wait_fleet_provisioning(instance: InstanceModel) -> bool:
     )
 
 
-def _patch_instance_configuration(instance: InstanceModel) -> InstanceConfiguration:
-    instance_configuration = get_instance_configuration(instance)
+def _get_instance_offer_for_instance(
+    instance_offer: InstanceOfferWithAvailability,
+    instance: InstanceModel,
+) -> InstanceOfferWithAvailability:
     if instance.fleet is None:
-        return instance_configuration
+        return instance_offer
 
     fleet = fleet_model_to_fleet(instance.fleet)
     master_instance = instance.fleet.instances[0]
     master_job_provisioning_data = get_instance_provisioning_data(master_instance)
+    instance_offer = instance_offer.copy()
     if (
         fleet.spec.configuration.placement == InstanceGroupPlacement.CLUSTER
         and master_job_provisioning_data is not None
+        and master_job_provisioning_data.availability_zone is not None
     ):
-        instance_configuration.availability_zone = master_job_provisioning_data.availability_zone
+        if instance_offer.availability_zones is None:
+            instance_offer.availability_zones = [master_job_provisioning_data.availability_zone]
+        instance_offer.availability_zones = [
+            z
+            for z in instance_offer.availability_zones
+            if instance_offer.availability_zones == master_job_provisioning_data.availability_zone
+        ]
 
-    return instance_configuration
+    return instance_offer
 
 
 def _create_placement_group_if_does_not_exist(
@@ -942,3 +977,7 @@ def _get_instance_timeout_interval(
     if backend_type == BackendType.VULTR and instance_type_name.startswith("vbm"):
         return timedelta(seconds=3300)
     return timedelta(seconds=600)
+
+
+def _ssh_keys_to_pkeys(ssh_keys: list[SSHKey]) -> list[PKey]:
+    return [pkey_from_str(sk.private) for sk in ssh_keys if sk.private is not None]

dstack/_internal/server/background/tasks/process_metrics.py

@@ -3,18 +3,19 @@ import json
 from typing import Dict, List, Optional
 
 from sqlalchemy import delete, select
-from sqlalchemy.orm import selectinload
+from sqlalchemy.orm import joinedload
 
 from dstack._internal.core.consts import DSTACK_RUNNER_HTTP_PORT
 from dstack._internal.core.models.runs import JobStatus
 from dstack._internal.server import settings
 from dstack._internal.server.db import get_session_ctx
-from dstack._internal.server.models import JobMetricsPoint, JobModel
+from dstack._internal.server.models import InstanceModel, JobMetricsPoint, JobModel
 from dstack._internal.server.schemas.runner import MetricsResponse
 from dstack._internal.server.services.jobs import get_job_provisioning_data, get_job_runtime_data
+from dstack._internal.server.services.pools import get_instance_ssh_private_keys
 from dstack._internal.server.services.runner import client
 from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel
-from dstack._internal.utils.common import batched, get_current_datetime, run_async
+from dstack._internal.utils.common import batched, get_current_datetime, get_or_error, run_async
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
@@ -29,14 +30,12 @@ async def collect_metrics():
     async with get_session_ctx() as session:
         res = await session.execute(
             select(JobModel)
-            .where(
-                JobModel.status.in_([JobStatus.RUNNING]),
-            )
-            .options(selectinload(JobModel.project))
+            .where(JobModel.status.in_([JobStatus.RUNNING]))
+            .options(joinedload(JobModel.instance).joinedload(InstanceModel.project))
             .order_by(JobModel.last_processed_at.asc())
             .limit(MAX_JOBS_FETCHED)
         )
-        job_models = res.scalars().all()
+        job_models = res.unique().scalars().all()
 
     for batch in batched(job_models, BATCH_SIZE):
         await _collect_jobs_metrics(batch)
@@ -87,6 +86,7 @@ def _get_recently_collected_metric_cutoff() -> int:
 
 
 async def _collect_job_metrics(job_model: JobModel) -> Optional[JobMetricsPoint]:
+    ssh_private_keys = get_instance_ssh_private_keys(get_or_error(job_model.instance))
     jpd = get_job_provisioning_data(job_model)
     jrd = get_job_runtime_data(job_model)
     if jpd is None:
@@ -94,7 +94,7 @@ async def _collect_job_metrics(job_model: JobModel) -> Optional[JobMetricsPoint]
     try:
         res = await run_async(
             _pull_runner_metrics,
-            job_model.project.ssh_private_key,
+            ssh_private_keys,
             jpd,
             jrd,
         )