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: