[SSH fleets] Support in-place updates of SSH keys

mkdocs/docs/concepts/fleets.md

@@ -530,8 +530,13 @@ If approved, `dstack` updates the fleet without re-creating it. For example, if
 
     SSH fleets support adding and removing `hosts`. Added hosts are provisioned, and removed hosts are terminated.
 
+    SSH fleets also support updating `identity_file` and `proxy_jump.identity_file`, including per-host overrides.
+
     Changes to other properties require re-creating the fleet. In this case, `dstack apply` offers to delete the fleet and create it again.
 
+??? info "SSH key rotation"
+    Updating `identity_file` and `proxy_jump.identity_file` does not require stopping runs. Allow about 20 seconds for changes to propagate before removing the old public key from the host.
+
 ### Delete fleets
 
 When a fleet isn't used by a run, you can delete it by passing the fleet configuration to `dstack delete`:

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

@@ -4,6 +4,7 @@
 from pathlib import Path
 from typing import Optional
 
+from rich.markup import escape
 from rich.table import Table
 
 from dstack._internal.cli.services.configurators.base import (
@@ -33,6 +34,9 @@
     FleetPlan,
     FleetSpec,
     InstanceGroupPlacement,
+    SSHHostParams,
+    SSHParams,
+    SSHProxyParams,
 )
 from dstack._internal.core.models.instances import InstanceStatus, SSHKey
 from dstack._internal.core.services.diff import copy_model, diff_models
@@ -378,12 +382,17 @@ def _render_fleet_spec_diff(old_spec: FleetSpec, new_spec: FleetSpec) -> Optiona
         if spec_field == "merged_profile":
             continue
         if spec_field == "configuration":
+            fields = []
+            for field in diff_models(old_spec.configuration, new_spec.configuration):
+                old_ssh = old_spec.configuration.ssh_config
+                new_ssh = new_spec.configuration.ssh_config
+                if field == "ssh_config" and old_ssh is not None and new_ssh is not None:
+                    fields.extend(_ssh_config_diff_fields(old_ssh, new_ssh))
+                else:
+                    fields.append(field)
             item = NestedListItem(
                 "Configuration properties:",
-                children=[
-                    NestedListItem(field)
-                    for field in diff_models(old_spec.configuration, new_spec.configuration)
-                ],
+                children=[NestedListItem(escape(field)) for field in fields],
             )
         elif spec_field == "profile":
             item = NestedListItem(
@@ -404,6 +413,37 @@ def _render_fleet_spec_diff(old_spec: FleetSpec, new_spec: FleetSpec) -> Optiona
     return nested_list.render()
 
 
+def _ssh_config_diff_fields(
+    old: SSHParams | SSHHostParams | SSHProxyParams,
+    new: SSHParams | SSHHostParams | SSHProxyParams,
+    prefix: str = "ssh_config",
+) -> list[str]:
+    fields = []
+    for field in diff_models(old, new):
+        old_value = getattr(old, field)
+        new_value = getattr(new, field)
+        path = f"{prefix}.{field}"
+        if field == "ssh_key":
+            continue
+        elif field == "hosts":
+            old_hosts = {h if isinstance(h, str) else h.hostname: h for h in old_value}
+            new_names = [h if isinstance(h, str) else h.hostname for h in new_value]
+            if list(old_hosts) != new_names:
+                fields.append(path)
+            for index, host in enumerate(new_value):
+                hostname = host if isinstance(host, str) else host.hostname
+                previous = old_hosts.get(hostname)
+                if isinstance(previous, SSHHostParams) and isinstance(host, SSHHostParams):
+                    fields.extend(_ssh_config_diff_fields(previous, host, f"{path}[{index}]"))
+                elif previous is not None and previous != host:
+                    fields.append(f"{path}[{index}]")
+        elif isinstance(old_value, SSHProxyParams) and isinstance(new_value, SSHProxyParams):
+            fields.extend(_ssh_config_diff_fields(old_value, new_value, path))
+        else:
+            fields.append(path)
+    return fields
+
+
 def _print_plan_header(plan: FleetPlan):
     def th(s: str) -> str:
         return f"[bold]{s}[/bold]"

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

@@ -189,6 +189,26 @@ def open_command(self) -> List[str]:
     def close_command(self) -> List[str]:
         return [*self._control_command_prefix(), "-O", "exit", self.destination]
 
+    def update_proxy_identity(self, index: int, private_key: str) -> None:
+        """Replace a proxy-hop key for the next open, keeping the current tunnel open.
+
+        ServiceConnection calls this when a fleet's `proxy_jump` key changes. Replace
+        the key file atomically so concurrent SSH startup never reads a partial key.
+        """
+        params, _ = self.ssh_proxies[index]
+        identity_path = os.path.join(self.temp_dir.name, f"proxy_identity_{index}")
+        # A concurrent SSH startup must see either complete key, never a partial write.
+        with tempfile.NamedTemporaryFile(mode="w", dir=self.temp_dir.name, delete=False) as f:
+            temporary_path = f.name
+            try:
+                f.write(private_key)
+                f.close()
+                os.replace(temporary_path, identity_path)
+            finally:
+                if os.path.exists(temporary_path):
+                    os.unlink(temporary_path)
+        self.ssh_proxies[index] = (params, identity_path)
+
     def check_command(self) -> List[str]:
         return [*self._control_command_prefix(), "-O", "check", self.destination]
 

src/dstack/_internal/proxy/gateway/repo/repo.py

@@ -1,7 +1,8 @@
+import tempfile
 from contextlib import asynccontextmanager
 from itertools import chain
 from pathlib import Path
-from typing import Optional
+from typing import Iterable, Optional
 
 from aiorwlock import RWLock
 from pydantic import BaseModel
@@ -53,6 +54,20 @@ async def set_service(self, service: Service) -> None:
         async with self.writer():
             self._state.services.setdefault(service.project_name, {})[service.run_name] = service
 
+    async def set_services(self, services: Iterable[Service]) -> None:
+        """Save a batch before publishing it, so failed saves remain retryable."""
+        async with self._lock.writer:
+            services_by_project = {
+                name: entries.copy() for name, entries in self._state.services.items()
+            }
+            for service in services:
+                services_by_project.setdefault(service.project_name, {})[service.run_name] = (
+                    service
+                )
+            state = self._state.model_copy(update={"services": services_by_project})
+            await run_async(self.save, state)
+            self._state = state
+
     async def delete_service(self, project_name: str, run_name: str) -> None:
         async with self.writer():
             project_services = self._state.services.get(project_name, {})
@@ -124,6 +139,17 @@ def load(state_file: Path) -> "GatewayProxyRepo":
             state = None
         return GatewayProxyRepo(state=state, file=state_file)
 
-    def save(self) -> None:
-        if self._file is not None:
-            self._file.write_text(self._state.model_dump_json())
+    def save(self, state: Optional[State] = None) -> None:
+        """Write the given state, or the current state, atomically."""
+        if self._file is None:
+            return
+        if state is None:
+            state = self._state
+        with tempfile.NamedTemporaryFile(mode="w", dir=self._file.parent, delete=False) as f:
+            temporary_path = Path(f.name)
+            try:
+                f.write(state.model_dump_json())
+                f.close()
+                temporary_path.replace(self._file)
+            finally:
+                temporary_path.unlink(missing_ok=True)

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

@@ -10,15 +10,17 @@
     RegisterReplicaRequest,
     RegisterServiceRequest,
     SetServiceIdRequest,
+    UpdateReplicaProxyJumpKeysRequest,
+    UpdateReplicaProxyJumpKeysResponse,
 )
 from dstack._internal.proxy.gateway.services.nginx import Nginx
 from dstack._internal.proxy.lib.deps import get_service_connection_pool
 from dstack._internal.proxy.lib.services.service_connection import ServiceConnectionPool
 
-router = APIRouter(prefix="/{project_name}")
+router = APIRouter()
 
 
-@router.post("/services/register")
+@router.post("/{project_name}/services/register")
 async def register_service(
     project_name: str,
     body: RegisterServiceRequest,
@@ -46,7 +48,7 @@ async def register_service(
     return OkResponse()
 
 
-@router.post("/services/{run_name}/unregister")
+@router.post("/{project_name}/services/{run_name}/unregister")
 async def unregister_service(
     project_name: str,
     run_name: str,
@@ -64,7 +66,7 @@ async def unregister_service(
     return OkResponse()
 
 
-@router.post("/services/{run_name}/set_id")
+@router.post("/{project_name}/services/{run_name}/set_id")
 async def set_service_id(
     project_name: str,
     run_name: str,
@@ -81,7 +83,7 @@ async def set_service_id(
     return OkResponse()
 
 
-@router.post("/services/{run_name}/replicas/register")
+@router.post("/{project_name}/services/{run_name}/replicas/register")
 async def register_replica(
     project_name: str,
     run_name: str,
@@ -109,7 +111,19 @@ async def register_replica(
     return OkResponse()
 
 
-@router.post("/services/{run_name}/replicas/{job_id}/unregister")
+@router.post("/replicas/proxy_jump_keys")
+async def update_replica_proxy_jump_keys(
+    body: UpdateReplicaProxyJumpKeysRequest,
+    repo: Annotated[GatewayProxyRepo, Depends(get_gateway_proxy_repo)],
+    service_conn_pool: Annotated[ServiceConnectionPool, Depends(get_service_connection_pool)],
+) -> UpdateReplicaProxyJumpKeysResponse:
+    failures = await registry_services.update_replica_proxy_jump_keys(
+        updates=body.replicas, repo=repo, service_conn_pool=service_conn_pool
+    )
+    return UpdateReplicaProxyJumpKeysResponse(failures=failures)
+
+
+@router.post("/{project_name}/services/{run_name}/replicas/{job_id}/unregister")
 async def unregister_replica(
     project_name: str,
     run_name: str,
@@ -129,7 +143,7 @@ async def unregister_replica(
     return OkResponse()
 
 
-@router.post("/entrypoints/register")
+@router.post("/{project_name}/entrypoints/register")
 async def register_entrypoint(
     project_name: str,
     body: RegisterEntrypointRequest,

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

@@ -62,6 +62,7 @@ class RegisterReplicaRequest(BaseModel):
     ssh_port: int
     ssh_proxy: Optional[SSHConnectionParams] = None
     ssh_proxy_private_key: Optional[str] = None
+    # TODO: Rename ssh_head_proxy* to ssh_proxy_jump* while preserving registration API compatibility.
     ssh_head_proxy: Optional[SSHConnectionParams] = None
     ssh_head_proxy_private_key: Optional[str] = None
     internal_ip: Optional[str] = None
@@ -71,3 +72,20 @@ class RegisterEntrypointRequest(BaseModel):
     domain: str
     https: bool
     read_timeout: int = DEFAULT_SERVICE_READ_TIMEOUT
+
+
+class ReplicaProxyJumpKeyUpdate(BaseModel):
+    project_name: str
+    run_name: str
+    job_id: str
+    ssh_proxy_jump_private_key: str
+    """Private key for the SSH fleet's `proxy_jump`."""
+
+
+class UpdateReplicaProxyJumpKeysRequest(BaseModel):
+    replicas: list[ReplicaProxyJumpKeyUpdate]
+
+
+class UpdateReplicaProxyJumpKeysResponse(BaseModel):
+    failures: dict[str, str]
+    """Job IDs whose key could not be updated; successful updates are saved together."""

src/dstack/_internal/proxy/gateway/schemas/services.py

@@ -3,6 +3,13 @@
 
 class ServiceListReplicaItem(CoreModel):
     id: str
+    ssh_proxy_jump_public_key: str | None = None
+    """Public key matching the SSH fleet's `proxy_jump` key saved for this replica on the gateway.
+
+    The server compares this with the instance's current `proxy_jump` key during gateway
+    synchronization and sends a replacement private key when they differ.
+    `None` means no `proxy_jump` private key is saved for this replica.
+    """
 
 
 class ServiceListItem(CoreModel):

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

@@ -212,6 +212,71 @@ async def register_replica(
     logger.info("Replica %s in service %s is registered now", replica.id, service.fmt())
 
 
+async def update_replica_proxy_jump_keys(
+    updates: list[schemas.ReplicaProxyJumpKeyUpdate],
+    repo: GatewayProxyRepo,
+    service_conn_pool: ServiceConnectionPool,
+) -> dict[str, str]:
+    """Replace replicas' SSH fleet `proxy_jump` keys in one gateway synchronization batch.
+
+    Keep existing tunnels, registrations, and Nginx configuration. Save successful
+    replacements once, including across services/projects. Return failures by job ID
+    so one removed replica or failed key-file write cannot block the other updates.
+    Repeating the same keys is a no-op; the next sync retries any remaining differences.
+    """
+    by_service: dict[tuple[str, str], dict[str, schemas.ReplicaProxyJumpKeyUpdate]] = {}
+    for item in updates:
+        by_service.setdefault((item.project_name.lower(), item.run_name.lower()), {})[
+            item.job_id
+        ] = item
+    failures = {}
+    changed_services = []
+    async with lock:
+        for (project_name, run_name), pending in by_service.items():
+            service = await repo.get_service(project_name, run_name)
+            project = await repo.get_project(project_name)
+            if service is None or project is None:
+                for job_id in pending:
+                    failures[job_id] = f"Service {project_name}/{run_name} not found"
+                continue
+            replicas = list(service.replicas)
+            changed = False
+            for index, replica in enumerate(replicas):
+                item = pending.pop(replica.id, None)
+                if item is None:
+                    continue
+                if replica.ssh_head_proxy is None:
+                    failures[replica.id] = "Replica has no proxy_jump"
+                    continue
+                if replica.ssh_head_proxy_private_key == item.ssh_proxy_jump_private_key:
+                    continue
+                # A fresh Replica discards its cached old public key. Preserve the
+                # unchanged Replica objects so polling does not reparse their keys.
+                updated_replica = models.Replica(
+                    **{
+                        **replica.model_dump(),
+                        "ssh_head_proxy_private_key": item.ssh_proxy_jump_private_key,
+                    }
+                )
+                try:
+                    await service_conn_pool.get_or_add(project, service, updated_replica)
+                except Exception as exc:
+                    logger.warning(
+                        "Could not update replica %s proxy_jump SSH key: %r", replica.id, exc
+                    )
+                    failures[replica.id] = "Could not update proxy_jump key"
+                    continue
+                replicas[index] = updated_replica
+                changed = True
+            for job_id in pending:
+                failures[job_id] = f"Replica not found in service {project_name}/{run_name}"
+            if changed:
+                changed_services.append(service.model_copy(update={"replicas": tuple(replicas)}))
+        if changed_services:
+            await repo.set_services(changed_services)
+    return failures
+
+
 async def unregister_replica(
     project_name: str,
     run_name: str,

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

@@ -4,17 +4,30 @@
     ServiceListReplicaItem,
     ServiceListResponse,
 )
+from dstack._internal.proxy.lib.models import Service
+from dstack._internal.utils.common import run_async
 
 
 async def list_services(repo: GatewayProxyRepo) -> ServiceListResponse:
     services = await repo.list_services()
+    # Parsing uncached proxy_jump keys must not block other gateway API requests.
+    return await run_async(_get_service_list_response, services)
+
+
+def _get_service_list_response(services: list[Service]) -> ServiceListResponse:
     return ServiceListResponse(
         services=[
             ServiceListItem(
                 id=service.id,
                 project_name=service.project_name,
                 run_name=service.run_name,
-                replicas=[ServiceListReplicaItem(id=replica.id) for replica in service.replicas],
+                replicas=[
+                    ServiceListReplicaItem(
+                        id=replica.id,
+                        ssh_proxy_jump_public_key=replica.ssh_proxy_jump_public_key,
+                    )
+                    for replica in service.replicas
+                ],
             )
             for service in services
         ]