From a936e73cce2a913d70031c41bb053d149c3c0dba Mon Sep 17 00:00:00 2001 From: Jvst Me Date: Tue, 11 Aug 2026 00:21:28 +0200 Subject: [PATCH] Gateway state sync mechanism Introduce a way to reconcile gateway replica state based on the current list of services. This allows newly started or newly recovered gateway replicas to learn about existing services. It also improves fault tolerance, since an unavailable gateway replica no longer blocks the service lifecycle - services and service replicas can still be registered on the remaining healthy gateway replicas. Behavioral details: - Run submission no longer waits for service registration before the run starts provisioning. This has both advantages (provisioning and service registration can happen in parallel) and disadvantages (compute may be provisioned for a service that will be terminated a moment later due to a registration error). This behavior may be revised in the future. - A service or a service replica is terminated with the `gateway_error` termination reason if it failed to be registered on all gateway replicas. - If registration succeeded at least on one gateway replica, registration on the remaining gateway replicas will be retried indefinitely. - Rolling deployments only proceed with terminating a service replica once the replacement replica is successfully registered on all gateway replicas. - Registration and unregistration tracking, including per-gateway-replica registration error messages, is available in events. --- mkdocs/docs/concepts/gateways.md | 1 - src/dstack/_internal/core/models/runs.py | 5 + .../pipeline_tasks/gateway_replicas.py | 911 ++++++++- .../background/pipeline_tasks/jobs_running.py | 172 +- .../pipeline_tasks/jobs_terminating.py | 77 +- .../pipeline_tasks/runs/__init__.py | 39 +- .../background/pipeline_tasks/runs/active.py | 36 +- .../pipeline_tasks/runs/terminating.py | 64 - ...10528_add_serviceregistrationmodel_and_.py | 123 ++ src/dstack/_internal/server/models.py | 85 +- .../server/services/gateways/client.py | 15 +- .../server/services/runs/replicas.py | 60 +- .../server/services/services/__init__.py | 80 +- .../pipeline_tasks/test_gateway_replicas.py | 1674 +++++++++++++++-- .../pipeline_tasks/test_running_jobs.py | 379 +++- .../pipeline_tasks/test_runs/test_active.py | 460 ++++- .../_internal/server/routers/test_runs.py | 67 +- .../server/services/services/test_services.py | 14 +- 18 files changed, 3619 insertions(+), 643 deletions(-) create mode 100644 src/dstack/_internal/server/migrations/versions/2026/08_13_0052_3d4f69210528_add_serviceregistrationmodel_and_.py diff --git a/mkdocs/docs/concepts/gateways.md b/mkdocs/docs/concepts/gateways.md index c011486036..e704e8a5d6 100644 --- a/mkdocs/docs/concepts/gateways.md +++ b/mkdocs/docs/concepts/gateways.md @@ -174,7 +174,6 @@ $ dstack gateway list - Changing the number of replicas or redeploying replicas is not supported. - HTTPS is only supported for AWS gateways with the `acm` [certificate type](#certificate). For other gateways, use an external load balancer for TLS termination. - - An unavailable gateway replica prevents any new services or service replicas from being added. - All replicas are bound to the same backend and region. - At most 3 replicas are allowed per gateway. diff --git a/src/dstack/_internal/core/models/runs.py b/src/dstack/_internal/core/models/runs.py index 9951696a44..7c2ec42600 100644 --- a/src/dstack/_internal/core/models/runs.py +++ b/src/dstack/_internal/core/models/runs.py @@ -94,6 +94,7 @@ class RunTerminationReason(str, Enum): STOPPED_BY_USER = "stopped_by_user" ABORTED_BY_USER = "aborted_by_user" SERVER_ERROR = "server_error" + GATEWAY_ERROR = "gateway_error" def to_job_termination_reason(self) -> "JobTerminationReason": """ @@ -107,6 +108,7 @@ def to_job_termination_reason(self) -> "JobTerminationReason": self.STOPPED_BY_USER: JobTerminationReason.TERMINATED_BY_USER, self.ABORTED_BY_USER: JobTerminationReason.ABORTED_BY_USER, self.SERVER_ERROR: JobTerminationReason.TERMINATED_BY_SERVER, + self.GATEWAY_ERROR: JobTerminationReason.TERMINATED_BY_SERVER, } return mapping[self] @@ -118,6 +120,7 @@ def to_status(self) -> "RunStatus": self.STOPPED_BY_USER: RunStatus.TERMINATED, self.ABORTED_BY_USER: RunStatus.TERMINATED, self.SERVER_ERROR: RunStatus.FAILED, + self.GATEWAY_ERROR: RunStatus.FAILED, } return mapping[self] @@ -126,6 +129,8 @@ def to_error(self) -> Optional[str]: return "retry limit exceeded" elif self == RunTerminationReason.SERVER_ERROR: return "server error" + elif self == RunTerminationReason.GATEWAY_ERROR: + return "gateway error" else: return None 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 d089b6eb6d..82fa6c8136 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py @@ -1,19 +1,25 @@ import asyncio import uuid -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import timedelta from typing import Any, Optional, Sequence -from sqlalchemy import and_, or_, select, update -from sqlalchemy.orm import InstrumentedAttribute, joinedload, load_only +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.sql.base import ExecutableOption from dstack._internal.core.backends.base.compute import ( ComputeWithGatewayLoadBalancerSupport, ComputeWithGatewaySupport, ) -from dstack._internal.core.errors import BackendError, BackendNotAvailable +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.runs import JobSpec, JobStatus, RunStatus, ServiceSpec +from dstack._internal.proxy.gateway.schemas.services import ServiceListItem +from dstack._internal.server import settings from dstack._internal.server.background.pipeline_tasks.base import ( Fetcher, Heartbeater, @@ -32,21 +38,36 @@ BackendModel, GatewayComputeModel, GatewayModel, + InstanceModel, + JobModel, ProjectModel, + RunModel, + ServiceRegistrationModel, + ServiceReplicaRegistrationModel, ) from dstack._internal.server.services import backends as backends_services +from dstack._internal.server.services import events from dstack._internal.server.services import gateways as gateways_services from dstack._internal.server.services.gateways import ( get_gateway_compute_configuration, get_gateway_configuration, get_gateway_lb_configuration, ) +from dstack._internal.server.services.gateways.client import GatewayClient +from dstack._internal.server.services.gateways.connection import GatewayConnection from dstack._internal.server.services.gateways.pool import gateway_connections_pool +from dstack._internal.server.services.instances import get_instance_remote_connection_info +from dstack._internal.server.services.jobs import job_model_to_job_submission from dstack._internal.server.services.locking import get_locker from dstack._internal.server.services.logging import fmt from dstack._internal.server.services.pipelines import PipelineHinterProtocol +from dstack._internal.server.services.runs import get_run_spec +from dstack._internal.server.services.services import ( + get_gateway_https, + should_configure_service_https_on_gateway, +) from dstack._internal.server.utils import tracing -from dstack._internal.utils.common import get_current_datetime, run_async +from dstack._internal.utils.common import get_current_datetime, get_or_error, run_async from dstack._internal.utils.logging import get_logger logger = get_logger(__name__) @@ -158,19 +179,10 @@ async def fetch(self, limit: int) -> list[GatewayReplicaPipelineItem]: [ GatewayReplicaStatus.SUBMITTED, GatewayReplicaStatus.PROVISIONING, + GatewayReplicaStatus.RUNNING, GatewayReplicaStatus.TERMINATING, ] ), - and_( - GatewayComputeModel.status == GatewayReplicaStatus.RUNNING, - or_( - GatewayComputeModel.scale_in == True, - GatewayModel.to_be_deleted == True, - GatewayModel.status == GatewayStatus.FAILED, - # Gateway was hard-deleted (unexpected, fetch to log an error) - GatewayModel.id.is_(None), - ), - ), ), or_( GatewayComputeModel.last_processed_at @@ -352,28 +364,52 @@ def _mark_terminating_if_needed( return update_map +# TODO: Consider refactoring the pipeline for consistency with other pipelines - split into process +# and apply phases instead of calling the `_commit_update()` helper from everywhere async def _commit_update( item: GatewayReplicaPipelineItem, replica_model: GatewayComputeModel, update_map: _GatewayReplicaUpdateMap, ) -> None: + async with get_session_ctx() as session: + await _apply_update(session, item, replica_model, update_map) + + +async def _apply_update( + session: AsyncSession, + item: GatewayReplicaPipelineItem, + replica_model: GatewayComputeModel, + update_map: _GatewayReplicaUpdateMap, +) -> bool: set_processed_update_map_fields(update_map) set_unlock_update_map_fields(update_map) - async with get_session_ctx() as session: - now = get_current_datetime() - resolve_now_placeholders(update_map, now=now) - res = await session.execute( - update(GatewayComputeModel) - .where( - GatewayComputeModel.id == replica_model.id, - GatewayComputeModel.lock_token == replica_model.lock_token, + now = get_current_datetime() + resolve_now_placeholders(update_map, now=now) + res = await session.execute( + update(GatewayComputeModel) + .where( + GatewayComputeModel.id == replica_model.id, + GatewayComputeModel.lock_token == replica_model.lock_token, + ) + .values(**update_map) + .returning(GatewayComputeModel.id) + ) + updated_ids = list(res.scalars().all()) + if len(updated_ids) == 0: + log_lock_token_changed_after_processing(logger, item) + return False + if update_map.get("deleted"): + await session.execute( + delete(ServiceRegistrationModel).where( + ServiceRegistrationModel.gateway_replica_id == replica_model.id ) - .values(**update_map) - .returning(GatewayComputeModel.id) ) - updated_ids = list(res.scalars().all()) - if len(updated_ids) == 0: - log_lock_token_changed_after_processing(logger, item) + await session.execute( + delete(ServiceReplicaRegistrationModel).where( + ServiceReplicaRegistrationModel.gateway_replica_id == replica_model.id + ) + ) + return True async def _process_submitted_item(item: GatewayReplicaPipelineItem): @@ -654,8 +690,17 @@ async def _process_running_item(item: GatewayReplicaPipelineItem): replica_fields=_REPLICA_FIELDS_MIN + [ GatewayComputeModel.scale_in, + GatewayComputeModel.ip_address, + GatewayComputeModel.ssh_private_key, ], - gateway_fields=_GATEWAY_FIELDS_MIN, + gateway_fields=_GATEWAY_FIELDS_MIN + + [ + GatewayModel.project_id, + GatewayModel.configuration, + GatewayModel.region, + GatewayModel.wildcard_domain, + ], + load_gateway_backend_type=True, ) if replica_model is None: return @@ -666,12 +711,814 @@ async def _process_running_item(item: GatewayReplicaPipelineItem): if update_map := _mark_terminating_if_needed(gateway_model, replica_model): await _commit_update(item, replica_model, update_map=update_map) return - logger.warning( - "%s replica %d: nothing to do in this pipeline tick", + try: + connection = await gateway_connections_pool.get_or_add( + hostname=get_or_error(replica_model.ip_address), + id_rsa=replica_model.ssh_private_key, + ) + except Exception as e: + logger.warning( + "%s replica %d: failed to connect to gateway: %s", + fmt(gateway_model), + replica_model.replica_num, + e, + ) + await _commit_update(item, replica_model, update_map={}) + return + async with connection.client() as client: + try: + currently_registered_services = await client.list_services() + except Exception as e: + if isinstance(e, HTTPStatusError) and e.response.status_code == 404: + logger.warning( + ( + "%s replica %d: got error 404 when listing services, which indicates a" + " pre-0.21.0 gateway. Skipping state sync until the gateway is updated" + ), + fmt(gateway_model), + replica_model.replica_num, + ) + else: + logger.warning( + "%s replica %d: failed to list services: %r", + fmt(gateway_model), + replica_model.replica_num, + e, + ) + await _commit_update(item, replica_model, update_map={}) + return + stmt = ( + select(RunModel) + .where( + RunModel.gateway_id == gateway_model.id, + RunModel.deleted == False, + RunModel.status.not_in(RunStatus.finished_statuses() + [RunStatus.TERMINATING]), + ) + .options( + load_only(RunModel.id), + joinedload(RunModel.jobs).load_only(JobModel.id), + with_loader_criteria( + JobModel, + and_( + JobModel.status == JobStatus.RUNNING, + JobModel.registered == True, + ), + ), + ) + ) + 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, + ) + run_models_by_id, job_models_by_id = await _load_runs_and_jobs_for_state_sync( + session, plan + ) + + sync_result = await _perform_state_sync( + connection, gateway_model, replica_model, run_models_by_id, job_models_by_id, plan + ) + + async with get_session_ctx() as session: + if not await _apply_update(session, item, replica_model, update_map={}): + return + reconcile_records_result = await _reconcile_registration_records( + session, replica_model, currently_registered_services, sync_result + ) + await _emit_state_sync_events( + session, + gateway_model, + replica_model, + run_models_by_id, + job_models_by_id, + sync_result, + reconcile_records_result, + ) + + +async def _perform_state_sync( + connection: GatewayConnection, + gateway_model: GatewayModel, + replica_model: GatewayComputeModel, + run_models_by_id: dict[uuid.UUID, RunModel], + job_models_by_id: dict[uuid.UUID, JobModel], + plan: "_StateSyncPlan", +) -> "_StateSyncResult": + result = _StateSyncResult() + for service_ref, run_id in plan.set_run_ids.items(): + logger.debug( + "%s replica %d: setting id %s for service %s/%s", + fmt(gateway_model), + replica_model.replica_num, + run_id, + service_ref.project_name, + service_ref.run_name, + ) + try: + async with connection.client() as client: + await client.set_service_id( + project=service_ref.project_name, + run_name=service_ref.run_name, + run_id=run_id, + ) + except Exception: + logger.exception( + "%s replica %d: failed to set id %s for service %s/%s", + fmt(gateway_model), + replica_model.replica_num, + run_id, + service_ref.project_name, + service_ref.run_name, + ) + continue + logger.info( + "%s replica %d: id %s set for service %s/%s", + fmt(gateway_model), + replica_model.replica_num, + run_id, + service_ref.project_name, + service_ref.run_name, + ) + for service_ref in plan.unregister_services: + logger.debug( + "%s replica %d: unregistering service %s/%s", + fmt(gateway_model), + replica_model.replica_num, + service_ref.project_name, + service_ref.run_name, + ) + try: + async with connection.client() as client: + await client.unregister_service( + project=service_ref.project_name, + run_name=service_ref.run_name, + ) + except GatewayError as e: + if service_ref.id is not None: + result.failed_service_unregistrations[service_ref.id] = e.msg + else: + logger.warning( + "%s replica %d: failed to unregister legacy service %s/%s with unknown ID: %s", + fmt(gateway_model), + replica_model.replica_num, + service_ref.project_name, + service_ref.run_name, + e.msg, + ) + continue + except Exception: + logger.exception( + "%s replica %d: failed to unregister service %s/%s", + fmt(gateway_model), + replica_model.replica_num, + service_ref.project_name, + service_ref.run_name, + ) + if service_ref.id is not None: + result.failed_service_unregistrations[service_ref.id] = "Unexpected error" + continue + if service_ref.id is not None: + result.unregistered_services.add(service_ref.id) + else: + logger.warning( + "%s replica %d: unregistered legacy service %s/%s with unknown ID", + fmt(gateway_model), + replica_model.replica_num, + service_ref.project_name, + service_ref.run_name, + ) + for service_ref, replica_ids in plan.unregister_replicas.items(): + for replica_id in replica_ids: + logger.debug( + "%s replica %d: unregistering replica %s for service %s/%s", + fmt(gateway_model), + replica_model.replica_num, + replica_id, + service_ref.project_name, + service_ref.run_name, + ) + try: + async with connection.client() as client: + await client.unregister_replica( + project=service_ref.project_name, + run_name=service_ref.run_name, + job_id=replica_id, + ) + except GatewayError as e: + result.failed_replica_unregistrations[replica_id] = e.msg + continue + except Exception: + logger.exception( + "%s replica %d: failed to unregister replica %s for service %s/%s", + fmt(gateway_model), + replica_model.replica_num, + replica_id, + service_ref.project_name, + service_ref.run_name, + ) + result.failed_replica_unregistrations[replica_id] = "Unexpected error" + continue + result.unregistered_replicas.add(replica_id) + for run_id in plan.register_services: + run_model = run_models_by_id.get(run_id) + if run_model is None: + error_message = "Run not found" + logger.error( + "%s replica %d: run %s not found, cannot register service", + fmt(gateway_model), + replica_model.replica_num, + run_id, + ) + result.failed_service_registrations[run_id] = error_message + continue + try: + async with connection.client() as client: + await _register_service(client, gateway_model, replica_model, run_model) + except GatewayError as e: + result.failed_service_registrations[run_id] = e.msg + continue + except Exception: + logger.exception( + "%s replica %d: failed to register service for run %s", + fmt(gateway_model), + replica_model.replica_num, + run_id, + ) + result.failed_service_registrations[run_id] = "Unexpected error" + continue + result.registered_services.add(run_id) + for run_id, replica_ids in plan.register_replicas.items(): + if run_id in result.failed_service_registrations: + continue + run_model = run_models_by_id.get(run_id) + if run_model is None: + logger.error( + "%s replica %d: run %s not found, cannot register replicas", + fmt(gateway_model), + replica_model.replica_num, + run_id, + ) + for replica_id in replica_ids: + result.failed_replica_registrations[replica_id] = "Run not found" + continue + for replica_id in replica_ids: + job_model = job_models_by_id.get(replica_id) + if job_model is None: + logger.error( + "%s replica %d: job %s not found, cannot register replica", + fmt(gateway_model), + replica_model.replica_num, + replica_id, + ) + result.failed_replica_registrations[replica_id] = "Job not found" + continue + try: + async with connection.client() as client: + await _register_replica( + client, gateway_model, replica_model, run_model, job_model + ) + except GatewayError as e: + result.failed_replica_registrations[replica_id] = e.msg + continue + except Exception: + logger.exception( + "%s replica %d: failed to register replica %s", + fmt(gateway_model), + replica_model.replica_num, + replica_id, + ) + result.failed_replica_registrations[replica_id] = "Unexpected error" + continue + result.registered_replicas.add(replica_id) + return result + + +def _get_or_create_service_registration( + session: AsyncSession, + existing_by_id: dict[uuid.UUID, ServiceRegistrationModel], + run_id: uuid.UUID, + gateway_replica_id: uuid.UUID, +) -> ServiceRegistrationModel: + if registration := existing_by_id.get(run_id): + return registration + registration = ServiceRegistrationModel( + run_id=run_id, + gateway_replica_id=gateway_replica_id, + register_attempt=0, + register_status_message=None, + unregister_attempt=0, + unregister_status_message=None, + ) + session.add(registration) + return registration + + +def _get_or_create_service_replica_registration( + session: AsyncSession, + existing_by_id: dict[uuid.UUID, ServiceReplicaRegistrationModel], + job_id: uuid.UUID, + gateway_replica_id: uuid.UUID, +) -> ServiceReplicaRegistrationModel: + if registration := existing_by_id.get(job_id): + return registration + registration = ServiceReplicaRegistrationModel( + job_id=job_id, + gateway_replica_id=gateway_replica_id, + register_attempt=0, + register_status_message=None, + unregister_attempt=0, + unregister_status_message=None, + ) + session.add(registration) + return registration + + +@dataclass +class _ReconcileRegistrationRecordsResult: + services_with_new_registration_error: set[uuid.UUID] = field(default_factory=set) + replicas_with_new_registration_error: set[uuid.UUID] = field(default_factory=set) + services_with_new_unregistration_error: set[uuid.UUID] = field(default_factory=set) + replicas_with_new_unregistration_error: set[uuid.UUID] = field(default_factory=set) + + +async def _reconcile_registration_records( + session: AsyncSession, + replica_model: GatewayComputeModel, + initially_registered: list[ServiceListItem], + sync_result: "_StateSyncResult", +) -> _ReconcileRegistrationRecordsResult: + result = _ReconcileRegistrationRecordsResult() + initially_registered_run_ids = { + uuid.UUID(s.id) for s in initially_registered if s.id is not None + } + initially_registered_replica_ids = { + uuid.UUID(r.id) for s in initially_registered for r in s.replicas + } + registered_run_ids = ( + initially_registered_run_ids - sync_result.unregistered_services + ) | sync_result.registered_services + registered_replica_ids = ( + initially_registered_replica_ids - sync_result.unregistered_replicas + ) | sync_result.registered_replicas + + keep_run_ids = registered_run_ids | sync_result.failed_service_registrations.keys() + keep_replica_ids = registered_replica_ids | sync_result.failed_replica_registrations.keys() + + await session.execute( + delete(ServiceRegistrationModel).where( + ServiceRegistrationModel.gateway_replica_id == replica_model.id, + ServiceRegistrationModel.run_id.not_in(keep_run_ids), + ) + ) + await session.execute( + delete(ServiceReplicaRegistrationModel).where( + ServiceReplicaRegistrationModel.gateway_replica_id == replica_model.id, + ServiceReplicaRegistrationModel.job_id.not_in(keep_replica_ids), + ) + ) + + service_registrations_by_run_id: dict[uuid.UUID, ServiceRegistrationModel] = {} + if keep_run_ids: + res = await session.execute( + select(ServiceRegistrationModel).where( + ServiceRegistrationModel.gateway_replica_id == replica_model.id, + ) + ) + service_registrations_by_run_id = {r.run_id: r for r in res.scalars().all()} + replica_registrations_by_job_id: dict[uuid.UUID, ServiceReplicaRegistrationModel] = {} + if keep_replica_ids: + res = await session.execute( + select(ServiceReplicaRegistrationModel).where( + ServiceReplicaRegistrationModel.gateway_replica_id == replica_model.id, + ) + ) + replica_registrations_by_job_id = {r.job_id: r for r in res.scalars().all()} + + for run_id in registered_run_ids: + registration = _get_or_create_service_registration( + session=session, + existing_by_id=service_registrations_by_run_id, + run_id=run_id, + gateway_replica_id=replica_model.id, + ) + registration.is_registered = True + registration.register_attempt = 0 + registration.register_status_message = None + unregister_error_message = sync_result.failed_service_unregistrations.get(run_id) + if unregister_error_message is None: + registration.unregister_attempt = 0 + registration.unregister_status_message = None + else: + registration.unregister_attempt += 1 + if unregister_error_message != registration.unregister_status_message: + registration.unregister_status_message = unregister_error_message + result.services_with_new_unregistration_error.add(run_id) + for job_id in registered_replica_ids: + registration = _get_or_create_service_replica_registration( + session=session, + existing_by_id=replica_registrations_by_job_id, + job_id=job_id, + gateway_replica_id=replica_model.id, + ) + registration.is_registered = True + registration.register_attempt = 0 + registration.register_status_message = None + unregister_error_message = sync_result.failed_replica_unregistrations.get(job_id) + if unregister_error_message is None: + registration.unregister_attempt = 0 + registration.unregister_status_message = None + else: + registration.unregister_attempt += 1 + if unregister_error_message != registration.unregister_status_message: + registration.unregister_status_message = unregister_error_message + result.replicas_with_new_unregistration_error.add(job_id) + for run_id, error_message in sync_result.failed_service_registrations.items(): + registration = _get_or_create_service_registration( + session=session, + existing_by_id=service_registrations_by_run_id, + run_id=run_id, + gateway_replica_id=replica_model.id, + ) + registration.is_registered = False + registration.register_attempt += 1 + if error_message != registration.register_status_message: + registration.register_status_message = error_message + result.services_with_new_registration_error.add(run_id) + for job_id, error_message in sync_result.failed_replica_registrations.items(): + registration = _get_or_create_service_replica_registration( + session=session, + existing_by_id=replica_registrations_by_job_id, + job_id=job_id, + gateway_replica_id=replica_model.id, + ) + registration.is_registered = False + registration.register_attempt += 1 + if error_message != registration.register_status_message: + registration.register_status_message = error_message + result.replicas_with_new_registration_error.add(job_id) + return result + + +async def _emit_state_sync_events( + session, + gateway_model: GatewayModel, + replica_model: GatewayComputeModel, + run_models_by_id: dict[uuid.UUID, RunModel], + job_models_by_id: dict[uuid.UUID, JobModel], + sync_result: "_StateSyncResult", + reconcile_records_result: _ReconcileRegistrationRecordsResult, +) -> None: + # TODO: once gateway replica event targets are supported, link events to gateway replicas + # instead of gateways, and remove gateway replica nums from messages. + for run_id in sync_result.unregistered_services: + run_model = run_models_by_id.get(run_id) + if run_model is None: + logger.error( + "%s replica %d: run %s not found, cannot emit service unregistration event", + fmt(gateway_model), + replica_model.replica_num, + run_id, + ) + continue + events.emit( + session, + f"Service unregistered from gateway replica {replica_model.replica_num}", + actor=events.SystemActor(), + targets=[events.Target.from_model(run_model), events.Target.from_model(gateway_model)], + ) + for job_id in sync_result.unregistered_replicas: + job_model = job_models_by_id.get(job_id) + if job_model is None: + logger.error( + "%s replica %d: job %s not found, cannot emit replica unregistration event", + fmt(gateway_model), + replica_model.replica_num, + job_id, + ) + continue + events.emit( + session, + f"Service replica unregistered from gateway replica {replica_model.replica_num}", + actor=events.SystemActor(), + targets=[events.Target.from_model(job_model), events.Target.from_model(gateway_model)], + ) + for run_id in sync_result.registered_services: + run_model = run_models_by_id.get(run_id) + if run_model is None: + logger.error( + "%s replica %d: run %s not found, cannot emit service registration event", + fmt(gateway_model), + replica_model.replica_num, + run_id, + ) + continue + events.emit( + session, + f"Service registered on gateway replica {replica_model.replica_num}", + actor=events.SystemActor(), + targets=[events.Target.from_model(run_model), events.Target.from_model(gateway_model)], + ) + for job_id in sync_result.registered_replicas: + job_model = job_models_by_id.get(job_id) + if job_model is None: + logger.error( + "%s replica %d: job %s not found, cannot emit replica registration event", + fmt(gateway_model), + replica_model.replica_num, + job_id, + ) + continue + events.emit( + session, + f"Service replica registered on gateway replica {replica_model.replica_num}", + actor=events.SystemActor(), + targets=[events.Target.from_model(job_model), events.Target.from_model(gateway_model)], + ) + for run_id, error_message in sync_result.failed_service_registrations.items(): + if run_id not in reconcile_records_result.services_with_new_registration_error: + continue # same error as before, do not emit duplicate event + run_model = run_models_by_id.get(run_id) + if run_model is None: + logger.error( + "%s replica %d: run %s not found, cannot emit service registration event", + fmt(gateway_model), + replica_model.replica_num, + run_id, + ) + continue + events.emit( + session, + f"Encountered service registration error on gateway replica {replica_model.replica_num}: {error_message}", + actor=events.SystemActor(), + targets=[events.Target.from_model(run_model), events.Target.from_model(gateway_model)], + ) + for job_id, error_message in sync_result.failed_replica_registrations.items(): + if job_id not in reconcile_records_result.replicas_with_new_registration_error: + continue # same error as before, do not emit duplicate event + job_model = job_models_by_id.get(job_id) + if job_model is None: + logger.error( + "%s replica %d: job %s not found, cannot emit replica registration event", + fmt(gateway_model), + replica_model.replica_num, + job_id, + ) + continue + events.emit( + session, + f"Encountered service replica registration error on gateway replica {replica_model.replica_num}: {error_message}", + actor=events.SystemActor(), + targets=[events.Target.from_model(job_model), events.Target.from_model(gateway_model)], + ) + for run_id, error_message in sync_result.failed_service_unregistrations.items(): + if run_id not in reconcile_records_result.services_with_new_unregistration_error: + continue # same error as before, do not emit duplicate event + run_model = run_models_by_id.get(run_id) + if run_model is None: + logger.error( + "%s replica %d: run %s not found, cannot emit service unregistration event", + fmt(gateway_model), + replica_model.replica_num, + run_id, + ) + continue + events.emit( + session, + f"Encountered service unregistration error on gateway replica {replica_model.replica_num}: {error_message}", + actor=events.SystemActor(), + targets=[events.Target.from_model(run_model), events.Target.from_model(gateway_model)], + ) + for job_id, error_message in sync_result.failed_replica_unregistrations.items(): + if job_id not in reconcile_records_result.replicas_with_new_unregistration_error: + continue # same error as before, do not emit duplicate event + job_model = job_models_by_id.get(job_id) + if job_model is None: + logger.error( + "%s replica %d: job %s not found, cannot emit replica unregistration event", + fmt(gateway_model), + replica_model.replica_num, + job_id, + ) + continue + events.emit( + session, + f"Encountered service replica unregistration error on gateway replica {replica_model.replica_num}: {error_message}", + actor=events.SystemActor(), + targets=[events.Target.from_model(job_model), events.Target.from_model(gateway_model)], + ) + + +async def _load_runs_and_jobs_for_state_sync( + session: AsyncSession, + plan: "_StateSyncPlan", +) -> tuple[dict[uuid.UUID, RunModel], dict[uuid.UUID, JobModel]]: + run_ids = ( + plan.register_services + | plan.register_replicas.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)) + ) + run_models_by_id = {run.id: run for run in res.unique().scalars().all()} + + job_ids: set[uuid.UUID] = set() + for replica_ids in plan.register_replicas.values(): + job_ids |= replica_ids + for replica_ids in plan.unregister_replicas.values(): + job_ids |= replica_ids + job_models_by_id: dict[uuid.UUID, JobModel] = {} + if job_ids: + res = await session.execute( + select(JobModel) + .where(JobModel.id.in_(job_ids)) + .options( + joinedload(JobModel.instance).joinedload(InstanceModel.project), + joinedload(JobModel.project).load_only(ProjectModel.id, ProjectModel.name), + ) + ) + job_models_by_id = {job.id: job for job in res.unique().scalars().all()} + + return run_models_by_id, job_models_by_id + + +async def _register_service( + client: GatewayClient, + gateway_model: GatewayModel, + replica_model: GatewayComputeModel, + run_model: RunModel, +) -> None: + run_spec = get_run_spec(run_model) + if run_spec.configuration.type != "service": + message = f"Run {run_model.id} is not a service, cannot register" + logger.error("%s replica %d: %s", fmt(gateway_model), replica_model.replica_num, message) + raise RuntimeError(message) + if run_model.service_spec is None: + message = f"Run {run_model.id} has no service spec, cannot register" + logger.error("%s replica %d: %s", fmt(gateway_model), replica_model.replica_num, message) + raise RuntimeError(message) + service_spec = validate_json_extra_ignore(ServiceSpec, run_model.service_spec) + domain = service_spec.get_domain() + if domain is None: + message = f"Run {run_model.id} service spec has no domain, cannot register" + logger.error("%s replica %d: %s", fmt(gateway_model), replica_model.replica_num, message) + raise RuntimeError(message) + + gateway_configuration = get_gateway_configuration(gateway_model) + has_replica_group_router = any( + g.router is not None for g in run_spec.configuration.replica_groups + ) + logger.debug( + "%s replica %d: registering service %s/%s", + fmt(gateway_model), + replica_model.replica_num, + run_model.project.name, + run_model.run_name, + ) + await client.register_service( + project=run_model.project.name, + run_id=run_model.id, + run_name=run_model.run_name, + domain=domain, + service_https=should_configure_service_https_on_gateway(run_spec, gateway_configuration), + gateway_https=get_gateway_https(gateway_configuration), + auth=run_spec.configuration.auth, + client_max_body_size=settings.DEFAULT_SERVICE_CLIENT_MAX_BODY_SIZE, + options=service_spec.options, + rate_limits=run_spec.configuration.rate_limits, + ssh_private_key=run_model.project.ssh_private_key, + has_router_replica=has_replica_group_router, + ) + + +async def _register_replica( + client: GatewayClient, + gateway_model: GatewayModel, + replica_model: GatewayComputeModel, + run_model: RunModel, + job_model: JobModel, +) -> None: + run_spec = get_run_spec(run_model) + if run_spec.configuration.type != "service": + message = f"Run {run_model.id} is not a service, cannot register replica" + logger.error("%s replica %d: %s", fmt(gateway_model), replica_model.replica_num, message) + raise RuntimeError(message) + instance = job_model.instance + if instance is None: + message = f"Job {job_model.id} has no instance, cannot register replica" + logger.error("%s replica %d: %s", fmt(gateway_model), replica_model.replica_num, message) + raise RuntimeError(message) + job_spec = validate_json_extra_ignore(JobSpec, job_model.job_spec_data) + job_submission = job_model_to_job_submission(job_model) + + instance_project_ssh_private_key = None + if job_model.project_id != instance.project_id: + instance_project_ssh_private_key = instance.project.ssh_private_key + ssh_head_proxy = None + ssh_head_proxy_private_key = None + rci = get_instance_remote_connection_info(instance) + if rci is not None and rci.ssh_proxy is not None: + ssh_head_proxy = rci.ssh_proxy + ssh_head_proxy_private_key = get_or_error(rci.ssh_proxy_keys)[0].private + + logger.debug( + "%s replica %d: registering replica %s for service %s/%s", fmt(gateway_model), replica_model.replica_num, + job_model.id, + run_model.project.name, + run_model.run_name, ) - await _commit_update(item, replica_model, update_map={}) + await client.register_replica( + project=run_model.project.name, + run_name=run_model.run_name, + configuration=run_spec.configuration, + job_spec=job_spec, + job_submission=job_submission, + instance_project_ssh_private_key=instance_project_ssh_private_key, + ssh_head_proxy=ssh_head_proxy, + ssh_head_proxy_private_key=ssh_head_proxy_private_key, + ) + + +@dataclass(frozen=True) +class _ServiceRef: + id: uuid.UUID | None + project_name: str + run_name: str + + +@dataclass +class _StateSyncPlan: + register_services: set[uuid.UUID] = field(default_factory=set) + unregister_services: set[_ServiceRef] = field(default_factory=set) + # run ID -> set[job ID] + 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) + + +@dataclass +class _StateSyncResult: + registered_services: set[uuid.UUID] = field(default_factory=set) + registered_replicas: set[uuid.UUID] = field(default_factory=set) + unregistered_services: set[uuid.UUID] = field(default_factory=set) + unregistered_replicas: set[uuid.UUID] = field(default_factory=set) + + # run ID -> error message + failed_service_registrations: dict[uuid.UUID, str] = field(default_factory=dict) + failed_replica_registrations: dict[uuid.UUID, str] = field(default_factory=dict) + failed_service_unregistrations: dict[uuid.UUID, str] = field(default_factory=dict) + failed_replica_unregistrations: dict[uuid.UUID, str] = field(default_factory=dict) + + +def _plan_state_sync( + currently_registered: list[ServiceListItem], expected: Sequence[RunModel] +) -> _StateSyncPlan: + plan = _StateSyncPlan() + expected_run_id_to_run = {run.id: run for run in expected} + expected_job_id_to_run = {job.id: run for run in expected for job in run.jobs} + expected_run_ids = {run.id for run in expected} + currently_registered_run_ids: set[uuid.UUID] = set() + + for service in currently_registered: + service_ref = _ServiceRef( + id=uuid.UUID(service.id) if service.id is not None else None, + project_name=service.project_name, + run_name=service.run_name, + ) + if service.id is not None: + run_id = uuid.UUID(service.id) + else: + # Try to recover ID for legacy pre-0.21.0 service + for replica in service.replicas: + if run := expected_job_id_to_run.get(uuid.UUID(replica.id)): + run_id = run.id + plan.set_run_ids[service_ref] = run_id + break + else: + # Could not recover ID, and none of the current replicas are relevant - unregister. + # If the service is relevant, we'll re-register it with ID. + plan.unregister_services.add(service_ref) + continue + currently_registered_run_ids.add(run_id) + if run := expected_run_id_to_run.get(run_id): + currently_registered_job_ids = {uuid.UUID(replica.id) for replica in service.replicas} + 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 + else: + plan.unregister_services.add(service_ref) + + for run_id in expected_run_ids - currently_registered_run_ids: + plan.register_services.add(run_id) + plan.register_replicas[run_id] = {job.id for job in expected_run_id_to_run[run_id].jobs} + + return plan async def _process_terminating_item(item: GatewayReplicaPipelineItem): diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py index 99192af163..77bf8576fa 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py @@ -6,13 +6,11 @@ from datetime import datetime, timedelta from typing import Dict, Iterable, Literal, Optional, Sequence, Union -import httpx from sqlalchemy import and_, exists, false, func, or_, select, true, update from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.orm import aliased, contains_eager, joinedload, load_only +from sqlalchemy.orm import aliased, contains_eager, joinedload, load_only, selectinload from dstack._internal.core.consts import DSTACK_RUNNER_HTTP_PORT, DSTACK_SHIM_HTTP_PORT -from dstack._internal.core.errors import GatewayError, SSHError from dstack._internal.core.models.common import ( NetworkMode, RegistryAuth, @@ -22,7 +20,8 @@ DevEnvironmentConfiguration, ) from dstack._internal.core.models.files import FileArchiveMapping -from dstack._internal.core.models.instances import InstanceStatus, SSHConnectionParams +from dstack._internal.core.models.gateways import GatewayReplicaStatus +from dstack._internal.core.models.instances import InstanceStatus from dstack._internal.core.models.metrics import Metric from dstack._internal.core.models.profiles import StartupOrder from dstack._internal.core.models.repos import RemoteRepoCreds @@ -60,6 +59,8 @@ from dstack._internal.server.models import ( ExportedFleetModel, FleetModel, + GatewayComputeModel, + GatewayModel, ImportModel, InstanceModel, JobModel, @@ -78,7 +79,7 @@ get_instance_specific_mounts, resolve_provisioning_image, ) -from dstack._internal.server.services.gateways import get_or_add_gateway_connections +from dstack._internal.server.services.gateways import get_gateway_compute_models from dstack._internal.server.services.instances import ( get_instance_remote_connection_info, get_instance_ssh_private_keys, @@ -370,16 +371,10 @@ class _JobUpdateMap(ItemUpdateMap, total=False): skip_min_processing_interval: bool -@dataclass -class _RegisterReplicaResult: - gateway_target: Optional[events.Target] # None = no gateway - - @dataclass class _ProcessResult: job_update_map: _JobUpdateMap = field(default_factory=_JobUpdateMap) new_probe_models: list[ProbeModel] = field(default_factory=list) - replica_registration: Optional[_RegisterReplicaResult] = None # None = not registered yet @dataclass @@ -403,7 +398,9 @@ async def _load_process_context(item: JobRunningPipelineItem) -> Optional[_Proce return None if item.status == JobStatus.RUNNING: # RUNNING jobs don't access run.jobs — skip loading sibling jobs entirely. - run_model = await _fetch_run_model(session=session, run_id=job_model.run_id) + run_model = await _fetch_run_model( + session=session, run_id=job_model.run_id, include_gateway=True + ) run = run_model_to_run(run_model, include_sensitive=True, include_jobs=False) job = Job( job_spec=get_job_spec(job_model), @@ -496,6 +493,7 @@ async def _process_running_job(context: _ProcessContext) -> _ProcessResult: ) await _maybe_register_replica(context=context, result=result) await _check_gpu_utilization(context=context, result=result) + _check_service_registration(context=context, result=result) elif _server_access_enabled(context) and context.job_model.status == JobStatus.RUNNING: # Removing on PROVISIONING/PULLING iterations would reset the failure time # tracked by the pool, breaking retry_timed_out() @@ -634,6 +632,7 @@ async def _refetch_locked_job_model( .options( joinedload(JobModel.run).load_only(RunModel.id, RunModel.run_spec, RunModel.status) ) + .options(selectinload(JobModel.service_replica_registrations)) .execution_options(populate_existing=True) ) return res.unique().scalar_one_or_none() @@ -644,6 +643,7 @@ async def _fetch_run_model( run_id: uuid.UUID, replica_num: Optional[int] = None, run_spec: Optional[RunSpec] = None, + include_gateway: bool = False, ) -> RunModel: """Fetch run model with related project, user, repo, and fleet. @@ -668,6 +668,16 @@ async def _fetch_run_model( .options(joinedload(RunModel.repo)) .options(joinedload(RunModel.fleet).load_only(FleetModel.id, FleetModel.name)) ) + if include_gateway: + query = query.options( + joinedload(RunModel.gateway) + .selectinload(GatewayModel.gateway_computes) + .load_only(GatewayComputeModel.id, GatewayComputeModel.status), + ).options( + joinedload(RunModel.gateway) + .joinedload(GatewayModel.gateway_compute) + .load_only(GatewayComputeModel.id, GatewayComputeModel.status), + ) if replica_num is not None: assert run_spec is not None, "run_spec must be provided when replica_num is set" router_group = get_router_replica_group(run_spec) @@ -1088,16 +1098,6 @@ def _emit_result_events( old_ready=job_model.ready, new_ready=result.job_update_map.get("ready", job_model.ready), ) - if result.replica_registration is not None: - targets = [events.Target.from_model(job_model)] - if result.replica_registration.gateway_target is not None: - targets.append(result.replica_registration.gateway_target) - events.emit( - session, - "Service replica registered to receive requests", - actor=events.SystemActor(), - targets=targets, - ) def _wait_for_instance_provisioning_data( @@ -1206,92 +1206,7 @@ async def _maybe_register_replica( if not is_ready or _get_result_registered(context.job_model, result): return - ssh_head_proxy: Optional[SSHConnectionParams] = None - ssh_head_proxy_private_key: Optional[str] = None - instance = get_or_error(context.job_model.instance) - rci = get_instance_remote_connection_info(instance) - if rci is not None and rci.ssh_proxy is not None: - ssh_head_proxy = rci.ssh_proxy - ssh_head_proxy_keys = get_or_error(rci.ssh_proxy_keys) - ssh_head_proxy_private_key = ssh_head_proxy_keys[0].private - - try: - gateway_target = await _register_service_replica( - context=context, - result=result, - ssh_head_proxy=ssh_head_proxy, - ssh_head_proxy_private_key=ssh_head_proxy_private_key, - ) - except GatewayError as e: - logger.warning("%s: failed to register service replica: %s", fmt(context.job_model), e) - _terminate_job( - job_model=context.job_model, - job_update_map=result.job_update_map, - termination_reason=JobTerminationReason.GATEWAY_ERROR, - termination_reason_message="Failed to register service replica", - ) - return - result.job_update_map["registered"] = True - result.replica_registration = _RegisterReplicaResult(gateway_target=gateway_target) - - -async def _register_service_replica( - context: _ProcessContext, - result: _ProcessResult, - ssh_head_proxy: Optional[SSHConnectionParams], - ssh_head_proxy_private_key: Optional[str], -) -> Optional[events.Target]: - if context.run_model.gateway_id is None: - return None - async with get_session_ctx() as session: - gateway_model, connections = await get_or_add_gateway_connections( - session, context.run_model.gateway_id - ) - gateway_target = events.Target.from_model(gateway_model) - assert context.job_model.instance is not None - instance_project_ssh_private_key = None - if context.job_model.project_id != context.job_model.instance.project_id: - instance_project_ssh_private_key = context.job_model.instance.project.ssh_private_key - # JobRuntimeData might change on PULLING -> RUNNING path - # so we must update job_submission with the result value. - job_submission = context.job_submission.model_copy(deep=True) - job_submission.job_runtime_data = _get_result_job_runtime_data(context.job_model, result) - for conn in connections: - try: - logger.debug( - "%s: registering replica for service %s on gateway replica %s", - fmt(context.job_model), - context.run.id.hex, - conn.ip_address, - ) - async with conn.client() as gateway_client: - await gateway_client.register_replica( - run=context.run, - job_spec=context.job.job_spec, - job_submission=job_submission, - instance_project_ssh_private_key=instance_project_ssh_private_key, - ssh_head_proxy=ssh_head_proxy, - ssh_head_proxy_private_key=ssh_head_proxy_private_key, - ) - except (httpx.RequestError, SSHError) as e: - logger.debug("Gateway request failed", exc_info=True) - raise GatewayError(repr(e)) - except GatewayError as e: - if "already exists in service" in e.msg: - logger.warning( - ( - "%s: could not register replica in gateway %s: %s." - " NOTE: if you just updated dstack from pre-0.19.25 to 0.19.25+," - " expect to see this warning once for every running service replica" - ), - fmt(context.job_model), - conn.ip_address, - e.msg, - ) - else: - raise - return gateway_target async def _check_gpu_utilization( @@ -1329,6 +1244,49 @@ async def _check_gpu_utilization( logger.debug("%s: GPU utilization check: OK", fmt(context.job_model)) +def _job_gateway_registration_failed(gateway: GatewayModel | None, job_model: JobModel) -> bool: + if gateway is None: + return False + running_gateway_replica_ids = { + replica.id + for replica in get_gateway_compute_models(gateway) + if replica.status == GatewayReplicaStatus.RUNNING + } + if not running_gateway_replica_ids: + return False + registration_by_replica_id = { + r.gateway_replica_id: r for r in job_model.service_replica_registrations + } + for replica_id in running_gateway_replica_ids: + registration = registration_by_replica_id.get(replica_id) + if ( + registration is None + or registration.is_registered + or registration.register_attempt == 0 + ): + return False + return True + + +def _check_service_registration( + context: _ProcessContext, + result: _ProcessResult, +) -> None: + # `run_model.gateway` is only loaded for jobs that were already RUNNING + if context.job_model.status != JobStatus.RUNNING: + return + if _get_result_status(context.job_model, result) != JobStatus.RUNNING: + return + if _job_gateway_registration_failed(context.run_model.gateway, context.job_model): + logger.debug("%s: service registration check: terminating", fmt(context.job_model)) + _terminate_job( + job_model=context.job_model, + job_update_map=result.job_update_map, + termination_reason=JobTerminationReason.GATEWAY_ERROR, + termination_reason_message="Failed to register service replica with the gateway", + ) + + def _should_terminate_due_to_low_gpu_util( min_util: int, gpus_util: Iterable[Iterable[int]] ) -> bool: diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py index b78916ad07..0a530a714d 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py @@ -5,7 +5,6 @@ from datetime import datetime, timedelta from typing import Optional, Sequence, TypedDict -import httpx from sqlalchemy import and_, delete, or_, select, update from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import joinedload, load_only @@ -13,7 +12,7 @@ from dstack._internal.core.backends.base.backend import Backend from dstack._internal.core.backends.base.compute import ComputeWithVolumeSupport from dstack._internal.core.consts import DSTACK_SHIM_HTTP_PORT -from dstack._internal.core.errors import BackendError, GatewayError, SSHError +from dstack._internal.core.errors import BackendError from dstack._internal.core.models.instances import InstanceStatus, InstanceTerminationReason from dstack._internal.core.models.runs import ( JobProvisioningData, @@ -51,7 +50,6 @@ ) from dstack._internal.server.services import backends as backends_services from dstack._internal.server.services import events -from dstack._internal.server.services.gateways import get_or_add_gateway_connections from dstack._internal.server.services.instances import ( emit_instance_status_change_event, get_instance_ssh_private_keys, @@ -323,11 +321,6 @@ class _VolumeUpdateRow(TypedDict): last_job_processed_at: UpdateMapDateTime -@dataclass -class _UnregisterReplicaResult: - gateway_target: Optional[events.Target] # None = no gateway - - @dataclass class _ProcessResult: job_update_map: _JobUpdateMap = field(default_factory=_JobUpdateMap) @@ -337,9 +330,6 @@ class _ProcessResult: detached_volume_ids: set[uuid.UUID] = field(default_factory=set) unassign_event_message: Optional[str] = None graceful_stop_event_message: Optional[str] = None - replica_unregistration: Optional[_UnregisterReplicaResult] = ( - None # None = not unregistered yet - ) @dataclass @@ -618,17 +608,6 @@ async def _apply_process_result( targets=[events.Target.from_model(job_model)], ) - if result.replica_unregistration is not None: - targets = [events.Target.from_model(job_model)] - if result.replica_unregistration.gateway_target is not None: - targets.append(result.replica_unregistration.gateway_target) - events.emit( - session, - "Service replica unregistered from receiving requests", - actor=events.SystemActor(), - targets=targets, - ) - async def _unlock_related_instance( session: AsyncSession, @@ -665,7 +644,7 @@ async def _process_terminating_job( result = _ProcessResult(instance_update_map=instance_update_map) if instance_model is None: - await _unregister_replica_and_update_result(result=result, job_model=job_model) + await _unset_registered(result=result, job_model=job_model) result.job_update_map["status"] = _get_job_termination_status(job_model) return result @@ -676,7 +655,7 @@ async def _process_terminating_job( result.instance_update_map = None result.delete_instance = True result.job_update_map["instance_id"] = None - await _unregister_replica_and_update_result(result=result, job_model=job_model) + await _unset_registered(result=result, job_model=job_model) result.job_update_map["status"] = _get_job_termination_status(job_model) return result @@ -733,7 +712,7 @@ async def _process_terminating_job( f" Instance blocks: {busy_blocks}/{instance_model.total_blocks} busy" ) - await _unregister_replica_and_update_result(result=result, job_model=job_model) + await _unset_registered(result=result, job_model=job_model) if detach_result.all_detached: result.job_update_map["status"] = _get_job_termination_status(job_model) return result @@ -804,55 +783,9 @@ async def _detach_job_volumes( return volume_update_rows, detach_result -async def _unregister_replica_and_update_result( - result: _ProcessResult, job_model: JobModel -) -> None: - gateway_target = await _unregister_replica(job_model=job_model) +async def _unset_registered(result: _ProcessResult, job_model: JobModel) -> None: if job_model.registered: result.job_update_map["registered"] = False - result.replica_unregistration = _UnregisterReplicaResult(gateway_target=gateway_target) - - -async def _unregister_replica( - job_model: JobModel, -) -> Optional[events.Target]: - if not job_model.registered: - return None - gateway_target = None - run_model = job_model.run - if run_model.gateway_id is not None: - async with get_session_ctx() as session: - gateway, connections = await get_or_add_gateway_connections( - session, run_model.gateway_id - ) - gateway_target = events.Target.from_model(gateway) - for conn in connections: - try: - logger.debug( - "%s: unregistering replica from service %s on gateway replica %s", - fmt(job_model), - job_model.run_id.hex, - conn.ip_address, - ) - async with conn.client() as client: - await client.unregister_replica( - project=run_model.project.name, - run_name=run_model.run_name, - job_id=job_model.id, - ) - except GatewayError as e: - logger.warning( - "%s: unregistering replica from service on gateway replica %s: %s", - fmt(job_model), - conn.ip_address, - e, - ) - except (httpx.RequestError, SSHError) as e: - logger.debug("Gateway request failed", exc_info=True) - # FIXME: Unhandled exception raised. - # Handle and retry unregister with timeout. - raise GatewayError(repr(e)) - return gateway_target def _get_job_termination_status(job_model: JobModel) -> JobStatus: diff --git a/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py b/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py index 727535f5bc..587fad33a9 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py @@ -6,7 +6,7 @@ from sqlalchemy import and_, func, or_, select, update from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.orm import aliased, contains_eager, joinedload, load_only +from sqlalchemy.orm import aliased, contains_eager, joinedload, load_only, selectinload import dstack._internal.server.background.pipeline_tasks.runs.active as active import dstack._internal.server.background.pipeline_tasks.runs.pending as pending @@ -29,7 +29,14 @@ delete_superseded_no_capacity_job_submissions, ) from dstack._internal.server.db import get_db, get_session_ctx -from dstack._internal.server.models import InstanceModel, JobModel, ProjectModel, RunModel +from dstack._internal.server.models import ( + GatewayComputeModel, + GatewayModel, + InstanceModel, + JobModel, + ProjectModel, + RunModel, +) from dstack._internal.server.services import events from dstack._internal.server.services.gateways import get_combined_gateway_stats from dstack._internal.server.services.jobs import emit_job_status_change_event @@ -547,6 +554,22 @@ async def _refetch_locked_run_for_active( .joinedload(JobModel.instance) .load_only(InstanceModel.fleet_id), ) + .options( + contains_eager(RunModel.jobs, alias=job_alias).selectinload( + JobModel.service_replica_registrations + ), + ) + .options(selectinload(RunModel.service_registrations)) + .options( + joinedload(RunModel.gateway) + .selectinload(GatewayModel.gateway_computes) + .load_only(GatewayComputeModel.id, GatewayComputeModel.status), + ) + .options( + joinedload(RunModel.gateway) + .joinedload(GatewayModel.gateway_compute) + .load_only(GatewayComputeModel.id, GatewayComputeModel.status), + ) .execution_options(populate_existing=True) ) return res.unique().scalar_one_or_none() @@ -851,7 +874,6 @@ async def _apply_terminating_result( context: terminating.TerminatingContext, result: terminating.TerminatingResult, ) -> None: - run_model = context.run_model set_processed_update_map_fields(result.run_update_map) set_unlock_update_map_fields(result.run_update_map) @@ -889,17 +911,6 @@ async def _apply_terminating_result( if job_update_rows: await session.execute(update(JobModel), job_update_rows) - if result.service_unregistration is not None: - targets = [events.Target.from_model(run_model)] - if result.service_unregistration.gateway_target is not None: - targets.append(result.service_unregistration.gateway_target) - events.emit( - session, - result.service_unregistration.event_message, - actor=events.SystemActor(), - targets=targets, - ) - _emit_terminating_job_status_change_events( session=session, context=context, diff --git a/src/dstack/_internal/server/background/pipeline_tasks/runs/active.py b/src/dstack/_internal/server/background/pipeline_tasks/runs/active.py index 1ad52e05ff..c533ca8b34 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/runs/active.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/runs/active.py @@ -9,6 +9,7 @@ from dstack._internal.core.errors import ServerError from dstack._internal.core.models.configurations import ServiceConfiguration +from dstack._internal.core.models.gateways import GatewayReplicaStatus from dstack._internal.core.models.profiles import RetryEvent, StopCriteria from dstack._internal.core.models.runs import ( JobStatus, @@ -26,6 +27,7 @@ ) from dstack._internal.server.db import get_session_ctx from dstack._internal.server.models import JobModel, RunModel +from dstack._internal.server.services.gateways import get_gateway_compute_models from dstack._internal.server.services.jobs import ( get_job_spec, get_job_specs_from_run_spec, @@ -366,6 +368,28 @@ def _should_stop_on_master_done(run_spec: RunSpec, run_model: RunModel) -> bool: return False +def _gateway_registration_failed(run_model: RunModel) -> bool: + if run_model.gateway is None: + return False + running_replica_ids = { + replica.id + for replica in get_gateway_compute_models(run_model.gateway) + if replica.status == GatewayReplicaStatus.RUNNING + } + if not running_replica_ids: + return False + registration_by_replica_id = {r.gateway_replica_id: r for r in run_model.service_registrations} + for replica_id in running_replica_ids: + registration = registration_by_replica_id.get(replica_id) + if ( + registration is None + or registration.is_registered + or registration.register_attempt == 0 + ): + return False + return True + + def _get_active_run_transition( run_spec: RunSpec, run_model: RunModel, @@ -384,6 +408,12 @@ def _get_active_run_transition( termination_reason=termination_reason, ) + if _gateway_registration_failed(run_model): + return _ActiveRunTransition( + new_status=RunStatus.TERMINATING, + termination_reason=RunTerminationReason.GATEWAY_ERROR, + ) + if _should_stop_on_master_done(run_spec, run_model): return _ActiveRunTransition( new_status=RunStatus.TERMINATING, @@ -686,11 +716,11 @@ async def _build_rolling_deployment_maps( max_new = max(j.replica_num for j in new_jobs) next_replica_num = max(next_replica_num, max_new + 1) - # Scale down: terminate unready out-of-date + excess ready replicas - replicas_to_stop = state.unready_out_of_date_replica_count + # Scale down: terminate not-receiving-traffic out-of-date + excess receiving-traffic replicas + replicas_to_stop = state.not_receiving_traffic_out_of_date_replica_count replicas_to_stop += max( 0, - state.ready_non_terminating_replica_count - group_desired, + state.receiving_traffic_non_terminating_replica_count - group_desired, ) if replicas_to_stop > 0: scale_down_maps = _build_scale_down_job_update_maps( diff --git a/src/dstack/_internal/server/background/pipeline_tasks/runs/terminating.py b/src/dstack/_internal/server/background/pipeline_tasks/runs/terminating.py index c9a75e3c71..a55f7741e7 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/runs/terminating.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/runs/terminating.py @@ -3,9 +3,6 @@ from datetime import datetime from typing import Optional -import httpx - -from dstack._internal.core.errors import GatewayError, SSHError from dstack._internal.core.models.runs import ( JobStatus, JobTerminationReason, @@ -14,10 +11,6 @@ ) from dstack._internal.server import models from dstack._internal.server.background.pipeline_tasks.base import ItemUpdateMap -from dstack._internal.server.db import get_session_ctx -from dstack._internal.server.services import events -from dstack._internal.server.services.gateways import get_or_add_gateway_connections -from dstack._internal.server.services.logging import fmt from dstack._internal.server.services.runs import _get_next_triggered_at, get_run_spec from dstack._internal.utils.common import get_or_error from dstack._internal.utils.logging import get_logger @@ -39,12 +32,6 @@ class TerminatingRunJobUpdateMap(ItemUpdateMap, total=False): skip_min_processing_interval: bool -@dataclass -class ServiceUnregistration: - event_message: str - gateway_target: Optional[events.Target] - - @dataclass class TerminatingContext: run_model: models.RunModel @@ -55,7 +42,6 @@ class TerminatingContext: class TerminatingResult: run_update_map: TerminatingRunUpdateMap = field(default_factory=TerminatingRunUpdateMap) job_id_to_update_map: dict[uuid.UUID, TerminatingRunJobUpdateMap] = field(default_factory=dict) - service_unregistration: Optional[ServiceUnregistration] = None async def process_terminating_run(context: TerminatingContext) -> TerminatingResult: @@ -92,16 +78,8 @@ async def process_terminating_run(context: TerminatingContext) -> TerminatingRes if any(not job_model.status.is_finished() for job_model in run_model.jobs): return TerminatingResult() - service_unregistration = None - if run_model.service_spec is not None: - try: - service_unregistration = await _unregister_service(run_model) - except Exception as e: - logger.warning("%s: failed to unregister service: %s", fmt(run_model), repr(e)) - return TerminatingResult( run_update_map=_get_run_update_map(run_model), - service_unregistration=service_unregistration, ) @@ -141,45 +119,3 @@ def _get_run_update_map(run_model: models.RunModel) -> TerminatingRunUpdateMap: resubmission_attempt=0, ) return TerminatingRunUpdateMap(status=termination_reason.to_status()) - - -async def _unregister_service(run_model: models.RunModel) -> Optional[ServiceUnregistration]: - if run_model.gateway_id is None: # in-server proxy - return None - - async with get_session_ctx() as session: - gateway, connections = await get_or_add_gateway_connections(session, run_model.gateway_id) - gateway_target = events.Target.from_model(gateway) - - gateway_errors = [] - for conn in connections: - try: - logger.debug( - "%s: unregistering service on gateway replica %s", fmt(run_model), conn.ip_address - ) - async with conn.client() as client: - await client.unregister_service( - project=run_model.project.name, - run_name=run_model.run_name, - ) - except GatewayError as e: - # Ignore if the service is not registered on this replica. - logger.warning( - "%s: unregistering service on gateway replica %s: %s", - fmt(run_model), - conn.ip_address, - e, - ) - gateway_errors.append(str(e)) - except (httpx.RequestError, SSHError) as e: - logger.debug("Gateway request failed", exc_info=True) - raise GatewayError(repr(e)) - - if gateway_errors: - event_message = f"Gateway error when unregistering service: {'; '.join(gateway_errors)}" - else: - event_message = "Service unregistered from gateway" - return ServiceUnregistration( - event_message=event_message, - gateway_target=gateway_target, - ) diff --git a/src/dstack/_internal/server/migrations/versions/2026/08_13_0052_3d4f69210528_add_serviceregistrationmodel_and_.py b/src/dstack/_internal/server/migrations/versions/2026/08_13_0052_3d4f69210528_add_serviceregistrationmodel_and_.py new file mode 100644 index 0000000000..c317039ad6 --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/08_13_0052_3d4f69210528_add_serviceregistrationmodel_and_.py @@ -0,0 +1,123 @@ +"""Add ServiceRegistrationModel and ServiceReplicaRegistrationModel + +Revision ID: 3d4f69210528 +Revises: 72cfa56364ad +Create Date: 2026-08-13 00:52:02.969244+00:00 + +""" + +import sqlalchemy as sa +import sqlalchemy_utils +from alembic import op + +# revision identifiers, used by Alembic. +revision = "3d4f69210528" +down_revision = "72cfa56364ad" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "service_registrations", + sa.Column("id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False), + sa.Column("run_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False), + sa.Column( + "gateway_replica_id", + sqlalchemy_utils.types.uuid.UUIDType(binary=False), + nullable=False, + ), + sa.Column("is_registered", sa.Boolean(), nullable=False), + sa.Column("register_attempt", sa.Integer(), nullable=False), + sa.Column("register_status_message", sa.Text(), nullable=True), + sa.Column("unregister_attempt", sa.Integer(), nullable=False), + sa.Column("unregister_status_message", sa.Text(), nullable=True), + sa.ForeignKeyConstraint( + ["gateway_replica_id"], + ["gateway_computes.id"], + name=op.f("fk_service_registrations_gateway_replica_id_gateway_computes"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["run_id"], + ["runs.id"], + name=op.f("fk_service_registrations_run_id_runs"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_service_registrations")), + sa.UniqueConstraint( + "run_id", + "gateway_replica_id", + name="uq_service_registrations_run_id_gateway_replica_id", + ), + ) + with op.batch_alter_table("service_registrations", schema=None) as batch_op: + batch_op.create_index( + batch_op.f("ix_service_registrations_gateway_replica_id"), + ["gateway_replica_id"], + unique=False, + ) + batch_op.create_index( + batch_op.f("ix_service_registrations_run_id"), ["run_id"], unique=False + ) + + op.create_table( + "service_replica_registrations", + sa.Column("id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False), + sa.Column("job_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False), + sa.Column( + "gateway_replica_id", + sqlalchemy_utils.types.uuid.UUIDType(binary=False), + nullable=False, + ), + sa.Column("is_registered", sa.Boolean(), nullable=False), + sa.Column("register_attempt", sa.Integer(), nullable=False), + sa.Column("register_status_message", sa.Text(), nullable=True), + sa.Column("unregister_attempt", sa.Integer(), nullable=False), + sa.Column("unregister_status_message", sa.Text(), nullable=True), + sa.ForeignKeyConstraint( + ["gateway_replica_id"], + ["gateway_computes.id"], + name=op.f("fk_service_replica_registrations_gateway_replica_id_gateway_computes"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["job_id"], + ["jobs.id"], + name=op.f("fk_service_replica_registrations_job_id_jobs"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_service_replica_registrations")), + sa.UniqueConstraint( + "job_id", + "gateway_replica_id", + name="uq_service_replica_registrations_job_id_gateway_replica_id", + ), + ) + with op.batch_alter_table("service_replica_registrations", schema=None) as batch_op: + batch_op.create_index( + batch_op.f("ix_service_replica_registrations_gateway_replica_id"), + ["gateway_replica_id"], + unique=False, + ) + batch_op.create_index( + batch_op.f("ix_service_replica_registrations_job_id"), ["job_id"], unique=False + ) + + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("service_replica_registrations", schema=None) as batch_op: + batch_op.drop_index(batch_op.f("ix_service_replica_registrations_job_id")) + batch_op.drop_index(batch_op.f("ix_service_replica_registrations_gateway_replica_id")) + + op.drop_table("service_replica_registrations") + with op.batch_alter_table("service_registrations", schema=None) as batch_op: + batch_op.drop_index(batch_op.f("ix_service_registrations_run_id")) + batch_op.drop_index(batch_op.f("ix_service_registrations_gateway_replica_id")) + + op.drop_table("service_registrations") + # ### end Alembic commands ### diff --git a/src/dstack/_internal/server/models.py b/src/dstack/_internal/server/models.py index 23fb31ed71..70f28eefdd 100644 --- a/src/dstack/_internal/server/models.py +++ b/src/dstack/_internal/server/models.py @@ -463,6 +463,9 @@ class RunModel(PipelineModelMixin, BaseModel): service_router_worker_sync: Mapped[Optional["ServiceRouterWorkerSyncModel"]] = relationship( back_populates="run", uselist=False ) + service_registrations: Mapped[List["ServiceRegistrationModel"]] = relationship( + back_populates="run" + ) __table_args__ = ( Index("ix_submitted_at_id", submitted_at.desc(), id), @@ -583,7 +586,8 @@ class JobModel(PipelineModelMixin, BaseModel): Always `False` for non-service runs. """ registered: Mapped[bool] = mapped_column(Boolean, server_default=false()) - """Whether the replica is registered to receive service requests from dstack-proxy. + """Whether the replica should be registered to receive service requests from dstack-proxy. + Registration on the gateway can happen with a delay after this field is flipped to `True`. Always `False` for non-service runs or jobs that shouldn't be registered (e.g., non-router replicas for services with routers). """ @@ -595,6 +599,10 @@ class JobModel(PipelineModelMixin, BaseModel): """ image_pull_progress: Mapped[Optional[str]] = mapped_column(Text) + service_replica_registrations: Mapped[List["ServiceReplicaRegistrationModel"]] = relationship( + back_populates="job" + ) + __table_args__ = ( Index( "ix_jobs_pipeline_fetch_q", @@ -746,6 +754,13 @@ class GatewayComputeModel(PipelineModelMixin, BaseModel): deleted: Mapped[bool] = mapped_column(Boolean, server_default=false()) app_updated_at: Mapped[datetime] = mapped_column(NaiveDateTime, default=get_current_datetime) + service_registrations: Mapped[List["ServiceRegistrationModel"]] = relationship( + back_populates="gateway_replica" + ) + service_replica_registrations: Mapped[List["ServiceReplicaRegistrationModel"]] = relationship( + back_populates="gateway_replica" + ) + __table_args__ = ( Index( "ix_gateway_computes_pipeline_fetch_q", @@ -756,6 +771,74 @@ class GatewayComputeModel(PipelineModelMixin, BaseModel): ) +class ServiceRegistrationModel(BaseModel): + """Many-to-many association between services and gateway replicas""" + + __tablename__ = "service_registrations" + + id: Mapped[uuid.UUID] = mapped_column( + UUIDType(binary=False), primary_key=True, default=uuid.uuid4 + ) + run_id: Mapped[uuid.UUID] = mapped_column( + ForeignKey("runs.id", ondelete="CASCADE"), index=True + ) + run: Mapped["RunModel"] = relationship(back_populates="service_registrations") + gateway_replica_id: Mapped[uuid.UUID] = mapped_column( + ForeignKey("gateway_computes.id", ondelete="CASCADE"), index=True + ) + gateway_replica: Mapped["GatewayComputeModel"] = relationship( + back_populates="service_registrations" + ) + is_registered: Mapped[bool] = mapped_column(Boolean, default=False) + """Whether the service is successfully registered on this gateway replica""" + register_attempt: Mapped[int] = mapped_column(Integer, default=0) + register_status_message: Mapped[Optional[str]] = mapped_column(Text) + unregister_attempt: Mapped[int] = mapped_column(Integer, default=0) + unregister_status_message: Mapped[Optional[str]] = mapped_column(Text) + + __table_args__ = ( + UniqueConstraint( + "run_id", + "gateway_replica_id", + name="uq_service_registrations_run_id_gateway_replica_id", + ), + ) + + +class ServiceReplicaRegistrationModel(BaseModel): + """Many-to-many association between service replicas and gateway replicas""" + + __tablename__ = "service_replica_registrations" + + id: Mapped[uuid.UUID] = mapped_column( + UUIDType(binary=False), primary_key=True, default=uuid.uuid4 + ) + job_id: Mapped[uuid.UUID] = mapped_column( + ForeignKey("jobs.id", ondelete="CASCADE"), index=True + ) + job: Mapped["JobModel"] = relationship(back_populates="service_replica_registrations") + gateway_replica_id: Mapped[uuid.UUID] = mapped_column( + ForeignKey("gateway_computes.id", ondelete="CASCADE"), index=True + ) + gateway_replica: Mapped["GatewayComputeModel"] = relationship( + back_populates="service_replica_registrations" + ) + is_registered: Mapped[bool] = mapped_column(Boolean, default=False) + """Whether the service replica is successfully registered on this gateway replica""" + register_attempt: Mapped[int] = mapped_column(Integer, default=0) + register_status_message: Mapped[Optional[str]] = mapped_column(Text) + unregister_attempt: Mapped[int] = mapped_column(Integer, default=0) + unregister_status_message: Mapped[Optional[str]] = mapped_column(Text) + + __table_args__ = ( + UniqueConstraint( + "job_id", + "gateway_replica_id", + name="uq_service_replica_registrations_job_id_gateway_replica_id", + ), + ) + + # TODO: Drop after the release without pools class PoolModel(BaseModel): __tablename__ = "pools" diff --git a/src/dstack/_internal/server/services/gateways/client.py b/src/dstack/_internal/server/services/gateways/client.py index 10f2558327..ed6c8d78b9 100644 --- a/src/dstack/_internal/server/services/gateways/client.py +++ b/src/dstack/_internal/server/services/gateways/client.py @@ -8,9 +8,9 @@ from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT from dstack._internal.core.errors import GatewayError from dstack._internal.core.models.common import validate_json_extra_ignore -from dstack._internal.core.models.configurations import RateLimit +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, Run, get_service_port +from dstack._internal.core.models.runs import JobSpec, JobSubmission, get_service_port 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 @@ -85,17 +85,18 @@ async def unregister_service(self, project: str, run_name: str): async def register_replica( self, - run: Run, + project: str, + run_name: str, + configuration: ServiceConfiguration, job_spec: JobSpec, job_submission: JobSubmission, instance_project_ssh_private_key: Optional[str], ssh_head_proxy: Optional[SSHConnectionParams], ssh_head_proxy_private_key: Optional[str], ): - assert run.run_spec.configuration.type == "service" payload = { "job_id": job_submission.id.hex, - "app_port": get_service_port(job_spec, run.run_spec.configuration), + "app_port": get_service_port(job_spec, configuration), "ssh_head_proxy": ssh_head_proxy.model_dump() if ssh_head_proxy is not None else None, "ssh_head_proxy_private_key": ssh_head_proxy_private_key, } @@ -130,9 +131,7 @@ async def register_replica( } ) resp = await self._client.post( - self._url( - f"/api/registry/{run.project_name}/services/{run.run_spec.run_name}/replicas/register" - ), + self._url(f"/api/registry/{project}/services/{run_name}/replicas/register"), json=payload, ) if resp.status_code == 400: diff --git a/src/dstack/_internal/server/services/runs/replicas.py b/src/dstack/_internal/server/services/runs/replicas.py index 633b959992..ddbdf7c4f7 100644 --- a/src/dstack/_internal/server/services/runs/replicas.py +++ b/src/dstack/_internal/server/services/runs/replicas.py @@ -3,9 +3,11 @@ from typing import Dict, List, Optional, Tuple, Union from dstack._internal.core.models.configurations import ReplicaGroup, ServiceConfiguration +from dstack._internal.core.models.gateways import GatewayReplicaStatus from dstack._internal.core.models.routers import RouterType from dstack._internal.core.models.runs import JobStatus, JobTerminationReason, RunSpec from dstack._internal.server.models import JobModel, RunModel +from dstack._internal.server.services.gateways import get_gateway_compute_models from dstack._internal.server.services.jobs import ( get_job_provisioning_data, get_job_spec, @@ -19,8 +21,8 @@ class GroupRolloutState: inactive_replicas: List[Tuple[int, bool, int, List[JobModel]]] has_out_of_date_replicas: bool non_terminated_replica_count: int - unready_out_of_date_replica_count: int - ready_non_terminating_replica_count: int + not_receiving_traffic_out_of_date_replica_count: int + receiving_traffic_non_terminating_replica_count: int class RouterEnvStatus(str, Enum): @@ -77,11 +79,11 @@ def build_replica_lists( elif {JobStatus.PROVISIONING, JobStatus.PULLING} & statuses: # if there are any provisioning or pulling jobs, the replica is active and has the importance of 1 active_replicas.append((1, is_out_of_date, replica_num, replica_jobs)) - elif not is_replica_ready(replica_jobs): + elif not is_replica_receiving_traffic(run_model, replica_jobs): # all jobs are running, but not receiving traffic, the replica is active and has the importance of 2 active_replicas.append((2, is_out_of_date, replica_num, replica_jobs)) else: - # all jobs are running and ready, the replica is active and has the importance of 3 + # all jobs are running and receiving traffic, the replica is active and has the importance of 3 active_replicas.append((3, is_out_of_date, replica_num, replica_jobs)) # Sort by is_out_of_date (up-to-date first), importance (desc), and replica_num (asc) @@ -98,8 +100,8 @@ def get_group_rollout_state(run_model: RunModel, group: ReplicaGroup) -> GroupRo ) non_terminated_replica_nums = set() - unready_out_of_date_replica_count = 0 - ready_non_terminating_replica_count = 0 + not_receiving_traffic_out_of_date_replica_count = 0 + receiving_traffic_non_terminating_replica_count = 0 for _, jobs in group_jobs_by_replica_latest(run_model.jobs): if not job_belongs_to_group(jobs[0], group.name): @@ -108,26 +110,32 @@ def get_group_rollout_state(run_model: RunModel, group: ReplicaGroup) -> GroupRo if any(not j.status.is_finished() for j in jobs): non_terminated_replica_nums.add(jobs[0].replica_num) + receiving_traffic = is_replica_receiving_traffic(run_model, jobs) + if ( any(j.deployment_num < run_model.deployment_num for j in jobs) and any( j.status not in [JobStatus.TERMINATING] + JobStatus.finished_statuses() for j in jobs ) - and not is_replica_ready(jobs) + and not receiving_traffic ): - unready_out_of_date_replica_count += 1 + not_receiving_traffic_out_of_date_replica_count += 1 - if is_replica_ready(jobs) and all(j.status != JobStatus.TERMINATING for j in jobs): - ready_non_terminating_replica_count += 1 + if receiving_traffic and all(j.status != JobStatus.TERMINATING for j in jobs): + receiving_traffic_non_terminating_replica_count += 1 return GroupRolloutState( active_replicas=active_replicas, inactive_replicas=inactive_replicas, has_out_of_date_replicas=has_out_of_date_replicas(run_model, group_filter=group.name), non_terminated_replica_count=len(non_terminated_replica_nums), - unready_out_of_date_replica_count=unready_out_of_date_replica_count, - ready_non_terminating_replica_count=ready_non_terminating_replica_count, + not_receiving_traffic_out_of_date_replica_count=( + not_receiving_traffic_out_of_date_replica_count + ), + receiving_traffic_non_terminating_replica_count=( + receiving_traffic_non_terminating_replica_count + ), ) @@ -149,9 +157,33 @@ def has_out_of_date_replicas(run: RunModel, group_filter: Optional[str] = None) return False -def is_replica_ready(jobs: list[JobModel]) -> bool: +def is_replica_receiving_traffic(run_model: RunModel, jobs: list[JobModel]) -> bool: # Only job_num=0 is supposed to receive service requests - return jobs[0].ready + job = jobs[0] + if not job.ready: + # waiting for probes to pass + return False + if not job.registered: + # served by the service's router replica + return True + if run_model.gateway is None: + # served by the in-server proxy + return True + running_gateway_replica_ids = { + replica.id + for replica in get_gateway_compute_models(run_model.gateway) + if replica.status == GatewayReplicaStatus.RUNNING + } + if not running_gateway_replica_ids: + return False + registration_by_replica_id = { + r.gateway_replica_id: r for r in job.service_replica_registrations + } + for replica_id in running_gateway_replica_ids: + registration = registration_by_replica_id.get(replica_id) + if registration is None or not registration.is_registered: + return False + return True def get_router_replica_group(run_spec: RunSpec) -> Optional[ReplicaGroup]: diff --git a/src/dstack/_internal/server/services/services/__init__.py b/src/dstack/_internal/server/services/services/__init__.py index eaf878260e..4421b60556 100644 --- a/src/dstack/_internal/server/services/services/__init__.py +++ b/src/dstack/_internal/server/services/services/__init__.py @@ -2,16 +2,12 @@ Application logic related to `type: service` runs. """ -from functools import partial - -import httpx from sqlalchemy.ext.asyncio import AsyncSession from dstack._internal.core.errors import ( GatewayError, ResourceNotExistsError, ServerClientError, - SSHError, ) from dstack._internal.core.models.configurations import ( SERVICE_HTTPS_DEFAULT, @@ -21,18 +17,14 @@ from dstack._internal.core.models.gateways import GatewayConfiguration, GatewayStatus from dstack._internal.core.models.runs import RunSpec, ServiceModelSpec, ServiceSpec from dstack._internal.core.models.services import OpenAIChatModel -from dstack._internal.proxy.gateway.const import SERVICE_ALREADY_REGISTERED_ERROR_TEMPLATE from dstack._internal.server import settings from dstack._internal.server.models import GatewayModel, RunModel -from dstack._internal.server.services import events from dstack._internal.server.services.gateways import ( get_gateway_compute_models, get_gateway_configuration, - get_or_add_gateway_connections, get_project_default_gateway_model, get_project_gateway_model_by_reference, ) -from dstack._internal.server.services.logging import fmt from dstack._internal.server.services.services.options import get_service_options from dstack._internal.utils.common import interpolate_gateway_domain from dstack._internal.utils.logging import get_logger @@ -106,13 +98,6 @@ async def _register_service_in_gateway( gateway_configuration = get_gateway_configuration(gateway) - has_replica_group_router = any( - g.router is not None for g in run_spec.configuration.replica_groups - ) - - configure_service_https = _should_configure_service_https_on_gateway( - run_spec, gateway_configuration - ) show_service_https = _should_show_service_https(run_spec, gateway_configuration) service_protocol = "https" if show_service_https else "http" @@ -131,7 +116,7 @@ async def _register_service_in_gateway( "Cannot run HTTPS service on gateway with no SSL certificates configured" ) - gateway_https = _get_gateway_https(gateway_configuration) + gateway_https = get_gateway_https(gateway_configuration) gateway_protocol = "https" if gateway_https else "http" wildcard_domain = gateway.wildcard_domain.lstrip("*.") if gateway.wildcard_domain else None @@ -152,65 +137,6 @@ async def _register_service_in_gateway( service_url=service_url, model_url=model_url, ) - - domain = service_spec.get_domain() - assert domain is not None - - _, connections = await get_or_add_gateway_connections(session, gateway.id) - for conn in connections: - try: - logger.debug("%s: registering service as %s", fmt(run_model), service_spec.url) - async with conn.client() as client: - do_register = partial( - client.register_service, - project=run_model.project.name, - run_id=run_model.id, - run_name=run_model.run_name, - domain=domain, - service_https=configure_service_https, - gateway_https=gateway_https, - auth=run_spec.configuration.auth, - client_max_body_size=settings.DEFAULT_SERVICE_CLIENT_MAX_BODY_SIZE, - options=service_spec.options, - rate_limits=run_spec.configuration.rate_limits, - ssh_private_key=run_model.project.ssh_private_key, - has_router_replica=has_replica_group_router, - ) - try: - await do_register() - except GatewayError as e: - if e.msg == SERVICE_ALREADY_REGISTERED_ERROR_TEMPLATE.format( - ref=f"{run_model.project.name}/{run_model.run_name}" - ): - # Happens if there was a communication issue with the gateway when last (un)registering - logger.warning( - "Service %s/%s is dangling on gateway replica %s, unregistering and re-registering", - run_model.project.name, - run_model.run_name, - conn.ip_address, - ) - await client.unregister_service( - project=run_model.project.name, - run_name=run_model.run_name, - ) - await do_register() - else: - raise - except SSHError: - raise ServerClientError("Gateway tunnel is not working") - except httpx.RequestError as e: - logger.debug("Gateway request failed", exc_info=True) - raise GatewayError(f"Gateway is not working: {e!r}") - - events.emit( - session, - "Service registered in gateway", - actor=events.SystemActor(), - targets=[ - events.Target.from_model(run_model), - events.Target.from_model(gateway), - ], - ) return service_spec @@ -265,7 +191,7 @@ def _get_service_spec( return service_spec -def _should_configure_service_https_on_gateway( +def should_configure_service_https_on_gateway( run_spec: RunSpec, configuration: GatewayConfiguration ) -> bool: """ @@ -304,7 +230,7 @@ def _should_show_service_https(run_spec: RunSpec, configuration: GatewayConfigur return https -def _get_gateway_https(configuration: GatewayConfiguration) -> bool: +def get_gateway_https(configuration: GatewayConfiguration) -> bool: if configuration.certificate is not None and configuration.certificate.type == "acm": return False if configuration.certificate is not None and configuration.certificate.type == "lets-encrypt": 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 e3af1c699a..678af6370b 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 @@ -1,34 +1,53 @@ import asyncio import uuid from datetime import datetime, timedelta, timezone -from unittest.mock import MagicMock, Mock, patch +from typing import Optional +from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch import pytest +from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from dstack._internal.core.backends.base.compute import ComputeWithGatewaySupport -from dstack._internal.core.errors import BackendError +from dstack._internal.core.errors import BackendError, GatewayError +from dstack._internal.core.models.configurations import ServiceConfiguration from dstack._internal.core.models.gateways import ( ACMGatewayCertificate, GatewayProvisioningData, GatewayReplicaStatus, GatewayStatus, ) +from dstack._internal.core.models.instances import InstanceStatus +from dstack._internal.core.models.runs import JobStatus, RunStatus, ServiceSpec +from dstack._internal.proxy.gateway.schemas.services import ServiceListItem, ServiceListReplicaItem from dstack._internal.server.background.pipeline_tasks.gateway_replicas import ( GatewayReplicaFetcher, GatewayReplicaPipeline, GatewayReplicaPipelineItem, GatewayReplicaWorker, ) -from dstack._internal.server.models import GatewayComputeModel +from dstack._internal.server.models import ( + GatewayComputeModel, + ServiceRegistrationModel, + ServiceReplicaRegistrationModel, +) from dstack._internal.server.testing.common import ( AsyncContextManager, ComputeMockSpec, create_backend, + create_fleet, create_gateway, create_gateway_compute, + create_instance, + create_job, create_project, + create_repo, + create_run, + create_user, get_gateway_compute_configuration, + get_job_provisioning_data, + get_run_spec, + list_events, ) from dstack._internal.utils.common import get_current_datetime @@ -152,23 +171,28 @@ async def test_fetch_selects_eligible_replicas_and_sets_lock_fields( items = await fetcher.fetch(limit=10) - assert {item.id for item in items} == {submitted.id, provisioning.id, terminating.id} + assert {item.id for item in items} == { + submitted.id, + provisioning.id, + terminating.id, + running.id, + } assert {(item.id, item.status) for item in items} == { (submitted.id, GatewayReplicaStatus.SUBMITTED), (provisioning.id, GatewayReplicaStatus.PROVISIONING), (terminating.id, GatewayReplicaStatus.TERMINATING), + (running.id, GatewayReplicaStatus.RUNNING), } for compute in [submitted, provisioning, terminating, running, terminated, recent, locked]: await session.refresh(compute) - fetched = [submitted, provisioning, terminating] + fetched = [submitted, provisioning, terminating, running] assert all(c.lock_owner == GatewayReplicaPipeline.__name__ for c in fetched) assert all(c.lock_expires_at is not None for c in fetched) assert all(c.lock_token is not None for c in fetched) assert len({c.lock_token for c in fetched}) == 1 - assert running.lock_owner is None assert terminated.lock_owner is None assert recent.lock_owner is None assert locked.lock_owner == "OtherPipeline" @@ -246,13 +270,15 @@ async def test_fetch_includes_running_replica_with_hard_deleted_gateway( assert items[0].status == GatewayReplicaStatus.RUNNING @pytest.mark.parametrize("legacy_compute", [False, True]) - async def test_fetch_excludes_running_replica_with_healthy_gateway( + async def test_fetch_includes_running_replica_with_healthy_gateway( self, test_db, session: AsyncSession, fetcher: GatewayReplicaFetcher, legacy_compute: bool, ): + # Healthy running replicas are still fetched periodically so the worker + # can run gateway state sync (see _process_running_item). project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) gateway = await create_gateway( @@ -270,7 +296,7 @@ async def test_fetch_excludes_running_replica_with_healthy_gateway( ) gateway.gateway_compute_id = compute.id else: - await create_gateway_compute( + compute = await create_gateway_compute( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, @@ -280,7 +306,9 @@ async def test_fetch_excludes_running_replica_with_healthy_gateway( items = await fetcher.fetch(limit=10) - assert len(items) == 0 + assert len(items) == 1 + assert items[0].id == compute.id + assert items[0].status == GatewayReplicaStatus.RUNNING async def test_fetch_includes_running_replica_marked_for_scale_in( self, @@ -667,265 +695,1523 @@ async def test_running_to_terminating_when_scaled_in( assert compute.status_message == "Scaled in" +def _get_client_mock(mock_gateway_connection: AsyncMock) -> AsyncMock: + return mock_gateway_connection.return_value.client.return_value.__aenter__.return_value + + @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) -class TestGatewayReplicaWorkerProvisioning: - @pytest.mark.parametrize("legacy_compute", [False, True]) - @pytest.mark.parametrize("populate_configuration", [True, False]) - async def test_provisioning_to_running( +class TestGatewayReplicaWorkerRunningStateSync: + """ + Covers `_process_running_item`'s gateway state sync: registering/unregistering + services and replicas so the gateway matches expected DB state, and recording + the outcome in `ServiceRegistrationModel`/`ServiceReplicaRegistrationModel`. + """ + + pytestmark = pytest.mark.usefixtures("image_config_mock") + + async def _create_service_run_and_job( + self, + session: AsyncSession, + project, + repo, + user, + gateway, + run_name: str, + run_status: RunStatus = RunStatus.RUNNING, + job_status: JobStatus = JobStatus.RUNNING, + job_registered: bool = True, + replica_num: int = 0, + service_url: Optional[str] = None, + ): + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_name=run_name, + status=run_status, + run_spec=get_run_spec( + run_name=run_name, + repo_id=repo.name, + configuration=ServiceConfiguration(port=80, image="ubuntu"), + ), + gateway=gateway, + ) + run.service_spec = ServiceSpec( + url=service_url or f"https://{run_name}.example.com" + ).model_dump_json() + await session.commit() + fleet = await create_fleet(session=session, project=project) + instance = await create_instance( + session=session, project=project, status=InstanceStatus.BUSY, fleet=fleet + ) + job = await create_job( + session=session, + run=run, + status=job_status, + job_provisioning_data=get_job_provisioning_data(dockerized=True), + instance=instance, + instance_assigned=True, + registered=job_registered, + ready=job_registered, + replica_num=replica_num, + ) + return run, job + + async def test_registers_new_service_and_replica( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker, - legacy_compute: bool, - populate_configuration: bool, + mock_gateway_connection: AsyncMock, ): project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) backend = await create_backend(session=session, project_id=project.id) gateway = await create_gateway( session=session, project_id=project.id, backend_id=backend.id, - status=GatewayStatus.PROVISIONING, - populate_configuration=populate_configuration, + status=GatewayStatus.RUNNING, + ) + compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + backend_id=backend.id, + status=GatewayReplicaStatus.RUNNING, + ssh_private_key="replica-private-key", + ) + run, job = await self._create_service_run_and_job( + session, project, repo, user, gateway, run_name="test-service" ) - if legacy_compute: - compute = await create_gateway_compute( - session=session, - status=GatewayReplicaStatus.PROVISIONING, - populate_configuration=populate_configuration, - ) - gateway.gateway_compute_id = compute.id - else: - compute = await create_gateway_compute( - session=session, - gateway_id=gateway.id, - status=GatewayReplicaStatus.PROVISIONING, - populate_configuration=populate_configuration, - ) _lock_compute(compute) await session.commit() - with patch( - "dstack._internal.server.services.gateways.gateway_connections_pool.get_or_add" - ) as pool_add: - pool_add.return_value = MagicMock() - pool_add.return_value.client.return_value = MagicMock(AsyncContextManager()) - await worker.process(_compute_to_pipeline_item(compute)) - pool_add.assert_called_once() + client_mock = _get_client_mock(mock_gateway_connection) + client_mock.list_services.return_value = [] - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.RUNNING - assert compute.active is True + await worker.process(_compute_to_pipeline_item(compute)) - async def test_provisioning_to_running_registers_with_load_balancer( - self, test_db, session: AsyncSession, worker: GatewayReplicaWorker + mock_gateway_connection.assert_called_once_with( + hostname=compute.ip_address, id_rsa="replica-private-key" + ) + client_mock.register_service.assert_called_once_with( + project=project.name, + run_id=run.id, + run_name="test-service", + domain="test-service.example.com", + service_https=ANY, + gateway_https=ANY, + auth=ANY, + client_max_body_size=ANY, + options={}, + rate_limits=[], + ssh_private_key=project.ssh_private_key, + has_router_replica=False, + ) + client_mock.register_replica.assert_called_once_with( + project=project.name, + run_name="test-service", + configuration=ANY, + job_spec=ANY, + job_submission=ANY, + instance_project_ssh_private_key=None, + ssh_head_proxy=None, + ssh_head_proxy_private_key=None, + ) + assert client_mock.register_replica.call_args.kwargs["job_submission"].id == job.id + client_mock.unregister_service.assert_not_called() + client_mock.unregister_replica.assert_not_called() + + service_registration = ( + await session.execute( + select(ServiceRegistrationModel).where( + ServiceRegistrationModel.run_id == run.id, + ServiceRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ).scalar_one() + assert service_registration.is_registered is True + assert service_registration.register_attempt == 0 + assert service_registration.register_status_message is None + + replica_registration = ( + await session.execute( + select(ServiceReplicaRegistrationModel).where( + ServiceReplicaRegistrationModel.job_id == job.id, + ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ).scalar_one() + assert replica_registration.is_registered is True + + events = await list_events(session) + assert {e.message for e in events} == { + f"Service registered on gateway replica {compute.replica_num}", + f"Service replica registered on gateway replica {compute.replica_num}", + } + + async def test_unregisters_dangling_service_and_stale_replica( + self, + test_db, + session: AsyncSession, + worker: GatewayReplicaWorker, + mock_gateway_connection: AsyncMock, ): project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) backend = await create_backend(session=session, project_id=project.id) gateway = await create_gateway( session=session, project_id=project.id, backend_id=backend.id, - status=GatewayStatus.PROVISIONING, - certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), - hostname="gateway-lb.example.com", - backend_data="lb-backend-data", + status=GatewayStatus.RUNNING, ) compute = await create_gateway_compute( session=session, gateway_id=gateway.id, backend_id=backend.id, - status=GatewayReplicaStatus.PROVISIONING, + status=GatewayReplicaStatus.RUNNING, + ) + # A live, still-expected service with one live replica and one stale + # replica (of the same run) that the gateway still thinks is registered. + run1, job1 = await self._create_service_run_and_job( + session, project, repo, user, gateway, run_name="live-service" + ) + stale_job = await create_job( + session=session, + run=run1, + status=JobStatus.TERMINATED, + registered=False, + replica_num=1, + ) + # A finished run whose service is still (erroneously) registered on the gateway. + run2, _ = await self._create_service_run_and_job( + session, + project, + repo, + user, + gateway, + run_name="dangling-service", + run_status=RunStatus.TERMINATED, + job_status=JobStatus.TERMINATED, + job_registered=False, ) _lock_compute(compute) await session.commit() - with ( - patch( - "dstack._internal.server.services.gateways.gateway_connections_pool.get_or_add" - ) as pool_add, - patch( - "dstack._internal.server.services.backends.get_project_backends_with_models" - ) as get_backends_mock, - ): - pool_add.return_value = MagicMock() - pool_add.return_value.client.return_value = MagicMock(AsyncContextManager()) - backend_mock = Mock() - backend_mock.compute.return_value = Mock(spec=ComputeMockSpec) - get_backends_mock.return_value = [(backend, backend_mock)] + client_mock = _get_client_mock(mock_gateway_connection) + client_mock.list_services.return_value = [ + ServiceListItem( + id=run1.id.hex, + project_name=project.name, + run_name="live-service", + replicas=[ + ServiceListReplicaItem(id=job1.id.hex), + ServiceListReplicaItem(id=stale_job.id.hex), + ], + ), + ServiceListItem( + id=run2.id.hex, + project_name=project.name, + run_name="dangling-service", + replicas=[], + ), + ] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_compute_to_pipeline_item(compute)) - register_mock = ( - backend_mock.compute.return_value.register_gateway_replica_with_load_balancer + client_mock.register_service.assert_not_called() + client_mock.register_replica.assert_not_called() + client_mock.unregister_service.assert_called_once_with( + project=project.name, run_name="dangling-service" + ) + client_mock.unregister_replica.assert_called_once_with( + project=project.name, run_name="live-service", job_id=stale_job.id + ) + + run1_registration = ( + ( + await session.execute( + select(ServiceRegistrationModel).where( + ServiceRegistrationModel.gateway_replica_id == compute.id, + ) + ) ) - register_mock.assert_called_once() - call_args = register_mock.call_args.args - assert call_args[0] == compute.instance_id - assert call_args[1].gateway_name == gateway.name - assert call_args[2] == "lb-backend-data" + .scalars() + .all() + ) + assert {r.run_id for r in run1_registration} == {run1.id} + assert all(r.is_registered for r in run1_registration) + + replica_registrations = ( + ( + await session.execute( + select(ServiceReplicaRegistrationModel).where( + ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ) + .scalars() + .all() + ) + assert {r.job_id for r in replica_registrations} == {job1.id} + assert all(r.is_registered for r in replica_registrations) - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.RUNNING - assert compute.active is True + events = await list_events(session) + assert {e.message for e in events} == { + f"Service unregistered from gateway replica {compute.replica_num}", + f"Service replica unregistered from gateway replica {compute.replica_num}", + } - async def test_provisioning_skips_load_balancer_registration_without_hostname( - self, test_db, session: AsyncSession, worker: GatewayReplicaWorker + async def test_deletes_registration_models_for_unregistered_service_and_replica( + self, + test_db, + session: AsyncSession, + worker: GatewayReplicaWorker, + mock_gateway_connection: AsyncMock, ): project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) backend = await create_backend(session=session, project_id=project.id) gateway = await create_gateway( session=session, project_id=project.id, backend_id=backend.id, - status=GatewayStatus.PROVISIONING, + status=GatewayStatus.RUNNING, ) compute = await create_gateway_compute( session=session, gateway_id=gateway.id, backend_id=backend.id, - status=GatewayReplicaStatus.PROVISIONING, + status=GatewayReplicaStatus.RUNNING, + ) + # A live, still-expected service with one live replica and one stale + # replica (of the same run) that the gateway still thinks is registered. + run1, job1 = await self._create_service_run_and_job( + session, project, repo, user, gateway, run_name="live-service" + ) + stale_job = await create_job( + session=session, + run=run1, + status=JobStatus.TERMINATED, + registered=False, + replica_num=1, + ) + # A finished run whose service is still (erroneously) registered on the gateway. + run2, _ = await self._create_service_run_and_job( + session, + project, + repo, + user, + gateway, + run_name="dangling-service", + run_status=RunStatus.TERMINATED, + job_status=JobStatus.TERMINATED, + job_registered=False, + ) + # Pre-existing registration records for everything currently on the + # gateway, including the ones about to be unregistered. + live_service_registration = ServiceRegistrationModel( + run_id=run1.id, gateway_replica_id=compute.id, is_registered=True + ) + live_replica_registration = ServiceReplicaRegistrationModel( + job_id=job1.id, gateway_replica_id=compute.id, is_registered=True + ) + stale_replica_registration = ServiceReplicaRegistrationModel( + job_id=stale_job.id, gateway_replica_id=compute.id, is_registered=True + ) + dangling_service_registration = ServiceRegistrationModel( + run_id=run2.id, gateway_replica_id=compute.id, is_registered=True + ) + session.add_all( + [ + live_service_registration, + live_replica_registration, + stale_replica_registration, + dangling_service_registration, + ] ) _lock_compute(compute) await session.commit() - with ( - patch( - "dstack._internal.server.services.gateways.gateway_connections_pool.get_or_add" - ) as pool_add, - patch( - "dstack._internal.server.services.backends.get_project_backends_with_models" - ) as get_backends_mock, - ): - pool_add.return_value = MagicMock() - pool_add.return_value.client.return_value = MagicMock(AsyncContextManager()) + client_mock = _get_client_mock(mock_gateway_connection) + client_mock.list_services.return_value = [ + ServiceListItem( + id=run1.id.hex, + project_name=project.name, + run_name="live-service", + replicas=[ + ServiceListReplicaItem(id=job1.id.hex), + ServiceListReplicaItem(id=stale_job.id.hex), + ], + ), + ServiceListItem( + id=run2.id.hex, + project_name=project.name, + run_name="dangling-service", + replicas=[], + ), + ] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_compute_to_pipeline_item(compute)) - get_backends_mock.assert_not_called() + client_mock.unregister_service.assert_called_once_with( + project=project.name, run_name="dangling-service" + ) + client_mock.unregister_replica.assert_called_once_with( + project=project.name, run_name="live-service", job_id=stale_job.id + ) - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.RUNNING - assert compute.active is True + remaining_service_registrations = ( + ( + await session.execute( + select(ServiceRegistrationModel).where( + ServiceRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ) + .scalars() + .all() + ) + assert {r.run_id for r in remaining_service_registrations} == {run1.id} + assert {r.id for r in remaining_service_registrations} == {live_service_registration.id} + + remaining_replica_registrations = ( + ( + await session.execute( + select(ServiceReplicaRegistrationModel).where( + ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ) + .scalars() + .all() + ) + assert {r.job_id for r in remaining_replica_registrations} == {job1.id} + assert {r.id for r in remaining_replica_registrations} == {live_replica_registration.id} + + # The registration records for the now-unregistered service and replica + # are gone + assert ( + await session.execute( + select(ServiceRegistrationModel).where( + ServiceRegistrationModel.id == dangling_service_registration.id, + ) + ) + ).scalar_one_or_none() is None + assert ( + await session.execute( + select(ServiceReplicaRegistrationModel).where( + ServiceReplicaRegistrationModel.id == stale_replica_registration.id, + ) + ) + ).scalar_one_or_none() is None - async def test_provisioning_to_terminating_when_load_balancer_registration_fails( - self, test_db, session: AsyncSession, worker: GatewayReplicaWorker + async def test_no_gateway_calls_when_state_already_in_sync( + self, + test_db, + session: AsyncSession, + worker: GatewayReplicaWorker, + mock_gateway_connection: AsyncMock, ): project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) backend = await create_backend(session=session, project_id=project.id) gateway = await create_gateway( session=session, project_id=project.id, backend_id=backend.id, - status=GatewayStatus.PROVISIONING, - certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), - hostname="gateway-lb.example.com", - backend_data="lb-backend-data", + status=GatewayStatus.RUNNING, ) compute = await create_gateway_compute( session=session, gateway_id=gateway.id, backend_id=backend.id, - status=GatewayReplicaStatus.PROVISIONING, + status=GatewayReplicaStatus.RUNNING, + ) + run, job = await self._create_service_run_and_job( + session, project, repo, user, gateway, run_name="synced-service" ) _lock_compute(compute) await session.commit() - with ( - patch( - "dstack._internal.server.services.gateways.gateway_connections_pool.get_or_add" - ) as pool_add, - patch( - "dstack._internal.server.services.backends.get_project_backends_with_models" - ) as get_backends_mock, - ): - pool_add.return_value = MagicMock() - pool_add.return_value.client.return_value = MagicMock(AsyncContextManager()) - backend_mock = Mock() - backend_mock.compute.return_value = Mock(spec=ComputeMockSpec) - backend_mock.compute.return_value.register_gateway_replica_with_load_balancer.side_effect = Exception( - "boom" - ) - get_backends_mock.return_value = [(backend, backend_mock)] + client_mock = _get_client_mock(mock_gateway_connection) + client_mock.list_services.return_value = [ + ServiceListItem( + id=run.id.hex, + project_name=project.name, + run_name="synced-service", + replicas=[ServiceListReplicaItem(id=job.id.hex)], + ), + ] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_compute_to_pipeline_item(compute)) - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATING - assert compute.active is False - assert compute.status_message == "Error registering with load balancer" + client_mock.register_service.assert_not_called() + client_mock.register_replica.assert_not_called() + client_mock.unregister_service.assert_not_called() + client_mock.unregister_replica.assert_not_called() + client_mock.set_service_id.assert_not_called() + + # Registration records are still (re)created to reflect the confirmed state, + # even though no gateway calls were needed this tick. + service_registration = ( + await session.execute( + select(ServiceRegistrationModel).where( + ServiceRegistrationModel.run_id == run.id, + ServiceRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ).scalar_one() + assert service_registration.is_registered is True + replica_registration = ( + await session.execute( + select(ServiceReplicaRegistrationModel).where( + ServiceReplicaRegistrationModel.job_id == job.id, + ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ).scalar_one() + assert replica_registration.is_registered is True - async def test_provisioning_to_terminating_when_backend_does_not_support_load_balancer( - self, test_db, session: AsyncSession, worker: GatewayReplicaWorker + # No events since nothing actually changed from the gateway's perspective. + assert not await list_events(session) + + async def test_recovers_legacy_service_id_by_matching_replica( + self, + test_db, + session: AsyncSession, + worker: GatewayReplicaWorker, + mock_gateway_connection: AsyncMock, ): project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) backend = await create_backend(session=session, project_id=project.id) gateway = await create_gateway( session=session, project_id=project.id, backend_id=backend.id, - status=GatewayStatus.PROVISIONING, - certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), - hostname="gateway-lb.example.com", - backend_data="lb-backend-data", + status=GatewayStatus.RUNNING, ) compute = await create_gateway_compute( session=session, gateway_id=gateway.id, backend_id=backend.id, - status=GatewayReplicaStatus.PROVISIONING, + status=GatewayReplicaStatus.RUNNING, + ) + run, job = await self._create_service_run_and_job( + session, project, repo, user, gateway, run_name="legacy-service" ) _lock_compute(compute) await session.commit() - with ( - patch( - "dstack._internal.server.services.gateways.gateway_connections_pool.get_or_add" - ) as pool_add, - patch( - "dstack._internal.server.services.backends.get_project_backends_with_models" - ) as get_backends_mock, - ): - pool_add.return_value = MagicMock() - pool_add.return_value.client.return_value = MagicMock(AsyncContextManager()) - backend_mock = Mock() - backend_mock.compute.return_value = Mock(spec=ComputeWithGatewaySupport) - get_backends_mock.return_value = [(backend, backend_mock)] + client_mock = _get_client_mock(mock_gateway_connection) + client_mock.list_services.return_value = [ + # Pre-0.21.0 gateways report services without an id. + ServiceListItem( + id=None, + project_name=project.name, + run_name="legacy-service", + replicas=[ServiceListReplicaItem(id=job.id.hex)], + ), + ] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_compute_to_pipeline_item(compute)) - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATING - assert compute.active is False - assert compute.status_message == "Backend does not support load balancer operations" + client_mock.set_service_id.assert_called_once_with( + project=project.name, run_name="legacy-service", run_id=run.id + ) + client_mock.register_service.assert_not_called() + client_mock.register_replica.assert_not_called() + client_mock.unregister_service.assert_not_called() + client_mock.unregister_replica.assert_not_called() - async def test_provisioning_waits_for_pending_acm_gateway_migration( - self, test_db, session: AsyncSession, worker: GatewayReplicaWorker + async def test_unregisters_and_reregisters_legacy_service_without_id_and_replicas( + self, + test_db, + session: AsyncSession, + worker: GatewayReplicaWorker, + mock_gateway_connection: AsyncMock, ): project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) backend = await create_backend(session=session, project_id=project.id) gateway = await create_gateway( session=session, project_id=project.id, backend_id=backend.id, status=GatewayStatus.RUNNING, - certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), - hostname=None, # migration not yet performed ) compute = await create_gateway_compute( session=session, gateway_id=gateway.id, backend_id=backend.id, - status=GatewayReplicaStatus.PROVISIONING, - hostname_deprecated_readonly="legacy-lb.example.com", + status=GatewayReplicaStatus.RUNNING, + ) + run, job = await self._create_service_run_and_job( + session, project, repo, user, gateway, run_name="legacy-service" ) _lock_compute(compute) - original_last_processed_at = compute.last_processed_at await session.commit() - with patch( - "dstack._internal.server.services.gateways.gateway_connections_pool.get_or_add" + client_mock = _get_client_mock(mock_gateway_connection) + client_mock.list_services.return_value = [ + # A legacy entry with no ID and no replicas: there's nothing to match + # it against, so it cannot be told apart from genuine garbage and must + # be unregistered. The real, still-expected run gets registered fresh. + ServiceListItem( + id=None, + project_name=project.name, + run_name="legacy-service", + replicas=[], + ), + ] + + await worker.process(_compute_to_pipeline_item(compute)) + + client_mock.set_service_id.assert_not_called() + client_mock.unregister_service.assert_called_once_with( + project=project.name, run_name="legacy-service" + ) + client_mock.register_service.assert_called_once_with( + project=project.name, + run_id=run.id, + run_name="legacy-service", + domain=ANY, + service_https=ANY, + gateway_https=ANY, + auth=ANY, + client_max_body_size=ANY, + options={}, + rate_limits=[], + ssh_private_key=project.ssh_private_key, + has_router_replica=False, + ) + client_mock.register_replica.assert_called_once_with( + project=project.name, + run_name="legacy-service", + configuration=ANY, + job_spec=ANY, + job_submission=ANY, + instance_project_ssh_private_key=None, + ssh_head_proxy=None, + ssh_head_proxy_private_key=None, + ) + + service_registration = ( + await session.execute( + select(ServiceRegistrationModel).where( + ServiceRegistrationModel.run_id == run.id, + ServiceRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ).scalar_one() + assert service_registration.is_registered is True + replica_registration = ( + await session.execute( + select(ServiceReplicaRegistrationModel).where( + ServiceReplicaRegistrationModel.job_id == job.id, + ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ).scalar_one() + assert replica_registration.is_registered is True + + # No unregistration event: the dangling legacy entry has no ID, so it can't + # be tied to a run for event targeting (only the fresh registration is). + events = await list_events(session) + assert {e.message for e in events} == { + f"Service registered on gateway replica {compute.replica_num}", + f"Service replica registered on gateway replica {compute.replica_num}", + } + + async def test_does_nothing_when_in_sync_and_registrations_already_exist( + self, + test_db, + session: AsyncSession, + worker: GatewayReplicaWorker, + mock_gateway_connection: AsyncMock, + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + backend_id=backend.id, + status=GatewayReplicaStatus.RUNNING, + ) + run, job = await self._create_service_run_and_job( + session, project, repo, user, gateway, run_name="synced-service" + ) + existing_service_registration = ServiceRegistrationModel( + run_id=run.id, + gateway_replica_id=compute.id, + is_registered=True, + register_attempt=0, + ) + existing_replica_registration = ServiceReplicaRegistrationModel( + job_id=job.id, + gateway_replica_id=compute.id, + is_registered=True, + register_attempt=0, + ) + session.add_all([existing_service_registration, existing_replica_registration]) + _lock_compute(compute) + await session.commit() + + client_mock = _get_client_mock(mock_gateway_connection) + client_mock.list_services.return_value = [ + ServiceListItem( + id=run.id.hex, + project_name=project.name, + run_name="synced-service", + replicas=[ServiceListReplicaItem(id=job.id.hex)], + ), + ] + + await worker.process(_compute_to_pipeline_item(compute)) + + client_mock.register_service.assert_not_called() + client_mock.register_replica.assert_not_called() + client_mock.unregister_service.assert_not_called() + client_mock.unregister_replica.assert_not_called() + client_mock.set_service_id.assert_not_called() + + service_registrations = ( + ( + await session.execute( + select(ServiceRegistrationModel).where( + ServiceRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ) + .scalars() + .all() + ) + assert len(service_registrations) == 1 + assert service_registrations[0].id == existing_service_registration.id + assert service_registrations[0].is_registered is True + assert service_registrations[0].register_attempt == 0 + + replica_registrations = ( + ( + await session.execute( + select(ServiceReplicaRegistrationModel).where( + ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ) + .scalars() + .all() + ) + assert len(replica_registrations) == 1 + assert replica_registrations[0].id == existing_replica_registration.id + assert replica_registrations[0].is_registered is True + assert replica_registrations[0].register_attempt == 0 + + assert not await list_events(session) + + async def test_reconciles_out_of_sync_registration_models( + self, + test_db, + session: AsyncSession, + worker: GatewayReplicaWorker, + mock_gateway_connection: AsyncMock, + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + backend_id=backend.id, + status=GatewayReplicaStatus.RUNNING, + ) + run, job = await self._create_service_run_and_job( + session, project, repo, user, gateway, run_name="drifted-service" + ) + # Local bookkeeping incorrectly thinks these failed to register, but the + # gateway actually already has them registered (e.g. the server crashed + # right after a successful registration, before it could record that). + stale_service_registration = ServiceRegistrationModel( + run_id=run.id, + gateway_replica_id=compute.id, + is_registered=False, + register_attempt=3, + register_status_message="stale error", + ) + stale_replica_registration = ServiceReplicaRegistrationModel( + job_id=job.id, + gateway_replica_id=compute.id, + is_registered=False, + register_attempt=2, + register_status_message="stale replica error", + ) + session.add_all([stale_service_registration, stale_replica_registration]) + _lock_compute(compute) + await session.commit() + + client_mock = _get_client_mock(mock_gateway_connection) + client_mock.list_services.return_value = [ + ServiceListItem( + id=run.id.hex, + project_name=project.name, + run_name="drifted-service", + replicas=[ServiceListReplicaItem(id=job.id.hex)], + ), + ] + + await worker.process(_compute_to_pipeline_item(compute)) + + # The gateway already reports it registered, so nothing needs to be + # (re)registered - only the local bookkeeping needs correcting. + client_mock.register_service.assert_not_called() + client_mock.register_replica.assert_not_called() + client_mock.unregister_service.assert_not_called() + client_mock.unregister_replica.assert_not_called() + + await session.refresh(stale_service_registration) + await session.refresh(stale_replica_registration) + assert stale_service_registration.is_registered is True + assert stale_service_registration.register_attempt == 0 + assert stale_service_registration.register_status_message is None + assert stale_replica_registration.is_registered is True + assert stale_replica_registration.register_attempt == 0 + assert stale_replica_registration.register_status_message is None + + # No events: as far as the gateway is concerned nothing changed, we only + # corrected stale local state. + assert not await list_events(session) + + @pytest.mark.parametrize( + ( + "make_error", + "expected_service_register_status_message", + "expected_replica_register_status_message", + ), + [ + pytest.param( + GatewayError, + "boom service", + "boom replica", + id="gateway_error", + ), + pytest.param( + Exception, + "Unexpected error", + "Unexpected error", + id="unexpected_error", + ), + ], + ) + async def test_propagates_registration_errors_and_increments_register_attempt( + self, + test_db, + session: AsyncSession, + worker: GatewayReplicaWorker, + mock_gateway_connection: AsyncMock, + make_error: type[Exception], + expected_service_register_status_message: str, + expected_replica_register_status_message: str, + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + backend_id=backend.id, + status=GatewayReplicaStatus.RUNNING, + ) + run_replica_fails, job_replica_fails = await self._create_service_run_and_job( + session, project, repo, user, gateway, run_name="replica-fails" + ) + run_service_fails, job_service_fails = await self._create_service_run_and_job( + session, project, repo, user, gateway, run_name="service-fails" + ) + # Simulate two earlier failed attempts to register this service. + existing_registration = ServiceRegistrationModel( + run_id=run_service_fails.id, + gateway_replica_id=compute.id, + is_registered=False, + register_attempt=2, + register_status_message="earlier error", + ) + session.add(existing_registration) + _lock_compute(compute) + await session.commit() + + client_mock = _get_client_mock(mock_gateway_connection) + client_mock.list_services.return_value = [] + + async def register_service_side_effect(**kwargs): + if kwargs["run_name"] == "service-fails": + raise make_error("boom service") + + async def register_replica_side_effect(**kwargs): + if kwargs["run_name"] == "replica-fails": + raise make_error("boom replica") + + client_mock.register_service.side_effect = register_service_side_effect + client_mock.register_replica.side_effect = register_replica_side_effect + + await worker.process(_compute_to_pipeline_item(compute)) + + client_mock.unregister_service.assert_not_called() + client_mock.unregister_replica.assert_not_called() + assert client_mock.register_service.call_count == 2 + # register_replica is only attempted for the run whose service registration + # succeeded. + client_mock.register_replica.assert_called_once_with( + project=project.name, + run_name="replica-fails", + configuration=ANY, + job_spec=ANY, + job_submission=ANY, + instance_project_ssh_private_key=None, + ssh_head_proxy=None, + ssh_head_proxy_private_key=None, + ) + + replica_fails_registration = ( + await session.execute( + select(ServiceRegistrationModel).where( + ServiceRegistrationModel.run_id == run_replica_fails.id, + ServiceRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ).scalar_one() + assert replica_fails_registration.is_registered is True + assert replica_fails_registration.register_attempt == 0 + assert replica_fails_registration.register_status_message is None + + replica_fails_replica_registration = ( + await session.execute( + select(ServiceReplicaRegistrationModel).where( + ServiceReplicaRegistrationModel.job_id == job_replica_fails.id, + ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ).scalar_one() + assert replica_fails_replica_registration.is_registered is False + assert replica_fails_replica_registration.register_attempt == 1 + assert ( + replica_fails_replica_registration.register_status_message + == expected_replica_register_status_message + ) + + await session.refresh(existing_registration) + assert existing_registration.is_registered is False + assert existing_registration.register_attempt == 3 + assert ( + existing_registration.register_status_message + == expected_service_register_status_message + ) + + # A service's replicas are never attempted once its own registration failed. + service_fails_replica_registration = ( + await session.execute( + select(ServiceReplicaRegistrationModel).where( + ServiceReplicaRegistrationModel.job_id == job_service_fails.id, + ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ).scalar_one_or_none() + assert service_fails_replica_registration is None + + events = await list_events(session) + assert {e.message for e in events} == { + f"Service registered on gateway replica {compute.replica_num}", + ( + f"Encountered service registration error on gateway replica " + f"{compute.replica_num}: {expected_service_register_status_message}" + ), + ( + f"Encountered service replica registration error on gateway replica " + f"{compute.replica_num}: {expected_replica_register_status_message}" + ), + } + + async def test_does_not_emit_duplicate_registration_error_event_for_unchanged_error( + self, + test_db, + session: AsyncSession, + worker: GatewayReplicaWorker, + mock_gateway_connection: AsyncMock, + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + backend_id=backend.id, + status=GatewayReplicaStatus.RUNNING, + ) + run, job = await self._create_service_run_and_job( + session, project, repo, user, gateway, run_name="replica-fails" + ) + + client_mock = _get_client_mock(mock_gateway_connection) + client_mock.list_services.return_value = [] + client_mock.register_replica.side_effect = GatewayError("boom replica") + + # First tick: the replica registration fails, recording the error and + # emitting one error event. + _lock_compute(compute) + await session.commit() + await worker.process(_compute_to_pipeline_item(compute)) + + replica_registration = ( + await session.execute( + select(ServiceReplicaRegistrationModel).where( + ServiceReplicaRegistrationModel.job_id == job.id, + ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ).scalar_one() + assert replica_registration.register_attempt == 1 + assert replica_registration.register_status_message == "boom replica" + + error_message = ( + f"Encountered service replica registration error on gateway replica " + f"{compute.replica_num}: boom replica" + ) + events_after_first_tick = await list_events(session) + assert [e.message for e in events_after_first_tick].count(error_message) == 1 + + # Second tick: the replica registration fails again with the exact same + # error. `register_attempt` keeps incrementing, but no duplicate event is + # emitted since nothing new happened from the user's perspective. + await session.refresh(compute) + _lock_compute(compute) + await session.commit() + await worker.process(_compute_to_pipeline_item(compute)) + + await session.refresh(replica_registration) + assert replica_registration.register_attempt == 2 + assert replica_registration.register_status_message == "boom replica" + + events_after_second_tick = await list_events(session) + assert [e.message for e in events_after_second_tick].count(error_message) == 1 + + @pytest.mark.parametrize( + ("make_error", "expected_service_status_message", "expected_replica_status_message"), + [ + pytest.param( + GatewayError, + "boom service", + "boom replica", + id="gateway_error", + ), + pytest.param( + Exception, + "Unexpected error", + "Unexpected error", + id="unexpected_error", + ), + ], + ) + async def test_propagates_unregistration_errors_and_increments_unregister_attempt( + self, + test_db, + session: AsyncSession, + worker: GatewayReplicaWorker, + mock_gateway_connection: AsyncMock, + make_error: type[Exception], + expected_service_status_message: str, + expected_replica_status_message: str, + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + backend_id=backend.id, + status=GatewayReplicaStatus.RUNNING, + ) + # A live, still-expected service with one live replica and one stale + # replica (of the same run) that the gateway still thinks is registered. + run, job = await self._create_service_run_and_job( + session, project, repo, user, gateway, run_name="live-service" + ) + stale_job = await create_job( + session=session, + run=run, + status=JobStatus.TERMINATED, + registered=False, + replica_num=1, + ) + # A finished run whose service is still (erroneously) registered on the gateway. + dangling_run, _ = await self._create_service_run_and_job( + session, + project, + repo, + user, + gateway, + run_name="dangling-service", + run_status=RunStatus.TERMINATED, + job_status=JobStatus.TERMINATED, + job_registered=False, + ) + _lock_compute(compute) + await session.commit() + + client_mock = _get_client_mock(mock_gateway_connection) + client_mock.list_services.return_value = [ + ServiceListItem( + id=run.id.hex, + project_name=project.name, + run_name="live-service", + replicas=[ + ServiceListReplicaItem(id=job.id.hex), + ServiceListReplicaItem(id=stale_job.id.hex), + ], + ), + ServiceListItem( + id=dangling_run.id.hex, + project_name=project.name, + run_name="dangling-service", + replicas=[], + ), + ] + client_mock.unregister_service.side_effect = make_error("boom service") + client_mock.unregister_replica.side_effect = make_error("boom replica") + + await worker.process(_compute_to_pipeline_item(compute)) + + # Both remain registered as far as the gateway is concerned, since we + # failed to remove them - only the unregister bookkeeping changes. + dangling_registration = ( + await session.execute( + select(ServiceRegistrationModel).where( + ServiceRegistrationModel.run_id == dangling_run.id, + ServiceRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ).scalar_one() + assert dangling_registration.is_registered is True + assert dangling_registration.unregister_attempt == 1 + assert dangling_registration.unregister_status_message == expected_service_status_message + + stale_replica_registration = ( + await session.execute( + select(ServiceReplicaRegistrationModel).where( + ServiceReplicaRegistrationModel.job_id == stale_job.id, + ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ).scalar_one() + assert stale_replica_registration.is_registered is True + assert stale_replica_registration.unregister_attempt == 1 + assert ( + stale_replica_registration.unregister_status_message == expected_replica_status_message + ) + + events = await list_events(session) + assert {e.message for e in events} == { + ( + f"Encountered service unregistration error on gateway replica " + f"{compute.replica_num}: {expected_service_status_message}" + ), + ( + f"Encountered service replica unregistration error on gateway replica " + f"{compute.replica_num}: {expected_replica_status_message}" + ), + } + + async def test_does_not_emit_duplicate_unregistration_error_event_for_unchanged_error( + self, + test_db, + session: AsyncSession, + worker: GatewayReplicaWorker, + mock_gateway_connection: AsyncMock, + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + backend_id=backend.id, + status=GatewayReplicaStatus.RUNNING, + ) + run, job = await self._create_service_run_and_job( + session, project, repo, user, gateway, run_name="live-service" + ) + stale_job = await create_job( + session=session, + run=run, + status=JobStatus.TERMINATED, + registered=False, + replica_num=1, + ) + _lock_compute(compute) + await session.commit() + + client_mock = _get_client_mock(mock_gateway_connection) + client_mock.list_services.return_value = [ + ServiceListItem( + id=run.id.hex, + project_name=project.name, + run_name="live-service", + replicas=[ + ServiceListReplicaItem(id=job.id.hex), + ServiceListReplicaItem(id=stale_job.id.hex), + ], + ), + ] + client_mock.unregister_replica.side_effect = GatewayError("boom replica") + + # First tick: unregistering the stale replica fails, recording the error + # and emitting one error event. + await worker.process(_compute_to_pipeline_item(compute)) + + replica_registration = ( + await session.execute( + select(ServiceReplicaRegistrationModel).where( + ServiceReplicaRegistrationModel.job_id == stale_job.id, + ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ) + ) + ).scalar_one() + assert replica_registration.unregister_attempt == 1 + assert replica_registration.unregister_status_message == "boom replica" + + error_message = ( + f"Encountered service replica unregistration error on gateway replica " + f"{compute.replica_num}: boom replica" + ) + events_after_first_tick = await list_events(session) + assert [e.message for e in events_after_first_tick].count(error_message) == 1 + + # Second tick: unregistering fails again with the exact same error. + # `unregister_attempt` keeps incrementing, but no duplicate event is + # emitted since nothing new happened from the user's perspective. + await session.refresh(compute) + _lock_compute(compute) + await session.commit() + await worker.process(_compute_to_pipeline_item(compute)) + + await session.refresh(replica_registration) + assert replica_registration.unregister_attempt == 2 + assert replica_registration.unregister_status_message == "boom replica" + + events_after_second_tick = await list_events(session) + assert [e.message for e in events_after_second_tick].count(error_message) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) +class TestGatewayReplicaWorkerProvisioning: + @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("populate_configuration", [True, False]) + async def test_provisioning_to_running( + self, + test_db, + session: AsyncSession, + worker: GatewayReplicaWorker, + legacy_compute: bool, + populate_configuration: bool, + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.PROVISIONING, + populate_configuration=populate_configuration, + ) + if legacy_compute: + compute = await create_gateway_compute( + session=session, + status=GatewayReplicaStatus.PROVISIONING, + populate_configuration=populate_configuration, + ) + gateway.gateway_compute_id = compute.id + else: + compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + status=GatewayReplicaStatus.PROVISIONING, + populate_configuration=populate_configuration, + ) + _lock_compute(compute) + await session.commit() + + with patch( + "dstack._internal.server.services.gateways.gateway_connections_pool.get_or_add" + ) as pool_add: + pool_add.return_value = MagicMock() + pool_add.return_value.client.return_value = MagicMock(AsyncContextManager()) + await worker.process(_compute_to_pipeline_item(compute)) + pool_add.assert_called_once() + + await session.refresh(compute) + assert compute.status == GatewayReplicaStatus.RUNNING + assert compute.active is True + + async def test_provisioning_to_running_registers_with_load_balancer( + self, test_db, session: AsyncSession, worker: GatewayReplicaWorker + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.PROVISIONING, + certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), + hostname="gateway-lb.example.com", + backend_data="lb-backend-data", + ) + compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + backend_id=backend.id, + status=GatewayReplicaStatus.PROVISIONING, + ) + _lock_compute(compute) + await session.commit() + + with ( + patch( + "dstack._internal.server.services.gateways.gateway_connections_pool.get_or_add" + ) as pool_add, + patch( + "dstack._internal.server.services.backends.get_project_backends_with_models" + ) as get_backends_mock, + ): + pool_add.return_value = MagicMock() + pool_add.return_value.client.return_value = MagicMock(AsyncContextManager()) + backend_mock = Mock() + backend_mock.compute.return_value = Mock(spec=ComputeMockSpec) + get_backends_mock.return_value = [(backend, backend_mock)] + + await worker.process(_compute_to_pipeline_item(compute)) + + register_mock = ( + backend_mock.compute.return_value.register_gateway_replica_with_load_balancer + ) + register_mock.assert_called_once() + call_args = register_mock.call_args.args + assert call_args[0] == compute.instance_id + assert call_args[1].gateway_name == gateway.name + assert call_args[2] == "lb-backend-data" + + await session.refresh(compute) + assert compute.status == GatewayReplicaStatus.RUNNING + assert compute.active is True + + async def test_provisioning_skips_load_balancer_registration_without_hostname( + self, test_db, session: AsyncSession, worker: GatewayReplicaWorker + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.PROVISIONING, + ) + compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + backend_id=backend.id, + status=GatewayReplicaStatus.PROVISIONING, + ) + _lock_compute(compute) + await session.commit() + + with ( + patch( + "dstack._internal.server.services.gateways.gateway_connections_pool.get_or_add" + ) as pool_add, + patch( + "dstack._internal.server.services.backends.get_project_backends_with_models" + ) as get_backends_mock, + ): + pool_add.return_value = MagicMock() + pool_add.return_value.client.return_value = MagicMock(AsyncContextManager()) + + await worker.process(_compute_to_pipeline_item(compute)) + + get_backends_mock.assert_not_called() + + await session.refresh(compute) + assert compute.status == GatewayReplicaStatus.RUNNING + assert compute.active is True + + async def test_provisioning_to_terminating_when_load_balancer_registration_fails( + self, test_db, session: AsyncSession, worker: GatewayReplicaWorker + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.PROVISIONING, + certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), + hostname="gateway-lb.example.com", + backend_data="lb-backend-data", + ) + compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + backend_id=backend.id, + status=GatewayReplicaStatus.PROVISIONING, + ) + _lock_compute(compute) + await session.commit() + + with ( + patch( + "dstack._internal.server.services.gateways.gateway_connections_pool.get_or_add" + ) as pool_add, + patch( + "dstack._internal.server.services.backends.get_project_backends_with_models" + ) as get_backends_mock, + ): + pool_add.return_value = MagicMock() + pool_add.return_value.client.return_value = MagicMock(AsyncContextManager()) + backend_mock = Mock() + backend_mock.compute.return_value = Mock(spec=ComputeMockSpec) + backend_mock.compute.return_value.register_gateway_replica_with_load_balancer.side_effect = Exception( + "boom" + ) + get_backends_mock.return_value = [(backend, backend_mock)] + + await worker.process(_compute_to_pipeline_item(compute)) + + await session.refresh(compute) + assert compute.status == GatewayReplicaStatus.TERMINATING + assert compute.active is False + assert compute.status_message == "Error registering with load balancer" + + async def test_provisioning_to_terminating_when_backend_does_not_support_load_balancer( + self, test_db, session: AsyncSession, worker: GatewayReplicaWorker + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.PROVISIONING, + certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), + hostname="gateway-lb.example.com", + backend_data="lb-backend-data", + ) + compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + backend_id=backend.id, + status=GatewayReplicaStatus.PROVISIONING, + ) + _lock_compute(compute) + await session.commit() + + with ( + patch( + "dstack._internal.server.services.gateways.gateway_connections_pool.get_or_add" + ) as pool_add, + patch( + "dstack._internal.server.services.backends.get_project_backends_with_models" + ) as get_backends_mock, + ): + pool_add.return_value = MagicMock() + pool_add.return_value.client.return_value = MagicMock(AsyncContextManager()) + backend_mock = Mock() + backend_mock.compute.return_value = Mock(spec=ComputeWithGatewaySupport) + get_backends_mock.return_value = [(backend, backend_mock)] + + await worker.process(_compute_to_pipeline_item(compute)) + + await session.refresh(compute) + assert compute.status == GatewayReplicaStatus.TERMINATING + assert compute.active is False + assert compute.status_message == "Backend does not support load balancer operations" + + async def test_provisioning_waits_for_pending_acm_gateway_migration( + self, test_db, session: AsyncSession, worker: GatewayReplicaWorker + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), + hostname=None, # migration not yet performed + ) + compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + backend_id=backend.id, + status=GatewayReplicaStatus.PROVISIONING, + hostname_deprecated_readonly="legacy-lb.example.com", + ) + _lock_compute(compute) + original_last_processed_at = compute.last_processed_at + await session.commit() + + with patch( + "dstack._internal.server.services.gateways.gateway_connections_pool.get_or_add" ) as pool_add: await worker.process(_compute_to_pipeline_item(compute)) pool_add.assert_not_called() @@ -1136,6 +2422,88 @@ async def test_terminating_to_terminated( assert compute.active is False assert compute.deleted is True + async def test_terminating_to_terminated_deletes_only_own_registration_records( + self, test_db, session: AsyncSession, worker: GatewayReplicaWorker + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + compute_to_terminate = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + backend_id=backend.id, + status=GatewayReplicaStatus.TERMINATING, + active=False, + replica_num=0, + ) + other_compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + backend_id=backend.id, + status=GatewayReplicaStatus.RUNNING, + ip_address="2.2.2.2", + instance_id="i-eeeeeeeeee", + replica_num=1, + ) + run = await create_run(session=session, project=project, repo=repo, user=user) + job = await create_job(session=session, run=run) + terminated_service_registration = ServiceRegistrationModel( + run_id=run.id, gateway_replica_id=compute_to_terminate.id, is_registered=True + ) + terminated_replica_registration = ServiceReplicaRegistrationModel( + job_id=job.id, gateway_replica_id=compute_to_terminate.id, is_registered=True + ) + other_service_registration = ServiceRegistrationModel( + run_id=run.id, gateway_replica_id=other_compute.id, is_registered=True + ) + other_replica_registration = ServiceReplicaRegistrationModel( + job_id=job.id, gateway_replica_id=other_compute.id, is_registered=True + ) + session.add_all( + [ + terminated_service_registration, + terminated_replica_registration, + other_service_registration, + other_replica_registration, + ] + ) + _lock_compute(compute_to_terminate) + await session.commit() + + with ( + patch( + "dstack._internal.server.services.backends.get_project_backends_with_models" + ) as get_backends_mock, + patch( + "dstack._internal.server.background.pipeline_tasks.gateway_replicas.gateway_connections_pool.remove" + ), + ): + backend_mock = Mock() + backend_mock.compute.return_value = Mock(spec=ComputeMockSpec) + get_backends_mock.return_value = [(backend, backend_mock)] + + await worker.process(_compute_to_pipeline_item(compute_to_terminate)) + + await session.refresh(compute_to_terminate) + assert compute_to_terminate.status == GatewayReplicaStatus.TERMINATED + assert compute_to_terminate.deleted is True + + remaining_service_registration = ( + (await session.execute(select(ServiceRegistrationModel))).scalars().one() + ) + assert remaining_service_registration.gateway_replica_id == other_compute.id + remaining_replica_registration = ( + (await session.execute(select(ServiceReplicaRegistrationModel))).scalars().one() + ) + assert remaining_replica_registration.gateway_replica_id == other_compute.id + async def test_terminating_deregisters_from_load_balancer_before_terminating( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker ): diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py index 6762237937..19dfd4219a 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py @@ -5,7 +5,7 @@ from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Optional -from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest from freezegun import freeze_time @@ -24,7 +24,7 @@ TaskConfiguration, ) from dstack._internal.core.models.duration import Duration -from dstack._internal.core.models.gateways import GatewayStatus +from dstack._internal.core.models.gateways import GatewayReplicaStatus, GatewayStatus from dstack._internal.core.models.instances import InstanceStatus from dstack._internal.core.models.profiles import StartupOrder, UtilizationPolicy from dstack._internal.core.models.runs import ( @@ -54,7 +54,7 @@ _SubmitJobToRunnerResult, ) from dstack._internal.server.background.pipeline_tasks.runs import RunPipeline -from dstack._internal.server.models import JobModel, ProbeModel +from dstack._internal.server.models import JobModel, ProbeModel, ServiceReplicaRegistrationModel from dstack._internal.server.schemas.runner import ( HealthcheckResponse, JobInfoResponse, @@ -1865,6 +1865,353 @@ async def test_gpu_utilization( assert job.termination_reason is None assert job.termination_reason_message is None + async def test_terminates_job_on_gateway_registration_failure( + self, + test_db, + session: AsyncSession, + worker: JobRunningWorker, + ) -> None: + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + gateway_compute = await create_gateway_compute(session=session, gateway_id=gateway.id) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + status=RunStatus.RUNNING, + run_spec=get_run_spec( + run_name="test-run", + repo_id=repo.name, + configuration=ServiceConfiguration(port=80, image="ubuntu"), + ), + gateway=gateway, + ) + instance = await create_instance( + session=session, + project=project, + status=InstanceStatus.BUSY, + ) + job = await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + job_provisioning_data=get_job_provisioning_data(), + instance=instance, + instance_assigned=True, + registered=True, + ready=True, + ) + session.add( + ServiceReplicaRegistrationModel( + job_id=job.id, + gateway_replica_id=gateway_compute.id, + is_registered=False, + register_attempt=2, + register_status_message="Connection refused", + ) + ) + await session.commit() + + with ( + patch("dstack._internal.server.services.runner.pool.SSHTunnel") as ssh_tunnel_cls, + patch( + "dstack._internal.server.services.runner.client.RunnerClient.from_address" + ) as runner_client_cls, + ): + runner_client_mock = runner_client_cls.return_value + runner_client_mock.pull.return_value = PullResponse( + job_states=[], + job_logs=[], + runner_logs=[], + last_updated=0, + no_connections_secs=0, + ) + await _process_job(session, worker, job) + ssh_tunnel_cls.assert_called_once() + runner_client_mock.pull.assert_called_once() + + await session.refresh(job) + assert job.status == JobStatus.TERMINATING + assert job.termination_reason == JobTerminationReason.GATEWAY_ERROR + + async def test_does_not_terminate_job_when_registered_with_at_least_one_gateway_replica( + self, + test_db, + session: AsyncSession, + worker: JobRunningWorker, + ) -> None: + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + gateway_compute_1 = await create_gateway_compute( + session=session, gateway_id=gateway.id, replica_num=0 + ) + gateway_compute_2 = await create_gateway_compute( + session=session, gateway_id=gateway.id, replica_num=1 + ) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + status=RunStatus.RUNNING, + run_spec=get_run_spec( + run_name="test-run", + repo_id=repo.name, + configuration=ServiceConfiguration(port=80, image="ubuntu"), + ), + gateway=gateway, + ) + instance = await create_instance( + session=session, + project=project, + status=InstanceStatus.BUSY, + ) + job = await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + job_provisioning_data=get_job_provisioning_data(), + instance=instance, + instance_assigned=True, + registered=True, + ready=True, + ) + session.add( + ServiceReplicaRegistrationModel( + job_id=job.id, + gateway_replica_id=gateway_compute_1.id, + is_registered=False, + register_attempt=2, + register_status_message="Connection refused", + ) + ) + session.add( + ServiceReplicaRegistrationModel( + job_id=job.id, + gateway_replica_id=gateway_compute_2.id, + is_registered=True, + register_attempt=0, + ) + ) + await session.commit() + + with ( + patch("dstack._internal.server.services.runner.pool.SSHTunnel") as ssh_tunnel_cls, + patch( + "dstack._internal.server.services.runner.client.RunnerClient.from_address" + ) as runner_client_cls, + ): + runner_client_mock = runner_client_cls.return_value + runner_client_mock.pull.return_value = PullResponse( + job_states=[], + job_logs=[], + runner_logs=[], + last_updated=0, + no_connections_secs=0, + ) + await _process_job(session, worker, job) + ssh_tunnel_cls.assert_called_once() + runner_client_mock.pull.assert_called_once() + + await session.refresh(job) + assert job.status == JobStatus.RUNNING + assert job.termination_reason is None + + async def test_does_not_terminate_job_when_gateway_replica_has_not_attempted_registration( + self, + test_db, + session: AsyncSession, + worker: JobRunningWorker, + ) -> None: + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + gateway_compute_1 = await create_gateway_compute( + session=session, gateway_id=gateway.id, replica_num=0 + ) + # Second running replica has not attempted registration yet (e.g. just came up). + await create_gateway_compute(session=session, gateway_id=gateway.id, replica_num=1) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + status=RunStatus.RUNNING, + run_spec=get_run_spec( + run_name="test-run", + repo_id=repo.name, + configuration=ServiceConfiguration(port=80, image="ubuntu"), + ), + gateway=gateway, + ) + instance = await create_instance( + session=session, + project=project, + status=InstanceStatus.BUSY, + ) + job = await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + job_provisioning_data=get_job_provisioning_data(), + instance=instance, + instance_assigned=True, + registered=True, + ready=True, + ) + session.add( + ServiceReplicaRegistrationModel( + job_id=job.id, + gateway_replica_id=gateway_compute_1.id, + is_registered=False, + register_attempt=2, + register_status_message="Connection refused", + ) + ) + await session.commit() + + with ( + patch("dstack._internal.server.services.runner.pool.SSHTunnel") as ssh_tunnel_cls, + patch( + "dstack._internal.server.services.runner.client.RunnerClient.from_address" + ) as runner_client_cls, + ): + runner_client_mock = runner_client_cls.return_value + runner_client_mock.pull.return_value = PullResponse( + job_states=[], + job_logs=[], + runner_logs=[], + last_updated=0, + no_connections_secs=0, + ) + await _process_job(session, worker, job) + ssh_tunnel_cls.assert_called_once() + runner_client_mock.pull.assert_called_once() + + await session.refresh(job) + assert job.status == JobStatus.RUNNING + assert job.termination_reason is None + + async def test_terminates_job_ignoring_registration_on_non_running_replica( + self, + test_db, + session: AsyncSession, + worker: JobRunningWorker, + ) -> None: + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + gateway_compute_running = await create_gateway_compute( + session=session, gateway_id=gateway.id, replica_num=0 + ) + # Terminated replica successfully registered before going away — should be ignored, + # since only currently running replicas count towards the predicate. + gateway_compute_terminating = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + replica_num=1, + status=GatewayReplicaStatus.TERMINATING, + ) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + status=RunStatus.RUNNING, + run_spec=get_run_spec( + run_name="test-run", + repo_id=repo.name, + configuration=ServiceConfiguration(port=80, image="ubuntu"), + ), + gateway=gateway, + ) + instance = await create_instance( + session=session, + project=project, + status=InstanceStatus.BUSY, + ) + job = await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + job_provisioning_data=get_job_provisioning_data(), + instance=instance, + instance_assigned=True, + registered=True, + ready=True, + ) + session.add( + ServiceReplicaRegistrationModel( + job_id=job.id, + gateway_replica_id=gateway_compute_running.id, + is_registered=False, + register_attempt=2, + register_status_message="Connection refused", + ) + ) + session.add( + ServiceReplicaRegistrationModel( + job_id=job.id, + gateway_replica_id=gateway_compute_terminating.id, + is_registered=True, + register_attempt=0, + ) + ) + await session.commit() + + with ( + patch("dstack._internal.server.services.runner.pool.SSHTunnel") as ssh_tunnel_cls, + patch( + "dstack._internal.server.services.runner.client.RunnerClient.from_address" + ) as runner_client_cls, + ): + runner_client_mock = runner_client_cls.return_value + runner_client_mock.pull.return_value = PullResponse( + job_states=[], + job_logs=[], + runner_logs=[], + last_updated=0, + no_connections_secs=0, + ) + await _process_job(session, worker, job) + ssh_tunnel_cls.assert_called_once() + runner_client_mock.pull.assert_called_once() + + await session.refresh(job) + assert job.status == JobStatus.TERMINATING + assert job.termination_reason == JobTerminationReason.GATEWAY_ERROR + @pytest.mark.parametrize("probe_count", [1, 2]) async def test_creates_probe_models_and_not_registers_service_replica( self, @@ -1971,7 +2318,6 @@ async def test_registers_service_replica_immediately_if_no_probes( assert {event.message for event in events} == { "Job status changed PULLING -> RUNNING", "Service replica ready to receive requests", - "Service replica registered to receive requests", } @pytest.mark.parametrize( @@ -2060,10 +2406,9 @@ async def test_registers_service_replica_only_after_probes_pass( events = await list_events(session) if expect_to_register: assert job.registered - assert len(events) == 2 + assert len(events) == 1 assert {event.message for event in events} == { "Service replica ready to receive requests", - "Service replica registered to receive requests", } else: assert not job.registered @@ -2077,7 +2422,6 @@ async def test_registers_service_replica_in_gateway( ssh_tunnel_mock: Mock, shim_client_mock: Mock, runner_client_mock: Mock, - mock_gateway_connection: AsyncMock, ): user = await create_user(session=session) project = await create_project(session=session, owner=user) @@ -2136,16 +2480,7 @@ async def test_registers_service_replica_in_gateway( assert {event.message for event in events} == { "Job status changed PULLING -> RUNNING", "Service replica ready to receive requests", - "Service replica registered to receive requests", } - mock_gateway_connection.return_value.client.return_value.__aenter__.return_value.register_replica.assert_called_once_with( - run=ANY, - job_spec=ANY, - job_submission=ANY, - instance_project_ssh_private_key=None, - ssh_head_proxy=None, - ssh_head_proxy_private_key=None, - ) async def test_registers_service_replica_in_gateway_when_running_on_imported_instance( self, @@ -2155,7 +2490,6 @@ async def test_registers_service_replica_in_gateway_when_running_on_imported_ins ssh_tunnel_mock: Mock, shim_client_mock: Mock, runner_client_mock: Mock, - mock_gateway_connection: AsyncMock, ): user = await create_user(session=session) exporter_project = await create_project( @@ -2223,16 +2557,7 @@ async def test_registers_service_replica_in_gateway_when_running_on_imported_ins assert {event.message for event in events} == { "Job status changed PULLING -> RUNNING", "Service replica ready to receive requests", - "Service replica registered to receive requests", } - mock_gateway_connection.return_value.client.return_value.__aenter__.return_value.register_replica.assert_called_once_with( - run=ANY, - job_spec=ANY, - job_submission=ANY, - instance_project_ssh_private_key="exporter-private-key", - ssh_head_proxy=None, - ssh_head_proxy_private_key=None, - ) @pytest.mark.parametrize("job_status", [JobStatus.RUNNING, JobStatus.PULLING]) async def test_terminates_job_when_instance_access_revoked( @@ -2470,7 +2795,6 @@ async def test_registers_router_replica_but_not_worker_replica_in_gateway( worker: JobRunningWorker, ssh_tunnel_mock: Mock, runner_client_mock: Mock, - mock_gateway_connection: AsyncMock, ): user = await create_user(session=session) project = await create_project(session=session, owner=user) @@ -2538,7 +2862,6 @@ async def test_registers_router_replica_but_not_worker_replica_in_gateway( events = await list_events(session) assert {event.message for event in events} == { "Service replica ready to receive requests", - "Service replica registered to receive requests", } await clear_events(session) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py b/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py index ed085e09d4..fb8f8aa8fd 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py @@ -13,6 +13,7 @@ TaskConfiguration, ) from dstack._internal.core.models.duration import Duration +from dstack._internal.core.models.gateways import GatewayReplicaStatus, GatewayStatus from dstack._internal.core.models.instances import InstanceStatus from dstack._internal.core.models.profiles import ( Profile, @@ -29,10 +30,17 @@ RunTerminationReason, ) from dstack._internal.server.background.pipeline_tasks.runs import RunWorker -from dstack._internal.server.models import JobModel +from dstack._internal.server.models import ( + JobModel, + ServiceRegistrationModel, + ServiceReplicaRegistrationModel, +) from dstack._internal.server.services.jobs import get_job_spec from dstack._internal.server.testing.common import ( + create_backend, create_fleet, + create_gateway, + create_gateway_compute, create_instance, create_job, create_project, @@ -174,6 +182,253 @@ async def test_terminates_run_on_job_failure( assert run.termination_reason == RunTerminationReason.JOB_FAILED assert run.lock_token is None + async def test_terminates_run_on_gateway_registration_failure( + self, test_db, session: AsyncSession, worker: RunWorker + ) -> None: + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + gateway_compute = await create_gateway_compute(session=session, gateway_id=gateway.id) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + status=RunStatus.RUNNING, + run_spec=get_run_spec( + run_name="test-run", + repo_id=repo.name, + configuration=ServiceConfiguration(port=80, image="ubuntu"), + ), + gateway=gateway, + ) + await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + registered=True, + ready=True, + ) + session.add( + ServiceRegistrationModel( + run_id=run.id, + gateway_replica_id=gateway_compute.id, + is_registered=False, + register_attempt=3, + register_status_message="Connection refused", + ) + ) + lock_run(run) + await session.commit() + + await worker.process(run_to_pipeline_item(run)) + + await session.refresh(run) + assert run.status == RunStatus.TERMINATING + assert run.termination_reason == RunTerminationReason.GATEWAY_ERROR + assert run.lock_token is None + + async def test_does_not_terminate_run_when_service_registered_despite_failed_attempt( + self, test_db, session: AsyncSession, worker: RunWorker + ) -> None: + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + gateway_compute_1 = await create_gateway_compute( + session=session, gateway_id=gateway.id, replica_num=0 + ) + gateway_compute_2 = await create_gateway_compute( + session=session, gateway_id=gateway.id, replica_num=1 + ) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + status=RunStatus.RUNNING, + run_spec=get_run_spec( + run_name="test-run", + repo_id=repo.name, + configuration=ServiceConfiguration(port=80, image="ubuntu"), + ), + gateway=gateway, + ) + await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + registered=True, + ready=True, + ) + session.add( + ServiceRegistrationModel( + run_id=run.id, + gateway_replica_id=gateway_compute_1.id, + is_registered=False, + register_attempt=3, + register_status_message="Connection refused", + ) + ) + session.add( + ServiceRegistrationModel( + run_id=run.id, + gateway_replica_id=gateway_compute_2.id, + is_registered=True, + register_attempt=0, + ) + ) + lock_run(run) + await session.commit() + + await worker.process(run_to_pipeline_item(run)) + + await session.refresh(run) + assert run.status == RunStatus.RUNNING + assert run.termination_reason is None + assert run.lock_token is None + + async def test_does_not_terminate_run_when_one_running_replica_has_not_attempted_registration( + self, test_db, session: AsyncSession, worker: RunWorker + ) -> None: + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + gateway_compute_1 = await create_gateway_compute( + session=session, gateway_id=gateway.id, replica_num=0 + ) + # Second running replica has not attempted registration yet (e.g. just came up). + await create_gateway_compute(session=session, gateway_id=gateway.id, replica_num=1) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + status=RunStatus.RUNNING, + run_spec=get_run_spec( + run_name="test-run", + repo_id=repo.name, + configuration=ServiceConfiguration(port=80, image="ubuntu"), + ), + gateway=gateway, + ) + await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + registered=True, + ready=True, + ) + session.add( + ServiceRegistrationModel( + run_id=run.id, + gateway_replica_id=gateway_compute_1.id, + is_registered=False, + register_attempt=3, + register_status_message="Connection refused", + ) + ) + lock_run(run) + await session.commit() + + await worker.process(run_to_pipeline_item(run)) + + await session.refresh(run) + assert run.status == RunStatus.RUNNING + assert run.termination_reason is None + assert run.lock_token is None + + async def test_terminates_run_ignoring_registration_on_non_running_replica( + self, test_db, session: AsyncSession, worker: RunWorker + ) -> None: + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + gateway_compute_running = await create_gateway_compute( + session=session, gateway_id=gateway.id, replica_num=0 + ) + # Terminated replica successfully registered before going away — should be ignored, + # since only currently running replicas count towards the predicate. + gateway_compute_terminating = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + replica_num=1, + status=GatewayReplicaStatus.TERMINATING, + ) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + status=RunStatus.RUNNING, + run_spec=get_run_spec( + run_name="test-run", + repo_id=repo.name, + configuration=ServiceConfiguration(port=80, image="ubuntu"), + ), + gateway=gateway, + ) + await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + registered=True, + ready=True, + ) + session.add( + ServiceRegistrationModel( + run_id=run.id, + gateway_replica_id=gateway_compute_running.id, + is_registered=False, + register_attempt=3, + register_status_message="Connection refused", + ) + ) + session.add( + ServiceRegistrationModel( + run_id=run.id, + gateway_replica_id=gateway_compute_terminating.id, + is_registered=True, + register_attempt=0, + ) + ) + lock_run(run) + await session.commit() + + await worker.process(run_to_pipeline_item(run)) + + await session.refresh(run) + assert run.status == RunStatus.TERMINATING + assert run.termination_reason == RunTerminationReason.GATEWAY_ERROR + assert run.lock_token is None + async def test_retries_failed_replica_within_retry_duration( self, test_db, session: AsyncSession, worker: RunWorker ) -> None: @@ -1072,6 +1327,209 @@ async def test_service_rolling_deployment_scale_down_old_unregistered( assert old_job.status == JobStatus.TERMINATING assert old_job.termination_reason == JobTerminationReason.SCALED_DOWN + @pytest.mark.parametrize("failed_registration", [False, True]) + async def test_service_rolling_deployment_keeps_old_replica_until_new_replica_registered_with_gateway( + self, test_db, session: AsyncSession, worker: RunWorker, failed_registration: bool + ) -> None: + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + gateway_compute_1 = await create_gateway_compute( + session=session, gateway_id=gateway.id, replica_num=0 + ) + gateway_compute_2 = await create_gateway_compute( + session=session, gateway_id=gateway.id, replica_num=1 + ) + run_spec = get_run_spec( + repo_id=repo.name, + run_name="service-run", + configuration=ServiceConfiguration( + port=8080, + commands=["echo new!"], + ), + ) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_name="service-run", + run_spec=run_spec, + status=RunStatus.RUNNING, + deployment_num=1, + gateway=gateway, + ) + old_job = await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + deployment_num=0, + registered=True, + ready=True, + replica_num=0, + ) + old_spec = get_job_spec(old_job) + old_spec.commands = ["echo old!"] + old_job.job_spec_data = old_spec.model_dump_json() + new_job = await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + deployment_num=1, + registered=True, + ready=True, + replica_num=1, + ) + # Old replica is fully registered — receiving traffic on every running gateway replica. + session.add( + ServiceReplicaRegistrationModel( + job_id=old_job.id, + gateway_replica_id=gateway_compute_1.id, + is_registered=True, + register_attempt=0, + ) + ) + session.add( + ServiceReplicaRegistrationModel( + job_id=old_job.id, + gateway_replica_id=gateway_compute_2.id, + is_registered=True, + register_attempt=0, + ) + ) + # New replica is only confirmed registered on the first gateway replica. + session.add( + ServiceReplicaRegistrationModel( + job_id=new_job.id, + gateway_replica_id=gateway_compute_1.id, + is_registered=True, + register_attempt=0, + ) + ) + if failed_registration: + # Registration on the second gateway replica was attempted and failed. + session.add( + ServiceReplicaRegistrationModel( + job_id=new_job.id, + gateway_replica_id=gateway_compute_2.id, + is_registered=False, + register_attempt=2, + ) + ) + await session.commit() + + lock_run(run) + await session.commit() + + await worker.process(run_to_pipeline_item(run)) + + await session.refresh(run) + assert run.status == RunStatus.RUNNING + + await session.refresh(old_job) + await session.refresh(new_job) + assert old_job.status == JobStatus.RUNNING + assert new_job.status == JobStatus.RUNNING + + async def test_service_rolling_deployment_scales_down_old_replica_once_new_replica_registered_with_gateway( + self, test_db, session: AsyncSession, worker: RunWorker + ) -> None: + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + gateway_compute_1 = await create_gateway_compute( + session=session, gateway_id=gateway.id, replica_num=0 + ) + gateway_compute_2 = await create_gateway_compute( + session=session, gateway_id=gateway.id, replica_num=1 + ) + run_spec = get_run_spec( + repo_id=repo.name, + run_name="service-run", + configuration=ServiceConfiguration( + port=8080, + commands=["echo new!"], + ), + ) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_name="service-run", + run_spec=run_spec, + status=RunStatus.RUNNING, + deployment_num=1, + gateway=gateway, + ) + old_job = await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + deployment_num=0, + registered=True, + ready=True, + replica_num=0, + ) + old_spec = get_job_spec(old_job) + old_spec.commands = ["echo old!"] + old_job.job_spec_data = old_spec.model_dump_json() + new_job = await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + deployment_num=1, + registered=True, + ready=True, + replica_num=1, + ) + for gateway_compute in (gateway_compute_1, gateway_compute_2): + session.add( + ServiceReplicaRegistrationModel( + job_id=old_job.id, + gateway_replica_id=gateway_compute.id, + is_registered=True, + register_attempt=0, + ) + ) + session.add( + ServiceReplicaRegistrationModel( + job_id=new_job.id, + gateway_replica_id=gateway_compute.id, + is_registered=True, + register_attempt=0, + ) + ) + await session.commit() + + lock_run(run) + await session.commit() + + await worker.process(run_to_pipeline_item(run)) + + await session.refresh(run) + assert run.status == RunStatus.RUNNING + + await session.refresh(old_job) + await session.refresh(new_job) + assert old_job.status == JobStatus.TERMINATING + assert old_job.termination_reason == JobTerminationReason.SCALED_DOWN + assert new_job.status == JobStatus.RUNNING + async def test_service_router_rolling_deployment_surges_ready_worker_replica( self, test_db, session: AsyncSession, worker: RunWorker ) -> None: diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py index 08207c944c..db9e20d49a 100644 --- a/src/tests/_internal/server/routers/test_runs.py +++ b/src/tests/_internal/server/routers/test_runs.py @@ -11,7 +11,6 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from dstack._internal.core.errors import GatewayError from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.common import ApplyAction, EntityReference from dstack._internal.core.models.configurations import ( @@ -3997,8 +3996,9 @@ async def test_submit_to_correct_proxy( assert response.status_code == 200 assert response.json()["service"]["url"] == expected_service_url assert response.json()["service"]["model"]["base_url"] == expected_model_url - events = await list_events(session) - assert ("Service registered in gateway" in {e.message for e in events}) == is_gateway + res = await session.execute(select(RunModel)) + run = res.scalar_one() + assert (run.gateway_id is not None) == is_gateway @pytest.mark.asyncio @pytest.mark.parametrize("populate_configuration", [True, False]) @@ -4049,8 +4049,9 @@ async def test_submit_to_gateway_by_name( ) assert response.status_code == 200 assert response.json()["service"]["url"] == "https://test-service.my-gateway.example" - events = await list_events(session) - assert "Service registered in gateway" in {e.message for e in events} + res = await session.execute(select(RunModel)) + run = res.scalar_one() + assert run.gateway_id is not None @pytest.mark.asyncio async def test_return_error_if_specified_gateway_not_exists( @@ -4402,62 +4403,6 @@ async def test_returns_error_if_imported_gateway_domain_has_unknown_variable( ] } - @pytest.mark.asyncio - async def test_unregister_dangling_service( - self, - test_db, - session: AsyncSession, - client: AsyncClient, - mock_gateway_connection: AsyncMock, - ) -> None: - user = await create_user(session=session, global_role=GlobalRole.USER) - project = await create_project(session=session, owner=user, name="test-project") - await add_project_member( - session=session, project=project, user=user, project_role=ProjectRole.USER - ) - repo = await create_repo(session=session, project_id=project.id) - backend = await create_backend(session=session, project_id=project.id) - gateway = await create_gateway( - session=session, - project_id=project.id, - backend_id=backend.id, - status=GatewayStatus.RUNNING, - wildcard_domain="example.com", - ) - await create_gateway_compute(session=session, backend_id=backend.id, gateway_id=gateway.id) - project.default_gateway_id = gateway.id - await session.commit() - - client_mock = ( - mock_gateway_connection.return_value.client.return_value.__aenter__.return_value - ) - client_mock.register_service.side_effect = [ - GatewayError("Service test-project/test-service is already registered"), - None, # Second call succeeds - ] - - response = await client.post( - f"/api/project/{project.name}/runs/apply", - headers=get_auth_headers(user.token), - json={ - "plan": { - "run_spec": get_service_run_spec(repo_id=repo.name, run_name="test-service"), - "current_resource": None, - }, - "force": False, - }, - ) - - assert response.status_code == 200 - assert response.json()["service"]["url"] == "https://test-service.example.com" - # Verify that unregister_service was called to clean up the dangling service - client_mock.unregister_service.assert_called_once_with( - project=project.name, - run_name="test-service", - ) - # Verify that register_service was called twice (first failed, then succeeded) - assert client_mock.register_service.call_count == 2 - @pytest.mark.asyncio async def test_return_error_if_default_gateway_forbids_new_services( self, diff --git a/src/tests/_internal/server/services/services/test_services.py b/src/tests/_internal/server/services/services/test_services.py index 0be5bed9df..f65fdbdf8b 100644 --- a/src/tests/_internal/server/services/services/test_services.py +++ b/src/tests/_internal/server/services/services/test_services.py @@ -15,8 +15,8 @@ from dstack._internal.core.models.runs import RunSpec from dstack._internal.server.services.services import ( _register_service_in_server, - _should_configure_service_https_on_gateway, _should_show_service_https, + should_configure_service_https_on_gateway, ) from dstack._internal.server.testing.common import get_run_spec @@ -59,36 +59,36 @@ class TestShouldConfigureServiceHttpsOnGateway: def test_auto_resolves_to_true_with_lets_encrypt_gateway(self) -> None: run_spec = _service_run_spec(https="auto") gw = _gateway_config(certificate=LetsEncryptGatewayCertificate()) - assert _should_configure_service_https_on_gateway(run_spec, gw) is True + assert should_configure_service_https_on_gateway(run_spec, gw) is True def test_auto_resolves_to_false_when_gateway_has_no_certificate(self) -> None: run_spec = _service_run_spec(https="auto") gw = _gateway_config(certificate=None) - assert _should_configure_service_https_on_gateway(run_spec, gw) is False + assert should_configure_service_https_on_gateway(run_spec, gw) is False def test_auto_resolves_to_false_with_acm_gateway(self) -> None: run_spec = _service_run_spec(https="auto") gw = _gateway_config( certificate=ACMGatewayCertificate(arn="arn:aws:acm:us-east-1:123:cert/abc") ) - assert _should_configure_service_https_on_gateway(run_spec, gw) is False + assert should_configure_service_https_on_gateway(run_spec, gw) is False def test_true_enables_https_when_gateway_has_no_certificate(self) -> None: run_spec = _service_run_spec(https=True) gw = _gateway_config(certificate=None) - assert _should_configure_service_https_on_gateway(run_spec, gw) is True + assert should_configure_service_https_on_gateway(run_spec, gw) is True def test_false_disables_https_regardless_of_gateway_certificate(self) -> None: run_spec = _service_run_spec(https=False) gw = _gateway_config(certificate=LetsEncryptGatewayCertificate()) - assert _should_configure_service_https_on_gateway(run_spec, gw) is False + assert should_configure_service_https_on_gateway(run_spec, gw) is False def test_true_does_not_configure_https_on_acm_gateway(self) -> None: run_spec = _service_run_spec(https=True) gw = _gateway_config( certificate=ACMGatewayCertificate(arn="arn:aws:acm:us-east-1:123:cert/abc") ) - assert _should_configure_service_https_on_gateway(run_spec, gw) is False + assert should_configure_service_https_on_gateway(run_spec, gw) is False class TestShouldShowServiceHttps: