diff --git a/mkdocs/docs/concepts/fleets.md b/mkdocs/docs/concepts/fleets.md index 92aeb8e3b2..412ddf15bd 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. + Changes to other properties require re-creating the fleet. In this case, `dstack apply` offers to delete the fleet and create it again. +??? info "SSH key rotation" + Updating `identity_file` and `proxy_jump.identity_file` does not require stopping runs. Allow about 20 seconds for changes to propagate before removing the old public key from the host. + ### Delete fleets When a fleet isn't used by a run, you can delete it by passing the fleet configuration to `dstack delete`: diff --git a/src/dstack/_internal/cli/services/configurators/fleet.py b/src/dstack/_internal/cli/services/configurators/fleet.py index ce0db5192f..0a6be89844 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,37 @@ def _render_fleet_spec_diff(old_spec: FleetSpec, new_spec: FleetSpec) -> Optiona return nested_list.render() +def _ssh_config_diff_fields( + old: SSHParams | SSHHostParams | SSHProxyParams, + new: SSHParams | SSHHostParams | SSHProxyParams, + prefix: str = "ssh_config", +) -> list[str]: + fields = [] + for field in diff_models(old, new): + old_value = getattr(old, field) + new_value = getattr(new, field) + path = f"{prefix}.{field}" + if field == "ssh_key": + continue + elif field == "hosts": + old_hosts = {h if isinstance(h, str) else h.hostname: h for h in old_value} + new_names = [h if isinstance(h, str) else h.hostname for h in new_value] + if list(old_hosts) != new_names: + fields.append(path) + for index, host in enumerate(new_value): + hostname = host if isinstance(host, str) else host.hostname + previous = old_hosts.get(hostname) + if isinstance(previous, SSHHostParams) and isinstance(host, SSHHostParams): + fields.extend(_ssh_config_diff_fields(previous, host, f"{path}[{index}]")) + elif previous is not None and previous != host: + fields.append(f"{path}[{index}]") + elif isinstance(old_value, SSHProxyParams) and isinstance(new_value, SSHProxyParams): + fields.extend(_ssh_config_diff_fields(old_value, new_value, path)) + else: + fields.append(path) + return fields + + def _print_plan_header(plan: FleetPlan): def th(s: str) -> str: return f"[bold]{s}[/bold]" diff --git a/src/dstack/_internal/core/services/ssh/tunnel.py b/src/dstack/_internal/core/services/ssh/tunnel.py index 313eeea6e1..71b6e91214 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 `proxy_jump` key changes. Replace + the key file atomically so concurrent SSH startup never reads a partial key. + """ + params, _ = self.ssh_proxies[index] + identity_path = os.path.join(self.temp_dir.name, f"proxy_identity_{index}") + # A concurrent SSH startup must see either complete key, never a partial write. + with tempfile.NamedTemporaryFile(mode="w", dir=self.temp_dir.name, delete=False) as f: + temporary_path = f.name + try: + f.write(private_key) + f.close() + os.replace(temporary_path, identity_path) + finally: + if os.path.exists(temporary_path): + os.unlink(temporary_path) + self.ssh_proxies[index] = (params, identity_path) + def check_command(self) -> List[str]: return [*self._control_command_prefix(), "-O", "check", self.destination] diff --git a/src/dstack/_internal/proxy/gateway/repo/repo.py b/src/dstack/_internal/proxy/gateway/repo/repo.py index a3b69d4565..be0c7cbb45 100644 --- a/src/dstack/_internal/proxy/gateway/repo/repo.py +++ b/src/dstack/_internal/proxy/gateway/repo/repo.py @@ -1,7 +1,8 @@ +import tempfile from contextlib import asynccontextmanager from itertools import chain from pathlib import Path -from typing import Optional +from typing import Iterable, Optional from aiorwlock import RWLock from pydantic import BaseModel @@ -53,6 +54,20 @@ async def set_service(self, service: Service) -> None: async with self.writer(): self._state.services.setdefault(service.project_name, {})[service.run_name] = service + async def set_services(self, services: Iterable[Service]) -> None: + """Save a batch before publishing it, so failed saves remain retryable.""" + async with self._lock.writer: + services_by_project = { + name: entries.copy() for name, entries in self._state.services.items() + } + for service in services: + services_by_project.setdefault(service.project_name, {})[service.run_name] = ( + service + ) + state = self._state.model_copy(update={"services": services_by_project}) + await run_async(self.save, state) + self._state = state + async def delete_service(self, project_name: str, run_name: str) -> None: async with self.writer(): project_services = self._state.services.get(project_name, {}) @@ -124,6 +139,17 @@ def load(state_file: Path) -> "GatewayProxyRepo": state = None return GatewayProxyRepo(state=state, file=state_file) - def save(self) -> None: - if self._file is not None: - self._file.write_text(self._state.model_dump_json()) + def save(self, state: Optional[State] = None) -> None: + """Write the given state, or the current state, atomically.""" + if self._file is None: + return + if state is None: + state = self._state + with tempfile.NamedTemporaryFile(mode="w", dir=self._file.parent, delete=False) as f: + temporary_path = Path(f.name) + try: + f.write(state.model_dump_json()) + f.close() + temporary_path.replace(self._file) + finally: + temporary_path.unlink(missing_ok=True) diff --git a/src/dstack/_internal/proxy/gateway/routers/registry.py b/src/dstack/_internal/proxy/gateway/routers/registry.py index 954ae2933a..b4145bc721 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, + UpdateReplicaProxyJumpKeysRequest, + UpdateReplicaProxyJumpKeysResponse, ) from dstack._internal.proxy.gateway.services.nginx import Nginx from dstack._internal.proxy.lib.deps import get_service_connection_pool from dstack._internal.proxy.lib.services.service_connection import ServiceConnectionPool -router = APIRouter(prefix="/{project_name}") +router = APIRouter() -@router.post("/services/register") +@router.post("/{project_name}/services/register") async def register_service( project_name: str, body: RegisterServiceRequest, @@ -46,7 +48,7 @@ async def register_service( return OkResponse() -@router.post("/services/{run_name}/unregister") +@router.post("/{project_name}/services/{run_name}/unregister") async def unregister_service( project_name: str, run_name: str, @@ -64,7 +66,7 @@ async def unregister_service( return OkResponse() -@router.post("/services/{run_name}/set_id") +@router.post("/{project_name}/services/{run_name}/set_id") async def set_service_id( project_name: str, run_name: str, @@ -81,7 +83,7 @@ async def set_service_id( return OkResponse() -@router.post("/services/{run_name}/replicas/register") +@router.post("/{project_name}/services/{run_name}/replicas/register") async def register_replica( project_name: str, run_name: str, @@ -109,7 +111,19 @@ async def register_replica( return OkResponse() -@router.post("/services/{run_name}/replicas/{job_id}/unregister") +@router.post("/replicas/proxy_jump_keys") +async def update_replica_proxy_jump_keys( + body: UpdateReplicaProxyJumpKeysRequest, + repo: Annotated[GatewayProxyRepo, Depends(get_gateway_proxy_repo)], + service_conn_pool: Annotated[ServiceConnectionPool, Depends(get_service_connection_pool)], +) -> UpdateReplicaProxyJumpKeysResponse: + failures = await registry_services.update_replica_proxy_jump_keys( + updates=body.replicas, repo=repo, service_conn_pool=service_conn_pool + ) + return UpdateReplicaProxyJumpKeysResponse(failures=failures) + + +@router.post("/{project_name}/services/{run_name}/replicas/{job_id}/unregister") async def unregister_replica( project_name: str, run_name: str, @@ -129,7 +143,7 @@ async def unregister_replica( return OkResponse() -@router.post("/entrypoints/register") +@router.post("/{project_name}/entrypoints/register") async def register_entrypoint( project_name: str, body: RegisterEntrypointRequest, diff --git a/src/dstack/_internal/proxy/gateway/schemas/registry.py b/src/dstack/_internal/proxy/gateway/schemas/registry.py index 25f801252a..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 @@ -71,3 +72,20 @@ class RegisterEntrypointRequest(BaseModel): domain: str https: bool read_timeout: int = DEFAULT_SERVICE_READ_TIMEOUT + + +class ReplicaProxyJumpKeyUpdate(BaseModel): + project_name: str + run_name: str + job_id: str + ssh_proxy_jump_private_key: str + """Private key for the SSH fleet's `proxy_jump`.""" + + +class UpdateReplicaProxyJumpKeysRequest(BaseModel): + replicas: list[ReplicaProxyJumpKeyUpdate] + + +class UpdateReplicaProxyJumpKeysResponse(BaseModel): + failures: dict[str, str] + """Job IDs whose key could not be updated; successful updates are saved together.""" diff --git a/src/dstack/_internal/proxy/gateway/schemas/services.py b/src/dstack/_internal/proxy/gateway/schemas/services.py index cbba05b518..f145c393a2 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_proxy_jump_public_key: str | None = None + """Public key matching the SSH fleet's `proxy_jump` key saved for this replica on the gateway. + + The server compares this with the instance's current `proxy_jump` key during gateway + synchronization and sends a replacement private key when they differ. + `None` means no `proxy_jump` private key is saved for this replica. + """ class ServiceListItem(CoreModel): diff --git a/src/dstack/_internal/proxy/gateway/services/registry.py b/src/dstack/_internal/proxy/gateway/services/registry.py index 1bc47d7195..e9f05f27cd 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_proxy_jump_keys( + updates: list[schemas.ReplicaProxyJumpKeyUpdate], + repo: GatewayProxyRepo, + service_conn_pool: ServiceConnectionPool, +) -> dict[str, str]: + """Replace replicas' SSH fleet `proxy_jump` keys in one gateway synchronization batch. + + Keep existing tunnels, registrations, and Nginx configuration. Save successful + replacements once, including across services/projects. Return failures by job ID + so one removed replica or failed key-file write cannot block the other updates. + Repeating the same keys is a no-op; the next sync retries any remaining differences. + """ + by_service: dict[tuple[str, str], dict[str, schemas.ReplicaProxyJumpKeyUpdate]] = {} + for item in updates: + by_service.setdefault((item.project_name.lower(), item.run_name.lower()), {})[ + item.job_id + ] = item + failures = {} + changed_services = [] + async with lock: + for (project_name, run_name), pending in by_service.items(): + service = await repo.get_service(project_name, run_name) + project = await repo.get_project(project_name) + if service is None or project is None: + for job_id in pending: + failures[job_id] = f"Service {project_name}/{run_name} not found" + continue + replicas = list(service.replicas) + changed = False + for index, replica in enumerate(replicas): + item = pending.pop(replica.id, None) + if item is None: + continue + if replica.ssh_head_proxy is None: + failures[replica.id] = "Replica has no proxy_jump" + continue + if replica.ssh_head_proxy_private_key == item.ssh_proxy_jump_private_key: + continue + # A fresh Replica discards its cached old public key. Preserve the + # unchanged Replica objects so polling does not reparse their keys. + updated_replica = models.Replica( + **{ + **replica.model_dump(), + "ssh_head_proxy_private_key": item.ssh_proxy_jump_private_key, + } + ) + try: + await service_conn_pool.get_or_add(project, service, updated_replica) + except Exception as exc: + logger.warning( + "Could not update replica %s proxy_jump SSH key: %r", replica.id, exc + ) + failures[replica.id] = "Could not update proxy_jump key" + continue + replicas[index] = updated_replica + changed = True + for job_id in pending: + failures[job_id] = f"Replica not found in service {project_name}/{run_name}" + if changed: + changed_services.append(service.model_copy(update={"replicas": tuple(replicas)})) + if changed_services: + await repo.set_services(changed_services) + return failures + + async def unregister_replica( project_name: str, run_name: str, diff --git a/src/dstack/_internal/proxy/gateway/services/services.py b/src/dstack/_internal/proxy/gateway/services/services.py index 5543d92160..89b6d0950e 100644 --- a/src/dstack/_internal/proxy/gateway/services/services.py +++ b/src/dstack/_internal/proxy/gateway/services/services.py @@ -4,17 +4,30 @@ ServiceListReplicaItem, ServiceListResponse, ) +from dstack._internal.proxy.lib.models import Service +from dstack._internal.utils.common import run_async async def list_services(repo: GatewayProxyRepo) -> ServiceListResponse: services = await repo.list_services() + # Parsing uncached proxy_jump keys must not block other gateway API requests. + return await run_async(_get_service_list_response, services) + + +def _get_service_list_response(services: list[Service]) -> ServiceListResponse: return ServiceListResponse( services=[ ServiceListItem( id=service.id, project_name=service.project_name, run_name=service.run_name, - replicas=[ServiceListReplicaItem(id=replica.id) for replica in service.replicas], + replicas=[ + ServiceListReplicaItem( + id=replica.id, + ssh_proxy_jump_public_key=replica.ssh_proxy_jump_public_key, + ) + for replica in service.replicas + ], ) for service in services ] diff --git a/src/dstack/_internal/proxy/lib/models.py b/src/dstack/_internal/proxy/lib/models.py index 44d5878bbd..63952311d5 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 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 @@ -25,11 +27,18 @@ 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_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) + 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..06cf81db33 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_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 `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. + """ + 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_proxy_jump_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 @@ -235,7 +272,7 @@ 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: + 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( 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 63da7292ef..8dafff69c0 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, 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 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 ( @@ -67,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 @@ -677,32 +686,57 @@ 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_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: + 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 = await run_async( + ssh_utils.get_public_key_from_private_key, get_or_error(key.private) + ) + instance_proxy_jump_public_keys[instance_id] = public_key + plan = _plan_state_sync( + currently_registered=currently_registered_services, + expected=expected_runs, + 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( session, plan ) @@ -911,6 +945,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_proxy_jump_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( + ReplicaProxyJumpKeyUpdate( + project_name=run.project.name, + run_name=run.run_name, + job_id=job_id.hex, + ssh_proxy_jump_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_proxy_jump_keys(key_updates) + except Exception as exc: + logger.warning( + "%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 proxy_jump 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 @@ -1225,12 +1301,15 @@ async def _load_runs_and_jobs_for_state_sync( run_ids = ( plan.register_services | plan.register_replicas.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] = {} 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()} @@ -1239,6 +1318,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_proxy_jump_keys.values(): + job_ids |= replica_ids job_models_by_id: dict[uuid.UUID, JobModel] = {} if job_ids: res = await session.execute( @@ -1370,6 +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_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 @@ -1387,7 +1470,9 @@ class _StateSyncResult: def _plan_state_sync( - currently_registered: list[ServiceListItem], expected: Sequence[RunModel] + currently_registered: list[ServiceListItem], + expected: Sequence[RunModel], + instance_proxy_jump_public_keys: dict[uuid.UUID, str], ) -> _StateSyncPlan: plan = _StateSyncPlan() expected_run_id_to_run = {run.id: run for run in expected} @@ -1421,6 +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_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_proxy_jump_public_keys.get(job.instance_id) + if public_key is None: + continue + 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/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 7be4bded4f..8c23f95873 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 @@ -599,9 +601,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( @@ -1147,63 +1147,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: @@ -1288,8 +1423,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: @@ -1313,15 +1452,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 @@ -1369,10 +1510,15 @@ 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 + 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): 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): diff --git a/src/dstack/_internal/server/services/gateways/client.py b/src/dstack/_internal/server/services/gateways/client.py index 0c9defecca..bf92b7687d 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 ( + ReplicaProxyJumpKeyUpdate, + UpdateReplicaProxyJumpKeysResponse, +) 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_proxy_jump_keys( + self, updates: list[ReplicaProxyJumpKeyUpdate] + ) -> dict[str, str]: + """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/proxy_jump_keys"), + json={"replicas": [item.model_dump() for item in updates]}, + ) + resp.raise_for_status() + self.is_server_ready = True + 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( self._url( diff --git a/src/dstack/_internal/server/services/proxy/repo.py b/src/dstack/_internal/server/services/proxy/repo.py index 7ea7df43cd..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,12 @@ 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: return None return Project( 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 35c6a313f0..3bed304e2c 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,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, proxy_jump + ): + 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 proxy_jump: + 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] + assert isinstance(changed, SSHHostParams) + if proxy_jump: + changed = changed.proxy_jump + assert changed is not None + 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" + 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/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 a9b265d9fc..44ecb4ca0e 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,108 @@ def sample_model_options(name: str = "test-model") -> dict: } +@pytest.mark.asyncio +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 + ): + 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_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_proxy_jump_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() + 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/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 + 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"}) + original_update = ServiceConnection.update_proxy_jump_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_proxy_jump_key", fail_one_key_file) + response = await client.post( + "/api/registry/replicas/proxy_jump_keys", json={"replicas": batch} + ) + assert response.status_code == 200 + 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): + 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/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_proxy_jump_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..ade75946ab 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("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 == proxy_jump_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=proxy_jump_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}], + "replicas": [ + { + "id": replica_id, + "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 f12eddf0da..3e891b3f44 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_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}) + 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 5b2d0241b3..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 @@ -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 @@ -48,10 +53,12 @@ get_gateway_replica_configuration, get_job_provisioning_data, get_run_spec, + get_ssh_key, list_events, mock_lease, ) from dstack._internal.utils.common import get_current_datetime +from tests._internal.utils.test_ssh import PRIVATE_KEY, PUBLIC_KEY @pytest.fixture @@ -764,6 +771,108 @@ class TestGatewayReplicaWorkerRunningStateSync: pytestmark = pytest.mark.usefixtures("image_config_mock") + @pytest.mark.parametrize( + "failure,public_key", + [(None, PUBLIC_KEY), (None, ""), ("update", PUBLIC_KEY), ("changed-again", PUBLIC_KEY)], + ) + 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) + 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_proxy_jump_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_proxy_jump_public_key = " ".join( + PUBLIC_KEY.split()[:2] + ) + if failure == "changed-again": + assert job.instance is not None + rci.ssh_proxy_keys = [get_ssh_key()] + job.instance.remote_connection_info = rci.model_dump_json() + await session.commit() + + return {} + + client.update_replica_proxy_jump_keys.side_effect = update + await worker.process(_replica_to_pipeline_item(replica)) + 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_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_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_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 + async def _create_service_run_and_job( self, session: AsyncSession, 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() diff --git a/src/tests/_internal/server/services/test_fleets.py b/src/tests/_internal/server/services/test_fleets.py index ce207d2421..578b088c89 100644 --- a/src/tests/_internal/server/services/test_fleets.py +++ b/src/tests/_internal/server/services/test_fleets.py @@ -1,35 +1,290 @@ +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.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.fleets import ( + _check_can_update_ssh_config, + 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, mock_lease, ) +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.cache.lease_project_backends", + mock_lease([]), + ) + 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.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 + 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.cache.lease_project_backends", + mock_lease([]), + ) + 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.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 + 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 TestGetPlanSSHFleetHostsValidation: