From c53aa02eb15779306fa7cd2882f384b5001bed44 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Tue, 3 Nov 2026 13:50:00 +0100 Subject: [PATCH 01/10] Support in-place SSH key rotation for SSH fleets --- mkdocs/docs/concepts/fleets.md | 5 + .../cli/services/configurators/fleet.py | 51 ++- .../_internal/core/services/ssh/tunnel.py | 20 ++ .../_internal/proxy/gateway/repo/repo.py | 10 +- .../proxy/gateway/routers/registry.py | 28 +- .../proxy/gateway/schemas/registry.py | 17 + .../proxy/gateway/schemas/services.py | 7 + .../proxy/gateway/services/registry.py | 65 ++++ .../proxy/gateway/services/services.py | 8 +- src/dstack/_internal/proxy/lib/models.py | 14 + .../proxy/lib/services/service_connection.py | 57 +++- .../pipeline_tasks/gateway_replicas.py | 128 ++++++-- .../pipeline_tasks/instances/__init__.py | 8 +- .../_internal/server/services/fleets.py | 297 +++++++++++++----- .../server/services/gateways/client.py | 20 ++ .../_internal/server/services/proxy/repo.py | 2 + .../cli/services/configurators/test_fleet.py | 67 +++- .../core/services/ssh/test_tunnel.py | 27 ++ .../proxy/gateway/routers/test_registry.py | 90 ++++++ .../proxy/gateway/routers/test_services.py | 2 +- .../lib/services/test_service_connection.py | 81 +++++ .../pipeline_tasks/instances/test_pipeline.py | 22 +- .../pipeline_tasks/test_gateway_replicas.py | 107 ++++++- .../_internal/server/services/test_fleets.py | 287 ++++++++++++++++- 24 files changed, 1294 insertions(+), 126 deletions(-) diff --git a/mkdocs/docs/concepts/fleets.md b/mkdocs/docs/concepts/fleets.md index 92aeb8e3b2..2244dbd731 100644 --- a/mkdocs/docs/concepts/fleets.md +++ b/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 and key contents at the same path. + 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" + To rotate SSH keys, authorize the new public key and apply the updated fleet configuration. Runs and established SSH connections continue. For `proxy_jump`, keep the old public key authorized until [gateways](gateways.md) and the optional [SSH proxy](../guides/server-deployment.md#ssh-proxy) receive the replacement; `dstack apply` does not wait for them. + ### Delete fleets When a fleet isn't used by a run, you can delete it by passing the fleet configuration to `dstack delete`: diff --git a/src/dstack/_internal/cli/services/configurators/fleet.py b/src/dstack/_internal/cli/services/configurators/fleet.py index ce0db5192f..94e85b189e 100644 --- a/src/dstack/_internal/cli/services/configurators/fleet.py +++ b/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,40 @@ 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": + # Key material is an internal field. Refer to its user-configured file, + # including replacements whose path has not changed. Never print keys. + if old.identity_file == new.identity_file: + fields.append(f"{prefix}.identity_file (key contents changed)") + 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]" diff --git a/src/dstack/_internal/core/services/ssh/tunnel.py b/src/dstack/_internal/core/services/ssh/tunnel.py index 313eeea6e1..afcb092988 100644 --- a/src/dstack/_internal/core/services/ssh/tunnel.py +++ b/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 head-node 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] diff --git a/src/dstack/_internal/proxy/gateway/repo/repo.py b/src/dstack/_internal/proxy/gateway/repo/repo.py index a3b69d4565..dd52745687 100644 --- a/src/dstack/_internal/proxy/gateway/repo/repo.py +++ b/src/dstack/_internal/proxy/gateway/repo/repo.py @@ -1,7 +1,7 @@ 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 +53,14 @@ 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 of changed services with one state-file write.""" + async with self.writer(): + for service in services: + self._state.services.setdefault(service.project_name, {})[service.run_name] = ( + service + ) + async def delete_service(self, project_name: str, run_name: str) -> None: async with self.writer(): project_services = self._state.services.get(project_name, {}) diff --git a/src/dstack/_internal/proxy/gateway/routers/registry.py b/src/dstack/_internal/proxy/gateway/routers/registry.py index 954ae2933a..5365f4c866 100644 --- a/src/dstack/_internal/proxy/gateway/routers/registry.py +++ b/src/dstack/_internal/proxy/gateway/routers/registry.py @@ -10,15 +10,17 @@ RegisterReplicaRequest, RegisterServiceRequest, SetServiceIdRequest, + UpdateReplicaHeadNodeKeysRequest, + UpdateReplicaHeadNodeKeysResponse, ) 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/head_node_keys") +async def update_replica_head_node_keys( + body: UpdateReplicaHeadNodeKeysRequest, + repo: Annotated[GatewayProxyRepo, Depends(get_gateway_proxy_repo)], + service_conn_pool: Annotated[ServiceConnectionPool, Depends(get_service_connection_pool)], +) -> UpdateReplicaHeadNodeKeysResponse: + failures = await registry_services.update_replica_head_node_keys( + updates=body.replicas, repo=repo, service_conn_pool=service_conn_pool + ) + return UpdateReplicaHeadNodeKeysResponse(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, diff --git a/src/dstack/_internal/proxy/gateway/schemas/registry.py b/src/dstack/_internal/proxy/gateway/schemas/registry.py index 25f801252a..de3b08cc74 100644 --- a/src/dstack/_internal/proxy/gateway/schemas/registry.py +++ b/src/dstack/_internal/proxy/gateway/schemas/registry.py @@ -71,3 +71,20 @@ class RegisterEntrypointRequest(BaseModel): domain: str https: bool read_timeout: int = DEFAULT_SERVICE_READ_TIMEOUT + + +class ReplicaHeadNodeKeyUpdate(BaseModel): + project_name: str + run_name: str + job_id: str + ssh_head_proxy_private_key: str + """Private key for the head node configured by the fleet's `proxy_jump`.""" + + +class UpdateReplicaHeadNodeKeysRequest(BaseModel): + replicas: list[ReplicaHeadNodeKeyUpdate] + + +class UpdateReplicaHeadNodeKeysResponse(BaseModel): + failures: dict[str, str] + """Job IDs whose key could not be updated; successful updates are saved together.""" diff --git a/src/dstack/_internal/proxy/gateway/schemas/services.py b/src/dstack/_internal/proxy/gateway/schemas/services.py index cbba05b518..25b0e80d44 100644 --- a/src/dstack/_internal/proxy/gateway/schemas/services.py +++ b/src/dstack/_internal/proxy/gateway/schemas/services.py @@ -3,6 +3,13 @@ class ServiceListReplicaItem(CoreModel): id: str + ssh_head_proxy_public_key: str | None = None + """Public key matching the head-node private key saved for this replica on the gateway. + + The server compares this with the instance's current head-node key during gateway + synchronization and sends a replacement private key when they differ. + `None` means no head-node private key is saved for this replica. + """ class ServiceListItem(CoreModel): diff --git a/src/dstack/_internal/proxy/gateway/services/registry.py b/src/dstack/_internal/proxy/gateway/services/registry.py index 1bc47d7195..0bf13c27d2 100644 --- a/src/dstack/_internal/proxy/gateway/services/registry.py +++ b/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_head_node_keys( + updates: list[schemas.ReplicaHeadNodeKeyUpdate], + repo: GatewayProxyRepo, + service_conn_pool: ServiceConnectionPool, +) -> dict[str, str]: + """Replace registered replicas' head-node 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.ReplicaHeadNodeKeyUpdate]] = {} + 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 head node" + continue + if replica.ssh_head_proxy_private_key == item.ssh_head_proxy_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_head_proxy_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 head-node SSH key: %r", replica.id, exc + ) + failures[replica.id] = "Could not update head-node 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, diff --git a/src/dstack/_internal/proxy/gateway/services/services.py b/src/dstack/_internal/proxy/gateway/services/services.py index 5543d92160..292d9cd5eb 100644 --- a/src/dstack/_internal/proxy/gateway/services/services.py +++ b/src/dstack/_internal/proxy/gateway/services/services.py @@ -14,7 +14,13 @@ async def list_services(repo: GatewayProxyRepo) -> ServiceListResponse: 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_head_proxy_public_key=replica.ssh_head_proxy_public_key, + ) + for replica in service.replicas + ], ) for service in services ] diff --git a/src/dstack/_internal/proxy/lib/models.py b/src/dstack/_internal/proxy/lib/models.py index 44d5878bbd..ac4ed7828d 100644 --- a/src/dstack/_internal/proxy/lib/models.py +++ b/src/dstack/_internal/proxy/lib/models.py @@ -1,6 +1,7 @@ """Things stored in BaseProxyRepo implementations.""" from datetime import datetime +from functools import cached_property, lru_cache from typing import Any, Iterable, Literal, Optional, Union from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -9,6 +10,7 @@ from dstack._internal.core.models.instances import SSHConnectionParams from dstack._internal.proxy.lib.const import DEFAULT_SERVICE_READ_TIMEOUT from dstack._internal.proxy.lib.errors import UnexpectedProxyError +from dstack._internal.utils import ssh as ssh_utils # Models should be immutable so that they can be stored in memory and safely shared by @@ -30,6 +32,18 @@ class Replica(ImmutableModel): ssh_head_proxy_private_key: Optional[str] = None internal_ip: Optional[str] = None + @cached_property + def ssh_head_proxy_public_key(self) -> Optional[str]: + if self.ssh_head_proxy_private_key is None: + return None + return _get_head_node_public_key(self.ssh_head_proxy_private_key) + + +@lru_cache(maxsize=128) +def _get_head_node_public_key(private_key: str) -> str: + # Replicas sharing a head node often share a key; parse it once, not per replica. + return ssh_utils.generate_public_key(ssh_utils.pkey_from_str(private_key)) + class IPAddressPartitioningKey(ImmutableModel): type: Literal["ip_address"] = "ip_address" diff --git a/src/dstack/_internal/proxy/lib/services/service_connection.py b/src/dstack/_internal/proxy/lib/services/service_connection.py index 030b6cb4fd..db3a12ef2f 100644 --- a/src/dstack/_internal/proxy/lib/services/service_connection.py +++ b/src/dstack/_internal/proxy/lib/services/service_connection.py @@ -42,6 +42,9 @@ class ServiceConnection: using the same socket path. The in-server proxy opens connections on demand. """ + replica: Replica + """Replica configuration applied to this connection, including its key for future SSH opens.""" + def __init__( self, project: Project, @@ -49,6 +52,8 @@ def __init__( replica: Replica, reconnect_semaphore: asyncio.Semaphore, ) -> None: + self.replica = replica + self._auto_reconnect = service.domain is not None self._temp_dir = TemporaryDirectory() options = { **SSH_DEFAULT_OPTIONS, @@ -81,7 +86,7 @@ def __init__( ), ], options=options, - background=service.domain is None, + background=not (self._auto_reconnect or replica.ssh_head_proxy is not None), ) self._client = ServiceClient( transport=AsyncHTTPTransport(uds=str(self._app_socket_path)), @@ -94,7 +99,6 @@ def __init__( self._lifecycle_lock = asyncio.Lock() self._closed = False self._monitor_task: Optional[asyncio.Task] = None - self._auto_reconnect = service.domain is not None self._replica_id = replica.id self._reconnect_semaphore = reconnect_semaphore @@ -113,7 +117,7 @@ async def open(self) -> None: # Removal may have started while SSH was connecting. raise UnexpectedProxyError("Service connection was removed while opening") self._is_open.set() - if self._auto_reconnect: + if self._auto_reconnect or self.replica.ssh_head_proxy is not None: self._monitor_task = asyncio.create_task(self._monitor_tunnel()) async def close(self) -> None: @@ -134,6 +138,20 @@ async def client(self) -> ServiceClient: await asyncio.wait_for(self._is_open.wait(), timeout=OPEN_TUNNEL_TIMEOUT) return self._client + async def update_head_node_key(self, replica: Replica) -> None: + """Replace the head-node key without reopening the tunnel or HTTP client. + + The pool calls this only when the head-node private key is the sole change + from self.replica. Apply the key file before replacing that configuration, + so a failed write leaves the previous configuration available for retry. + The replacement key is used on the next SSH open; the socket path is unchanged. + """ + async with self._lifecycle_lock: + if self._closed: + raise UnexpectedProxyError("Cannot update a closed service connection") + self._tunnel.update_proxy_identity(0, get_or_error(replica.ssh_head_proxy_private_key)) + self.replica = replica + async def _close(self) -> None: async with self._lifecycle_lock: if self._monitor_task is not None: @@ -152,6 +170,13 @@ async def _monitor_tunnel(self) -> None: while True: started_at = loop.time() await self._tunnel.wait_closed() + if not self._auto_reconnect: + # The in-server pool retains stopped replicas. Reopen only when another + # request selects this replica, rather than retrying stopped workloads. + async with self._lifecycle_lock: + self._is_open.clear() + await self._tunnel.aclose() + return if loop.time() - started_at >= MAX_TUNNEL_RECONNECT_DELAY: retry_delay = 0 logger.warning("SSH tunnel to replica %s exited, reconnecting", self._replica_id) @@ -187,6 +212,18 @@ async def get_or_add( ) -> ServiceConnection: connection = self.connections.get(replica.id) if connection is not None: + if ( + replica.ssh_head_proxy is not None + and connection.replica.ssh_head_proxy_private_key + != replica.ssh_head_proxy_private_key + and connection.replica.model_copy( + update={"ssh_head_proxy_private_key": replica.ssh_head_proxy_private_key} + ) + == replica + ): + await connection.update_head_node_key(replica) + if service.domain is None and replica.ssh_head_proxy is not None: + await connection.open() return connection connection = ServiceConnection(project, service, replica, self._reconnect_semaphore) self.connections[replica.id] = connection @@ -234,12 +271,10 @@ async def get_service_replica_client( ) # Nginx not available, forward directly to the tunnel replica = random.choice(service.replicas) - connection = await service_conn_pool.get(replica.id) - if connection is None: - project = await repo.get_project(service.project_name) - if project is None: - raise UnexpectedProxyError( - f"Expected to find project {service.project_name} but could not" - ) - connection = await service_conn_pool.get_or_add(project, service, replica) + project = await repo.get_project(service.project_name) + if project is None: + raise UnexpectedProxyError( + f"Expected to find project {service.project_name} but could not" + ) + connection = await service_conn_pool.get_or_add(project, service, replica) return await connection.client() diff --git a/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py b/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py index 2530e10ec3..3ece85e70e 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py @@ -7,7 +7,13 @@ from httpx import HTTPStatusError from sqlalchemy import and_, delete, or_, select, update from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.orm import InstrumentedAttribute, joinedload, load_only, with_loader_criteria +from sqlalchemy.orm import ( + InstrumentedAttribute, + joinedload, + load_only, + raiseload, + with_loader_criteria, +) from sqlalchemy.sql.base import ExecutableOption from dstack._internal.core.backends.base.compute import ( @@ -17,7 +23,9 @@ from dstack._internal.core.errors import BackendError, BackendNotAvailable, GatewayError from dstack._internal.core.models.common import validate_json_extra_ignore from dstack._internal.core.models.gateways import GatewayReplicaStatus, GatewayStatus +from dstack._internal.core.models.instances import RemoteConnectionInfo from dstack._internal.core.models.runs import JobSpec, JobStatus, RunStatus, ServiceSpec +from dstack._internal.proxy.gateway.schemas.registry import ReplicaHeadNodeKeyUpdate from dstack._internal.proxy.gateway.schemas.services import ServiceListItem from dstack._internal.server import settings from dstack._internal.server.background.pipeline_tasks.base import ( @@ -678,32 +686,52 @@ async def _process_running_item(item: GatewayReplicaPipelineItem): logger.warning("%s: failed to list services: %r", fmt(replica_model), e) await _commit_update(item, replica_model, update_map={}) return + run_filters = ( + RunModel.gateway_id == gateway_model.id, + RunModel.deleted == False, + RunModel.status.not_in(RunStatus.finished_statuses() + [RunStatus.TERMINATING]), + ) + job_filters = (JobModel.status == JobStatus.RUNNING, JobModel.registered == True) stmt = ( select(RunModel) - .where( - RunModel.gateway_id == gateway_model.id, - RunModel.deleted == False, - RunModel.status.not_in(RunStatus.finished_statuses() + [RunStatus.TERMINATING]), - ) + .where(*run_filters) .options( load_only(RunModel.id), - joinedload(RunModel.jobs).load_only(JobModel.id), - with_loader_criteria( - JobModel, - and_( - JobModel.status == JobStatus.RUNNING, - JobModel.registered == True, - ), - ), + joinedload(RunModel.jobs).load_only(JobModel.id, JobModel.instance_id), + with_loader_criteria(JobModel, and_(*job_filters)), ) ) async with get_session_ctx() as session: res = await session.execute(stmt) expected_runs = res.scalars().unique().all() - plan = _plan_state_sync( - currently_registered=currently_registered_services, - expected=expected_runs, + # Fetch SSH connection data once per instance, not once per service replica. + # Cloud instances have no remote_connection_info and return no key data. + res = await session.execute( + select(InstanceModel.id, InstanceModel.remote_connection_info).where( + InstanceModel.remote_connection_info.is_not(None), + InstanceModel.id.in_( + select(JobModel.instance_id) + .join(RunModel, JobModel.run_id == RunModel.id) + .where(*run_filters, *job_filters) + ), + ) ) + instance_connection_info = res.all() + + # Parse keys and compare gateway state without holding a DB connection. + instance_head_node_public_keys = {} + for instance_id, connection_info in instance_connection_info: + rci = validate_json_extra_ignore(RemoteConnectionInfo, get_or_error(connection_info)) + if rci.ssh_proxy is not None: + instance_head_node_public_keys[instance_id] = " ".join( + get_or_error(rci.ssh_proxy_keys)[0].public.split()[:2] + ) + plan = _plan_state_sync( + currently_registered=currently_registered_services, + expected=expected_runs, + instance_head_node_public_keys=instance_head_node_public_keys, + ) + async with get_session_ctx() as session: run_models_by_id, job_models_by_id = await _load_runs_and_jobs_for_state_sync( session, plan ) @@ -912,6 +940,48 @@ async def _perform_state_sync( result.failed_replica_registrations[replica_id] = "Unexpected error" continue result.registered_replicas.add(replica_id) + # Send all replacements together. The gateway keeps working tunnels open and + # saves its state once, rather than rewriting the full state for each replica. + key_updates = [] + for run_id, job_ids in plan.update_replica_head_node_keys.items(): + run = run_models_by_id.get(run_id) + if run is None: + continue + for job_id in job_ids: + job = job_models_by_id.get(job_id) + if job is None or job.instance is None: + continue + rci = get_instance_remote_connection_info(job.instance) + if rci is None or rci.ssh_proxy is None: + continue + key_updates.append( + ReplicaHeadNodeKeyUpdate( + project_name=run.project.name, + run_name=run.run_name, + job_id=job_id.hex, + ssh_head_proxy_private_key=get_or_error( + get_or_error(rci.ssh_proxy_keys)[0].private + ), + ) + ) + if key_updates: + try: + async with connection.client() as client: + failures = await client.update_replica_head_node_keys(key_updates) + except Exception as exc: + logger.warning( + "%s: could not update service head-node SSH keys: %r", fmt(replica_model), exc + ) + else: + for job_id, error in failures.items(): + logger.warning( + "%s: could not update replica %s head-node SSH key: %s", + fmt(replica_model), + job_id, + error, + ) + # Poll actual gateway keys again next time, including after a lost response + # or a partial failure. No separate retry state is required. return result @@ -1226,12 +1296,15 @@ async def _load_runs_and_jobs_for_state_sync( run_ids = ( plan.register_services | plan.register_replicas.keys() + | plan.update_replica_head_node_keys.keys() | {s.id for s in plan.unregister_services if s.id is not None} ) run_models_by_id: dict[uuid.UUID, RunModel] = {} if run_ids: res = await session.execute( - select(RunModel).where(RunModel.id.in_(run_ids)).options(joinedload(RunModel.project)) + select(RunModel) + .where(RunModel.id.in_(run_ids)) + .options(joinedload(RunModel.project), raiseload(RunModel.jobs)) ) run_models_by_id = {run.id: run for run in res.unique().scalars().all()} @@ -1240,6 +1313,8 @@ async def _load_runs_and_jobs_for_state_sync( job_ids |= replica_ids for replica_ids in plan.unregister_replicas.values(): job_ids |= replica_ids + for replica_ids in plan.update_replica_head_node_keys.values(): + job_ids |= replica_ids job_models_by_id: dict[uuid.UUID, JobModel] = {} if job_ids: res = await session.execute( @@ -1371,6 +1446,8 @@ class _StateSyncPlan: register_replicas: dict[uuid.UUID, set[uuid.UUID]] = field(default_factory=dict) unregister_replicas: dict[_ServiceRef, set[uuid.UUID]] = field(default_factory=dict) set_run_ids: dict[_ServiceRef, uuid.UUID] = field(default_factory=dict) + update_replica_head_node_keys: dict[uuid.UUID, set[uuid.UUID]] = field(default_factory=dict) + """Run IDs mapped to registered job IDs whose saved gateway head-node key needs replacing.""" @dataclass @@ -1388,7 +1465,9 @@ class _StateSyncResult: def _plan_state_sync( - currently_registered: list[ServiceListItem], expected: Sequence[RunModel] + currently_registered: list[ServiceListItem], + expected: Sequence[RunModel], + instance_head_node_public_keys: dict[uuid.UUID, str], ) -> _StateSyncPlan: plan = _StateSyncPlan() expected_run_id_to_run = {run.id: run for run in expected} @@ -1422,6 +1501,17 @@ def _plan_state_sync( expected_job_ids = {job.id for job in run.jobs} plan.register_replicas[run.id] = expected_job_ids - currently_registered_job_ids plan.unregister_replicas[service_ref] = currently_registered_job_ids - expected_job_ids + gateway_head_node_public_keys = { + uuid.UUID(r.id): r.ssh_head_proxy_public_key for r in service.replicas + } + for job in run.jobs: + if job.id not in currently_registered_job_ids or job.instance_id is None: + continue + public_key = instance_head_node_public_keys.get(job.instance_id) + if public_key is None: + continue + if gateway_head_node_public_keys[job.id] != public_key: + plan.update_replica_head_node_keys.setdefault(run.id, set()).add(job.id) else: plan.unregister_services.add(service_ref) plan.unregister_replicas[service_ref] = { diff --git a/src/dstack/_internal/server/background/pipeline_tasks/instances/__init__.py b/src/dstack/_internal/server/background/pipeline_tasks/instances/__init__.py index 7b0cae9a1d..40d26a2f04 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/instances/__init__.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/instances/__init__.py @@ -216,7 +216,13 @@ async def fetch(self, limit: int) -> list[InstancePipelineItem]: # because the fleet pipeline cannot reclaim stale instance locks. or_( InstanceModel.fleet_id.is_(None), - FleetModel.lock_owner.is_(None), + and_( + FleetModel.lock_owner.is_(None), + or_( + FleetModel.lock_expires_at.is_(None), + FleetModel.lock_expires_at < now, + ), + ), ), InstanceModel.lock_expires_at.is_(None), ), diff --git a/src/dstack/_internal/server/services/fleets.py b/src/dstack/_internal/server/services/fleets.py index ac677b4948..6b27533682 100644 --- a/src/dstack/_internal/server/services/fleets.py +++ b/src/dstack/_internal/server/services/fleets.py @@ -1,11 +1,13 @@ import asyncio +import time import uuid from collections.abc import Callable -from datetime import datetime +from contextlib import asynccontextmanager +from datetime import datetime, timedelta from functools import wraps -from typing import List, Literal, Optional, Tuple, TypeVar, Union +from typing import AsyncIterator, List, Literal, Optional, Tuple, TypeVar, Union -from sqlalchemy import and_, exists, false, func, or_, select +from sqlalchemy import and_, exists, false, func, or_, select, update from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import aliased, joinedload, selectinload @@ -597,9 +599,7 @@ async def apply_plan( get_locker(get_db().dialect_name).lock_ctx(InstanceModel.__tablename__, instances_ids), ): # Refetch after lock - # TODO: Lock instances with FOR UPDATE? - # We do not respect InstanceModel.lock_* fields here because FleetPipeline does not update SSH instances. - # TODO: Respect InstanceModel.lock_* fields if FleetPipeline and apply update the same instances. + # _update_fleet also locks instances and checks their pipeline locks before changing keys. res = await session.execute( select(FleetModel) .where( @@ -1145,63 +1145,198 @@ async def _update_fleet( _check_can_update_fleet_spec(fleet_sensitive.spec, spec) - fleet_model.spec = spec.model_dump_json() - # Reset consolidation attempt so the next pipeline pass picks up the spec change promptly. - fleet_model.consolidation_attempt = 0 + async with _lock_ssh_key_update_instances( + session, fleet_model, spec.configuration.ssh_config + ) as key_updates: + for instance_model, connection_info in key_updates: + instance_model.remote_connection_info = connection_info - if ( - fleet_sensitive.spec.configuration.ssh_config is not None - and spec.configuration.ssh_config is not None - ): - added_hosts, removed_hosts, changed_hosts = _calculate_ssh_hosts_changes( - current=fleet_sensitive.spec.configuration.ssh_config.hosts, - new=spec.configuration.ssh_config.hosts, - ) - # `_check_can_update_fleet_spec` ensures hosts are not changed - assert not changed_hosts, changed_hosts - active_instance_nums: set[int] = set() - removed_instance_nums: list[int] = [] - if removed_hosts or added_hosts: - for instance_model in fleet_model.instances: - if instance_model.deleted: - continue - active_instance_nums.add(instance_model.instance_num) - rci = get_instance_remote_connection_info(instance_model) - if rci is None: - logger.error( - "Cloud instance %s in SSH fleet %s", - instance_model.id, - fleet_model.id, + fleet_model.spec = spec.model_dump_json() + # Reset consolidation attempt so the next pipeline pass picks up the spec change promptly. + fleet_model.consolidation_attempt = 0 + + if ( + fleet_sensitive.spec.configuration.ssh_config is not None + and spec.configuration.ssh_config is not None + ): + added_hosts, removed_hosts, changed_hosts = _calculate_ssh_hosts_changes( + current=fleet_sensitive.spec.configuration.ssh_config.hosts, + new=spec.configuration.ssh_config.hosts, + ) + # `_check_can_update_fleet_spec` ensures hosts are not changed + assert not changed_hosts, changed_hosts + active_instance_nums: set[int] = set() + removed_instance_nums: list[int] = [] + if removed_hosts or added_hosts: + for instance_model in fleet_model.instances: + if instance_model.deleted: + continue + active_instance_nums.add(instance_model.instance_num) + rci = get_instance_remote_connection_info(instance_model) + if rci is None: + logger.error( + "Cloud instance %s in SSH fleet %s", + instance_model.id, + fleet_model.id, + ) + continue + if rci.host in removed_hosts: + removed_instance_nums.append(instance_model.instance_num) + if added_hosts: + await _check_ssh_hosts_not_yet_added(session, spec, fleet.id) + for host in added_hosts.values(): + instance_num = get_next_instance_num(active_instance_nums) + instance_model = await create_fleet_ssh_instance_model( + project=project, + spec=spec, + ssh_params=spec.configuration.ssh_config, + env=spec.configuration.env, + blocks=spec.configuration.blocks, + instance_num=instance_num, + host=host, ) - continue - if rci.host in removed_hosts: - removed_instance_nums.append(instance_model.instance_num) - if added_hosts: - await _check_ssh_hosts_not_yet_added(session, spec, fleet.id) - for host in added_hosts.values(): - instance_num = get_next_instance_num(active_instance_nums) - instance_model = await create_fleet_ssh_instance_model( - project=project, - spec=spec, - ssh_params=spec.configuration.ssh_config, - env=spec.configuration.env, - blocks=spec.configuration.blocks, - instance_num=instance_num, - host=host, + fleet_model.instances.append(instance_model) + events.emit( + session, + f"Instance created on fleet update. Status: {instance_model.status.upper()}", + actor=events.UserActor.from_user(user), + targets=[events.Target.from_model(instance_model)], + ) + active_instance_nums.add(instance_num) + if removed_instance_nums: + _terminate_fleet_instances(session, fleet_model, removed_instance_nums, actor=user) + + await session.commit() + return fleet_model_to_fleet(fleet_model) + + +_SSH_KEY_UPDATE_WAIT_SECONDS = 30 + + +@asynccontextmanager +async def _lock_ssh_key_update_instances( + session: AsyncSession, fleet_model: FleetModel, ssh_config: Optional[SSHParams] +) -> AsyncIterator[list[tuple[InstanceModel, str]]]: + """Let instance processing finish before atomically updating SSH keys. + + `_update_fleet` already holds the fleet row lock. If keys change, reserve the + fleet using its existing lease fields so instance fetchers stop taking new + locks. Release the DB connection between attempts, and leave the keys unchanged + on timeout or cancellation. The lease expires after a crashed API process. + """ + if ssh_config is None or not _get_ssh_instance_key_updates(fleet_model.instances, ssh_config): + yield [] + return + + fleet_id, previous_spec = fleet_model.id, fleet_model.spec + lock_token = uuid.uuid4() + deadline = time.monotonic() + _SSH_KEY_UPDATE_WAIT_SECONDS + # Leave enough time after the wait for the final DB-only update. No connection + # is held while sleeping, and an abandoned lease is recovered by the pipelines. + lease_duration = timedelta(seconds=_SSH_KEY_UPDATE_WAIT_SECONDS + 30) + await session.commit() + try: + res = await session.execute( + update(FleetModel) + .where( + FleetModel.id == fleet_id, + FleetModel.spec == previous_spec, + FleetModel.deleted == False, + FleetModel.lock_expires_at.is_(None), + ) + .values(lock_token=lock_token, lock_expires_at=get_current_datetime() + lease_duration) + .returning(FleetModel.id) + .execution_options(synchronize_session=False) + ) + if res.scalar_one_or_none() is None: + raise ServerClientError( + "Failed to update SSH keys: fleet changed or is being processed. Try again later." + ) + await session.commit() + + while True: + # Start with a write so SQLite cannot acquire a stale read snapshot + # before renewing the lease. The same guard detects a lost lease on Postgres. + res = await session.execute( + update(FleetModel) + .where( + FleetModel.id == fleet_id, + FleetModel.lock_token == lock_token, + FleetModel.lock_expires_at > get_current_datetime(), ) - fleet_model.instances.append(instance_model) - events.emit( - session, - f"Instance created on fleet update. Status: {instance_model.status.upper()}", - actor=events.UserActor.from_user(user), - targets=[events.Target.from_model(instance_model)], + .values(lock_expires_at=get_current_datetime() + lease_duration) + .returning(FleetModel.id) + .execution_options(synchronize_session=False) + ) + if res.scalar_one_or_none() is None: + raise ServerClientError( + "Failed to update SSH keys: fleet lock was lost. Try again later." + ) + res = await session.execute( + select(InstanceModel) + .where(InstanceModel.fleet_id == fleet_id, InstanceModel.deleted == False) + .order_by(InstanceModel.id) + .with_for_update(key_share=True) + .execution_options(populate_existing=True) + ) + key_updates = _get_ssh_instance_key_updates(list(res.scalars()), ssh_config) + locked = [i for i, _ in key_updates if i.lock_expires_at is not None] + if not locked: + # Instance processing may have completed or deleted instances during + # the wait. Refetch the relationship before applying host-list changes. + await session.execute( + select(FleetModel) + .where(FleetModel.id == fleet_id) + .options(selectinload(FleetModel.instances)) + .execution_options(populate_existing=True) + ) + fleet_model.lock_token = None + fleet_model.lock_expires_at = None + yield key_updates + return + await session.commit() + if time.monotonic() >= deadline: + raise ServerClientError( + f"Timed out waiting to update SSH keys: instance {locked[0].name!r} is still being processed. Try again later." ) - active_instance_nums.add(instance_num) - if removed_instance_nums: - _terminate_fleet_instances(session, fleet_model, removed_instance_nums, actor=user) + await asyncio.sleep(0.1) + finally: + # Roll back before releasing the lease so exceptions cannot persist only + # part of a fleet update. A stale caller must never unlock another owner. + await session.rollback() + await session.execute( + update(FleetModel) + .where(FleetModel.id == fleet_id, FleetModel.lock_token == lock_token) + .values(lock_token=None, lock_expires_at=None) + .execution_options(synchronize_session=False) + ) + await session.commit() - await session.commit() - return fleet_model_to_fleet(fleet_model) + +def _get_ssh_instance_key_updates( + instances: list[InstanceModel], ssh_config: SSHParams +) -> list[tuple[InstanceModel, str]]: + hosts = {h if isinstance(h, str) else h.hostname: h for h in ssh_config.hosts} + updates = [] + for instance in instances: + rci = get_instance_remote_connection_info(instance) + if instance.deleted or rci is None or (host := hosts.get(rci.host)) is None: + continue + ssh_key, proxy_jump = ssh_config.ssh_key, ssh_config.proxy_jump + if not isinstance(host, str): + ssh_key, proxy_jump = host.ssh_key or ssh_key, host.proxy_jump or proxy_jump + assert ssh_key is not None + if proxy_jump is not None: + assert proxy_jump.ssh_key is not None + updated = rci.model_copy( + update={ + "ssh_keys": [ssh_key], + "ssh_proxy_keys": [proxy_jump.ssh_key] if proxy_jump is not None else None, + } + ) + if updated != rci: + updates.append((instance, updated.model_dump_json())) + return updates def _can_update_fleet_spec(current_fleet_spec: FleetSpec, new_fleet_spec: FleetSpec) -> bool: @@ -1286,8 +1421,12 @@ def _check_can_update_fleet_configuration(current: FleetConfiguration, new: Flee _check_can_update_ssh_config(current_ssh_config, new_ssh_config) -@_check_can_update("hosts") +@_check_can_update("hosts", "identity_file", "ssh_key", "proxy_jump") def _check_can_update_ssh_config(current: SSHParams, new: SSHParams, diff: ModelDiff): + if "proxy_jump" in diff: + if current.proxy_jump is None or new.proxy_jump is None: + raise ServerClientError("Adding or removing proxy_jump cannot be updated in-place") + _check_can_update_inner(current.proxy_jump, new.proxy_jump, ("identity_file", "ssh_key")) if "hosts" in diff: _, _, changed_hosts = _calculate_ssh_hosts_changes(current.hosts, new.hosts) if changed_hosts: @@ -1311,15 +1450,17 @@ def _calculate_ssh_hosts_changes( if current_host != new_host: changed_hosts.add(host) else: - current_host = copy_model(current_host, reset={"identity_file"}) - new_host = copy_model(new_host, reset={"identity_file"}) + current_host = copy_model(current_host, reset={"identity_file", "ssh_key"}) + new_host = copy_model(new_host, reset={"identity_file", "ssh_key"}) # XXX: cannot use copy_model() or diff_models() with # `reset={..., "proxy_jump": {"identity_file"}}` # as SSHProxyParams.identity_file has no default value if current_host.proxy_jump is not None: current_host.proxy_jump.identity_file = "" + current_host.proxy_jump.ssh_key = None if new_host.proxy_jump is not None: new_host.proxy_jump.identity_file = "" + new_host.proxy_jump.ssh_key = None if diff_models(current_host, new_host): changed_hosts.add(host) return added_hosts, removed_hosts, changed_hosts @@ -1367,10 +1508,17 @@ async def _check_ssh_hosts_not_yet_added( def _remove_fleet_spec_sensitive_info(spec: FleetSpec): if spec.configuration.ssh_config is not None: - spec.configuration.ssh_config.ssh_key = None + ssh_config = spec.configuration.ssh_config + if ssh_config.ssh_key is not None: + ssh_config.ssh_key.private = None + if ssh_config.proxy_jump is not None and ssh_config.proxy_jump.ssh_key is not None: + ssh_config.proxy_jump.ssh_key.private = None for host in spec.configuration.ssh_config.hosts: if not isinstance(host, str): - host.ssh_key = None + if host.ssh_key is not None: + host.ssh_key.private = None + if host.proxy_jump is not None and host.proxy_jump.ssh_key is not None: + host.proxy_jump.ssh_key.private = None def _validate_fleet_spec_and_set_defaults(spec: FleetSpec): @@ -1383,13 +1531,16 @@ def _validate_fleet_spec_and_set_defaults(spec: FleetSpec): validate_dstack_resource_name(spec.configuration.name) _validate_fleet_configuration_subtype_specific_fields(spec.configuration) if spec.configuration.ssh_config is not None: - _validate_all_ssh_params_specified(spec.configuration.ssh_config) - if spec.configuration.ssh_config.ssh_key is not None: - _validate_ssh_key(spec.configuration.ssh_config.ssh_key) - for host in spec.configuration.ssh_config.hosts: - if isinstance(host, SSHHostParams) and host.ssh_key is not None: - _validate_ssh_key(host.ssh_key) - _validate_internal_ips(spec.configuration.ssh_config) + ssh_config = spec.configuration.ssh_config + _validate_all_ssh_params_specified(ssh_config) + ssh_params: list[SSHParams | SSHHostParams] = [ssh_config] + ssh_params.extend(host for host in ssh_config.hosts if isinstance(host, SSHHostParams)) + for params in ssh_params: + if params.ssh_key is not None: + _validate_ssh_key(params.ssh_key) + if params.proxy_jump is not None and params.proxy_jump.ssh_key is not None: + _validate_ssh_key(params.proxy_jump.ssh_key) + _validate_internal_ips(ssh_config) _set_fleet_spec_defaults(spec) @@ -1447,12 +1598,16 @@ def _validate_ssh_key(ssh_key: SSHKey): if ssh_key.private is None: raise ServerClientError("Private key not provided") try: - ssh_utils.pkey_from_str(ssh_key.private) + private_key = ssh_utils.pkey_from_str(ssh_key.private) except ValueError: raise ServerClientError( "Unsupported key type. " "The key type should be RSA, ECDSA, or Ed25519 and should not be encrypted with passphrase." ) + # Plans and gateway synchronization compare public keys. A stale .pub file must not + # hide replacement of the private key at the same path. + if " ".join(ssh_key.public.split()[:2]) != ssh_utils.generate_public_key(private_key): + raise ServerClientError("Public SSH key does not match the private key") def _validate_internal_ips(ssh_config: SSHParams): diff --git a/src/dstack/_internal/server/services/gateways/client.py b/src/dstack/_internal/server/services/gateways/client.py index 0c9defecca..2a176c69a1 100644 --- a/src/dstack/_internal/server/services/gateways/client.py +++ b/src/dstack/_internal/server/services/gateways/client.py @@ -11,6 +11,10 @@ from dstack._internal.core.models.configurations import RateLimit, ServiceConfiguration from dstack._internal.core.models.instances import SSHConnectionParams from dstack._internal.core.models.runs import JobSpec, JobSubmission, get_service_port +from dstack._internal.proxy.gateway.schemas.registry import ( + ReplicaHeadNodeKeyUpdate, + UpdateReplicaHeadNodeKeysResponse, +) from dstack._internal.proxy.gateway.schemas.services import ServiceListItem, ServiceListResponse from dstack._internal.proxy.gateway.schemas.stats import ServiceStats from dstack._internal.server import settings @@ -141,6 +145,22 @@ async def register_replica( resp.raise_for_status() self.is_server_ready = True + async def update_replica_head_node_keys( + self, updates: list[ReplicaHeadNodeKeyUpdate] + ) -> dict[str, str]: + """Send one synchronization batch of SSH-fleet head-node key replacements. + + The gateway preserves existing service tunnels and saves successful updates + once. Return failures by job ID; the next sync retries keys still different. + """ + resp = await self._client.post( + self._url("/api/registry/replicas/head_node_keys"), + json={"replicas": [item.model_dump() for item in updates]}, + ) + resp.raise_for_status() + self.is_server_ready = True + return UpdateReplicaHeadNodeKeysResponse.model_validate(resp.json()).failures + async def unregister_replica(self, project: str, run_name: str, job_id: uuid.UUID): resp = await self._client.post( self._url( diff --git a/src/dstack/_internal/server/services/proxy/repo.py b/src/dstack/_internal/server/services/proxy/repo.py index 7ea7df43cd..12f79f9537 100644 --- a/src/dstack/_internal/server/services/proxy/repo.py +++ b/src/dstack/_internal/server/services/proxy/repo.py @@ -182,6 +182,8 @@ async def get_model(self, project_name: str, name: str) -> Optional[ChatModel]: async def get_project(self, name: str) -> Optional[Project]: res = await self.session.execute(select(ProjectModel).where(ProjectModel.name == name)) project = res.scalar_one_or_none() + # Release the DB connection before opening/replacing a service SSH tunnel. + await self.session.commit() if project is None: return None return Project( diff --git a/src/tests/_internal/cli/services/configurators/test_fleet.py b/src/tests/_internal/cli/services/configurators/test_fleet.py index 35c6a313f0..62ae203bdd 100644 --- a/src/tests/_internal/cli/services/configurators/test_fleet.py +++ b/src/tests/_internal/cli/services/configurators/test_fleet.py @@ -26,6 +26,8 @@ FleetSpec, FleetStatus, InstanceGroupPlacement, + SSHHostParams, + SSHProxyParams, ) from dstack._internal.core.models.instances import SSHKey from dstack._internal.core.models.profiles import Profile @@ -179,6 +181,69 @@ def test_prints_in_place_update_diff(self, monkeypatch: pytest.MonkeyPatch): assert " - nodes" in output confirm_ask.assert_called_once_with("Update the fleet in-place?") + @pytest.mark.parametrize( + "per_host,head_node,same_path", + [(False, False, False), (False, True, True), (True, False, True), (True, True, False)], + ) + def test_identifies_the_changed_ssh_key_without_printing_key_material( + self, monkeypatch, per_host, head_node, same_path + ): + console, confirm_ask = patch_console_and_confirm(monkeypatch) + current = get_ssh_fleet_spec() + ssh = current.configuration.ssh_config + assert ssh is not None + target = ssh + path = "ssh_config" + if per_host: + target = SSHHostParams(hostname="10.0.0.100") + ssh.hosts = [target] + path += ".hosts[0]" + if head_node: + jump = SSHProxyParams(hostname="head", user="ubuntu", identity_file="old-key") + target.proxy_jump = jump + target = jump + path += ".proxy_jump" + target.identity_file = "old-key" + target.ssh_key = SSHKey(public="old-public-key", private="old-private-key") + new = current.model_copy(deep=True) + changed = new.configuration.ssh_config + assert changed is not None + if per_host: + changed = changed.hosts[0] + if head_node: + changed = changed.proxy_jump + assert changed is not None + if not same_path: + changed.identity_file = "new-key" + changed.ssh_key = SSHKey(public="new-public-key", private="new-private-key") + + plan = create_fleet_plan(current_spec=current, spec=new, action=ApplyAction.UPDATE) + FleetConfigurator(Mock())._apply_plan(plan, get_command_args()) + + output = console.export_text() + label = f"{path}.identity_file" + (" (key contents changed)" if same_path else "") + assert f" - {label}\n" in output + assert "ssh_key" not in output + assert "public-key" not in output + assert "private-key" not in output + confirm_ask.assert_called_once_with("Update the fleet in-place?") + + @pytest.mark.parametrize("field,value", [("user", "other-user"), ("port", 2222)]) + def test_identifies_ssh_route_changes_requiring_recreation(self, monkeypatch, field, value): + console, confirm_ask = patch_console_and_confirm(monkeypatch) + current = get_ssh_fleet_spec() + new = current.model_copy(deep=True) + assert new.configuration.ssh_config is not None + setattr(new.configuration.ssh_config, field, value) + plan = create_fleet_plan(current_spec=current, spec=new, action=ApplyAction.CREATE) + + FleetConfigurator(Mock())._apply_plan(plan, get_command_args()) + + output = console.export_text() + assert "Detected changes that cannot be updated in-place:" in output + assert f" - ssh_config.{field}\n" in output + confirm_ask.assert_called_once_with("Re-create the fleet?") + def test_prints_recreate_diff(self, monkeypatch: pytest.MonkeyPatch): console, confirm_ask = patch_console_and_confirm(monkeypatch) current_spec = get_cloud_fleet_spec(placement=InstanceGroupPlacement.ANY) @@ -240,7 +305,7 @@ def test_renders_ssh_hosts_change(self): == dedent( """ - Configuration properties: - - ssh_config + - ssh_config.hosts """ ).lstrip() ) diff --git a/src/tests/_internal/core/services/ssh/test_tunnel.py b/src/tests/_internal/core/services/ssh/test_tunnel.py index 4e3099ab6b..4f825af088 100644 --- a/src/tests/_internal/core/services/ssh/test_tunnel.py +++ b/src/tests/_internal/core/services/ssh/test_tunnel.py @@ -22,6 +22,33 @@ class TestSSHTunnel: + @pytest.mark.usefixtures("ssh_client_info") + def test_update_proxy_identity_keeps_sockets_and_other_keys(self, tmp_path): + external_key = tmp_path / "external-key" + external_key.write_text("original external key") + tunnel = SSHTunnel( + destination="ubuntu@host", + identity=FileContent("project key"), + ssh_proxies=[ + ( + SSHConnectionParams(hostname="jump", username="ubuntu", port=22), + FilePath(external_key), + ) + ], + ) + socket = tunnel.control_sock_path + tunnel.update_proxy_identity(0, "replacement") + path = Path(tunnel.ssh_proxies[0][1]) + assert path.read_text() == "replacement" + assert path.stat().st_mode & 0o777 == 0o600 + assert external_key.read_text() == "original external key" + assert Path(tunnel.identity_path).read_text() == "project key" + assert tunnel.control_sock_path == socket + assert str(path) in " ".join(tunnel.open_command()) + tunnel.update_proxy_identity(0, "next key") + assert Path(tunnel.ssh_proxies[0][1]) == path + assert path.read_text() == "next key" + @pytest.fixture def ssh_client_info(self, monkeypatch: pytest.MonkeyPatch) -> SSHClientInfo: ssh_client_info = SSHClientInfo.from_raw_version("OpenSSH_9.7p1", Path("/usr/bin/ssh")) diff --git a/src/tests/_internal/proxy/gateway/routers/test_registry.py b/src/tests/_internal/proxy/gateway/routers/test_registry.py index a9b265d9fc..1f6e35dc5e 100644 --- a/src/tests/_internal/proxy/gateway/routers/test_registry.py +++ b/src/tests/_internal/proxy/gateway/routers/test_registry.py @@ -3,6 +3,7 @@ from datetime import datetime, timezone from pathlib import Path from typing import Optional +from unittest.mock import patch import httpx import pytest @@ -14,7 +15,10 @@ from dstack._internal.proxy.gateway.services.nginx import Nginx from dstack._internal.proxy.gateway.testing.common import Mocks from dstack._internal.proxy.lib.models import ChatModel, OpenAIChatModelFormat +from dstack._internal.proxy.lib.services.service_connection import ServiceConnection from dstack._internal.proxy.lib.testing.common import make_project, make_service +from dstack._internal.server.testing.common import get_ssh_key +from tests._internal.utils.test_ssh import PRIVATE_KEY, PUBLIC_KEY def make_client(nginx_dir: Path, repo: Optional[GatewayProxyRepo] = None) -> httpx.AsyncClient: @@ -90,6 +94,92 @@ def sample_model_options(name: str = "test-model") -> dict: } +@pytest.mark.asyncio +class TestUpdateReplicaHeadNodeKeys: + @pytest.mark.parametrize("failure", ["missing", "key-file"]) + async def test_batch_saves_once_preserves_tunnels_and_retries_failures( + self, tmp_path, system_mocks, monkeypatch, failure + ): + state_file = tmp_path / "state.json" + repo = GatewayProxyRepo(file=state_file) + client = make_client(tmp_path, repo=repo) + replacement = get_ssh_key() + updates = [] + for n in range(2): + project, run, job = f"project-{n}", f"run-{n}", f"job-{n}" + response = await client.post( + f"/api/registry/{project}/services/register", + json=register_service_payload(run_name=run, domain=f"{run}.test"), + ) + assert response.status_code == 200 + payload = register_replica_payload_with_head_proxy(job_id=job) + payload["ssh_head_proxy_private_key"] = PRIVATE_KEY + response = await client.post( + f"/api/registry/{project}/services/{run}/replicas/register", json=payload + ) + assert response.status_code == 200 + updates.append( + dict( + project_name=project, + run_name=run, + job_id=job, + ssh_head_proxy_private_key=replacement.private, + ) + ) + # Populate the cached old public keys before replacing their private keys. + before = (await client.get("/api/services/list")).json() + assert all( + service["replicas"][0]["ssh_head_proxy_public_key"] == " ".join(PUBLIC_KEY.split()[:2]) + for service in before["services"] + ) + system_mocks.open_conn.reset_mock() + system_mocks.close_conn.reset_mock() + system_mocks.reload_nginx.reset_mock() + batch = list(updates) + if failure == "missing": + batch.append({**updates[0], "job_id": "missing"}) + original_update = ServiceConnection.update_head_node_key + + async def fail_one_key_file(connection, replica): + if replica.id == "job-1": + raise OSError("key file write failed") + return await original_update(connection, replica) + + with patch.object(repo, "save", wraps=repo.save) as save: + with monkeypatch.context() as scoped: + if failure == "key-file": + scoped.setattr(ServiceConnection, "update_head_node_key", fail_one_key_file) + response = await client.post( + "/api/registry/replicas/head_node_keys", json={"replicas": batch} + ) + assert response.status_code == 200 + assert set(response.json()["failures"]) == { + "missing" if failure == "missing" else "job-1" + } + assert save.call_count == 1 + restored = GatewayProxyRepo.load(state_file) + for n in range(2): + service = await restored.get_service(f"project-{n}", f"run-{n}") + assert service is not None and len(service.replicas) == 1 + expected = PRIVATE_KEY if failure == "key-file" and n == 1 else replacement.private + assert service.replicas[0].ssh_head_proxy_private_key == expected + # Retry a failed write, then repeat a successful batch as after a lost response. + for _ in range(2): + response = await client.post( + "/api/registry/replicas/head_node_keys", json={"replicas": updates} + ) + assert response.json() == {"failures": {}} + assert save.call_count == (2 if failure == "key-file" else 1) + after = (await client.get("/api/services/list")).json() + assert all( + service["replicas"][0]["ssh_head_proxy_public_key"] == replacement.public + for service in after["services"] + ) + system_mocks.open_conn.assert_not_awaited() + system_mocks.close_conn.assert_not_awaited() + system_mocks.reload_nginx.assert_not_called() + + @pytest.mark.asyncio class TestRegisterService: async def test_register(self, tmp_path: Path, system_mocks: Mocks) -> None: diff --git a/src/tests/_internal/proxy/gateway/routers/test_services.py b/src/tests/_internal/proxy/gateway/routers/test_services.py index 90c9238308..c4dd228f50 100644 --- a/src/tests/_internal/proxy/gateway/routers/test_services.py +++ b/src/tests/_internal/proxy/gateway/routers/test_services.py @@ -51,7 +51,7 @@ async def test_list(self): "id": service_id, "project_name": "test-proj", "run_name": "srv-1", - "replicas": [{"id": replica_id}], + "replicas": [{"id": replica_id, "ssh_head_proxy_public_key": None}], } ] } diff --git a/src/tests/_internal/proxy/lib/services/test_service_connection.py b/src/tests/_internal/proxy/lib/services/test_service_connection.py index f12eddf0da..444da6ec80 100644 --- a/src/tests/_internal/proxy/lib/services/test_service_connection.py +++ b/src/tests/_internal/proxy/lib/services/test_service_connection.py @@ -6,6 +6,7 @@ import pytest_asyncio from dstack._internal.core.errors import SSHError +from dstack._internal.core.models.instances import SSHConnectionParams from dstack._internal.core.services.ssh.tunnel import SSHTunnel from dstack._internal.proxy.lib.errors import UnexpectedProxyError from dstack._internal.proxy.lib.services import service_connection @@ -125,6 +126,31 @@ async def test_in_server_connection_does_not_monitor( tunnel_factory.return_value.wait_closed.assert_not_awaited() assert tunnel_factory.call_args.kwargs["background"] is True + async def test_in_server_jump_tunnel_reopens_only_on_a_new_request( + self, project, service, pool, tunnel_factory, monitor_clock + ): + service = service.model_copy(update={"domain": None}) + replica = service.replicas[0].model_copy( + update={ + "ssh_head_proxy": SSHConnectionParams(hostname="jump", username="ubuntu", port=22), + "ssh_head_proxy_private_key": "key", + } + ) + connection = await pool.get_or_add(project, service, replica) + tunnel = tunnel_factory.return_value + tunnel.exits.put_nowait(255) + await _yield_to_event_loop() + + tunnel.aopen.assert_awaited_once() + tunnel.aclose.assert_awaited_once() + assert not connection._is_open.is_set() + assert monitor_clock.pending.empty() + assert tunnel_factory.call_args.kwargs["background"] is False + + assert await pool.get_or_add(project, service, replica) is connection + assert tunnel.aopen.await_count == 2 + assert connection._is_open.is_set() + async def test_removal_cancels_reconnect_before_final_tunnel_cleanup( self, project, service, pool, tunnel_factory ): @@ -165,6 +191,61 @@ async def close_tunnel(): @pytest.mark.asyncio class TestServiceConnectionPool: + @pytest.mark.parametrize("gateway", [False, True]) + async def test_head_node_key_update_preserves_client_socket_and_reconnects_with_new_key( + self, project, service, pool, tunnel_factory, gateway + ): + service = service.model_copy(update={"domain": service.domain if gateway else None}) + replica = service.replicas[0].model_copy( + update={ + "ssh_head_proxy": SSHConnectionParams(hostname="jump", username="ubuntu", port=22), + "ssh_head_proxy_private_key": "old key", + } + ) + connection = await pool.get_or_add(project, service, replica) + client = await connection.client() + socket = connection.app_socket_path + changed = replica.model_copy(update={"ssh_head_proxy_private_key": "new key"}) + assert await pool.get_or_add(project, service, changed) is connection + assert await connection.client() is client + assert not client.is_closed + assert connection.app_socket_path == socket + tunnel = tunnel_factory.return_value + tunnel.update_proxy_identity.assert_called_once_with(0, "new key") + tunnel.aclose.assert_not_awaited() + tunnel.aopen.assert_awaited_once() + assert await pool.get_or_add(project, service, changed) is connection + tunnel.update_proxy_identity.assert_called_once() + tunnel.exits.put_nowait(255) + await _yield_to_event_loop() + if not gateway: + tunnel.aopen.assert_awaited_once() + assert await pool.get_or_add(project, service, changed) is connection + assert tunnel.aopen.await_count == 2 + assert await pool.get(replica.id) is connection + + async def test_failed_key_file_update_keeps_working_connection_and_can_retry( + self, project, service, pool, tunnel_factory + ): + replica = service.replicas[0].model_copy( + update={ + "ssh_head_proxy": SSHConnectionParams(hostname="jump", username="ubuntu", port=22), + "ssh_head_proxy_private_key": "old key", + } + ) + connection = await pool.get_or_add(project, service, replica) + client = await connection.client() + changed = replica.model_copy(update={"ssh_head_proxy_private_key": "new key"}) + tunnel = tunnel_factory.return_value + tunnel.update_proxy_identity.side_effect = OSError("write failed") + with pytest.raises(OSError, match="write failed"): + await pool.get_or_add(project, service, changed) + assert connection.replica == replica + assert not client.is_closed + tunnel.aclose.assert_not_awaited() + tunnel.update_proxy_identity.side_effect = None + assert await pool.get_or_add(project, service, changed) is connection + async def test_cancelled_removal_finishes_monitor_and_resource_cleanup( self, project, service, pool, tunnel_factory ): diff --git a/src/tests/_internal/server/background/pipeline_tasks/instances/test_pipeline.py b/src/tests/_internal/server/background/pipeline_tasks/instances/test_pipeline.py index cf1639617d..b1df098970 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/instances/test_pipeline.py +++ b/src/tests/_internal/server/background/pipeline_tasks/instances/test_pipeline.py @@ -240,16 +240,22 @@ async def test_fetch_allows_stale_instance_locks_if_fleet_is_waiting_for_instanc await session.refresh(instance) assert instance.lock_owner == InstancePipeline.__name__ - async def test_fetch_excludes_fresh_instances_when_fleet_is_waiting_for_instance_locks( - self, test_db, session: AsyncSession, fetcher: InstanceFetcher + @pytest.mark.parametrize("lease", ["pipeline", "key-update", "expired-key-update"]) + async def test_fetch_respects_fleet_processing_and_key_update_leases( + self, test_db, session: AsyncSession, fetcher: InstanceFetcher, lease ): project = await create_project(session=session) fleet = await create_fleet(session=session, project=project) stale = get_current_datetime() - dt.timedelta(minutes=1) - fleet.lock_owner = FleetPipeline.__name__ + fleet.lock_owner = FleetPipeline.__name__ if lease == "pipeline" else None fleet.lock_token = None - fleet.lock_expires_at = None + fleet.lock_expires_at = ( + None + if lease == "pipeline" + else get_current_datetime() + + dt.timedelta(seconds=-1 if lease == "expired-key-update" else 60) + ) instance = await create_instance( session=session, @@ -263,10 +269,14 @@ async def test_fetch_excludes_fresh_instances_when_fleet_is_waiting_for_instance items = await fetcher.fetch(limit=10) - assert items == [] + assert [item.id for item in items] == ( + [instance.id] if lease == "expired-key-update" else [] + ) await session.refresh(instance) - assert instance.lock_owner is None + assert instance.lock_owner == ( + InstancePipeline.__name__ if lease == "expired-key-update" else None + ) @pytest.mark.asyncio diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_gateway_replicas.py b/src/tests/_internal/server/background/pipeline_tasks/test_gateway_replicas.py index 1c1ca8bb4f..07a77e0ad1 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_gateway_replicas.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_gateway_replicas.py @@ -17,7 +17,12 @@ GatewayReplicaStatus, GatewayStatus, ) -from dstack._internal.core.models.instances import InstanceStatus +from dstack._internal.core.models.instances import ( + InstanceStatus, + RemoteConnectionInfo, + SSHConnectionParams, + SSHKey, +) from dstack._internal.core.models.runs import JobStatus, RunStatus, ServiceSpec from dstack._internal.proxy.gateway.schemas.services import ServiceListItem, ServiceListReplicaItem from dstack._internal.server import settings @@ -51,6 +56,7 @@ list_events, ) from dstack._internal.utils.common import get_current_datetime +from tests._internal.utils.test_ssh import PRIVATE_KEY, PUBLIC_KEY @pytest.fixture @@ -763,6 +769,105 @@ class TestGatewayReplicaWorkerRunningStateSync: pytestmark = pytest.mark.usefixtures("image_config_mock") + @pytest.mark.parametrize("failure", [None, "update", "changed-again"]) + async def test_updates_head_node_key_without_unregistering_or_changing_job( + self, test_db, session, worker, mock_gateway_connection, failure + ): + project = await create_project(session) + user = await create_user(session) + repo = await create_repo(session, project_id=project.id) + backend = await create_backend(session, project.id) + gateway = await create_gateway( + session, project.id, backend.id, status=GatewayStatus.RUNNING + ) + replica = await create_gateway_replica(session, backend, gateway_id=gateway.id) + run, job = await self._create_service_run_and_job( + session, project, repo, user, gateway, run_name="rotation" + ) + assert job.instance is not None + other_job = await create_job( + session, + run, + status=JobStatus.RUNNING, + instance=job.instance, + registered=True, + ready=True, + replica_num=1, + ) + key = SSHKey(public=PUBLIC_KEY, private=PRIVATE_KEY) + rci = RemoteConnectionInfo( + host="host", + port=22, + ssh_user="ubuntu", + ssh_keys=[key], + ssh_proxy=SSHConnectionParams(hostname="jump", port=22, username="ubuntu"), + ssh_proxy_keys=[key], + ) + job.instance.remote_connection_info = rci.model_dump_json() + for service_job in (job, other_job): + session.add( + ServiceReplicaRegistrationModel( + job_id=service_job.id, gateway_replica_id=replica.id, is_registered=True + ) + ) + _lock_replica(replica) + await session.commit() + reported = { + service_job.id: ServiceListReplicaItem( + id=service_job.id.hex, ssh_head_proxy_public_key="old key" + ) + for service_job in (job, other_job) + } + service = ServiceListItem( + id=run.id.hex, + project_name=project.name, + run_name=run.run_name, + replicas=list(reported.values()), + ) + client = _get_client_mock(mock_gateway_connection) + client.list_services.return_value = [service] + + async def update(updates): + if failure == "update": + raise GatewayError("update failed") + for item in updates: + reported[uuid.UUID(item.job_id)].ssh_head_proxy_public_key = " ".join( + PUBLIC_KEY.split()[:2] + ) + if failure == "changed-again": + assert job.instance is not None + rci.ssh_proxy_keys = [SSHKey(public="ssh-rsa another", private="another private")] + job.instance.remote_connection_info = rci.model_dump_json() + await session.commit() + + return {} + + client.update_replica_head_node_keys.side_effect = update + await worker.process(_replica_to_pipeline_item(replica)) + client.update_replica_head_node_keys.assert_awaited_once() + updates = client.update_replica_head_node_keys.call_args.args[0] + assert {item.job_id for item in updates} == {job.id.hex, other_job.id.hex} + assert all( + item.project_name == project.name + and item.run_name == run.run_name + and item.ssh_head_proxy_private_key == PRIVATE_KEY + for item in updates + ) + client.register_replica.assert_not_called() + client.unregister_replica.assert_not_called() + client.register_service.assert_not_called() + client.unregister_service.assert_not_called() + client.update_replica_head_node_keys.reset_mock(side_effect=True) + client.update_replica_head_node_keys.return_value = {} + await session.refresh(replica) + _lock_replica(replica) + await session.commit() + await worker.process(_replica_to_pipeline_item(replica)) + assert client.update_replica_head_node_keys.call_count == (0 if failure is None else 1) + await session.refresh(job) + assert job.status == JobStatus.RUNNING + assert job.submission_num == 0 + async def _create_service_run_and_job( self, session: AsyncSession, diff --git a/src/tests/_internal/server/services/test_fleets.py b/src/tests/_internal/server/services/test_fleets.py index 19ae011d65..94d73cc6eb 100644 --- a/src/tests/_internal/server/services/test_fleets.py +++ b/src/tests/_internal/server/services/test_fleets.py @@ -1,36 +1,319 @@ +import asyncio +import uuid +from datetime import timedelta from typing import Optional, Union from unittest.mock import Mock import pytest +from sqlalchemy import update from sqlalchemy.ext.asyncio import AsyncSession from dstack._internal.core.backends.base.backend import Backend from dstack._internal.core.errors import ServerClientError from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.common import ApplyAction from dstack._internal.core.models.fleets import ( + ApplyFleetPlanInput, FleetConfiguration, FleetNodesSpec, FleetSpec, InstanceGroupPlacement, SSHHostParams, SSHParams, + SSHProxyParams, ) -from dstack._internal.core.models.instances import RemoteConnectionInfo -from dstack._internal.server.models import FleetModel, ProjectModel +from dstack._internal.core.models.instances import ( + InstanceStatus, + RemoteConnectionInfo, + SSHConnectionParams, + SSHKey, +) +from dstack._internal.core.models.runs import JobStatus, RunStatus +from dstack._internal.server.background.pipeline_tasks.base import Heartbeater +from dstack._internal.server.background.pipeline_tasks.instances import InstanceFetcher +from dstack._internal.server.models import FleetModel, InstanceModel, ProjectModel from dstack._internal.server.services.backends import get_project_backends from dstack._internal.server.services.fleets import ( + _check_can_update_ssh_config, + _validate_fleet_spec_and_set_defaults, + _validate_ssh_key, + apply_plan, get_fleet_master_instance_provisioning_data, get_plan, ) +from dstack._internal.server.services.instances import get_instance_remote_connection_info from dstack._internal.server.testing.common import ( create_fleet, create_instance, + create_job, create_project, + create_repo, + create_run, create_user, get_fleet_spec, get_job_provisioning_data, get_ssh_key, ) +from dstack._internal.utils.common import get_current_datetime +from tests._internal.utils.test_ssh import PRIVATE_KEY, PUBLIC_KEY + + +@pytest.mark.asyncio +@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) +class TestApplySSHFleetKeys: + @pytest.mark.parametrize( + "per_host,busy,rotation", + [(False, False, "host"), (False, True, "jump"), (True, True, "both")], + ) + async def test_updates_existing_instances_without_reprovisioning( + self, test_db, session, monkeypatch, per_host, busy, rotation + ): + monkeypatch.setattr( + "dstack._internal.server.services.backends.get_project_backends", Mock(return_value=[]) + ) + user = await create_user(session) + project = await create_project(session, owner=user) + old_key, new_key = get_ssh_key(), SSHKey(public=PUBLIC_KEY, private=PRIVATE_KEY) + jump = SSHProxyParams(hostname="jump", user="ubuntu", identity_file="key", ssh_key=old_key) + ssh_config = SSHParams( + user="ubuntu", + identity_file="key", + ssh_key=old_key, + proxy_jump=jump, + hosts=[ + SSHHostParams( + hostname="host", + identity_file="key", + ssh_key=old_key, + proxy_jump=jump.model_copy(deep=True), + ) + if per_host + else "host" + ], + ) + spec = get_fleet_spec(conf=FleetConfiguration(name="fleet", ssh_config=ssh_config)) + fleet = await create_fleet(session, project=project, spec=spec) + instance = await create_instance( + session, + project=project, + fleet=fleet, + backend=BackendType.REMOTE, + status=InstanceStatus.BUSY if busy else InstanceStatus.IDLE, + remote_connection_info=RemoteConnectionInfo( + host="host", + port=22, + ssh_user="ubuntu", + ssh_keys=[old_key], + ssh_proxy=SSHConnectionParams(hostname="jump", port=22, username="ubuntu"), + ssh_proxy_keys=[old_key], + ), + ) + job = None + if busy: + repo = await create_repo(session, project_id=project.id) + run = await create_run( + session, project, repo, user, fleet=fleet, status=RunStatus.RUNNING + ) + job = await create_job(session, run, instance=instance, status=JobStatus.RUNNING) + provisioning = instance.job_provisioning_data + new = spec.model_copy(deep=True) + keys = new.configuration.ssh_config + assert keys is not None + if per_host: + keys = keys.hosts[0] + assert isinstance(keys, SSHHostParams) + if rotation != "jump": + keys.ssh_key = new_key + assert keys.proxy_jump is not None + if rotation != "host": + keys.proxy_jump.ssh_key = new_key + plan = await get_plan(session, project, user, new) + assert plan.action == ApplyAction.UPDATE + assert plan.current_resource is not None and plan.effective_spec is not None + assert plan.current_resource.spec != plan.effective_spec # same file path, different key + assert PRIVATE_KEY not in plan.effective_spec.model_dump_json() + updated = await apply_plan( + session, + user, + project, + ApplyFleetPlanInput(spec=new, current_resource=plan.current_resource), + False, + Mock(), + ) + assert updated.id == fleet.id + await session.refresh(instance) + rci = get_instance_remote_connection_info(instance) + assert rci is not None + assert rci.ssh_keys == [old_key if rotation == "jump" else new_key] + assert rci.ssh_proxy_keys == [old_key if rotation == "host" else new_key] + assert instance.job_provisioning_data == provisioning + assert instance.status == (InstanceStatus.BUSY if busy else InstanceStatus.IDLE) + if job is not None: + await session.refresh(job) + assert job.status == JobStatus.RUNNING + assert job.submission_num == 0 + + @pytest.mark.parametrize("outcome", ["success", "timeout", "cancel", "lost-lease"]) + async def test_waits_for_instance_processing_without_holding_db_connection( + self, test_db, session, monkeypatch, outcome + ): + monkeypatch.setattr( + "dstack._internal.server.services.backends.get_project_backends", Mock(return_value=[]) + ) + user = await create_user(session) + project = await create_project(session, owner=user) + old_key = get_ssh_key() + spec = get_fleet_spec( + conf=FleetConfiguration( + name="fleet", ssh_config=SSHParams(user="ubuntu", hosts=["host"], ssh_key=old_key) + ) + ) + fleet = await create_fleet(session, project=project, spec=spec) + instance = await create_instance( + session, + project=project, + fleet=fleet, + backend=BackendType.REMOTE, + remote_connection_info=RemoteConnectionInfo( + host="host", port=22, ssh_user="ubuntu", ssh_keys=[old_key] + ), + ) + instance.lock_expires_at = get_current_datetime() + timedelta(seconds=30) + await session.commit() + new = spec.model_copy(deep=True) + assert new.configuration.ssh_config is not None + new.configuration.ssh_config.ssh_key = SSHKey(public=PUBLIC_KEY, private=PRIVATE_KEY) + plan = await get_plan(session, project, user, new) + fleet_id, instance_id = fleet.id, instance.id + original_sleep = asyncio.sleep + waited = False + replacement_token = None + + async def finish_processing(delay): + nonlocal waited, replacement_token + if delay != 0.1: + return await original_sleep(0) + waited = True + if hasattr(test_db.engine.pool, "checkedout"): + assert test_db.engine.pool.checkedout() == 0 + # Exercise the actual fetcher while apply owns the fleet lease. + fetcher = InstanceFetcher( + queue=asyncio.Queue(), + queue_desired_minsize=1, + min_processing_interval=timedelta(seconds=7), + lock_timeout=timedelta(seconds=30), + heartbeater=Heartbeater( + model_type=InstanceModel, + lock_timeout=timedelta(seconds=30), + heartbeat_trigger=timedelta(seconds=15), + ), + ) + assert await fetcher.fetch(1) == [] + if outcome == "cancel": + raise asyncio.CancelledError + async with test_db.get_session() as processing: + if outcome == "lost-lease": + replacement_token = uuid.uuid4() + await processing.execute( + update(FleetModel) + .where(FleetModel.id == fleet_id) + .values(lock_token=replacement_token) + ) + else: + await processing.execute( + update(InstanceModel) + .where(InstanceModel.id == instance_id) + .values(lock_expires_at=None) + ) + await processing.commit() + + monkeypatch.setattr(asyncio, "sleep", finish_processing) + if outcome == "timeout": + monkeypatch.setattr( + "dstack._internal.server.services.fleets._SSH_KEY_UPDATE_WAIT_SECONDS", 0 + ) + request = ApplyFleetPlanInput(spec=new, current_resource=plan.current_resource) + if outcome == "success": + await apply_plan(session, user, project, request, False, Mock()) + elif outcome == "cancel": + with pytest.raises(asyncio.CancelledError): + await apply_plan(session, user, project, request, False, Mock()) + else: + with pytest.raises( + ServerClientError, match="Timed out" if outcome == "timeout" else "lock was lost" + ): + await apply_plan(session, user, project, request, False, Mock()) + assert waited == (outcome != "timeout") + await session.refresh(instance) + await session.refresh(fleet) + expected = ( + SSHKey(public=PUBLIC_KEY, private=PRIVATE_KEY) if outcome == "success" else old_key + ) + rci = get_instance_remote_connection_info(instance) + assert rci is not None and rci.ssh_keys == [expected] + stored = FleetSpec.model_validate_json(fleet.spec).configuration.ssh_config + assert stored is not None and stored.ssh_key == expected + if outcome == "lost-lease": + assert fleet.lock_token == replacement_token + assert fleet.lock_expires_at is not None + else: + assert fleet.lock_token is None + assert fleet.lock_expires_at is None + + +class TestCheckCanUpdateSSHConfig: + @pytest.mark.parametrize( + "field,value", [("hostname", "other"), ("port", 2222), ("user", "other")] + ) + @pytest.mark.parametrize("per_host", [False, True]) + def test_rejects_jump_route_changes(self, field, value, per_host): + jump = SSHProxyParams(hostname="jump", user="ubuntu", identity_file="key") + current = SSHParams( + hosts=[SSHHostParams(hostname="host", proxy_jump=jump)] if per_host else ["host"], + proxy_jump=None if per_host else jump, + ) + new = current.model_copy(deep=True) + if per_host: + host = new.hosts[0] + assert isinstance(host, SSHHostParams) + target = host.proxy_jump + else: + target = new.proxy_jump + assert target is not None + setattr(target, field, value) + with pytest.raises(ServerClientError): + _check_can_update_ssh_config(current, new) + + +class TestValidateSSHKey: + @pytest.mark.parametrize("per_host", [False, True]) + def test_rejects_stale_head_node_public_key(self, per_host): + head_node = SSHProxyParams( + hostname="head", + user="ubuntu", + identity_file="key", + ssh_key=SSHKey(public=get_ssh_key().public, private=PRIVATE_KEY), + ) + ssh_config = SSHParams( + user="ubuntu", + ssh_key=get_ssh_key(), + hosts=[SSHHostParams(hostname="host", proxy_jump=head_node)] if per_host else ["host"], + proxy_jump=None if per_host else head_node, + ) + spec = get_fleet_spec(conf=FleetConfiguration(name="fleet", ssh_config=ssh_config)) + with pytest.raises(ServerClientError, match="Public SSH key does not match"): + _validate_fleet_spec_and_set_defaults(spec) + + def test_rejects_stale_public_key(self): + key = get_ssh_key() + key.private = PRIVATE_KEY + with pytest.raises(ServerClientError, match="Public SSH key does not match"): + _validate_ssh_key(key) + + def test_ignores_public_key_comment(self): + key = get_ssh_key() + key.public += " rotated-key\n" + _validate_ssh_key(key) class TestGetPlanSSHFleetHostsValidation: From f3404f1c10322c9504aa5bdb311667a5c7f63c1e Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Tue, 3 Nov 2026 14:41:01 +0100 Subject: [PATCH 02/10] Narrow SSH host type in fleet CLI test --- src/tests/_internal/cli/services/configurators/test_fleet.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/tests/_internal/cli/services/configurators/test_fleet.py b/src/tests/_internal/cli/services/configurators/test_fleet.py index 62ae203bdd..9895b8d84c 100644 --- a/src/tests/_internal/cli/services/configurators/test_fleet.py +++ b/src/tests/_internal/cli/services/configurators/test_fleet.py @@ -210,6 +210,7 @@ def test_identifies_the_changed_ssh_key_without_printing_key_material( assert changed is not None if per_host: changed = changed.hosts[0] + assert isinstance(changed, SSHHostParams) if head_node: changed = changed.proxy_jump assert changed is not None From e14966541662207812fb95dd0720dace1fc4f982 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Tue, 3 Nov 2026 14:47:30 +0100 Subject: [PATCH 03/10] Preserve private-only SSH key API requests --- .../_internal/server/services/fleets.py | 6 ++++- .../_internal/server/routers/test_fleets.py | 23 +++++++++++++++---- .../_internal/server/services/test_fleets.py | 7 ++++++ 3 files changed, 31 insertions(+), 5 deletions(-) diff --git a/src/dstack/_internal/server/services/fleets.py b/src/dstack/_internal/server/services/fleets.py index 6b27533682..a3b9b907ca 100644 --- a/src/dstack/_internal/server/services/fleets.py +++ b/src/dstack/_internal/server/services/fleets.py @@ -1606,7 +1606,11 @@ def _validate_ssh_key(ssh_key: SSHKey): ) # Plans and gateway synchronization compare public keys. A stale .pub file must not # hide replacement of the private key at the same path. - if " ".join(ssh_key.public.split()[:2]) != ssh_utils.generate_public_key(private_key): + public_key = ssh_utils.generate_public_key(private_key) + if not ssh_key.public.strip(): + # Older API requests may provide only the private key. + ssh_key.public = public_key + elif " ".join(ssh_key.public.split()[:2]) != public_key: raise ServerClientError("Public SSH key does not match the private key") diff --git a/src/tests/_internal/server/routers/test_fleets.py b/src/tests/_internal/server/routers/test_fleets.py index 90c1c9bd3f..f920cc16a8 100644 --- a/src/tests/_internal/server/routers/test_fleets.py +++ b/src/tests/_internal/server/routers/test_fleets.py @@ -56,6 +56,7 @@ get_ssh_fleet_configuration, ) from dstack._internal.server.testing.matchers import SomeUUID4Str +from dstack._internal.utils import ssh as ssh_utils pytestmark = pytest.mark.usefixtures("image_config_mock") @@ -1077,7 +1078,12 @@ async def test_creates_ssh_fleet(self, test_db, session: AsyncSession, client: A "user": "ubuntu", "port": None, "identity_file": None, - "ssh_key": None, # should not return ssh_key + "ssh_key": { + "public": ssh_utils.generate_public_key( + ssh_utils.pkey_from_str(get_private_key_string()) + ), + "private": None, + }, "proxy_jump": None, "hosts": ["1.1.1.1"], "network": None, @@ -1299,7 +1305,12 @@ async def test_updates_ssh_fleet(self, test_db, session: AsyncSession, client: A "user": "ubuntu", "port": None, "identity_file": None, - "ssh_key": None, # should not return ssh_key + "ssh_key": { + "public": ssh_utils.generate_public_key( + ssh_utils.pkey_from_str(get_private_key_string()) + ), + "private": None, + }, "proxy_jump": None, "hosts": ["10.0.0.101"], "network": None, @@ -2359,7 +2370,9 @@ async def test_returns_update_plan_for_existing_fleet( spec = get_fleet_spec(conf=conf) effective_spec = spec.model_copy(deep=True) assert effective_spec.configuration.ssh_config is not None - effective_spec.configuration.ssh_config.ssh_key = None + effective_spec.configuration.ssh_config.ssh_key = SSHKey( + public=ssh_utils.generate_public_key(ssh_utils.pkey_from_str(get_private_key_string())) + ) current_spec = spec.model_copy(deep=True) # `hosts` can be updated in-place assert current_spec.configuration.ssh_config is not None @@ -2484,7 +2497,9 @@ async def test_returns_create_plan_for_existing_fleet( spec = get_fleet_spec(conf=conf) effective_spec = spec.model_copy(deep=True) assert effective_spec.configuration.ssh_config is not None - effective_spec.configuration.ssh_config.ssh_key = None + effective_spec.configuration.ssh_config.ssh_key = SSHKey( + public=ssh_utils.generate_public_key(ssh_utils.pkey_from_str(get_private_key_string())) + ) current_spec = spec.model_copy(deep=True) # `placement` cannot be updated in-place current_spec.configuration.placement = InstanceGroupPlacement.CLUSTER diff --git a/src/tests/_internal/server/services/test_fleets.py b/src/tests/_internal/server/services/test_fleets.py index 94d73cc6eb..aa5a3e4ff1 100644 --- a/src/tests/_internal/server/services/test_fleets.py +++ b/src/tests/_internal/server/services/test_fleets.py @@ -286,6 +286,13 @@ def test_rejects_jump_route_changes(self, field, value, per_host): class TestValidateSSHKey: + def test_derives_empty_public_key(self): + key = get_ssh_key() + public_key = key.public + key.public = "" + _validate_ssh_key(key) + assert key.public == public_key + @pytest.mark.parametrize("per_host", [False, True]) def test_rejects_stale_head_node_public_key(self, per_host): head_node = SSHProxyParams( From 1670e6682ca75427cdb8b645e0508998030c796e Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Wed, 4 Nov 2026 09:19:09 +0100 Subject: [PATCH 04/10] Clarify SSH fleet key rotation instructions --- mkdocs/docs/concepts/fleets.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mkdocs/docs/concepts/fleets.md b/mkdocs/docs/concepts/fleets.md index 2244dbd731..c81cde27e4 100644 --- a/mkdocs/docs/concepts/fleets.md +++ b/mkdocs/docs/concepts/fleets.md @@ -535,7 +535,9 @@ If approved, `dstack` updates the fleet without re-creating it. For example, if 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" - To rotate SSH keys, authorize the new public key and apply the updated fleet configuration. Runs and established SSH connections continue. For `proxy_jump`, keep the old public key authorized until [gateways](gateways.md) and the optional [SSH proxy](../guides/server-deployment.md#ssh-proxy) receive the replacement; `dstack apply` does not wait for them. + Authorize the new public key on the hosts or head node, update `identity_file` or `proxy_jump.identity_file`, and run `dstack apply`. Running workloads and existing SSH connections are not interrupted. + + For head-node keys, `dstack apply` does not wait for [gateways](gateways.md) to update their saved keys or the [SSH proxy](../guides/server-deployment.md#ssh-proxy)'s 10-second cache to expire. Revoking the old public key before these components use the new key can prevent new SSH connections or service-tunnel reconnects. ### Delete fleets From 1f0fa86fa24320657568eb5fae72efc6dec11771 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Wed, 4 Nov 2026 09:21:44 +0100 Subject: [PATCH 05/10] Explain when previous SSH fleet keys can be removed --- mkdocs/docs/concepts/fleets.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mkdocs/docs/concepts/fleets.md b/mkdocs/docs/concepts/fleets.md index c81cde27e4..24571722ab 100644 --- a/mkdocs/docs/concepts/fleets.md +++ b/mkdocs/docs/concepts/fleets.md @@ -535,9 +535,9 @@ If approved, `dstack` updates the fleet without re-creating it. For example, if 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" - Authorize the new public key on the hosts or head node, update `identity_file` or `proxy_jump.identity_file`, and run `dstack apply`. Running workloads and existing SSH connections are not interrupted. + In-place SSH key updates do not interrupt workloads or existing SSH connections. After `dstack apply` succeeds, you can remove the old public key from the fleet hosts. - For head-node keys, `dstack apply` does not wait for [gateways](gateways.md) to update their saved keys or the [SSH proxy](../guides/server-deployment.md#ssh-proxy)'s 10-second cache to expire. Revoking the old public key before these components use the new key can prevent new SSH connections or service-tunnel reconnects. + For a head node (`proxy_jump`), first wait for [gateways](gateways.md) serving the fleet's services to use the new key and the [SSH proxy](../guides/server-deployment.md#ssh-proxy)'s 10-second cache to expire. `dstack apply` does not wait for either. ### Delete fleets From afa5b6d29b2d22accc8100f576937967fe5cede4 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Wed, 4 Nov 2026 09:42:27 +0100 Subject: [PATCH 06/10] Limit SSH fleet key rotation to identity file updates --- mkdocs/docs/concepts/fleets.md | 6 +-- .../cli/services/configurators/fleet.py | 5 +-- src/dstack/_internal/proxy/lib/models.py | 10 +---- .../proxy/lib/services/service_connection.py | 14 +++--- .../pipeline_tasks/gateway_replicas.py | 10 +++-- .../_internal/server/services/fleets.py | 41 ++++++----------- src/dstack/_internal/utils/ssh.py | 7 +++ .../cli/services/configurators/test_fleet.py | 11 +++-- .../pipeline_tasks/test_gateway_replicas.py | 12 +++-- .../_internal/server/routers/test_fleets.py | 23 ++-------- .../_internal/server/services/test_fleets.py | 45 ++----------------- 11 files changed, 62 insertions(+), 122 deletions(-) diff --git a/mkdocs/docs/concepts/fleets.md b/mkdocs/docs/concepts/fleets.md index 24571722ab..9e1432ea89 100644 --- a/mkdocs/docs/concepts/fleets.md +++ b/mkdocs/docs/concepts/fleets.md @@ -530,14 +530,12 @@ 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 and key contents at the same path. + 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" - In-place SSH key updates do not interrupt workloads or existing SSH connections. After `dstack apply` succeeds, you can remove the old public key from the fleet hosts. - - For a head node (`proxy_jump`), first wait for [gateways](gateways.md) serving the fleet's services to use the new key and the [SSH proxy](../guides/server-deployment.md#ssh-proxy)'s 10-second cache to expire. `dstack apply` does not wait for either. + Updating `identity_file` and `proxy_jump.identity_file` does not require stopping runs. Allow up to 15 seconds for changes to propagate before removing the old public key from the host. ### Delete fleets diff --git a/src/dstack/_internal/cli/services/configurators/fleet.py b/src/dstack/_internal/cli/services/configurators/fleet.py index 94e85b189e..0a6be89844 100644 --- a/src/dstack/_internal/cli/services/configurators/fleet.py +++ b/src/dstack/_internal/cli/services/configurators/fleet.py @@ -424,10 +424,7 @@ def _ssh_config_diff_fields( new_value = getattr(new, field) path = f"{prefix}.{field}" if field == "ssh_key": - # Key material is an internal field. Refer to its user-configured file, - # including replacements whose path has not changed. Never print keys. - if old.identity_file == new.identity_file: - fields.append(f"{prefix}.identity_file (key contents changed)") + 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] diff --git a/src/dstack/_internal/proxy/lib/models.py b/src/dstack/_internal/proxy/lib/models.py index ac4ed7828d..7811760cb3 100644 --- a/src/dstack/_internal/proxy/lib/models.py +++ b/src/dstack/_internal/proxy/lib/models.py @@ -1,7 +1,7 @@ """Things stored in BaseProxyRepo implementations.""" from datetime import datetime -from functools import cached_property, lru_cache +from functools import cached_property from typing import Any, Iterable, Literal, Optional, Union from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -36,13 +36,7 @@ class Replica(ImmutableModel): def ssh_head_proxy_public_key(self) -> Optional[str]: if self.ssh_head_proxy_private_key is None: return None - return _get_head_node_public_key(self.ssh_head_proxy_private_key) - - -@lru_cache(maxsize=128) -def _get_head_node_public_key(private_key: str) -> str: - # Replicas sharing a head node often share a key; parse it once, not per replica. - return ssh_utils.generate_public_key(ssh_utils.pkey_from_str(private_key)) + return ssh_utils.get_public_key_from_private_key(self.ssh_head_proxy_private_key) class IPAddressPartitioningKey(ImmutableModel): diff --git a/src/dstack/_internal/proxy/lib/services/service_connection.py b/src/dstack/_internal/proxy/lib/services/service_connection.py index db3a12ef2f..5b85835d0a 100644 --- a/src/dstack/_internal/proxy/lib/services/service_connection.py +++ b/src/dstack/_internal/proxy/lib/services/service_connection.py @@ -271,10 +271,12 @@ async def get_service_replica_client( ) # Nginx not available, forward directly to the tunnel replica = random.choice(service.replicas) - project = await repo.get_project(service.project_name) - if project is None: - raise UnexpectedProxyError( - f"Expected to find project {service.project_name} but could not" - ) - connection = await service_conn_pool.get_or_add(project, service, replica) + connection = await service_conn_pool.get(replica.id) + if connection is None or replica.ssh_head_proxy is not None: + project = await repo.get_project(service.project_name) + if project is None: + raise UnexpectedProxyError( + f"Expected to find project {service.project_name} but could not" + ) + connection = await service_conn_pool.get_or_add(project, service, replica) return await connection.client() diff --git a/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py b/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py index 3ece85e70e..7f97f31f18 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py @@ -75,6 +75,7 @@ should_configure_service_https_on_gateway_replica, ) from dstack._internal.server.utils import tracing +from dstack._internal.utils import ssh as ssh_utils from dstack._internal.utils.common import get_current_datetime, get_or_error, run_async from dstack._internal.utils.logging import get_logger @@ -723,9 +724,12 @@ async def _process_running_item(item: GatewayReplicaPipelineItem): for instance_id, connection_info in instance_connection_info: rci = validate_json_extra_ignore(RemoteConnectionInfo, get_or_error(connection_info)) if rci.ssh_proxy is not None: - instance_head_node_public_keys[instance_id] = " ".join( - get_or_error(rci.ssh_proxy_keys)[0].public.split()[:2] - ) + key = get_or_error(rci.ssh_proxy_keys)[0] + public_key = " ".join(key.public.split()[:2]) + # Existing API requests may provide only the private key. + if not public_key: + public_key = ssh_utils.get_public_key_from_private_key(get_or_error(key.private)) + instance_head_node_public_keys[instance_id] = public_key plan = _plan_state_sync( currently_registered=currently_registered_services, expected=expected_runs, diff --git a/src/dstack/_internal/server/services/fleets.py b/src/dstack/_internal/server/services/fleets.py index a3b9b907ca..7bb5467191 100644 --- a/src/dstack/_internal/server/services/fleets.py +++ b/src/dstack/_internal/server/services/fleets.py @@ -1509,16 +1509,14 @@ async def _check_ssh_hosts_not_yet_added( def _remove_fleet_spec_sensitive_info(spec: FleetSpec): if spec.configuration.ssh_config is not None: ssh_config = spec.configuration.ssh_config - if ssh_config.ssh_key is not None: - ssh_config.ssh_key.private = None - if ssh_config.proxy_jump is not None and ssh_config.proxy_jump.ssh_key is not None: - ssh_config.proxy_jump.ssh_key.private = None + ssh_config.ssh_key = None + if ssh_config.proxy_jump is not None: + ssh_config.proxy_jump.ssh_key = None for host in spec.configuration.ssh_config.hosts: if not isinstance(host, str): - if host.ssh_key is not None: - host.ssh_key.private = None - if host.proxy_jump is not None and host.proxy_jump.ssh_key is not None: - host.proxy_jump.ssh_key.private = None + host.ssh_key = None + if host.proxy_jump is not None: + host.proxy_jump.ssh_key = None def _validate_fleet_spec_and_set_defaults(spec: FleetSpec): @@ -1531,16 +1529,13 @@ def _validate_fleet_spec_and_set_defaults(spec: FleetSpec): validate_dstack_resource_name(spec.configuration.name) _validate_fleet_configuration_subtype_specific_fields(spec.configuration) if spec.configuration.ssh_config is not None: - ssh_config = spec.configuration.ssh_config - _validate_all_ssh_params_specified(ssh_config) - ssh_params: list[SSHParams | SSHHostParams] = [ssh_config] - ssh_params.extend(host for host in ssh_config.hosts if isinstance(host, SSHHostParams)) - for params in ssh_params: - if params.ssh_key is not None: - _validate_ssh_key(params.ssh_key) - if params.proxy_jump is not None and params.proxy_jump.ssh_key is not None: - _validate_ssh_key(params.proxy_jump.ssh_key) - _validate_internal_ips(ssh_config) + _validate_all_ssh_params_specified(spec.configuration.ssh_config) + if spec.configuration.ssh_config.ssh_key is not None: + _validate_ssh_key(spec.configuration.ssh_config.ssh_key) + for host in spec.configuration.ssh_config.hosts: + if isinstance(host, SSHHostParams) and host.ssh_key is not None: + _validate_ssh_key(host.ssh_key) + _validate_internal_ips(spec.configuration.ssh_config) _set_fleet_spec_defaults(spec) @@ -1598,20 +1593,12 @@ def _validate_ssh_key(ssh_key: SSHKey): if ssh_key.private is None: raise ServerClientError("Private key not provided") try: - private_key = ssh_utils.pkey_from_str(ssh_key.private) + ssh_utils.pkey_from_str(ssh_key.private) except ValueError: raise ServerClientError( "Unsupported key type. " "The key type should be RSA, ECDSA, or Ed25519 and should not be encrypted with passphrase." ) - # Plans and gateway synchronization compare public keys. A stale .pub file must not - # hide replacement of the private key at the same path. - public_key = ssh_utils.generate_public_key(private_key) - if not ssh_key.public.strip(): - # Older API requests may provide only the private key. - ssh_key.public = public_key - elif " ".join(ssh_key.public.split()[:2]) != public_key: - raise ServerClientError("Public SSH key does not match the private key") def _validate_internal_ips(ssh_config: SSHParams): diff --git a/src/dstack/_internal/utils/ssh.py b/src/dstack/_internal/utils/ssh.py index b556700064..c8401c3878 100644 --- a/src/dstack/_internal/utils/ssh.py +++ b/src/dstack/_internal/utils/ssh.py @@ -9,6 +9,7 @@ import tempfile from contextlib import suppress from dataclasses import dataclass +from functools import lru_cache from pathlib import Path from typing import Dict, Optional, Union @@ -343,6 +344,12 @@ def generate_public_key(private_key: PKey) -> str: return public_key +@lru_cache(maxsize=128) +def get_public_key_from_private_key(private_key: str) -> str: + """Derive an OpenSSH public key, reusing results for identical private keys.""" + return generate_public_key(pkey_from_str(private_key)) + + def resolve_ssh_key( path: PathLike, ) -> Union[ diff --git a/src/tests/_internal/cli/services/configurators/test_fleet.py b/src/tests/_internal/cli/services/configurators/test_fleet.py index 9895b8d84c..b191d3582d 100644 --- a/src/tests/_internal/cli/services/configurators/test_fleet.py +++ b/src/tests/_internal/cli/services/configurators/test_fleet.py @@ -182,11 +182,11 @@ def test_prints_in_place_update_diff(self, monkeypatch: pytest.MonkeyPatch): confirm_ask.assert_called_once_with("Update the fleet in-place?") @pytest.mark.parametrize( - "per_host,head_node,same_path", - [(False, False, False), (False, True, True), (True, False, True), (True, True, False)], + "per_host,head_node", + [(False, False), (False, True), (True, False), (True, True)], ) def test_identifies_the_changed_ssh_key_without_printing_key_material( - self, monkeypatch, per_host, head_node, same_path + self, monkeypatch, per_host, head_node ): console, confirm_ask = patch_console_and_confirm(monkeypatch) current = get_ssh_fleet_spec() @@ -214,15 +214,14 @@ def test_identifies_the_changed_ssh_key_without_printing_key_material( if head_node: changed = changed.proxy_jump assert changed is not None - if not same_path: - changed.identity_file = "new-key" + changed.identity_file = "new-key" changed.ssh_key = SSHKey(public="new-public-key", private="new-private-key") plan = create_fleet_plan(current_spec=current, spec=new, action=ApplyAction.UPDATE) FleetConfigurator(Mock())._apply_plan(plan, get_command_args()) output = console.export_text() - label = f"{path}.identity_file" + (" (key contents changed)" if same_path else "") + label = f"{path}.identity_file" assert f" - {label}\n" in output assert "ssh_key" not in output assert "public-key" not in output diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_gateway_replicas.py b/src/tests/_internal/server/background/pipeline_tasks/test_gateway_replicas.py index 07a77e0ad1..050b9cc227 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_gateway_replicas.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_gateway_replicas.py @@ -53,6 +53,7 @@ get_gateway_replica_configuration, get_job_provisioning_data, get_run_spec, + get_ssh_key, list_events, ) from dstack._internal.utils.common import get_current_datetime @@ -769,9 +770,12 @@ class TestGatewayReplicaWorkerRunningStateSync: pytestmark = pytest.mark.usefixtures("image_config_mock") - @pytest.mark.parametrize("failure", [None, "update", "changed-again"]) + @pytest.mark.parametrize( + "failure,public_key", + [(None, PUBLIC_KEY), (None, ""), ("update", PUBLIC_KEY), ("changed-again", PUBLIC_KEY)], + ) async def test_updates_head_node_key_without_unregistering_or_changing_job( - self, test_db, session, worker, mock_gateway_connection, failure + self, test_db, session, worker, mock_gateway_connection, failure, public_key ): project = await create_project(session) user = await create_user(session) @@ -794,7 +798,7 @@ async def test_updates_head_node_key_without_unregistering_or_changing_job( ready=True, replica_num=1, ) - key = SSHKey(public=PUBLIC_KEY, private=PRIVATE_KEY) + key = SSHKey(public=public_key, private=PRIVATE_KEY) rci = RemoteConnectionInfo( host="host", port=22, @@ -836,7 +840,7 @@ async def update(updates): ) if failure == "changed-again": assert job.instance is not None - rci.ssh_proxy_keys = [SSHKey(public="ssh-rsa another", private="another private")] + rci.ssh_proxy_keys = [get_ssh_key()] job.instance.remote_connection_info = rci.model_dump_json() await session.commit() diff --git a/src/tests/_internal/server/routers/test_fleets.py b/src/tests/_internal/server/routers/test_fleets.py index f920cc16a8..90c1c9bd3f 100644 --- a/src/tests/_internal/server/routers/test_fleets.py +++ b/src/tests/_internal/server/routers/test_fleets.py @@ -56,7 +56,6 @@ get_ssh_fleet_configuration, ) from dstack._internal.server.testing.matchers import SomeUUID4Str -from dstack._internal.utils import ssh as ssh_utils pytestmark = pytest.mark.usefixtures("image_config_mock") @@ -1078,12 +1077,7 @@ async def test_creates_ssh_fleet(self, test_db, session: AsyncSession, client: A "user": "ubuntu", "port": None, "identity_file": None, - "ssh_key": { - "public": ssh_utils.generate_public_key( - ssh_utils.pkey_from_str(get_private_key_string()) - ), - "private": None, - }, + "ssh_key": None, # should not return ssh_key "proxy_jump": None, "hosts": ["1.1.1.1"], "network": None, @@ -1305,12 +1299,7 @@ async def test_updates_ssh_fleet(self, test_db, session: AsyncSession, client: A "user": "ubuntu", "port": None, "identity_file": None, - "ssh_key": { - "public": ssh_utils.generate_public_key( - ssh_utils.pkey_from_str(get_private_key_string()) - ), - "private": None, - }, + "ssh_key": None, # should not return ssh_key "proxy_jump": None, "hosts": ["10.0.0.101"], "network": None, @@ -2370,9 +2359,7 @@ async def test_returns_update_plan_for_existing_fleet( spec = get_fleet_spec(conf=conf) effective_spec = spec.model_copy(deep=True) assert effective_spec.configuration.ssh_config is not None - effective_spec.configuration.ssh_config.ssh_key = SSHKey( - public=ssh_utils.generate_public_key(ssh_utils.pkey_from_str(get_private_key_string())) - ) + effective_spec.configuration.ssh_config.ssh_key = None current_spec = spec.model_copy(deep=True) # `hosts` can be updated in-place assert current_spec.configuration.ssh_config is not None @@ -2497,9 +2484,7 @@ async def test_returns_create_plan_for_existing_fleet( spec = get_fleet_spec(conf=conf) effective_spec = spec.model_copy(deep=True) assert effective_spec.configuration.ssh_config is not None - effective_spec.configuration.ssh_config.ssh_key = SSHKey( - public=ssh_utils.generate_public_key(ssh_utils.pkey_from_str(get_private_key_string())) - ) + effective_spec.configuration.ssh_config.ssh_key = None current_spec = spec.model_copy(deep=True) # `placement` cannot be updated in-place current_spec.configuration.placement = InstanceGroupPlacement.CLUSTER diff --git a/src/tests/_internal/server/services/test_fleets.py b/src/tests/_internal/server/services/test_fleets.py index aa5a3e4ff1..0cb6ab6435 100644 --- a/src/tests/_internal/server/services/test_fleets.py +++ b/src/tests/_internal/server/services/test_fleets.py @@ -35,8 +35,6 @@ from dstack._internal.server.services.backends import get_project_backends from dstack._internal.server.services.fleets import ( _check_can_update_ssh_config, - _validate_fleet_spec_and_set_defaults, - _validate_ssh_key, apply_plan, get_fleet_master_instance_provisioning_data, get_plan, @@ -123,14 +121,16 @@ async def test_updates_existing_instances_without_reprovisioning( keys = keys.hosts[0] assert isinstance(keys, SSHHostParams) if rotation != "jump": + keys.identity_file = "new-key" keys.ssh_key = new_key assert keys.proxy_jump is not None if rotation != "host": + keys.proxy_jump.identity_file = "new-key" keys.proxy_jump.ssh_key = new_key plan = await get_plan(session, project, user, new) assert plan.action == ApplyAction.UPDATE assert plan.current_resource is not None and plan.effective_spec is not None - assert plan.current_resource.spec != plan.effective_spec # same file path, different key + assert plan.current_resource.spec != plan.effective_spec assert PRIVATE_KEY not in plan.effective_spec.model_dump_json() updated = await apply_plan( session, @@ -182,6 +182,7 @@ async def test_waits_for_instance_processing_without_holding_db_connection( await session.commit() new = spec.model_copy(deep=True) assert new.configuration.ssh_config is not None + new.configuration.ssh_config.identity_file = "new-key" new.configuration.ssh_config.ssh_key = SSHKey(public=PUBLIC_KEY, private=PRIVATE_KEY) plan = await get_plan(session, project, user, new) fleet_id, instance_id = fleet.id, instance.id @@ -285,44 +286,6 @@ def test_rejects_jump_route_changes(self, field, value, per_host): _check_can_update_ssh_config(current, new) -class TestValidateSSHKey: - def test_derives_empty_public_key(self): - key = get_ssh_key() - public_key = key.public - key.public = "" - _validate_ssh_key(key) - assert key.public == public_key - - @pytest.mark.parametrize("per_host", [False, True]) - def test_rejects_stale_head_node_public_key(self, per_host): - head_node = SSHProxyParams( - hostname="head", - user="ubuntu", - identity_file="key", - ssh_key=SSHKey(public=get_ssh_key().public, private=PRIVATE_KEY), - ) - ssh_config = SSHParams( - user="ubuntu", - ssh_key=get_ssh_key(), - hosts=[SSHHostParams(hostname="host", proxy_jump=head_node)] if per_host else ["host"], - proxy_jump=None if per_host else head_node, - ) - spec = get_fleet_spec(conf=FleetConfiguration(name="fleet", ssh_config=ssh_config)) - with pytest.raises(ServerClientError, match="Public SSH key does not match"): - _validate_fleet_spec_and_set_defaults(spec) - - def test_rejects_stale_public_key(self): - key = get_ssh_key() - key.private = PRIVATE_KEY - with pytest.raises(ServerClientError, match="Public SSH key does not match"): - _validate_ssh_key(key) - - def test_ignores_public_key_comment(self): - key = get_ssh_key() - key.public += " rotated-key\n" - _validate_ssh_key(key) - - class TestGetPlanSSHFleetHostsValidation: @pytest.fixture def get_project_backends_mock(self, monkeypatch: pytest.MonkeyPatch) -> list[Backend]: From 73f3c74266b9368361bfbc5533b658644dd148f4 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Wed, 4 Nov 2026 11:11:03 +0100 Subject: [PATCH 07/10] Move SSH key parsing off the event loop --- .../proxy/gateway/services/services.py | 7 ++++++ .../pipeline_tasks/gateway_replicas.py | 4 +++- .../proxy/gateway/routers/test_services.py | 22 +++++++++++++++++-- 3 files changed, 30 insertions(+), 3 deletions(-) diff --git a/src/dstack/_internal/proxy/gateway/services/services.py b/src/dstack/_internal/proxy/gateway/services/services.py index 292d9cd5eb..036d4ed739 100644 --- a/src/dstack/_internal/proxy/gateway/services/services.py +++ b/src/dstack/_internal/proxy/gateway/services/services.py @@ -4,10 +4,17 @@ 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 head-node 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( diff --git a/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py b/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py index 7f97f31f18..4475c170b5 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py @@ -728,7 +728,9 @@ async def _process_running_item(item: GatewayReplicaPipelineItem): public_key = " ".join(key.public.split()[:2]) # Existing API requests may provide only the private key. if not public_key: - public_key = ssh_utils.get_public_key_from_private_key(get_or_error(key.private)) + public_key = await run_async( + ssh_utils.get_public_key_from_private_key, get_or_error(key.private) + ) instance_head_node_public_keys[instance_id] = public_key plan = _plan_state_sync( currently_registered=currently_registered_services, diff --git a/src/tests/_internal/proxy/gateway/routers/test_services.py b/src/tests/_internal/proxy/gateway/routers/test_services.py index c4dd228f50..41f3de33af 100644 --- a/src/tests/_internal/proxy/gateway/routers/test_services.py +++ b/src/tests/_internal/proxy/gateway/routers/test_services.py @@ -1,3 +1,4 @@ +import threading import uuid import httpx @@ -7,6 +8,7 @@ from dstack._internal.proxy.gateway.repo.repo import GatewayProxyRepo from dstack._internal.proxy.lib.models import Replica from dstack._internal.proxy.lib.testing.common import make_project, make_service +from dstack._internal.utils import ssh as ssh_utils def make_client(repo: GatewayProxyRepo) -> httpx.AsyncClient: @@ -23,7 +25,17 @@ async def test_empty(self): assert resp.status_code == 200 assert resp.json() == {"services": []} - async def test_list(self): + @pytest.mark.parametrize("head_node_key", [None, "head-private-key"]) + async def test_list(self, monkeypatch, head_node_key): + event_loop_thread = threading.get_ident() + + def get_public_key(private_key): + assert private_key == head_node_key + # Uncached RSA parsing must run outside the gateway's event loop. + assert threading.get_ident() != event_loop_thread + return "public-key" + + monkeypatch.setattr(ssh_utils, "get_public_key_from_private_key", get_public_key) repo = GatewayProxyRepo() service_id = uuid.uuid4().hex replica_id = uuid.uuid4().hex @@ -37,6 +49,7 @@ async def test_list(self): ssh_destination="ubuntu@server", ssh_port=22, ssh_proxy=None, + ssh_head_proxy_private_key=head_node_key, ) ] ) @@ -51,7 +64,12 @@ async def test_list(self): "id": service_id, "project_name": "test-proj", "run_name": "srv-1", - "replicas": [{"id": replica_id, "ssh_head_proxy_public_key": None}], + "replicas": [ + { + "id": replica_id, + "ssh_head_proxy_public_key": "public-key" if head_node_key else None, + } + ], } ] } From 5c153e1b21f5f30426276a00e6720ca6db9d3237 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Wed, 4 Nov 2026 12:56:08 +0100 Subject: [PATCH 08/10] Fix gateway persistence during SSH key rotation --- mkdocs/docs/concepts/fleets.md | 2 +- .../_internal/proxy/gateway/repo/repo.py | 30 +++++++++++++---- .../_internal/proxy/gateway/repo/test_repo.py | 32 +++++++++++++++++++ .../proxy/gateway/routers/test_registry.py | 22 +++++++++++-- 4 files changed, 76 insertions(+), 10 deletions(-) diff --git a/mkdocs/docs/concepts/fleets.md b/mkdocs/docs/concepts/fleets.md index 9e1432ea89..412ddf15bd 100644 --- a/mkdocs/docs/concepts/fleets.md +++ b/mkdocs/docs/concepts/fleets.md @@ -535,7 +535,7 @@ If approved, `dstack` updates the fleet without re-creating it. For example, if 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 up to 15 seconds for changes to propagate before removing the old public key from the host. + 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 diff --git a/src/dstack/_internal/proxy/gateway/repo/repo.py b/src/dstack/_internal/proxy/gateway/repo/repo.py index dd52745687..be0c7cbb45 100644 --- a/src/dstack/_internal/proxy/gateway/repo/repo.py +++ b/src/dstack/_internal/proxy/gateway/repo/repo.py @@ -1,3 +1,4 @@ +import tempfile from contextlib import asynccontextmanager from itertools import chain from pathlib import Path @@ -54,12 +55,18 @@ async def set_service(self, service: Service) -> None: self._state.services.setdefault(service.project_name, {})[service.run_name] = service async def set_services(self, services: Iterable[Service]) -> None: - """Save a batch of changed services with one state-file write.""" - async with self.writer(): + """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: - self._state.services.setdefault(service.project_name, {})[service.run_name] = ( + 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(): @@ -132,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) diff --git a/src/tests/_internal/proxy/gateway/repo/test_repo.py b/src/tests/_internal/proxy/gateway/repo/test_repo.py index bf95c7060f..3db014b52c 100644 --- a/src/tests/_internal/proxy/gateway/repo/test_repo.py +++ b/src/tests/_internal/proxy/gateway/repo/test_repo.py @@ -39,3 +39,35 @@ async def test_persist_repo(tmp_path: Path) -> None: srv_2, ] assert await repo.list_models("proj-1") == [model_1] + + +@pytest.mark.asyncio +class TestSetServices: + async def test_failed_replace_preserves_state_and_retries( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + state_file = tmp_path / "state.json" + repo = GatewayProxyRepo.load(state_file) + old_service = make_service("project", "run") + new_service = old_service.model_copy(update={"auth": not old_service.auth}) + await repo.set_services([old_service]) + previous_content = state_file.read_bytes() + + def fail_replace(source: Path, target: Path) -> None: + # The replacement is complete, but cannot be installed over the saved state. + assert source.read_bytes() != previous_content + raise OSError("state replacement failed") + + with monkeypatch.context() as scoped: + scoped.setattr(Path, "replace", fail_replace) + with pytest.raises(OSError, match="state replacement failed"): + await repo.set_services([new_service]) + + assert await repo.get_service("project", "run") == old_service + assert state_file.read_bytes() == previous_content + assert await GatewayProxyRepo.load(state_file).get_service("project", "run") == old_service + assert list(tmp_path.iterdir()) == [state_file] + + await repo.set_services([new_service]) + assert await repo.get_service("project", "run") == new_service + assert await GatewayProxyRepo.load(state_file).get_service("project", "run") == new_service diff --git a/src/tests/_internal/proxy/gateway/routers/test_registry.py b/src/tests/_internal/proxy/gateway/routers/test_registry.py index 1f6e35dc5e..acb9ef4eef 100644 --- a/src/tests/_internal/proxy/gateway/routers/test_registry.py +++ b/src/tests/_internal/proxy/gateway/routers/test_registry.py @@ -96,7 +96,7 @@ def sample_model_options(name: str = "test-model") -> dict: @pytest.mark.asyncio class TestUpdateReplicaHeadNodeKeys: - @pytest.mark.parametrize("failure", ["missing", "key-file"]) + @pytest.mark.parametrize("failure", ["missing", "key-file", "state-save"]) async def test_batch_saves_once_preserves_tunnels_and_retries_failures( self, tmp_path, system_mocks, monkeypatch, failure ): @@ -135,6 +135,19 @@ async def test_batch_saves_once_preserves_tunnels_and_retries_failures( system_mocks.open_conn.reset_mock() system_mocks.close_conn.reset_mock() system_mocks.reload_nginx.reset_mock() + if failure == "state-save": + with patch.object(repo, "save", side_effect=OSError("state save failed")): + with pytest.raises(OSError, match="state save failed"): + await client.post( + "/api/registry/replicas/head_node_keys", json={"replicas": updates} + ) + # Keep reporting the old keys so the server retries the failed save. + assert (await client.get("/api/services/list")).json() == before + restored = GatewayProxyRepo.load(state_file) + for n in range(2): + service = await restored.get_service(f"project-{n}", f"run-{n}") + assert service is not None + assert service.replicas[0].ssh_head_proxy_private_key == PRIVATE_KEY batch = list(updates) if failure == "missing": batch.append({**updates[0], "job_id": "missing"}) @@ -153,9 +166,12 @@ async def fail_one_key_file(connection, replica): "/api/registry/replicas/head_node_keys", json={"replicas": batch} ) assert response.status_code == 200 - assert set(response.json()["failures"]) == { - "missing" if failure == "missing" else "job-1" + expected_failures = { + "missing": {"missing"}, + "key-file": {"job-1"}, + "state-save": set(), } + assert set(response.json()["failures"]) == expected_failures[failure] assert save.call_count == 1 restored = GatewayProxyRepo.load(state_file) for n in range(2): From 8a1be679cf1a92439deab248e77e470f8d208bde Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Wed, 4 Nov 2026 13:07:30 +0100 Subject: [PATCH 09/10] Reuse loaded project data in the in-server proxy --- .../_internal/server/services/proxy/repo.py | 11 ++- .../server/services/proxy/test_repo.py | 84 +++++++++++++++++++ 2 files changed, 92 insertions(+), 3 deletions(-) create mode 100644 src/tests/_internal/server/services/proxy/test_repo.py diff --git a/src/dstack/_internal/server/services/proxy/repo.py b/src/dstack/_internal/server/services/proxy/repo.py index 12f79f9537..9551541973 100644 --- a/src/dstack/_internal/server/services/proxy/repo.py +++ b/src/dstack/_internal/server/services/proxy/repo.py @@ -43,12 +43,13 @@ class ServerProxyRepo(BaseProxyRepo): """ - A repo implementation used by dstack-proxy running within dstack-server. + A request-scoped repo used by dstack-proxy running within dstack-server. Retrieves data from dstack-server's database. """ def __init__(self, session: AsyncSession) -> None: self.session = session + self._service_projects: dict[str, ProjectModel] = {} async def get_service(self, project_name: str, run_name: str) -> Optional[Service]: res = await self.session.execute( @@ -76,6 +77,8 @@ async def get_service(self, project_name: str, run_name: str) -> Optional[Servic run_spec = get_run_spec(run) if not isinstance(run_spec.configuration, ServiceConfiguration): return None + # Reuse the project loaded by this query when opening/updating service connections. + self._service_projects[project_name] = jobs[0].project router_group = next( (g for g in run_spec.configuration.replica_groups if g.router is not None), None, @@ -180,8 +183,10 @@ async def get_model(self, project_name: str, name: str) -> Optional[ChatModel]: return max(models, key=lambda m: m.created_at) async def get_project(self, name: str) -> Optional[Project]: - res = await self.session.execute(select(ProjectModel).where(ProjectModel.name == name)) - project = res.scalar_one_or_none() + project = self._service_projects.get(name) + if project is None: + res = await self.session.execute(select(ProjectModel).where(ProjectModel.name == name)) + project = res.scalar_one_or_none() # Release the DB connection before opening/replacing a service SSH tunnel. await self.session.commit() if project is None: diff --git a/src/tests/_internal/server/services/proxy/test_repo.py b/src/tests/_internal/server/services/proxy/test_repo.py new file mode 100644 index 0000000000..fd31ed6c7a --- /dev/null +++ b/src/tests/_internal/server/services/proxy/test_repo.py @@ -0,0 +1,84 @@ +from unittest.mock import patch + +import pytest +from sqlalchemy.ext.asyncio import AsyncSession + +from dstack._internal.core.models.configurations import ServiceConfiguration +from dstack._internal.core.models.instances import SSHConnectionParams +from dstack._internal.core.models.runs import JobStatus +from dstack._internal.proxy.lib.models import Project +from dstack._internal.server.services.proxy.repo import ServerProxyRepo +from dstack._internal.server.testing.common import ( + create_instance, + create_job, + create_project, + create_repo, + create_run, + create_user, + get_job_provisioning_data, + get_remote_connection_info, + get_run_spec, + get_ssh_key, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) +class TestServerProxyRepo: + async def test_service_project_needs_no_additional_query( + self, test_db, session: AsyncSession + ) -> None: + user = await create_user(session) + project = await create_project(session, owner=user, ssh_private_key="project key") + source_repo = await create_repo(session, project_id=project.id) + run = await create_run( + session, + project=project, + repo=source_repo, + user=user, + run_spec=get_run_spec( + repo_id=source_repo.name, + configuration=ServiceConfiguration(commands=["python serve.py"], port=8000), + ), + ) + jump_key = get_ssh_key() + instance = await create_instance( + session, + project=project, + remote_connection_info=get_remote_connection_info( + ssh_proxy=SSHConnectionParams(hostname="jump", username="ubuntu", port=22), + ssh_proxy_keys=[jump_key], + ), + ) + await create_job( + session, + run=run, + instance=instance, + status=JobStatus.RUNNING, + ready=True, + registered=True, + job_provisioning_data=get_job_provisioning_data(), + ) + + repo = ServerProxyRepo(session) + with patch.object(session, "execute", wraps=session.execute) as execute: + service = await repo.get_service(project.name, run.run_name) + assert service is not None + assert service.replicas[0].ssh_head_proxy_private_key == jump_key.private + assert await repo.get_project(project.name) == Project( + name=project.name, ssh_private_key="project key" + ) + # Reuse the project loaded by the service query; release its connection + # before the caller opens or updates an SSH tunnel. + assert execute.await_count == 1 + assert not session.in_transaction() + + async def test_get_project_without_loading_service( + self, test_db, session: AsyncSession + ) -> None: + project = await create_project(session, ssh_private_key="project key") + repo = ServerProxyRepo(session) + assert await repo.get_project(project.name) == Project( + name=project.name, ssh_private_key="project key" + ) + assert not session.in_transaction() From 41c54759b043692cd7d8adb2bc3ef34052163689 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Thu, 5 Nov 2026 10:21:43 +0100 Subject: [PATCH 10/10] Use proxy_jump terminology for SSH key updates --- .../_internal/core/services/ssh/tunnel.py | 2 +- .../proxy/gateway/routers/registry.py | 16 ++++---- .../proxy/gateway/schemas/registry.py | 13 +++--- .../proxy/gateway/schemas/services.py | 8 ++-- .../proxy/gateway/services/registry.py | 18 ++++----- .../proxy/gateway/services/services.py | 4 +- src/dstack/_internal/proxy/lib/models.py | 5 ++- .../proxy/lib/services/service_connection.py | 8 ++-- .../pipeline_tasks/gateway_replicas.py | 40 +++++++++---------- .../server/services/gateways/client.py | 14 +++---- .../cli/services/configurators/test_fleet.py | 8 ++-- .../proxy/gateway/routers/test_registry.py | 18 ++++----- .../proxy/gateway/routers/test_services.py | 10 ++--- .../lib/services/test_service_connection.py | 2 +- .../pipeline_tasks/test_gateway_replicas.py | 20 +++++----- 15 files changed, 94 insertions(+), 92 deletions(-) diff --git a/src/dstack/_internal/core/services/ssh/tunnel.py b/src/dstack/_internal/core/services/ssh/tunnel.py index afcb092988..71b6e91214 100644 --- a/src/dstack/_internal/core/services/ssh/tunnel.py +++ b/src/dstack/_internal/core/services/ssh/tunnel.py @@ -192,7 +192,7 @@ def close_command(self) -> List[str]: 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 head-node key changes. Replace + 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] diff --git a/src/dstack/_internal/proxy/gateway/routers/registry.py b/src/dstack/_internal/proxy/gateway/routers/registry.py index 5365f4c866..b4145bc721 100644 --- a/src/dstack/_internal/proxy/gateway/routers/registry.py +++ b/src/dstack/_internal/proxy/gateway/routers/registry.py @@ -10,8 +10,8 @@ RegisterReplicaRequest, RegisterServiceRequest, SetServiceIdRequest, - UpdateReplicaHeadNodeKeysRequest, - UpdateReplicaHeadNodeKeysResponse, + UpdateReplicaProxyJumpKeysRequest, + UpdateReplicaProxyJumpKeysResponse, ) from dstack._internal.proxy.gateway.services.nginx import Nginx from dstack._internal.proxy.lib.deps import get_service_connection_pool @@ -111,16 +111,16 @@ async def register_replica( return OkResponse() -@router.post("/replicas/head_node_keys") -async def update_replica_head_node_keys( - body: UpdateReplicaHeadNodeKeysRequest, +@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)], -) -> UpdateReplicaHeadNodeKeysResponse: - failures = await registry_services.update_replica_head_node_keys( +) -> UpdateReplicaProxyJumpKeysResponse: + failures = await registry_services.update_replica_proxy_jump_keys( updates=body.replicas, repo=repo, service_conn_pool=service_conn_pool ) - return UpdateReplicaHeadNodeKeysResponse(failures=failures) + return UpdateReplicaProxyJumpKeysResponse(failures=failures) @router.post("/{project_name}/services/{run_name}/replicas/{job_id}/unregister") diff --git a/src/dstack/_internal/proxy/gateway/schemas/registry.py b/src/dstack/_internal/proxy/gateway/schemas/registry.py index de3b08cc74..56c60e3add 100644 --- a/src/dstack/_internal/proxy/gateway/schemas/registry.py +++ b/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 @@ -73,18 +74,18 @@ class RegisterEntrypointRequest(BaseModel): read_timeout: int = DEFAULT_SERVICE_READ_TIMEOUT -class ReplicaHeadNodeKeyUpdate(BaseModel): +class ReplicaProxyJumpKeyUpdate(BaseModel): project_name: str run_name: str job_id: str - ssh_head_proxy_private_key: str - """Private key for the head node configured by the fleet's `proxy_jump`.""" + ssh_proxy_jump_private_key: str + """Private key for the SSH fleet's `proxy_jump`.""" -class UpdateReplicaHeadNodeKeysRequest(BaseModel): - replicas: list[ReplicaHeadNodeKeyUpdate] +class UpdateReplicaProxyJumpKeysRequest(BaseModel): + replicas: list[ReplicaProxyJumpKeyUpdate] -class UpdateReplicaHeadNodeKeysResponse(BaseModel): +class UpdateReplicaProxyJumpKeysResponse(BaseModel): failures: dict[str, str] """Job IDs whose key could not be updated; successful updates are saved together.""" diff --git a/src/dstack/_internal/proxy/gateway/schemas/services.py b/src/dstack/_internal/proxy/gateway/schemas/services.py index 25b0e80d44..f145c393a2 100644 --- a/src/dstack/_internal/proxy/gateway/schemas/services.py +++ b/src/dstack/_internal/proxy/gateway/schemas/services.py @@ -3,12 +3,12 @@ class ServiceListReplicaItem(CoreModel): id: str - ssh_head_proxy_public_key: str | None = None - """Public key matching the head-node private key saved for this replica on the gateway. + 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 head-node key during 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 head-node private key is saved for this replica. + `None` means no `proxy_jump` private key is saved for this replica. """ diff --git a/src/dstack/_internal/proxy/gateway/services/registry.py b/src/dstack/_internal/proxy/gateway/services/registry.py index 0bf13c27d2..e9f05f27cd 100644 --- a/src/dstack/_internal/proxy/gateway/services/registry.py +++ b/src/dstack/_internal/proxy/gateway/services/registry.py @@ -212,19 +212,19 @@ async def register_replica( logger.info("Replica %s in service %s is registered now", replica.id, service.fmt()) -async def update_replica_head_node_keys( - updates: list[schemas.ReplicaHeadNodeKeyUpdate], +async def update_replica_proxy_jump_keys( + updates: list[schemas.ReplicaProxyJumpKeyUpdate], repo: GatewayProxyRepo, service_conn_pool: ServiceConnectionPool, ) -> dict[str, str]: - """Replace registered replicas' head-node keys in one gateway synchronization batch. + """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.ReplicaHeadNodeKeyUpdate]] = {} + 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 @@ -246,25 +246,25 @@ async def update_replica_head_node_keys( if item is None: continue if replica.ssh_head_proxy is None: - failures[replica.id] = "Replica has no head node" + failures[replica.id] = "Replica has no proxy_jump" continue - if replica.ssh_head_proxy_private_key == item.ssh_head_proxy_private_key: + 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_head_proxy_private_key, + "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 head-node SSH key: %r", replica.id, exc + "Could not update replica %s proxy_jump SSH key: %r", replica.id, exc ) - failures[replica.id] = "Could not update head-node key" + failures[replica.id] = "Could not update proxy_jump key" continue replicas[index] = updated_replica changed = True diff --git a/src/dstack/_internal/proxy/gateway/services/services.py b/src/dstack/_internal/proxy/gateway/services/services.py index 036d4ed739..89b6d0950e 100644 --- a/src/dstack/_internal/proxy/gateway/services/services.py +++ b/src/dstack/_internal/proxy/gateway/services/services.py @@ -10,7 +10,7 @@ async def list_services(repo: GatewayProxyRepo) -> ServiceListResponse: services = await repo.list_services() - # Parsing uncached head-node keys must not block other gateway API requests. + # Parsing uncached proxy_jump keys must not block other gateway API requests. return await run_async(_get_service_list_response, services) @@ -24,7 +24,7 @@ def _get_service_list_response(services: list[Service]) -> ServiceListResponse: replicas=[ ServiceListReplicaItem( id=replica.id, - ssh_head_proxy_public_key=replica.ssh_head_proxy_public_key, + ssh_proxy_jump_public_key=replica.ssh_proxy_jump_public_key, ) for replica in service.replicas ], diff --git a/src/dstack/_internal/proxy/lib/models.py b/src/dstack/_internal/proxy/lib/models.py index 7811760cb3..63952311d5 100644 --- a/src/dstack/_internal/proxy/lib/models.py +++ b/src/dstack/_internal/proxy/lib/models.py @@ -27,13 +27,14 @@ class Replica(ImmutableModel): ssh_proxy: Optional[SSHConnectionParams] = None ssh_proxy_private_key: Optional[str] = None "`None` means same as service project's key" - # Optional outer proxy, a head node/bastion + # SSH fleet proxy_jump. + # TODO: Rename ssh_head_proxy* to ssh_proxy_jump* while preserving gateway state/API compatibility. ssh_head_proxy: Optional[SSHConnectionParams] = None ssh_head_proxy_private_key: Optional[str] = None internal_ip: Optional[str] = None @cached_property - def ssh_head_proxy_public_key(self) -> Optional[str]: + def ssh_proxy_jump_public_key(self) -> Optional[str]: if self.ssh_head_proxy_private_key is None: return None return ssh_utils.get_public_key_from_private_key(self.ssh_head_proxy_private_key) diff --git a/src/dstack/_internal/proxy/lib/services/service_connection.py b/src/dstack/_internal/proxy/lib/services/service_connection.py index 5b85835d0a..06cf81db33 100644 --- a/src/dstack/_internal/proxy/lib/services/service_connection.py +++ b/src/dstack/_internal/proxy/lib/services/service_connection.py @@ -138,10 +138,10 @@ async def client(self) -> ServiceClient: await asyncio.wait_for(self._is_open.wait(), timeout=OPEN_TUNNEL_TIMEOUT) return self._client - async def update_head_node_key(self, replica: Replica) -> None: - """Replace the head-node key without reopening the tunnel or HTTP client. + async def update_proxy_jump_key(self, replica: Replica) -> None: + """Replace the SSH fleet's `proxy_jump` key without reopening the tunnel or HTTP client. - The pool calls this only when the head-node private key is the sole change + The pool calls this only when the `proxy_jump` private key is the sole change from self.replica. Apply the key file before replacing that configuration, so a failed write leaves the previous configuration available for retry. The replacement key is used on the next SSH open; the socket path is unchanged. @@ -221,7 +221,7 @@ async def get_or_add( ) == replica ): - await connection.update_head_node_key(replica) + await connection.update_proxy_jump_key(replica) if service.domain is None and replica.ssh_head_proxy is not None: await connection.open() return connection diff --git a/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py b/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py index 8dbfdceb67..8dafff69c0 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py @@ -25,7 +25,7 @@ from dstack._internal.core.models.gateways import GatewayReplicaStatus, GatewayStatus from dstack._internal.core.models.instances import RemoteConnectionInfo from dstack._internal.core.models.runs import JobSpec, JobStatus, RunStatus, ServiceSpec -from dstack._internal.proxy.gateway.schemas.registry import ReplicaHeadNodeKeyUpdate +from dstack._internal.proxy.gateway.schemas.registry import ReplicaProxyJumpKeyUpdate from dstack._internal.proxy.gateway.schemas.services import ServiceListItem from dstack._internal.server import settings from dstack._internal.server.background.pipeline_tasks.base import ( @@ -719,7 +719,7 @@ async def _process_running_item(item: GatewayReplicaPipelineItem): instance_connection_info = res.all() # Parse keys and compare gateway state without holding a DB connection. - instance_head_node_public_keys = {} + instance_proxy_jump_public_keys = {} for instance_id, connection_info in instance_connection_info: rci = validate_json_extra_ignore(RemoteConnectionInfo, get_or_error(connection_info)) if rci.ssh_proxy is not None: @@ -730,11 +730,11 @@ async def _process_running_item(item: GatewayReplicaPipelineItem): public_key = await run_async( ssh_utils.get_public_key_from_private_key, get_or_error(key.private) ) - instance_head_node_public_keys[instance_id] = public_key + instance_proxy_jump_public_keys[instance_id] = public_key plan = _plan_state_sync( currently_registered=currently_registered_services, expected=expected_runs, - instance_head_node_public_keys=instance_head_node_public_keys, + instance_proxy_jump_public_keys=instance_proxy_jump_public_keys, ) async with get_session_ctx() as session: run_models_by_id, job_models_by_id = await _load_runs_and_jobs_for_state_sync( @@ -948,7 +948,7 @@ async def _perform_state_sync( # Send all replacements together. The gateway keeps working tunnels open and # saves its state once, rather than rewriting the full state for each replica. key_updates = [] - for run_id, job_ids in plan.update_replica_head_node_keys.items(): + for run_id, job_ids in plan.update_replica_proxy_jump_keys.items(): run = run_models_by_id.get(run_id) if run is None: continue @@ -960,11 +960,11 @@ async def _perform_state_sync( if rci is None or rci.ssh_proxy is None: continue key_updates.append( - ReplicaHeadNodeKeyUpdate( + ReplicaProxyJumpKeyUpdate( project_name=run.project.name, run_name=run.run_name, job_id=job_id.hex, - ssh_head_proxy_private_key=get_or_error( + ssh_proxy_jump_private_key=get_or_error( get_or_error(rci.ssh_proxy_keys)[0].private ), ) @@ -972,15 +972,15 @@ async def _perform_state_sync( if key_updates: try: async with connection.client() as client: - failures = await client.update_replica_head_node_keys(key_updates) + failures = await client.update_replica_proxy_jump_keys(key_updates) except Exception as exc: logger.warning( - "%s: could not update service head-node SSH keys: %r", fmt(replica_model), exc + "%s: could not update service proxy_jump SSH keys: %r", fmt(replica_model), exc ) else: for job_id, error in failures.items(): logger.warning( - "%s: could not update replica %s head-node SSH key: %s", + "%s: could not update replica %s proxy_jump SSH key: %s", fmt(replica_model), job_id, error, @@ -1301,7 +1301,7 @@ async def _load_runs_and_jobs_for_state_sync( run_ids = ( plan.register_services | plan.register_replicas.keys() - | plan.update_replica_head_node_keys.keys() + | plan.update_replica_proxy_jump_keys.keys() | {s.id for s in plan.unregister_services if s.id is not None} ) run_models_by_id: dict[uuid.UUID, RunModel] = {} @@ -1318,7 +1318,7 @@ async def _load_runs_and_jobs_for_state_sync( job_ids |= replica_ids for replica_ids in plan.unregister_replicas.values(): job_ids |= replica_ids - for replica_ids in plan.update_replica_head_node_keys.values(): + for replica_ids in plan.update_replica_proxy_jump_keys.values(): job_ids |= replica_ids job_models_by_id: dict[uuid.UUID, JobModel] = {} if job_ids: @@ -1451,8 +1451,8 @@ class _StateSyncPlan: register_replicas: dict[uuid.UUID, set[uuid.UUID]] = field(default_factory=dict) unregister_replicas: dict[_ServiceRef, set[uuid.UUID]] = field(default_factory=dict) set_run_ids: dict[_ServiceRef, uuid.UUID] = field(default_factory=dict) - update_replica_head_node_keys: dict[uuid.UUID, set[uuid.UUID]] = field(default_factory=dict) - """Run IDs mapped to registered job IDs whose saved gateway head-node key needs replacing.""" + update_replica_proxy_jump_keys: dict[uuid.UUID, set[uuid.UUID]] = field(default_factory=dict) + """Run IDs mapped to registered job IDs whose saved gateway proxy_jump key needs replacing.""" @dataclass @@ -1472,7 +1472,7 @@ class _StateSyncResult: def _plan_state_sync( currently_registered: list[ServiceListItem], expected: Sequence[RunModel], - instance_head_node_public_keys: dict[uuid.UUID, str], + instance_proxy_jump_public_keys: dict[uuid.UUID, str], ) -> _StateSyncPlan: plan = _StateSyncPlan() expected_run_id_to_run = {run.id: run for run in expected} @@ -1506,17 +1506,17 @@ def _plan_state_sync( expected_job_ids = {job.id for job in run.jobs} plan.register_replicas[run.id] = expected_job_ids - currently_registered_job_ids plan.unregister_replicas[service_ref] = currently_registered_job_ids - expected_job_ids - gateway_head_node_public_keys = { - uuid.UUID(r.id): r.ssh_head_proxy_public_key for r in service.replicas + gateway_proxy_jump_public_keys = { + uuid.UUID(r.id): r.ssh_proxy_jump_public_key for r in service.replicas } for job in run.jobs: if job.id not in currently_registered_job_ids or job.instance_id is None: continue - public_key = instance_head_node_public_keys.get(job.instance_id) + public_key = instance_proxy_jump_public_keys.get(job.instance_id) if public_key is None: continue - if gateway_head_node_public_keys[job.id] != public_key: - plan.update_replica_head_node_keys.setdefault(run.id, set()).add(job.id) + if gateway_proxy_jump_public_keys[job.id] != public_key: + plan.update_replica_proxy_jump_keys.setdefault(run.id, set()).add(job.id) else: plan.unregister_services.add(service_ref) plan.unregister_replicas[service_ref] = { diff --git a/src/dstack/_internal/server/services/gateways/client.py b/src/dstack/_internal/server/services/gateways/client.py index 2a176c69a1..bf92b7687d 100644 --- a/src/dstack/_internal/server/services/gateways/client.py +++ b/src/dstack/_internal/server/services/gateways/client.py @@ -12,8 +12,8 @@ from dstack._internal.core.models.instances import SSHConnectionParams from dstack._internal.core.models.runs import JobSpec, JobSubmission, get_service_port from dstack._internal.proxy.gateway.schemas.registry import ( - ReplicaHeadNodeKeyUpdate, - UpdateReplicaHeadNodeKeysResponse, + ReplicaProxyJumpKeyUpdate, + UpdateReplicaProxyJumpKeysResponse, ) from dstack._internal.proxy.gateway.schemas.services import ServiceListItem, ServiceListResponse from dstack._internal.proxy.gateway.schemas.stats import ServiceStats @@ -145,21 +145,21 @@ async def register_replica( resp.raise_for_status() self.is_server_ready = True - async def update_replica_head_node_keys( - self, updates: list[ReplicaHeadNodeKeyUpdate] + async def update_replica_proxy_jump_keys( + self, updates: list[ReplicaProxyJumpKeyUpdate] ) -> dict[str, str]: - """Send one synchronization batch of SSH-fleet head-node key replacements. + """Send one synchronization batch of SSH fleet `proxy_jump` key replacements. The gateway preserves existing service tunnels and saves successful updates once. Return failures by job ID; the next sync retries keys still different. """ resp = await self._client.post( - self._url("/api/registry/replicas/head_node_keys"), + self._url("/api/registry/replicas/proxy_jump_keys"), json={"replicas": [item.model_dump() for item in updates]}, ) resp.raise_for_status() self.is_server_ready = True - return UpdateReplicaHeadNodeKeysResponse.model_validate(resp.json()).failures + return UpdateReplicaProxyJumpKeysResponse.model_validate(resp.json()).failures async def unregister_replica(self, project: str, run_name: str, job_id: uuid.UUID): resp = await self._client.post( diff --git a/src/tests/_internal/cli/services/configurators/test_fleet.py b/src/tests/_internal/cli/services/configurators/test_fleet.py index b191d3582d..3bed304e2c 100644 --- a/src/tests/_internal/cli/services/configurators/test_fleet.py +++ b/src/tests/_internal/cli/services/configurators/test_fleet.py @@ -182,11 +182,11 @@ def test_prints_in_place_update_diff(self, monkeypatch: pytest.MonkeyPatch): confirm_ask.assert_called_once_with("Update the fleet in-place?") @pytest.mark.parametrize( - "per_host,head_node", + "per_host,proxy_jump", [(False, False), (False, True), (True, False), (True, True)], ) def test_identifies_the_changed_ssh_key_without_printing_key_material( - self, monkeypatch, per_host, head_node + self, monkeypatch, per_host, proxy_jump ): console, confirm_ask = patch_console_and_confirm(monkeypatch) current = get_ssh_fleet_spec() @@ -198,7 +198,7 @@ def test_identifies_the_changed_ssh_key_without_printing_key_material( target = SSHHostParams(hostname="10.0.0.100") ssh.hosts = [target] path += ".hosts[0]" - if head_node: + if proxy_jump: jump = SSHProxyParams(hostname="head", user="ubuntu", identity_file="old-key") target.proxy_jump = jump target = jump @@ -211,7 +211,7 @@ def test_identifies_the_changed_ssh_key_without_printing_key_material( if per_host: changed = changed.hosts[0] assert isinstance(changed, SSHHostParams) - if head_node: + if proxy_jump: changed = changed.proxy_jump assert changed is not None changed.identity_file = "new-key" diff --git a/src/tests/_internal/proxy/gateway/routers/test_registry.py b/src/tests/_internal/proxy/gateway/routers/test_registry.py index acb9ef4eef..44ecb4ca0e 100644 --- a/src/tests/_internal/proxy/gateway/routers/test_registry.py +++ b/src/tests/_internal/proxy/gateway/routers/test_registry.py @@ -95,7 +95,7 @@ def sample_model_options(name: str = "test-model") -> dict: @pytest.mark.asyncio -class TestUpdateReplicaHeadNodeKeys: +class TestUpdateReplicaProxyJumpKeys: @pytest.mark.parametrize("failure", ["missing", "key-file", "state-save"]) async def test_batch_saves_once_preserves_tunnels_and_retries_failures( self, tmp_path, system_mocks, monkeypatch, failure @@ -123,13 +123,13 @@ async def test_batch_saves_once_preserves_tunnels_and_retries_failures( project_name=project, run_name=run, job_id=job, - ssh_head_proxy_private_key=replacement.private, + ssh_proxy_jump_private_key=replacement.private, ) ) # Populate the cached old public keys before replacing their private keys. before = (await client.get("/api/services/list")).json() assert all( - service["replicas"][0]["ssh_head_proxy_public_key"] == " ".join(PUBLIC_KEY.split()[:2]) + service["replicas"][0]["ssh_proxy_jump_public_key"] == " ".join(PUBLIC_KEY.split()[:2]) for service in before["services"] ) system_mocks.open_conn.reset_mock() @@ -139,7 +139,7 @@ async def test_batch_saves_once_preserves_tunnels_and_retries_failures( with patch.object(repo, "save", side_effect=OSError("state save failed")): with pytest.raises(OSError, match="state save failed"): await client.post( - "/api/registry/replicas/head_node_keys", json={"replicas": updates} + "/api/registry/replicas/proxy_jump_keys", json={"replicas": updates} ) # Keep reporting the old keys so the server retries the failed save. assert (await client.get("/api/services/list")).json() == before @@ -151,7 +151,7 @@ async def test_batch_saves_once_preserves_tunnels_and_retries_failures( batch = list(updates) if failure == "missing": batch.append({**updates[0], "job_id": "missing"}) - original_update = ServiceConnection.update_head_node_key + original_update = ServiceConnection.update_proxy_jump_key async def fail_one_key_file(connection, replica): if replica.id == "job-1": @@ -161,9 +161,9 @@ async def fail_one_key_file(connection, replica): with patch.object(repo, "save", wraps=repo.save) as save: with monkeypatch.context() as scoped: if failure == "key-file": - scoped.setattr(ServiceConnection, "update_head_node_key", fail_one_key_file) + scoped.setattr(ServiceConnection, "update_proxy_jump_key", fail_one_key_file) response = await client.post( - "/api/registry/replicas/head_node_keys", json={"replicas": batch} + "/api/registry/replicas/proxy_jump_keys", json={"replicas": batch} ) assert response.status_code == 200 expected_failures = { @@ -182,13 +182,13 @@ async def fail_one_key_file(connection, replica): # Retry a failed write, then repeat a successful batch as after a lost response. for _ in range(2): response = await client.post( - "/api/registry/replicas/head_node_keys", json={"replicas": updates} + "/api/registry/replicas/proxy_jump_keys", json={"replicas": updates} ) assert response.json() == {"failures": {}} assert save.call_count == (2 if failure == "key-file" else 1) after = (await client.get("/api/services/list")).json() assert all( - service["replicas"][0]["ssh_head_proxy_public_key"] == replacement.public + service["replicas"][0]["ssh_proxy_jump_public_key"] == replacement.public for service in after["services"] ) system_mocks.open_conn.assert_not_awaited() diff --git a/src/tests/_internal/proxy/gateway/routers/test_services.py b/src/tests/_internal/proxy/gateway/routers/test_services.py index 41f3de33af..ade75946ab 100644 --- a/src/tests/_internal/proxy/gateway/routers/test_services.py +++ b/src/tests/_internal/proxy/gateway/routers/test_services.py @@ -25,12 +25,12 @@ async def test_empty(self): assert resp.status_code == 200 assert resp.json() == {"services": []} - @pytest.mark.parametrize("head_node_key", [None, "head-private-key"]) - async def test_list(self, monkeypatch, head_node_key): + @pytest.mark.parametrize("proxy_jump_key", [None, "head-private-key"]) + async def test_list(self, monkeypatch, proxy_jump_key): event_loop_thread = threading.get_ident() def get_public_key(private_key): - assert private_key == head_node_key + assert private_key == proxy_jump_key # Uncached RSA parsing must run outside the gateway's event loop. assert threading.get_ident() != event_loop_thread return "public-key" @@ -49,7 +49,7 @@ def get_public_key(private_key): ssh_destination="ubuntu@server", ssh_port=22, ssh_proxy=None, - ssh_head_proxy_private_key=head_node_key, + ssh_head_proxy_private_key=proxy_jump_key, ) ] ) @@ -67,7 +67,7 @@ def get_public_key(private_key): "replicas": [ { "id": replica_id, - "ssh_head_proxy_public_key": "public-key" if head_node_key else None, + "ssh_proxy_jump_public_key": "public-key" if proxy_jump_key else None, } ], } diff --git a/src/tests/_internal/proxy/lib/services/test_service_connection.py b/src/tests/_internal/proxy/lib/services/test_service_connection.py index 444da6ec80..3e891b3f44 100644 --- a/src/tests/_internal/proxy/lib/services/test_service_connection.py +++ b/src/tests/_internal/proxy/lib/services/test_service_connection.py @@ -192,7 +192,7 @@ async def close_tunnel(): @pytest.mark.asyncio class TestServiceConnectionPool: @pytest.mark.parametrize("gateway", [False, True]) - async def test_head_node_key_update_preserves_client_socket_and_reconnects_with_new_key( + async def test_proxy_jump_key_update_preserves_client_socket_and_reconnects_with_new_key( self, project, service, pool, tunnel_factory, gateway ): service = service.model_copy(update={"domain": service.domain if gateway else None}) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_gateway_replicas.py b/src/tests/_internal/server/background/pipeline_tasks/test_gateway_replicas.py index 506d3f060b..5e61a972a1 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_gateway_replicas.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_gateway_replicas.py @@ -775,7 +775,7 @@ class TestGatewayReplicaWorkerRunningStateSync: "failure,public_key", [(None, PUBLIC_KEY), (None, ""), ("update", PUBLIC_KEY), ("changed-again", PUBLIC_KEY)], ) - async def test_updates_head_node_key_without_unregistering_or_changing_job( + async def test_updates_proxy_jump_key_without_unregistering_or_changing_job( self, test_db, session, worker, mock_gateway_connection, failure, public_key ): project = await create_project(session) @@ -819,7 +819,7 @@ async def test_updates_head_node_key_without_unregistering_or_changing_job( await session.commit() reported = { service_job.id: ServiceListReplicaItem( - id=service_job.id.hex, ssh_head_proxy_public_key="old key" + id=service_job.id.hex, ssh_proxy_jump_public_key="old key" ) for service_job in (job, other_job) } @@ -836,7 +836,7 @@ async def update(updates): if failure == "update": raise GatewayError("update failed") for item in updates: - reported[uuid.UUID(item.job_id)].ssh_head_proxy_public_key = " ".join( + reported[uuid.UUID(item.job_id)].ssh_proxy_jump_public_key = " ".join( PUBLIC_KEY.split()[:2] ) if failure == "changed-again": @@ -847,28 +847,28 @@ async def update(updates): return {} - client.update_replica_head_node_keys.side_effect = update + client.update_replica_proxy_jump_keys.side_effect = update await worker.process(_replica_to_pipeline_item(replica)) - client.update_replica_head_node_keys.assert_awaited_once() - updates = client.update_replica_head_node_keys.call_args.args[0] + client.update_replica_proxy_jump_keys.assert_awaited_once() + updates = client.update_replica_proxy_jump_keys.call_args.args[0] assert {item.job_id for item in updates} == {job.id.hex, other_job.id.hex} assert all( item.project_name == project.name and item.run_name == run.run_name - and item.ssh_head_proxy_private_key == PRIVATE_KEY + and item.ssh_proxy_jump_private_key == PRIVATE_KEY for item in updates ) client.register_replica.assert_not_called() client.unregister_replica.assert_not_called() client.register_service.assert_not_called() client.unregister_service.assert_not_called() - client.update_replica_head_node_keys.reset_mock(side_effect=True) - client.update_replica_head_node_keys.return_value = {} + client.update_replica_proxy_jump_keys.reset_mock(side_effect=True) + client.update_replica_proxy_jump_keys.return_value = {} await session.refresh(replica) _lock_replica(replica) await session.commit() await worker.process(_replica_to_pipeline_item(replica)) - assert client.update_replica_head_node_keys.call_count == (0 if failure is None else 1) + assert client.update_replica_proxy_jump_keys.call_count == (0 if failure is None else 1) await session.refresh(job) assert job.status == JobStatus.RUNNING assert job.submission_num == 0