Add head node support for SSH fleets

docs/docs/concepts/fleets.md

@@ -300,6 +300,47 @@ ssh_config:
     - 3.255.177.52
 ```
 
+#### Head node
+
+If fleet nodes are behind a head node, configure [`proxy_jump`](../reference/dstack.yml/fleet.md#proxy_jump):
+
+<div editor-title="examples/misc/fleets/.dstack.yml">
+
+    ```yaml
+    type: fleet
+    name: my-fleet
+
+    ssh_config:
+      user: ubuntu
+      identity_file: ~/.ssh/worker_node_key
+      hosts:
+        - 3.255.177.51
+        - 3.255.177.52
+      proxy_jump:
+        hostname: 3.255.177.50
+        user: ubuntu
+        identity_file: ~/.ssh/head_node_key
+    ```
+
+</div>
+
+To be able to attach to runs, both explicitly with `dstack attach` and implicitly with `dstack apply`, you must either
+add a front node key (`~/.ssh/head_node_key`) to an SSH agent or configure a key path in `~/.ssh/config`:
+
+<div editor-title="~/.ssh/config">
+
+    ```
+    Host 3.255.177.50
+        IdentityFile ~/.ssh/head_node_key
+    ```
+
+</div>
+
+where `Host` must match `ssh_config.proxy_jump.hostname` or `ssh_config.hosts[n].proxy_jump.hostname` if you configure head nodes
+on a per-worker basis.
+
+> Currently, [services](services.md) do not work on instances with a head node setup.
+
 !!! info "Reference"
     For all SSH fleet configuration options, refer to the [reference](../reference/dstack.yml/fleet.md).
 

docs/docs/reference/dstack.yml/fleet.md

@@ -17,12 +17,26 @@ The `fleet` configuration type allows creating and updating fleets.
       show_root_heading: false
       item_id_prefix: ssh_config-
 
+#### `ssh_config.proxy_jump` { #ssh_config-proxy_jump data-toc-label="proxy_jump" }
+
+#SCHEMA# dstack._internal.core.models.fleets.SSHProxyParams
+    overrides:
+      show_root_heading: false
+      item_id_prefix: proxy_jump-
+
 #### `ssh_config.hosts[n]` { #ssh_config-hosts data-toc-label="hosts" }
 
 #SCHEMA# dstack._internal.core.models.fleets.SSHHostParams
     overrides:
       show_root_heading: false
 
+##### `ssh_config.hosts[n].proxy_jump` { #proxy_jump data-toc-label="hosts[n].proxy_jump" }
+
+#SCHEMA# dstack._internal.core.models.fleets.SSHProxyParams
+    overrides:
+      show_root_heading: false
+      item_id_prefix: hosts-proxy_jump-
+
 ### `resources`
 
 #SCHEMA# dstack._internal.core.models.resources.ResourcesSpecSchema

src/dstack/_internal/cli/services/configurators/fleet.py

@@ -201,13 +201,16 @@ def apply_args(self, conf: FleetConfiguration, args: argparse.Namespace, unknown
 
 
 def _preprocess_spec(spec: FleetSpec):
-    if spec.configuration.ssh_config is not None:
-        spec.configuration.ssh_config.ssh_key = _resolve_ssh_key(
-            spec.configuration.ssh_config.identity_file
-        )
-        for host in spec.configuration.ssh_config.hosts:
+    ssh_config = spec.configuration.ssh_config
+    if ssh_config is not None:
+        ssh_config.ssh_key = _resolve_ssh_key(ssh_config.identity_file)
+        if ssh_config.proxy_jump is not None:
+            ssh_config.proxy_jump.ssh_key = _resolve_ssh_key(ssh_config.proxy_jump.identity_file)
+        for host in ssh_config.hosts:
             if not isinstance(host, str):
                 host.ssh_key = _resolve_ssh_key(host.identity_file)
+                if host.proxy_jump is not None:
+                    host.proxy_jump.ssh_key = _resolve_ssh_key(host.proxy_jump.identity_file)
 
 
 def _resolve_ssh_key(ssh_key_path: Optional[str]) -> Optional[SSHKey]:

src/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 @@
     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

src/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(
@@ -96,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(

src/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()