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

dstack/_internal/core/backends/remote/provisioning.py

@@ -1,9 +1,9 @@
 import io
 import json
 import time
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from textwrap import dedent
-from typing import Any, Dict, Generator, List
+from typing import Any, Dict, Generator, List, Optional
 
 import paramiko
 from gpuhunt import AcceleratorVendor, correct_gpu_memory_gib
@@ -17,6 +17,7 @@ from dstack._internal.core.models.instances import (
     Gpu,
     InstanceType,
     Resources,
+    SSHConnectionParams,
 )
 from dstack._internal.utils.gpu import (
     convert_amd_gpu_name,
@@ -262,35 +263,72 @@ def host_info_to_instance_type(host_info: Dict[str, Any]) -> InstanceType:
 
 @contextmanager
 def get_paramiko_connection(
-    ssh_user: str, host: str, port: int, pkeys: List[paramiko.PKey]
+    ssh_user: str,
+    host: str,
+    port: int,
+    pkeys: List[paramiko.PKey],
+    proxy: Optional[SSHConnectionParams] = None,
+    proxy_pkeys: Optional[list[paramiko.PKey]] = None,
 ) -> Generator[paramiko.SSHClient, None, None]:
-    with paramiko.SSHClient() as client:
-        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-        for pkey in pkeys:
-            conn_url = f"{ssh_user}@{host}:{port}"
+    if proxy is not None:
+        if proxy_pkeys is None:
+            raise ProvisioningError("Missing proxy private keys")
+        proxy_ctx = get_paramiko_connection(
+            proxy.username, proxy.hostname, proxy.port, proxy_pkeys
+        )
+    else:
+        proxy_ctx = nullcontext()
+    conn_url = f"{ssh_user}@{host}:{port}"
+    with proxy_ctx as proxy_client, paramiko.SSHClient() as client:
+        proxy_channel: Optional[paramiko.Channel] = None
+        if proxy_client is not None:
             try:
-                logger.debug("Try to connect to %s with key %s", conn_url, pkey.fingerprint)
-                client.connect(
-                    username=ssh_user,
-                    hostname=host,
-                    port=port,
-                    pkey=pkey,
-                    look_for_keys=False,
-                    allow_agent=False,
-                    timeout=SSH_CONNECT_TIMEOUT,
+                proxy_channel = proxy_client.get_transport().open_channel(
+                    "direct-tcpip", (host, port), ("", 0)
                 )
-            except paramiko.AuthenticationException:
-                logger.debug(
-                    f'Authentication failed to connect to "{conn_url}" and {pkey.fingerprint}'
-                )
-                continue  # try next key
             except (paramiko.SSHException, OSError) as e:
-                raise ProvisioningError(f"Connect failed: {e}") from e
-            else:
+                raise ProvisioningError(f"Proxy channel failed: {e}") from e
+        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        for pkey in pkeys:
+            logger.debug("Try to connect to %s with key %s", conn_url, pkey.fingerprint)
+            connected = _paramiko_connect(client, ssh_user, host, port, pkey, proxy_channel)
+            if connected:
                 yield client
                 return
-        else:
-            keys_fp = ", ".join(f"{pk.fingerprint!r}" for pk in pkeys)
-            raise ProvisioningError(
-                f"SSH connection to the {conn_url} with keys [{keys_fp}] was unsuccessful"
+            logger.debug(
+                f'Authentication failed to connect to "{conn_url}" and {pkey.fingerprint}'
             )
+        keys_fp = ", ".join(f"{pk.fingerprint!r}" for pk in pkeys)
+        raise ProvisioningError(
+            f"SSH connection to the {conn_url} with keys [{keys_fp}] was unsuccessful"
+        )
+
+
+def _paramiko_connect(
+    client: paramiko.SSHClient,
+    user: str,
+    host: str,
+    port: int,
+    pkey: paramiko.PKey,
+    channel: Optional[paramiko.Channel] = None,
+) -> bool:
+    """
+    Returns `True` if connected, `False` if auth failed, and raises `ProvisioningError`
+    on other errors.
+    """
+    try:
+        client.connect(
+            username=user,
+            hostname=host,
+            port=port,
+            pkey=pkey,
+            look_for_keys=False,
+            allow_agent=False,
+            timeout=SSH_CONNECT_TIMEOUT,
+            sock=channel,
+        )
+        return True
+    except paramiko.AuthenticationException:
+        return False
+    except (paramiko.SSHException, OSError) as e:
+        raise ProvisioningError(f"Connect failed: {e}") from e

dstack/_internal/core/backends/runpod/compute.py

@@ -113,6 +113,7 @@ class RunpodCompute(Compute):
             bid_per_gpu=bid_per_gpu,
             network_volume_id=network_volume_id,
             volume_mount_path=volume_mount_path,
+            env={"RUNPOD_POD_USER": "0"},
         )
 
         instance_id = resp["id"]

dstack/_internal/core/models/backends/azure.py

@@ -11,6 +11,7 @@ class AzureConfigInfo(CoreModel):
     type: Literal["azure"] = "azure"
     tenant_id: str
     subscription_id: str
+    resource_group: Optional[str] = None
     locations: Optional[List[str]] = None
     vpc_ids: Optional[Dict[str, str]] = None
     public_ips: Optional[bool] = None
@@ -48,6 +49,7 @@ class AzureConfigInfoWithCredsPartial(CoreModel):
     creds: Optional[AnyAzureCreds]
     tenant_id: Optional[str]
     subscription_id: Optional[str]
+    resource_group: Optional[str]
     locations: Optional[List[str]]
     vpc_ids: Optional[Dict[str, str]]
     public_ips: Optional[bool]
@@ -63,4 +65,4 @@ class AzureConfigValues(CoreModel):
 
 
 class AzureStoredConfig(AzureConfigInfo):
-    resource_group: str
+    resource_group: str = ""

dstack/_internal/core/models/configurations.py

@@ -10,7 +10,7 @@ from dstack._internal.core.models.common import CoreModel, Duration, RegistryAut
 from dstack._internal.core.models.envs import Env
 from dstack._internal.core.models.fleets import FleetConfiguration
 from dstack._internal.core.models.gateways import GatewayConfiguration
-from dstack._internal.core.models.profiles import ProfileParams
+from dstack._internal.core.models.profiles import ProfileParams, parse_off_duration
 from dstack._internal.core.models.repos.base import Repo
 from dstack._internal.core.models.repos.virtual import VirtualRepo
 from dstack._internal.core.models.resources import Range, ResourcesSpec
@@ -212,6 +212,29 @@ class DevEnvironmentConfigurationParams(CoreModel):
     ide: Annotated[Literal["vscode"], Field(description="The IDE to run")]
     version: Annotated[Optional[str], Field(description="The version of the IDE")]
     init: Annotated[CommandsList, Field(description="The bash commands to run on startup")] = []
+    inactivity_duration: Annotated[
+        Optional[Union[Literal["off"], int, bool, str]],
+        Field(
+            description=(
+                "The maximum amount of time the dev environment can be inactive"
+                " (e.g., `2h`, `1d`, etc)."
+                " After it elapses, the dev environment is automatically stopped."
+                " Inactivity is defined as the absence of SSH connections to the"
+                " dev environment, including VS Code connections, `ssh <run name>`"
+                " shells, and attached `dstack apply` or `dstack attach` commands."
+                " Use `off` for unlimited duration. Defaults to `off`"
+            )
+        ),
+    ]
+
+    @validator("inactivity_duration", pre=True, allow_reuse=True)
+    def parse_inactivity_duration(
+        cls, v: Optional[Union[Literal["off"], int, bool, str]]
+    ) -> Optional[int]:
+        v = parse_off_duration(v)
+        if isinstance(v, int):
+            return v
+        return None
 
 
 class DevEnvironmentConfiguration(

dstack/_internal/core/models/fleets.py

@@ -39,6 +39,14 @@ class InstanceGroupPlacement(str, Enum):
     CLUSTER = "cluster"
 
 
+class SSHProxyParams(CoreModel):
+    hostname: Annotated[str, Field(description="The IP address or domain of proxy host")]
+    port: Annotated[Optional[int], Field(description="The SSH port of proxy host")] = None
+    user: Annotated[str, Field(description="The user to log in with for proxy host")]
+    identity_file: Annotated[str, Field(description="The private key to use for proxy host")]
+    ssh_key: Optional[SSHKey] = None
+
+
 class SSHHostParams(CoreModel):
     hostname: Annotated[str, Field(description="The IP address or domain to connect to")]
     port: Annotated[
@@ -50,6 +58,9 @@ class SSHHostParams(CoreModel):
     identity_file: Annotated[
         Optional[str], Field(description="The private key to use for this host")
     ] = None
+    proxy_jump: Annotated[
+        Optional[SSHProxyParams], Field(description="The SSH proxy configuration for this host")
+    ] = None
     internal_ip: Annotated[
         Optional[str],
         Field(
@@ -61,6 +72,19 @@ class SSHHostParams(CoreModel):
     ] = None
     ssh_key: Optional[SSHKey] = None
 
+    blocks: Annotated[
+        Union[Literal["auto"], int],
+        Field(
+            description=(
+                "The amount of blocks to split the instance into, a number or `auto`."
+                " `auto` means as many as possible."
+                " The number of GPUs and CPUs must be divisible by the number of blocks."
+                " Defaults to `1`, i.e. do not split"
+            ),
+            ge=1,
+        ),
+    ] = 1
+
     @validator("internal_ip")
     def validate_internal_ip(cls, value):
         if value is None:
@@ -83,6 +107,9 @@ class SSHParams(CoreModel):
         Optional[str], Field(description="The private key to use for all hosts")
     ] = None
     ssh_key: Optional[SSHKey] = None
+    proxy_jump: Annotated[
+        Optional[SSHProxyParams], Field(description="The SSH proxy configuration for all hosts")
+    ] = None
     hosts: Annotated[
         List[Union[SSHHostParams, str]],
         Field(
@@ -142,6 +169,19 @@ class InstanceGroupParams(CoreModel):
         Field(description="The resources requirements"),
     ] = ResourcesSpec()
 
+    blocks: Annotated[
+        Union[Literal["auto"], int],
+        Field(
+            description=(
+                "The amount of blocks to split the instance into, a number or `auto`."
+                " `auto` means as many as possible."
+                " The number of GPUs and CPUs must be divisible by the number of blocks."
+                " Defaults to `1`, i.e. do not split"
+            ),
+            ge=1,
+        ),
+    ] = 1
+
     backends: Annotated[
         Optional[List[BackendType]],
         Field(description="The backends to consider for provisioning (e.g., `[aws, gcp]`)"),
@@ -152,6 +192,12 @@ class InstanceGroupParams(CoreModel):
             description="The regions to consider for provisioning (e.g., `[eu-west-1, us-west4, westeurope]`)"
         ),
     ] = None
+    availability_zones: Annotated[
+        Optional[List[str]],
+        Field(
+            description="The availability zones to consider for provisioning (e.g., `[eu-west-1a, us-west4-a]`)"
+        ),
+    ] = None
     instance_types: Annotated[
         Optional[List[str]],
         Field(

dstack/_internal/core/models/instances.py

@@ -92,6 +92,8 @@ class RemoteConnectionInfo(CoreModel):
     port: int
     ssh_user: str
     ssh_keys: List[SSHKey]
+    ssh_proxy: Optional[SSHConnectionParams] = None
+    ssh_proxy_keys: Optional[list[SSHKey]] = None
     env: Env = Env()
 
 
@@ -101,7 +103,6 @@ class InstanceConfiguration(CoreModel):
     user: str  # dstack user name
     ssh_keys: List[SSHKey]
     instance_id: Optional[str] = None
-    availability_zone: Optional[str] = None
     placement_group_name: Optional[str] = None
     reservation: Optional[str] = None
     volumes: Optional[List[Volume]] = None
@@ -140,7 +141,10 @@ class InstanceOffer(CoreModel):
 
 class InstanceOfferWithAvailability(InstanceOffer):
     availability: InstanceAvailability
+    availability_zones: Optional[List[str]] = None
     instance_runtime: InstanceRuntime = InstanceRuntime.SHIM
+    blocks: int = 1
+    total_blocks: int = 1
 
 
 class InstanceStatus(str, Enum):

dstack/_internal/core/models/pools.py

@@ -25,14 +25,17 @@ class Instance(CoreModel):
     fleet_name: Optional[str] = None
     instance_num: int
     pool_name: Optional[str] = None
-    job_name: Optional[str] = None
+    job_name: Optional[str] = None  # deprecated, always None (instance can have more than one job)
     hostname: Optional[str] = None
     status: InstanceStatus
     unreachable: bool = False
     termination_reason: Optional[str] = None
     created: datetime.datetime
     region: Optional[str] = None
+    availability_zone: Optional[str] = None
     price: Optional[float] = None
+    total_blocks: Optional[int] = None
+    busy_blocks: int = 0
 
 
 class PoolInstances(CoreModel):

dstack/_internal/core/models/profiles.py

@@ -40,15 +40,15 @@ def parse_duration(v: Optional[Union[int, str]]) -> Optional[int]:
     return Duration.parse(v)
 
 
-def parse_max_duration(v: Optional[Union[int, str, bool]]) -> Optional[Union[str, int, bool]]:
+def parse_max_duration(v: Optional[Union[int, str, bool]]) -> Optional[Union[str, int]]:
     return parse_off_duration(v)
 
 
-def parse_stop_duration(v: Optional[Union[int, str, bool]]) -> Optional[Union[str, int, bool]]:
+def parse_stop_duration(v: Optional[Union[int, str, bool]]) -> Optional[Union[str, int]]:
     return parse_off_duration(v)
 
 
-def parse_off_duration(v: Optional[Union[int, str, bool]]) -> Optional[Union[str, int, bool]]:
+def parse_off_duration(v: Optional[Union[int, str, bool]]) -> Optional[Union[str, int]]:
     if v == "off" or v is False:
         return "off"
     if v is True:
@@ -123,6 +123,12 @@ class ProfileParams(CoreModel):
             description="The regions to consider for provisioning (e.g., `[eu-west-1, us-west4, westeurope]`)"
         ),
     ]
+    availability_zones: Annotated[
+        Optional[List[str]],
+        Field(
+            description="The availability zones to consider for provisioning (e.g., `[eu-west-1a, us-west4-a]`)"
+        ),
+    ] = None
     instance_types: Annotated[
         Optional[List[str]],
         Field(
@@ -162,7 +168,7 @@ class ProfileParams(CoreModel):
         Optional[Union[Literal["off"], str, int, bool]],
         Field(
             description=(
-                "The maximum duration of a run gracefull stopping."
+                "The maximum duration of a run graceful stopping."
                 " After it elapses, the run is automatically forced stopped."
                 " This includes force detaching volumes used by the run."
                 " Use `off` for unlimited duration. Defaults to `5m`"

dstack/_internal/core/models/runs.py

@@ -27,6 +27,7 @@ from dstack._internal.core.models.profiles import (
 from dstack._internal.core.models.repos import AnyRunRepoData
 from dstack._internal.core.models.resources import Memory, ResourcesSpec
 from dstack._internal.core.models.unix import UnixUser
+from dstack._internal.core.models.volumes import MountPoint
 from dstack._internal.utils import common as common_utils
 from dstack._internal.utils.common import format_pretty_duration
 
@@ -112,6 +113,7 @@ class JobTerminationReason(str, Enum):
     DONE_BY_RUNNER = "done_by_runner"
     ABORTED_BY_USER = "aborted_by_user"
     TERMINATED_BY_SERVER = "terminated_by_server"
+    INACTIVITY_DURATION_EXCEEDED = "inactivity_duration_exceeded"
     # Set by the runner
     CONTAINER_EXITED_WITH_ERROR = "container_exited_with_error"
     PORTS_BINDING_FAILED = "ports_binding_failed"
@@ -132,6 +134,7 @@ class JobTerminationReason(str, Enum):
             self.DONE_BY_RUNNER: JobStatus.DONE,
             self.ABORTED_BY_USER: JobStatus.ABORTED,
             self.TERMINATED_BY_SERVER: JobStatus.TERMINATED,
+            self.INACTIVITY_DURATION_EXCEEDED: JobStatus.TERMINATED,
             self.CONTAINER_EXITED_WITH_ERROR: JobStatus.FAILED,
             self.PORTS_BINDING_FAILED: JobStatus.FAILED,
             self.CREATING_CONTAINER_ERROR: JobStatus.FAILED,
@@ -190,6 +193,7 @@ class JobSpec(CoreModel):
     registry_auth: Optional[RegistryAuth]
     requirements: Requirements
     retry: Optional[Retry]
+    volumes: Optional[List[MountPoint]] = None
     # For backward compatibility with 0.18.x when retry_policy was required.
     # TODO: remove in 0.19
     retry_policy: ProfileRetryPolicy = ProfileRetryPolicy(retry=False)
@@ -231,6 +235,17 @@ class JobProvisioningData(CoreModel):
 
 
 class JobRuntimeData(CoreModel):
+    """
+    Holds various information only available after the job is submitted, such as:
+        * offer (depends on the instance)
+        * volumes used by the job
+        * resource constraints for container (depend on the instance)
+        * port mapping (reported by the shim only after the container is started)
+
+    Some fields are mutable, for example, `ports` only available when the shim starts
+    the container.
+    """
+
     network_mode: NetworkMode
     # GPU, CPU, memory resource shares. None means all available (no limit)
     gpu: Optional[int] = None
@@ -240,6 +255,10 @@ class JobRuntimeData(CoreModel):
     # None if data is not yet available (on vm-based backends and ssh instances)
     # or not applicable (container-based backends)
     ports: Optional[dict[int, int]] = None
+    # List of volumes used by the job
+    volume_names: Optional[list[str]] = None  # None for backward compalibility
+    # Virtual shared offer
+    offer: Optional[InstanceOfferWithAvailability] = None  # None for backward compalibility
 
 
 class ClusterInfo(CoreModel):
@@ -254,6 +273,7 @@ class JobSubmission(CoreModel):
     submitted_at: datetime
     last_processed_at: datetime
     finished_at: Optional[datetime]
+    inactivity_secs: Optional[int]
     status: JobStatus
     termination_reason: Optional[JobTerminationReason]
     termination_reason_message: Optional[str]

dstack/_internal/core/models/volumes.py

@@ -32,6 +32,9 @@ class VolumeConfiguration(CoreModel):
     name: Annotated[Optional[str], Field(description="The volume name")] = None
     backend: Annotated[BackendType, Field(description="The volume backend")]
     region: Annotated[str, Field(description="The volume region")]
+    availability_zone: Annotated[
+        Optional[str], Field(description="The volume availability zone")
+    ] = None
     size: Annotated[
         Optional[Memory],
         Field(description="The volume size. Must be specified when creating new volumes"),

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

@@ -2,7 +2,7 @@ import atexit
 import re
 import time
 from pathlib import Path
-from typing import Optional
+from typing import Optional, Union
 
 import psutil
 
@@ -14,6 +14,8 @@ from dstack._internal.core.services.ssh.ports import PortsLock
 from dstack._internal.core.services.ssh.tunnel import SSHTunnel, ports_to_forwarded_sockets
 from dstack._internal.utils.path import FilePath, PathLike
 from dstack._internal.utils.ssh import (
+    default_ssh_config_path,
+    get_host_config,
     include_ssh_config,
     normalize_path,
     update_ssh_config,
@@ -88,28 +90,63 @@ class SSHAttach:
             },
         )
         self.ssh_proxy = ssh_proxy
-        if ssh_proxy is None:
-            self.host_config = {
+
+        hosts: dict[str, dict[str, Union[str, int, FilePath]]] = {}
+        self.hosts = hosts
+
+        if local_backend:
+            hosts[run_name] = {
                 "HostName": hostname,
-                "Port": ssh_port,
-                "User": user if dockerized else container_user,
-                "IdentityFile": self.identity_file,
-                "IdentitiesOnly": "yes",
-                "StrictHostKeyChecking": "no",
-                "UserKnownHostsFile": "/dev/null",
-            }
-        else:
-            self.host_config = {
-                "HostName": ssh_proxy.hostname,
-                "Port": ssh_proxy.port,
-                "User": ssh_proxy.username,
+                "Port": container_ssh_port,
+                "User": container_user,
                 "IdentityFile": self.identity_file,
                 "IdentitiesOnly": "yes",
                 "StrictHostKeyChecking": "no",
                 "UserKnownHostsFile": "/dev/null",
             }
-        if dockerized and not local_backend:
-            self.container_config = {
+        elif dockerized:
+            if ssh_proxy is not None:
+                # SSH instance with jump host
+                # dstack has no IdentityFile for jump host, it must be either preconfigured
+                # in the ~/.ssh/config or loaded into ssh-agent
+                hosts[f"{run_name}-jump-host"] = {
+                    "HostName": ssh_proxy.hostname,
+                    "Port": ssh_proxy.port,
+                    "User": ssh_proxy.username,
+                    "StrictHostKeyChecking": "no",
+                    "UserKnownHostsFile": "/dev/null",
+                }
+                jump_host_config = get_host_config(ssh_proxy.hostname, default_ssh_config_path)
+                jump_host_identity_files = jump_host_config.get("identityfile")
+                if jump_host_identity_files:
+                    hosts[f"{run_name}-jump-host"].update(
+                        {
+                            "IdentityFile": jump_host_identity_files[0],
+                            "IdentitiesOnly": "yes",
+                        }
+                    )
+                hosts[f"{run_name}-host"] = {
+                    "HostName": hostname,
+                    "Port": ssh_port,
+                    "User": user,
+                    "IdentityFile": self.identity_file,
+                    "IdentitiesOnly": "yes",
+                    "StrictHostKeyChecking": "no",
+                    "UserKnownHostsFile": "/dev/null",
+                    "ProxyJump": f"{run_name}-jump-host",
+                }
+            else:
+                # Regular SSH instance or VM-based cloud instance
+                hosts[f"{run_name}-host"] = {
+                    "HostName": hostname,
+                    "Port": ssh_port,
+                    "User": user,
+                    "IdentityFile": self.identity_file,
+                    "IdentitiesOnly": "yes",
+                    "StrictHostKeyChecking": "no",
+                    "UserKnownHostsFile": "/dev/null",
+                }
+            hosts[run_name] = {
                 "HostName": "localhost",
                 "Port": container_ssh_port,
                 "User": container_user,
@@ -119,32 +156,41 @@ class SSHAttach:
                 "UserKnownHostsFile": "/dev/null",
                 "ProxyJump": f"{run_name}-host",
             }
-        elif ssh_proxy is not None:
-            self.container_config = {
-                "HostName": hostname,
-                "Port": ssh_port,
-                "User": container_user,
-                "IdentityFile": self.identity_file,
-                "IdentitiesOnly": "yes",
-                "StrictHostKeyChecking": "no",
-                "UserKnownHostsFile": "/dev/null",
-                "ProxyJump": f"{run_name}-jump-host",
-            }
         else:
-            self.container_config = None
-        if local_backend:
-            self.container_config = None
-            self.host_config = {
-                "HostName": hostname,
-                "Port": container_ssh_port,
-                "User": container_user,
-                "IdentityFile": self.identity_file,
-                "IdentitiesOnly": "yes",
-                "StrictHostKeyChecking": "no",
-                "UserKnownHostsFile": "/dev/null",
-            }
-        if self.container_config is not None and get_ssh_client_info().supports_multiplexing:
-            self.container_config.update(
+            if ssh_proxy is not None:
+                # Kubernetes
+                hosts[f"{run_name}-jump-host"] = {
+                    "HostName": ssh_proxy.hostname,
+                    "Port": ssh_proxy.port,
+                    "User": ssh_proxy.username,
+                    "IdentityFile": self.identity_file,
+                    "IdentitiesOnly": "yes",
+                    "StrictHostKeyChecking": "no",
+                    "UserKnownHostsFile": "/dev/null",
+                }
+                hosts[run_name] = {
+                    "HostName": hostname,
+                    "Port": ssh_port,
+                    "User": container_user,
+                    "IdentityFile": self.identity_file,
+                    "IdentitiesOnly": "yes",
+                    "StrictHostKeyChecking": "no",
+                    "UserKnownHostsFile": "/dev/null",
+                    "ProxyJump": f"{run_name}-jump-host",
+                }
+            else:
+                # Container-based backends
+                hosts[run_name] = {
+                    "HostName": hostname,
+                    "Port": ssh_port,
+                    "User": container_user,
+                    "IdentityFile": self.identity_file,
+                    "IdentitiesOnly": "yes",
+                    "StrictHostKeyChecking": "no",
+                    "UserKnownHostsFile": "/dev/null",
+                }
+        if get_ssh_client_info().supports_multiplexing:
+            hosts[run_name].update(
                 {
                     "ControlMaster": "auto",
                     "ControlPath": self.control_sock_path,
@@ -153,14 +199,8 @@ class SSHAttach:
 
     def attach(self):
         include_ssh_config(self.ssh_config_path)
-        if self.container_config is None:
-            update_ssh_config(self.ssh_config_path, self.run_name, self.host_config)
-        elif self.ssh_proxy is not None:
-            update_ssh_config(self.ssh_config_path, f"{self.run_name}-jump-host", self.host_config)
-            update_ssh_config(self.ssh_config_path, self.run_name, self.container_config)
-        else:
-            update_ssh_config(self.ssh_config_path, f"{self.run_name}-host", self.host_config)
-            update_ssh_config(self.ssh_config_path, self.run_name, self.container_config)
+        for host, options in self.hosts.items():
+            update_ssh_config(self.ssh_config_path, host, options)
 
         max_retries = 10
         self._ports_lock.release()
@@ -178,9 +218,8 @@ class SSHAttach:
 
     def detach(self):
         self.tunnel.close()
-        update_ssh_config(self.ssh_config_path, f"{self.run_name}-jump-host", {})
-        update_ssh_config(self.ssh_config_path, f"{self.run_name}-host", {})
-        update_ssh_config(self.ssh_config_path, self.run_name, {})
+        for host in self.hosts:
+            update_ssh_config(self.ssh_config_path, host, {})
 
     def __enter__(self):
         self.attach()