From d9292a7b6726913f6313b9fe01d5cbd5b7ea54c5 Mon Sep 17 00:00:00 2001 From: Jvst Me Date: Mon, 17 Aug 2026 23:13:11 +0200 Subject: [PATCH] [chore]: Gateway replica naming cleanup - Switch from "gateway compute" to "gateway replica" everywhere in the codebase (class and function names, variables, log messages, etc). - Switch from "gateway" to "gateway replica" where this term is more precise semantically. This change does not affect public `dstack` APIs. Database tables and columns preserve their original names to avoid migration. --- .../_internal/core/backends/aws/compute.py | 16 +- .../_internal/core/backends/azure/compute.py | 16 +- .../_internal/core/backends/base/compute.py | 28 +- .../_internal/core/backends/gcp/compute.py | 16 +- .../core/backends/kubernetes/compute.py | 20 +- src/dstack/_internal/core/models/gateways.py | 4 +- .../pipeline_tasks/gateway_replicas.py | 191 ++-- .../background/pipeline_tasks/gateways.py | 134 +-- .../background/pipeline_tasks/jobs_running.py | 14 +- .../pipeline_tasks/runs/__init__.py | 20 +- .../background/scheduled_tasks/gateways.py | 4 +- src/dstack/_internal/server/models.py | 46 +- .../server/services/gateways/__init__.py | 242 ++--- .../server/services/gateways/connection.py | 6 +- .../server/services/runs/__init__.py | 4 +- .../server/services/runs/replicas.py | 4 +- .../server/services/services/__init__.py | 8 +- src/dstack/_internal/server/testing/common.py | 24 +- .../core/backends/base/test_compute.py | 4 +- .../pipeline_tasks/test_gateway_replicas.py | 828 +++++++++--------- .../pipeline_tasks/test_gateways.py | 306 +++---- .../pipeline_tasks/test_running_jobs.py | 34 +- .../pipeline_tasks/test_runs/test_active.py | 50 +- .../pipeline_tasks/test_runs/test_pending.py | 6 +- .../_internal/server/routers/test_gateways.py | 116 +-- .../_internal/server/routers/test_runs.py | 18 +- .../server/services/gateways/test_gateways.py | 34 +- 27 files changed, 1106 insertions(+), 1087 deletions(-) diff --git a/src/dstack/_internal/core/backends/aws/compute.py b/src/dstack/_internal/core/backends/aws/compute.py index 39521cc7c6..78b956f5b2 100644 --- a/src/dstack/_internal/core/backends/aws/compute.py +++ b/src/dstack/_internal/core/backends/aws/compute.py @@ -57,10 +57,10 @@ from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.common import CoreModel, validate_json_extra_ignore from dstack._internal.core.models.gateways import ( - GatewayComputeConfiguration, GatewayLoadBalancerConfiguration, GatewayLoadBalancerData, - GatewayProvisioningData, + GatewayReplicaConfiguration, + GatewayReplicaProvisioningData, ) from dstack._internal.core.models.instances import ( InstanceAvailability, @@ -524,10 +524,10 @@ def is_suitable_placement_group( return False return placement_group.configuration.region == instance_offer.region - def create_gateway( + def create_gateway_replica( self, - configuration: GatewayComputeConfiguration, - ) -> GatewayProvisioningData: + configuration: GatewayReplicaConfiguration, + ) -> GatewayReplicaProvisioningData: ec2_resource = self.session.resource("ec2", region_name=configuration.region) ec2_client = self.session.client("ec2", region_name=configuration.region) @@ -587,7 +587,7 @@ def create_gateway( instance.wait_until_running() instance.reload() # populate instance.public_ip_address ip_address = _get_instance_ip(instance, configuration.public_ip) - return GatewayProvisioningData( + return GatewayReplicaProvisioningData( instance_id=instance.instance_id, region=configuration.region, availability_zone=availability_zone, @@ -716,10 +716,10 @@ def create_gateway_load_balancer( ).model_dump_json(), ) - def terminate_gateway( + def terminate_gateway_replica( self, instance_id: str, - configuration: GatewayComputeConfiguration, + configuration: GatewayReplicaConfiguration, backend_data: Optional[str] = None, ): self.terminate_instance( diff --git a/src/dstack/_internal/core/backends/azure/compute.py b/src/dstack/_internal/core/backends/azure/compute.py index 83a590a4b1..f13a7f4a19 100644 --- a/src/dstack/_internal/core/backends/azure/compute.py +++ b/src/dstack/_internal/core/backends/azure/compute.py @@ -61,8 +61,8 @@ from dstack._internal.core.errors import ComputeError, NoCapacityError from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.gateways import ( - GatewayComputeConfiguration, - GatewayProvisioningData, + GatewayReplicaConfiguration, + GatewayReplicaProvisioningData, ) from dstack._internal.core.models.instances import ( InstanceAvailability, @@ -233,10 +233,10 @@ def terminate_instance( instance_name=instance_id, ) - def create_gateway( + def create_gateway_replica( self, - configuration: GatewayComputeConfiguration, - ) -> GatewayProvisioningData: + configuration: GatewayReplicaConfiguration, + ) -> GatewayReplicaProvisioningData: if configuration.instance_type is not None: # TODO: support instance_type. Requires selecting a VM image to avoid errors like this: # > The selected VM size 'Standard_E4s_v6' cannot boot Hypervisor Generation '1' @@ -306,16 +306,16 @@ def create_gateway( resource_group=self.config.resource_group, vm=vm, ) - return GatewayProvisioningData( + return GatewayReplicaProvisioningData( instance_id=vm.name, ip_address=public_ip, region=configuration.region, ) - def terminate_gateway( + def terminate_gateway_replica( self, instance_id: str, - configuration: GatewayComputeConfiguration, + configuration: GatewayReplicaConfiguration, backend_data: Optional[str] = None, ): self.terminate_instance( diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py index 8aae80e5ec..673f396429 100644 --- a/src/dstack/_internal/core/backends/base/compute.py +++ b/src/dstack/_internal/core/backends/base/compute.py @@ -29,10 +29,10 @@ from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.compute_groups import ComputeGroup, ComputeGroupProvisioningData from dstack._internal.core.models.gateways import ( - GatewayComputeConfiguration, GatewayLoadBalancerConfiguration, GatewayLoadBalancerData, - GatewayProvisioningData, + GatewayReplicaConfiguration, + GatewayReplicaProvisioningData, ) from dstack._internal.core.models.instances import ( InstanceConfiguration, @@ -556,25 +556,25 @@ class ComputeWithGatewaySupport(ABC): """ @abstractmethod - def create_gateway( + def create_gateway_replica( self, - configuration: GatewayComputeConfiguration, - ) -> GatewayProvisioningData: + configuration: GatewayReplicaConfiguration, + ) -> GatewayReplicaProvisioningData: """ - Creates a gateway instance. + Creates a gateway replica instance. """ pass @abstractmethod - def terminate_gateway( + def terminate_gateway_replica( self, instance_id: str, - configuration: GatewayComputeConfiguration, + configuration: GatewayReplicaConfiguration, backend_data: Optional[str] = None, ): """ - Terminates a gateway instance. Generally, it passes the call to `terminate_instance()`, - but may perform additional work such as deleting a load balancer when a gateway has one. + Terminates a gateway replica instance. Generally, it passes the call to + `terminate_instance()`, but may perform additional work if necessary. """ pass @@ -631,7 +631,7 @@ def deregister_gateway_replica_from_load_balancer( class ComputeWithPrivateGatewaySupport: """ Must be subclassed to support private gateways. - `create_gateway()` must be able to create private gateways. + `create_gateway_replica()` must be able to create private gateways. """ pass @@ -751,15 +751,15 @@ def generate_unique_instance_name_for_job( def generate_unique_gateway_instance_name( - gateway_compute_configuration: GatewayComputeConfiguration, + gateway_replica_configuration: GatewayReplicaConfiguration, max_length: int = _DEFAULT_MAX_RESOURCE_NAME_LEN, ) -> str: """ Generates a unique gateway instance name valid across all backends. """ return generate_unique_backend_name( - resource_name=gateway_compute_configuration.instance_name, - project_name=gateway_compute_configuration.project_name, + resource_name=gateway_replica_configuration.instance_name, + project_name=gateway_replica_configuration.project_name, max_length=max_length, ) diff --git a/src/dstack/_internal/core/backends/gcp/compute.py b/src/dstack/_internal/core/backends/gcp/compute.py index 9c94ad2ad2..5b30226004 100644 --- a/src/dstack/_internal/core/backends/gcp/compute.py +++ b/src/dstack/_internal/core/backends/gcp/compute.py @@ -57,8 +57,8 @@ from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.common import CoreModel, validate_extra_ignore from dstack._internal.core.models.gateways import ( - GatewayComputeConfiguration, - GatewayProvisioningData, + GatewayReplicaConfiguration, + GatewayReplicaProvisioningData, ) from dstack._internal.core.models.instances import ( InstanceAvailability, @@ -561,10 +561,10 @@ def are_placement_groups_compatible_with_reservations(self, backend_type: Backen # Instead, we use the placement policy defined in reservation settings. return False - def create_gateway( + def create_gateway_replica( self, - configuration: GatewayComputeConfiguration, - ) -> GatewayProvisioningData: + configuration: GatewayReplicaConfiguration, + ) -> GatewayReplicaProvisioningData: if self.config.vpc_project_id is None: gcp_resources.create_gateway_firewall_rules( firewalls_client=self.firewalls_client, @@ -628,7 +628,7 @@ def create_gateway( instance = self.instances_client.get( project=self.config.project_id, zone=zone, instance=instance_name ) - return GatewayProvisioningData( + return GatewayReplicaProvisioningData( instance_id=instance_name, region=configuration.region, # used for instance termination availability_zone=zone, @@ -636,10 +636,10 @@ def create_gateway( backend_data=json.dumps({"zone": zone}), ) - def terminate_gateway( + def terminate_gateway_replica( self, instance_id: str, - configuration: GatewayComputeConfiguration, + configuration: GatewayReplicaConfiguration, backend_data: Optional[str] = None, ): self.terminate_instance( diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py index 5b21377cd1..f60c285bd8 100644 --- a/src/dstack/_internal/core/backends/kubernetes/compute.py +++ b/src/dstack/_internal/core/backends/kubernetes/compute.py @@ -79,8 +79,8 @@ from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.common import CoreModel, validate_json_extra_ignore from dstack._internal.core.models.gateways import ( - GatewayComputeConfiguration, - GatewayProvisioningData, + GatewayReplicaConfiguration, + GatewayReplicaProvisioningData, ) from dstack._internal.core.models.instances import ( InstanceOfferWithAvailability, @@ -462,10 +462,10 @@ def terminate_instance( if not all(deleted): raise ComputeError("Not all objects were deleted, check logs") - def create_gateway( + def create_gateway_replica( self, - configuration: GatewayComputeConfiguration, - ) -> GatewayProvisioningData: + configuration: GatewayReplicaConfiguration, + ) -> GatewayReplicaProvisioningData: cluster = self.region_cluster_map.get(configuration.region) if cluster is None: raise ComputeError(f"Unknown region: {configuration.region!r}") @@ -577,19 +577,19 @@ def create_gateway( if address is None: self.terminate_instance(instance_name, region=configuration.region) raise ComputeError( - "Failed to get gateway hostname. " + "Failed to get gateway replica hostname. " "Ensure the Kubernetes cluster supports Load Balancer services." ) - return GatewayProvisioningData( + return GatewayReplicaProvisioningData( instance_id=instance_name, ip_address=address, region=cluster.region, ) - def terminate_gateway( + def terminate_gateway_replica( self, instance_id: str, - configuration: GatewayComputeConfiguration, + configuration: GatewayReplicaConfiguration, backend_data: Optional[str] = None, ): region = configuration.region @@ -600,7 +600,7 @@ def terminate_gateway( if cluster is not None: logger.warning( ( - "Terminating gateway %s in unknown region %s." + "Terminating gateway replica %s in unknown region %s." " Assuming it was created before multi-cluster support was added" " and is located in cluster %s" ), diff --git a/src/dstack/_internal/core/models/gateways.py b/src/dstack/_internal/core/models/gateways.py index a070c2a92f..d9decdaea1 100644 --- a/src/dstack/_internal/core/models/gateways.py +++ b/src/dstack/_internal/core/models/gateways.py @@ -187,7 +187,7 @@ class ApplyGatewayPlanInput(CoreModel): ] = None -class GatewayComputeConfiguration(CoreModel): +class GatewayReplicaConfiguration(CoreModel): project_name: str instance_name: str backend: BackendType @@ -199,7 +199,7 @@ class GatewayComputeConfiguration(CoreModel): tags: Optional[Dict[str, str]] = None -class GatewayProvisioningData(CoreModel): +class GatewayReplicaProvisioningData(CoreModel): instance_id: str # TODO: rename `ip_address`; Kubernetes uses domain names here. ip_address: str 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 8271e60606..0ba64dc8ad 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py @@ -36,8 +36,8 @@ from dstack._internal.server.db import get_db, get_session_ctx from dstack._internal.server.models import ( BackendModel, - GatewayComputeModel, GatewayModel, + GatewayReplicaModel, InstanceModel, JobModel, ProjectModel, @@ -49,9 +49,9 @@ 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, + get_gateway_replica_configuration, ) from dstack._internal.server.services.gateways.client import GatewayClient from dstack._internal.server.services.gateways.connection import GatewayConnection @@ -99,7 +99,7 @@ def __init__( heartbeat_trigger=heartbeat_trigger, ) self.__heartbeater = Heartbeater[GatewayReplicaPipelineItem]( - model_type=GatewayComputeModel, + model_type=GatewayReplicaModel, lock_timeout=self._lock_timeout, heartbeat_trigger=self._heartbeat_trigger, ) @@ -121,7 +121,7 @@ def __init__( @property def hint_fetch_model_name(self) -> str: - return GatewayComputeModel.__name__ + return GatewayReplicaModel.__name__ @property def _heartbeater(self) -> Heartbeater[GatewayReplicaPipelineItem]: @@ -158,24 +158,24 @@ def __init__( @tracing.instrument_pipeline_task("GatewayReplicaFetcher.fetch") async def fetch(self, limit: int) -> list[GatewayReplicaPipelineItem]: replica_lock, _ = get_locker(get_db().dialect_name).get_lockset( - GatewayComputeModel.__tablename__ + GatewayReplicaModel.__tablename__ ) async with replica_lock: async with get_session_ctx() as session: now = get_current_datetime() res = await session.execute( - select(GatewayComputeModel) + select(GatewayReplicaModel) .outerjoin( GatewayModel, or_( - GatewayModel.id == GatewayComputeModel.gateway_id, - GatewayModel.gateway_compute_id == GatewayComputeModel.id, + GatewayModel.id == GatewayReplicaModel.gateway_id, + GatewayModel.gateway_replica_id == GatewayReplicaModel.id, ), ) .where( - GatewayComputeModel.deleted == False, + GatewayReplicaModel.deleted == False, or_( - GatewayComputeModel.status.in_( + GatewayReplicaModel.status.in_( [ GatewayReplicaStatus.SUBMITTED, GatewayReplicaStatus.PROVISIONING, @@ -185,29 +185,29 @@ async def fetch(self, limit: int) -> list[GatewayReplicaPipelineItem]: ), ), or_( - GatewayComputeModel.last_processed_at + GatewayReplicaModel.last_processed_at <= now - self._min_processing_interval, - GatewayComputeModel.last_processed_at - == GatewayComputeModel.created_at, + GatewayReplicaModel.last_processed_at + == GatewayReplicaModel.created_at, ), or_( - GatewayComputeModel.lock_expires_at.is_(None), - GatewayComputeModel.lock_expires_at < now, + GatewayReplicaModel.lock_expires_at.is_(None), + GatewayReplicaModel.lock_expires_at < now, ), or_( - GatewayComputeModel.lock_owner.is_(None), - GatewayComputeModel.lock_owner == GatewayReplicaPipeline.__name__, + GatewayReplicaModel.lock_owner.is_(None), + GatewayReplicaModel.lock_owner == GatewayReplicaPipeline.__name__, ), ) - .order_by(GatewayComputeModel.last_processed_at.asc()) + .order_by(GatewayReplicaModel.last_processed_at.asc()) .limit(limit) - .with_for_update(skip_locked=True, key_share=True, of=GatewayComputeModel) + .with_for_update(skip_locked=True, key_share=True, of=GatewayReplicaModel) .options( load_only( - GatewayComputeModel.id, - GatewayComputeModel.lock_token, - GatewayComputeModel.lock_expires_at, - GatewayComputeModel.status, + GatewayReplicaModel.id, + GatewayReplicaModel.lock_token, + GatewayReplicaModel.lock_expires_at, + GatewayReplicaModel.status, ) ) ) @@ -222,7 +222,7 @@ async def fetch(self, limit: int) -> list[GatewayReplicaPipelineItem]: replica_model.lock_owner = GatewayReplicaPipeline.__name__ items.append( GatewayReplicaPipelineItem( - __tablename__=GatewayComputeModel.__tablename__, + __tablename__=GatewayReplicaModel.__tablename__, id=replica_model.id, lock_expires_at=lock_expires_at, lock_token=lock_token, @@ -271,10 +271,10 @@ class _GatewayReplicaUpdateMap(ItemUpdateMap, total=False): _REPLICA_FIELDS_MIN: list[InstrumentedAttribute[Any]] = [ - GatewayComputeModel.id, - GatewayComputeModel.lock_token, - GatewayComputeModel.status, - GatewayComputeModel.replica_num, + GatewayReplicaModel.id, + GatewayReplicaModel.lock_token, + GatewayReplicaModel.status, + GatewayReplicaModel.replica_num, ] _GATEWAY_FIELDS_MIN: list[InstrumentedAttribute[Any]] = [ @@ -291,7 +291,7 @@ async def _load_gateway_replica( gateway_fields: list[InstrumentedAttribute[Any]], load_backends: bool = False, load_gateway_backend_type: bool = False, -) -> Optional[GatewayComputeModel]: +) -> Optional[GatewayReplicaModel]: def build_gateway_options( gateway_attr: InstrumentedAttribute[GatewayModel | None], ) -> list[ExecutableOption]: @@ -309,15 +309,15 @@ def build_gateway_options( async with get_session_ctx() as session: stmt = ( - select(GatewayComputeModel) + select(GatewayReplicaModel) .where( - GatewayComputeModel.id == item.id, - GatewayComputeModel.lock_token == item.lock_token, + GatewayReplicaModel.id == item.id, + GatewayReplicaModel.lock_token == item.lock_token, ) .options( load_only(*replica_fields), - *build_gateway_options(GatewayComputeModel.gateway), - *build_gateway_options(GatewayComputeModel.legacy_gateway), + *build_gateway_options(GatewayReplicaModel.gateway), + *build_gateway_options(GatewayReplicaModel.legacy_gateway), ) ) res = await session.execute(stmt) @@ -329,7 +329,7 @@ def build_gateway_options( return replica_model -def _get_loaded_gateway_model(replica_model: GatewayComputeModel) -> Optional[GatewayModel]: +def _get_loaded_gateway_model(replica_model: GatewayReplicaModel) -> Optional[GatewayModel]: gateway_model = replica_model.gateway or replica_model.legacy_gateway if gateway_model is None: logger.error("Gateway replica %s is not attached to a gateway", replica_model.id) @@ -337,7 +337,7 @@ def _get_loaded_gateway_model(replica_model: GatewayComputeModel) -> Optional[Ga def _mark_terminating_if_needed( - gateway_model: GatewayModel, replica_model: GatewayComputeModel + gateway_model: GatewayModel, replica_model: GatewayReplicaModel ) -> Optional[_GatewayReplicaUpdateMap]: if gateway_model.to_be_deleted or gateway_model.status == GatewayStatus.FAILED: status_message = None @@ -368,7 +368,7 @@ def _mark_terminating_if_needed( # and apply phases instead of calling the `_commit_update()` helper from everywhere async def _commit_update( item: GatewayReplicaPipelineItem, - replica_model: GatewayComputeModel, + replica_model: GatewayReplicaModel, update_map: _GatewayReplicaUpdateMap, ) -> None: async with get_session_ctx() as session: @@ -378,7 +378,7 @@ async def _commit_update( async def _apply_update( session: AsyncSession, item: GatewayReplicaPipelineItem, - replica_model: GatewayComputeModel, + replica_model: GatewayReplicaModel, update_map: _GatewayReplicaUpdateMap, ) -> bool: set_processed_update_map_fields(update_map) @@ -386,13 +386,13 @@ async def _apply_update( now = get_current_datetime() resolve_now_placeholders(update_map, now=now) res = await session.execute( - update(GatewayComputeModel) + update(GatewayReplicaModel) .where( - GatewayComputeModel.id == replica_model.id, - GatewayComputeModel.lock_token == replica_model.lock_token, + GatewayReplicaModel.id == replica_model.id, + GatewayReplicaModel.lock_token == replica_model.lock_token, ) .values(**update_map) - .returning(GatewayComputeModel.id) + .returning(GatewayReplicaModel.id) ) updated_ids = list(res.scalars().all()) if len(updated_ids) == 0: @@ -417,10 +417,10 @@ async def _process_submitted_item(item: GatewayReplicaPipelineItem): item, replica_fields=_REPLICA_FIELDS_MIN + [ - GatewayComputeModel.backend_id, - GatewayComputeModel.configuration, - GatewayComputeModel.ssh_public_key, - GatewayComputeModel.scale_in, + GatewayReplicaModel.backend_id, + GatewayReplicaModel.configuration, + GatewayReplicaModel.ssh_public_key, + GatewayReplicaModel.scale_in, ], gateway_fields=_GATEWAY_FIELDS_MIN + [ @@ -446,7 +446,7 @@ async def _process_submitted_item(item: GatewayReplicaPipelineItem): async def _provision_gateway_replica( gateway_model: GatewayModel, - replica_model: GatewayComputeModel, + replica_model: GatewayReplicaModel, ) -> _GatewayReplicaUpdateMap: try: if replica_model.backend_id is None: # unexpected @@ -468,21 +468,21 @@ async def _provision_gateway_replica( compute = backend.compute() assert isinstance(compute, ComputeWithGatewaySupport) - compute_configuration = get_gateway_compute_configuration(replica_model, gateway_model) + replica_configuration = get_gateway_replica_configuration(replica_model, gateway_model) logger.debug( - "%s replica %d: creating gateway compute", + "%s replica %d: creating gateway replica", fmt(gateway_model), replica_model.replica_num, ) try: - gpd = await run_async(compute.create_gateway, compute_configuration) + gpd = await run_async(compute.create_gateway_replica, replica_configuration) except BackendError as e: status_message = f"Backend error: {repr(e)}" if len(e.args) > 0: status_message = str(e.args[0]) logger.warning( - "%s replica %d: failed to create gateway compute: %s", + "%s replica %d: failed to create gateway replica: %s", fmt(gateway_model), replica_model.replica_num, status_message, @@ -495,7 +495,7 @@ async def _provision_gateway_replica( ) except Exception: logger.exception( - "%s replica %d: unexpected error when creating gateway compute", + "%s replica %d: unexpected error when creating gateway replica", fmt(gateway_model), replica_model.replica_num, ) @@ -507,7 +507,7 @@ async def _provision_gateway_replica( ) logger.info( - "%s replica %d: gateway compute created", + "%s replica %d: gateway replica created", fmt(gateway_model), replica_model.replica_num, ) @@ -526,12 +526,12 @@ async def _process_provisioning_item(item: GatewayReplicaPipelineItem): item, replica_fields=_REPLICA_FIELDS_MIN + [ - GatewayComputeModel.ip_address, - GatewayComputeModel.ssh_private_key, - GatewayComputeModel.scale_in, - GatewayComputeModel.instance_id, - GatewayComputeModel.backend_id, - GatewayComputeModel.configuration, + GatewayReplicaModel.ip_address, + GatewayReplicaModel.ssh_private_key, + GatewayReplicaModel.scale_in, + GatewayReplicaModel.instance_id, + GatewayReplicaModel.backend_id, + GatewayReplicaModel.configuration, ], gateway_fields=_GATEWAY_FIELDS_MIN + [ @@ -603,7 +603,7 @@ async def _process_provisioning_item(item: GatewayReplicaPipelineItem): async def _register_replica_with_load_balancer( gateway_model: GatewayModel, - replica_model: GatewayComputeModel, + replica_model: GatewayReplicaModel, ) -> Optional[str]: """Registers the replica instance with the gateway's load balancer. Returns an error message on failure, None on success. @@ -646,40 +646,41 @@ async def _register_replica_with_load_balancer( async def _connect_and_configure_gateway_replica( gateway_model: GatewayModel, - gateway_compute: GatewayComputeModel, + gateway_replica: GatewayReplicaModel, ) -> Optional[str]: """Returns an error message on failure, None on success.""" logger.debug( - "%s replica %d: connecting to gateway compute", + "%s replica %d: connecting to gateway replica", fmt(gateway_model), - gateway_compute.replica_num, + gateway_replica.replica_num, ) # TODO: do only one connection/configuration attempt per pipeline tick. - # Blocking on connect_to_gateway_with_retry and configure_gateway now has these cons: + # Blocking on connect_to_gateway_replica_with_retry and configure_gateway_replica now has + # these cons: # - cannot terminate the gateway replica before it is provisioned because the DB model is locked # - connection retry counter is reset on server restart # - only one server replica is processing the gateway replica - connection = await gateways_services.connect_to_gateway_with_retry(gateway_compute) + connection = await gateways_services.connect_to_gateway_replica_with_retry(gateway_replica) if connection is None: logger.warning( - "%s replica %d: failed to connect to gateway compute", + "%s replica %d: failed to connect to gateway replica", fmt(gateway_model), - gateway_compute.replica_num, + gateway_replica.replica_num, ) - return "Failed to connect to gateway" + return "Failed to connect to gateway replica" try: - await gateways_services.configure_gateway(connection) + await gateways_services.configure_gateway_replica(connection) except Exception: logger.exception( - "%s replica %d: failed to configure gateway", + "%s replica %d: failed to configure gateway replica", fmt(gateway_model), - gateway_compute.replica_num, + gateway_replica.replica_num, ) - return "Failed to configure gateway" + return "Failed to configure gateway replica" logger.info( - "%s replica %d: gateway compute connected and configured", + "%s replica %d: gateway replica connected and configured", fmt(gateway_model), - gateway_compute.replica_num, + gateway_replica.replica_num, ) return None @@ -689,9 +690,9 @@ async def _process_running_item(item: GatewayReplicaPipelineItem): item, replica_fields=_REPLICA_FIELDS_MIN + [ - GatewayComputeModel.scale_in, - GatewayComputeModel.ip_address, - GatewayComputeModel.ssh_private_key, + GatewayReplicaModel.scale_in, + GatewayReplicaModel.ip_address, + GatewayReplicaModel.ssh_private_key, ], gateway_fields=_GATEWAY_FIELDS_MIN + [ @@ -801,7 +802,7 @@ async def _process_running_item(item: GatewayReplicaPipelineItem): async def _perform_state_sync( connection: GatewayConnection, gateway_model: GatewayModel, - replica_model: GatewayComputeModel, + replica_model: GatewayReplicaModel, run_models_by_id: dict[uuid.UUID, RunModel], job_models_by_id: dict[uuid.UUID, JobModel], plan: "_StateSyncPlan", @@ -1049,7 +1050,7 @@ class _ReconcileRegistrationRecordsResult: async def _reconcile_registration_records( session: AsyncSession, - replica_model: GatewayComputeModel, + replica_model: GatewayReplicaModel, initially_registered: list[ServiceListItem], sync_result: "_StateSyncResult", ) -> _ReconcileRegistrationRecordsResult: @@ -1168,7 +1169,7 @@ async def _reconcile_registration_records( async def _emit_state_sync_events( session, gateway_model: GatewayModel, - replica_model: GatewayComputeModel, + replica_model: GatewayReplicaModel, run_models_by_id: dict[uuid.UUID, RunModel], job_models_by_id: dict[uuid.UUID, JobModel], sync_result: "_StateSyncResult", @@ -1353,7 +1354,7 @@ async def _load_runs_and_jobs_for_state_sync( async def _register_service( client: GatewayClient, gateway_model: GatewayModel, - replica_model: GatewayComputeModel, + replica_model: GatewayReplicaModel, run_model: RunModel, ) -> None: run_spec = get_run_spec(run_model) @@ -1402,7 +1403,7 @@ async def _register_service( async def _register_replica( client: GatewayClient, gateway_model: GatewayModel, - replica_model: GatewayComputeModel, + replica_model: GatewayReplicaModel, run_model: RunModel, job_model: JobModel, ) -> None: @@ -1533,12 +1534,12 @@ async def _process_terminating_item(item: GatewayReplicaPipelineItem): item, replica_fields=_REPLICA_FIELDS_MIN + [ - GatewayComputeModel.instance_id, - GatewayComputeModel.ip_address, - GatewayComputeModel.backend_id, - GatewayComputeModel.configuration, - GatewayComputeModel.backend_data, - GatewayComputeModel.ssh_public_key, + GatewayReplicaModel.instance_id, + GatewayReplicaModel.ip_address, + GatewayReplicaModel.backend_id, + GatewayReplicaModel.configuration, + GatewayReplicaModel.backend_data, + GatewayReplicaModel.ssh_public_key, ], gateway_fields=_GATEWAY_FIELDS_MIN + [ @@ -1577,7 +1578,7 @@ async def _process_terminating_item(item: GatewayReplicaPipelineItem): return compute = backend.compute() assert isinstance(compute, ComputeWithGatewaySupport) - compute_configuration = get_gateway_compute_configuration(replica_model, gateway_model) + replica_configuration = get_gateway_replica_configuration(replica_model, gateway_model) if replica_model.instance_id is None: logger.warning( "%s replica %d: instance_id is None, skipping gateway replica termination", @@ -1594,20 +1595,20 @@ async def _process_terminating_item(item: GatewayReplicaPipelineItem): await _deregister_gateway_replica_from_load_balancer(compute, gateway_model, replica_model) logger.debug( - "%s replica %d: terminating gateway compute", + "%s replica %d: terminating gateway replica", fmt(gateway_model), replica_model.replica_num, ) try: await run_async( - compute.terminate_gateway, + compute.terminate_gateway_replica, replica_model.instance_id, - compute_configuration, + replica_configuration, replica_model.backend_data, ) except Exception: logger.exception( - "%s replica %d: error when terminating gateway compute", + "%s replica %d: error when terminating gateway replica", fmt(gateway_model), replica_model.replica_num, ) @@ -1615,7 +1616,7 @@ async def _process_terminating_item(item: GatewayReplicaPipelineItem): return logger.info( - "%s replica %d: gateway compute terminated", + "%s replica %d: gateway replica terminated", fmt(gateway_model), replica_model.replica_num, ) @@ -1629,7 +1630,7 @@ async def _process_terminating_item(item: GatewayReplicaPipelineItem): async def _deregister_gateway_replica_from_load_balancer( compute: ComputeWithGatewaySupport, gateway_model: GatewayModel, - replica_model: GatewayComputeModel, + replica_model: GatewayReplicaModel, ) -> None: if not isinstance(compute, ComputeWithGatewayLoadBalancerSupport): logger.error( diff --git a/src/dstack/_internal/server/background/pipeline_tasks/gateways.py b/src/dstack/_internal/server/background/pipeline_tasks/gateways.py index 55eb6b800e..0a879fd09b 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/gateways.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/gateways.py @@ -34,8 +34,8 @@ from dstack._internal.server.db import get_db, get_session_ctx from dstack._internal.server.models import ( BackendModel, - GatewayComputeModel, GatewayModel, + GatewayReplicaModel, ProjectModel, ) from dstack._internal.server.services import backends as backends_services @@ -43,8 +43,8 @@ from dstack._internal.server.services import gateways as gateways_services from dstack._internal.server.services.gateways import ( emit_gateway_status_change_event, - get_gateway_compute_models, get_gateway_lb_configuration, + get_gateway_replica_models, ) from dstack._internal.server.services.locking import get_locker from dstack._internal.server.services.logging import fmt @@ -146,11 +146,11 @@ async def fetch(self, limit: int) -> list[GatewayPipelineItem]: async with get_session_ctx() as session: now = get_current_datetime() active_replica_count_subquery = ( - select(func.count(GatewayComputeModel.id)) + select(func.count(GatewayReplicaModel.id)) .where( or_( - GatewayComputeModel.gateway_id == GatewayModel.id, - GatewayComputeModel.id == GatewayModel.gateway_compute_id, + GatewayReplicaModel.gateway_id == GatewayModel.id, + GatewayReplicaModel.id == GatewayModel.gateway_replica_id, ), *_get_active_replica_filters(), ) @@ -158,13 +158,13 @@ async def fetch(self, limit: int) -> list[GatewayPipelineItem]: .scalar_subquery() ) unmigrated_hostname_exists_subquery = ( - select(GatewayComputeModel.id) + select(GatewayReplicaModel.id) .where( or_( - GatewayComputeModel.gateway_id == GatewayModel.id, - GatewayComputeModel.id == GatewayModel.gateway_compute_id, + GatewayReplicaModel.gateway_id == GatewayModel.id, + GatewayReplicaModel.id == GatewayModel.gateway_replica_id, ), - GatewayComputeModel.hostname_deprecated_readonly.is_not(None), + GatewayReplicaModel.hostname_deprecated_readonly.is_not(None), ) .correlate(GatewayModel) .exists() @@ -286,7 +286,7 @@ class _GatewayUpdateMap(ItemUpdateMap, total=False): @dataclass class _ReplicaScalingResult: needs_more_replicas: bool = False - new_gateway_compute_models: list[GatewayComputeModel] = field(default_factory=list) + new_gateway_replica_models: list[GatewayReplicaModel] = field(default_factory=list) scale_in_replica_ids: list[uuid.UUID] = field(default_factory=list) gateway_update_map: _GatewayUpdateMap = field(default_factory=_GatewayUpdateMap) limit_reached: bool = False @@ -427,16 +427,16 @@ async def _process_provisioning_item(item: GatewayPipelineItem): ) .options(joinedload(GatewayModel.project).load_only(ProjectModel.name)) .options(joinedload(GatewayModel.backend).load_only(BackendModel.type)) - .options(joinedload(GatewayModel.gateway_compute)) + .options(joinedload(GatewayModel.gateway_replica)) .options( - selectinload(GatewayModel.gateway_computes).load_only( - GatewayComputeModel.id, - GatewayComputeModel.status, - GatewayComputeModel.replica_num, - GatewayComputeModel.created_at, - GatewayComputeModel.scale_in, - GatewayComputeModel.hostname_deprecated_readonly, - GatewayComputeModel.backend_data, + selectinload(GatewayModel.gateway_replicas).load_only( + GatewayReplicaModel.id, + GatewayReplicaModel.status, + GatewayReplicaModel.replica_num, + GatewayReplicaModel.created_at, + GatewayReplicaModel.scale_in, + GatewayReplicaModel.hostname_deprecated_readonly, + GatewayReplicaModel.backend_data, ) ) ) @@ -483,18 +483,18 @@ class _ProvisioningResult: def _process_provisioning_gateway(gateway_model: GatewayModel) -> _ProvisioningResult: - gateway_computes = get_gateway_compute_models(gateway_model) - # Provisioning gateways must have compute. - assert len(gateway_computes) > 0 + gateway_replicas = get_gateway_replica_models(gateway_model) + # Provisioning gateways must have a replica. + assert len(gateway_replicas) > 0 - scale_result = _reconcile_gateway_replica_count(gateway_model, gateway_computes) + scale_result = _reconcile_gateway_replica_count(gateway_model, gateway_replicas) statuses = { - gc.status - for gc in gateway_computes - if not gc.scale_in and gc.id not in scale_result.scale_in_replica_ids + replica.status + for replica in gateway_replicas + if not replica.scale_in and replica.id not in scale_result.scale_in_replica_ids } update_map = _migrate_hostname_and_backend_data_from_legacy_replica( - gateway_model, gateway_computes + gateway_model, gateway_replicas ) if statuses & {GatewayReplicaStatus.TERMINATING, GatewayReplicaStatus.TERMINATED}: @@ -535,16 +535,16 @@ async def _process_running_item(item: GatewayPipelineItem): ) .options(joinedload(GatewayModel.project).load_only(ProjectModel.name)) .options(joinedload(GatewayModel.backend).load_only(BackendModel.type)) - .options(joinedload(GatewayModel.gateway_compute)) + .options(joinedload(GatewayModel.gateway_replica)) .options( - selectinload(GatewayModel.gateway_computes).load_only( - GatewayComputeModel.id, - GatewayComputeModel.status, - GatewayComputeModel.replica_num, - GatewayComputeModel.created_at, - GatewayComputeModel.scale_in, - GatewayComputeModel.hostname_deprecated_readonly, - GatewayComputeModel.backend_data, + selectinload(GatewayModel.gateway_replicas).load_only( + GatewayReplicaModel.id, + GatewayReplicaModel.status, + GatewayReplicaModel.replica_num, + GatewayReplicaModel.created_at, + GatewayReplicaModel.scale_in, + GatewayReplicaModel.hostname_deprecated_readonly, + GatewayReplicaModel.backend_data, ) ) ) @@ -553,12 +553,12 @@ async def _process_running_item(item: GatewayPipelineItem): log_lock_token_mismatch(logger, item) return - gateway_computes = get_gateway_compute_models(gateway_model) - scale_result = _reconcile_gateway_replica_count(gateway_model, gateway_computes) + gateway_replicas = get_gateway_replica_models(gateway_model) + scale_result = _reconcile_gateway_replica_count(gateway_model, gateway_replicas) update_map = _GatewayUpdateMap() update_map.update( - _migrate_hostname_and_backend_data_from_legacy_replica(gateway_model, gateway_computes) + _migrate_hostname_and_backend_data_from_legacy_replica(gateway_model, gateway_replicas) ) update_map.update(scale_result.gateway_update_map) set_processed_update_map_fields(update_map) @@ -592,13 +592,13 @@ async def _process_to_be_deleted_item(item: GatewayPipelineItem): ) .options(joinedload(GatewayModel.project).joinedload(ProjectModel.backends)) .options(joinedload(GatewayModel.backend).load_only(BackendModel.type)) - .options(joinedload(GatewayModel.gateway_compute)) + .options(joinedload(GatewayModel.gateway_replica)) .options( - selectinload(GatewayModel.gateway_computes).load_only( - GatewayComputeModel.id, - GatewayComputeModel.status, - GatewayComputeModel.hostname_deprecated_readonly, - GatewayComputeModel.backend_data, + selectinload(GatewayModel.gateway_replicas).load_only( + GatewayReplicaModel.id, + GatewayReplicaModel.status, + GatewayReplicaModel.hostname_deprecated_readonly, + GatewayReplicaModel.backend_data, ) ) ) @@ -661,16 +661,16 @@ class _ProcessToBeDeletedResult: async def _process_to_be_deleted_gateway(gateway_model: GatewayModel) -> _ProcessToBeDeletedResult: - gateway_computes = get_gateway_compute_models(gateway_model) + gateway_replicas = get_gateway_replica_models(gateway_model) if update_map := _migrate_hostname_and_backend_data_from_legacy_replica( - gateway_model, gateway_computes + gateway_model, gateway_replicas ): return _ProcessToBeDeletedResult( delete_gateway=False, update_map=update_map, ) all_replicas_terminated = all( - gc.status == GatewayReplicaStatus.TERMINATED for gc in gateway_computes + replica.status == GatewayReplicaStatus.TERMINATED for replica in gateway_replicas ) lb_terminated = True if all_replicas_terminated and gateway_model.hostname is not None: @@ -714,7 +714,7 @@ async def _terminate_gateway_load_balancer(gateway_model: GatewayModel) -> bool: } -def _is_replica_active(replica: GatewayComputeModel) -> bool: +def _is_replica_active(replica: GatewayReplicaModel) -> bool: # should match _get_active_replica_filters return not replica.scale_in and replica.status not in ( GatewayReplicaStatus.TERMINATING, @@ -725,8 +725,8 @@ def _is_replica_active(replica: GatewayComputeModel) -> bool: def _get_active_replica_filters() -> list[ColumnElement[bool]]: # should match _is_replica_active return [ - GatewayComputeModel.scale_in == False, - GatewayComputeModel.status.not_in( + GatewayReplicaModel.scale_in == False, + GatewayReplicaModel.status.not_in( [GatewayReplicaStatus.TERMINATING, GatewayReplicaStatus.TERMINATED] ), ] @@ -734,7 +734,7 @@ def _get_active_replica_filters() -> list[ColumnElement[bool]]: def _reconcile_gateway_replica_count( gateway_model: GatewayModel, - gateway_replicas: list[GatewayComputeModel], + gateway_replicas: list[GatewayReplicaModel], ) -> _ReplicaScalingResult: desired_replica_count = gateway_model.desired_replica_count if desired_replica_count is None: # pre-0.21.0 gateway @@ -769,8 +769,8 @@ def _reconcile_gateway_replica_count( r.replica_num for r in gateway_replicas if r.status != GatewayReplicaStatus.TERMINATED } new_nums = itertools.islice(get_lowest_unused_nums(used_nums), diff) - new_gateway_compute_models = [ - gateways_services.create_gateway_compute_model( + new_gateway_replica_models = [ + gateways_services.create_gateway_replica_model( project_name=gateway_model.project.name, configuration=configuration, replica_num=replica_num, @@ -788,7 +788,7 @@ def _reconcile_gateway_replica_count( 0 if reset_replica_scale_attempt else gateway_model.replica_scale_attempt ) + 1 return _ReplicaScalingResult( - new_gateway_compute_models=new_gateway_compute_models, + new_gateway_replica_models=new_gateway_replica_models, gateway_update_map={ "replica_scale_attempt": new_attempt, "last_replica_scale_attempt_at": NOW_PLACEHOLDER, @@ -863,14 +863,14 @@ async def _apply_replica_scaling( gateway_model: GatewayModel, scale_result: _ReplicaScalingResult, ) -> None: - for gateway_compute_model in scale_result.new_gateway_compute_models: - session.add(gateway_compute_model) + for gateway_replica_model in scale_result.new_gateway_replica_models: + session.add(gateway_replica_model) if scale_result.scale_in_replica_ids: # The gateway pipeline does not need to lock gateway replicas — it only mutates `scale_in`, # which can only ever be flipped from False to True, so no races are expected. await session.execute( - update(GatewayComputeModel) - .where(GatewayComputeModel.id.in_(scale_result.scale_in_replica_ids)) + update(GatewayReplicaModel) + .where(GatewayReplicaModel.id.in_(scale_result.scale_in_replica_ids)) .values(scale_in=True) ) if scale_result.limit_reached: @@ -888,24 +888,24 @@ async def _apply_replica_scaling( def _migrate_hostname_and_backend_data_from_legacy_replica( gateway_model: GatewayModel, - gateway_computes: list[GatewayComputeModel], + gateway_replicas: list[GatewayReplicaModel], ) -> _GatewayUpdateMap: """ - Move `hostname` and `backend_data` from pre-0.21.0 GatewayComputeModel onto GatewayModel. + Move `hostname` and `backend_data` from pre-0.21.0 GatewayReplicaModel onto GatewayModel. Alembic migration ecc9e8a0bfac does the same thing. This function is a fallback in case any gateways are created by an older server replica after the migration passes. """ if gateway_model.hostname is not None: return {} - for gateway_compute in gateway_computes: - if gateway_compute.hostname_deprecated_readonly is not None: + for gateway_replica in gateway_replicas: + if gateway_replica.hostname_deprecated_readonly is not None: update_map: _GatewayUpdateMap = { - "hostname": gateway_compute.hostname_deprecated_readonly, - # Pre-0.21.0 AWS ACM gateways used GatewayComputeModel.backend_data exclusively + "hostname": gateway_replica.hostname_deprecated_readonly, + # Pre-0.21.0 AWS ACM gateways used GatewayReplicaModel.backend_data exclusively # for load-balancer related fields, and not for gateway replica instance fields. - # So GatewayComputeModel.backend_data is copied entirely. - "backend_data": gateway_compute.backend_data, + # So GatewayReplicaModel.backend_data is copied entirely. + "backend_data": gateway_replica.backend_data, } logger.info( "%s: migrating hostname and backend_data onto GatewayModel", fmt(gateway_model) 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 d93ed2f2b0..01ec6ba50b 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py @@ -59,8 +59,8 @@ from dstack._internal.server.models import ( ExportedFleetModel, FleetModel, - GatewayComputeModel, GatewayModel, + GatewayReplicaModel, ImportModel, InstanceModel, JobModel, @@ -79,7 +79,7 @@ get_instance_specific_mounts, resolve_provisioning_image, ) -from dstack._internal.server.services.gateways import get_gateway_compute_models +from dstack._internal.server.services.gateways import get_gateway_replica_models from dstack._internal.server.services.instances import ( get_instance_remote_connection_info, get_instance_ssh_private_keys, @@ -731,12 +731,12 @@ async def _fetch_run_model( if include_gateway: query = query.options( joinedload(RunModel.gateway) - .selectinload(GatewayModel.gateway_computes) - .load_only(GatewayComputeModel.id, GatewayComputeModel.status), + .selectinload(GatewayModel.gateway_replicas) + .load_only(GatewayReplicaModel.id, GatewayReplicaModel.status), ).options( joinedload(RunModel.gateway) - .joinedload(GatewayModel.gateway_compute) - .load_only(GatewayComputeModel.id, GatewayComputeModel.status), + .joinedload(GatewayModel.gateway_replica) + .load_only(GatewayReplicaModel.id, GatewayReplicaModel.status), ) if replica_num is not None: assert run_spec is not None, "run_spec must be provided when replica_num is set" @@ -1309,7 +1309,7 @@ def _job_gateway_registration_failed(gateway: GatewayModel | None, job_model: Jo return False running_gateway_replica_ids = { replica.id - for replica in get_gateway_compute_models(gateway) + for replica in get_gateway_replica_models(gateway) if replica.status == GatewayReplicaStatus.RUNNING } if not running_gateway_replica_ids: 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 2e34816dc9..20461f0814 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py @@ -37,8 +37,8 @@ ) from dstack._internal.server.db import get_db, get_session_ctx from dstack._internal.server.models import ( - GatewayComputeModel, GatewayModel, + GatewayReplicaModel, InstanceModel, JobModel, ProjectModel, @@ -47,7 +47,7 @@ from dstack._internal.server.services import events from dstack._internal.server.services.gateways import ( get_combined_gateway_stats, - get_gateway_compute_models, + get_gateway_replica_models, ) from dstack._internal.server.services.jobs import emit_job_status_change_event from dstack._internal.server.services.locking import get_locker @@ -334,7 +334,7 @@ async def _load_pending_context( gateway_stats = None if run_spec.configuration.type == "service" and run_model.gateway is not None: gateway_stats = await get_combined_gateway_stats( - get_gateway_compute_models(run_model.gateway), + get_gateway_replica_models(run_model.gateway), run_model.project.name, run_model.run_name, ) @@ -383,12 +383,12 @@ async def _refetch_locked_run_for_pending( ) .options(selectinload(RunModel.service_registrations)) .options( - joinedload(RunModel.gateway).selectinload(GatewayModel.gateway_computes), + joinedload(RunModel.gateway).selectinload(GatewayModel.gateway_replicas), ) .options( - joinedload(RunModel.gateway).joinedload(GatewayModel.gateway_compute), + joinedload(RunModel.gateway).joinedload(GatewayModel.gateway_replica), ) - .options(with_loader_criteria(GatewayComputeModel, GatewayComputeModel.deleted == False)) + .options(with_loader_criteria(GatewayReplicaModel, GatewayReplicaModel.deleted == False)) .execution_options(populate_existing=True) ) return res.unique().scalar_one_or_none() @@ -536,7 +536,7 @@ async def _load_active_context( gateway_stats = None if run_spec.configuration.type == "service" and run_model.gateway is not None: gateway_stats = await get_combined_gateway_stats( - get_gateway_compute_models(run_model.gateway), + get_gateway_replica_models(run_model.gateway), run_model.project.name, run_model.run_name, ) @@ -590,12 +590,12 @@ async def _refetch_locked_run_for_active( ) .options(selectinload(RunModel.service_registrations)) .options( - joinedload(RunModel.gateway).selectinload(GatewayModel.gateway_computes), + joinedload(RunModel.gateway).selectinload(GatewayModel.gateway_replicas), ) .options( - joinedload(RunModel.gateway).joinedload(GatewayModel.gateway_compute), + joinedload(RunModel.gateway).joinedload(GatewayModel.gateway_replica), ) - .options(with_loader_criteria(GatewayComputeModel, GatewayComputeModel.deleted == False)) + .options(with_loader_criteria(GatewayReplicaModel, GatewayReplicaModel.deleted == False)) .execution_options(populate_existing=True) ) return res.unique().scalar_one_or_none() diff --git a/src/dstack/_internal/server/background/scheduled_tasks/gateways.py b/src/dstack/_internal/server/background/scheduled_tasks/gateways.py index 2d7dbb93fd..4bcd90f42e 100644 --- a/src/dstack/_internal/server/background/scheduled_tasks/gateways.py +++ b/src/dstack/_internal/server/background/scheduled_tasks/gateways.py @@ -6,7 +6,7 @@ from dstack._internal.core.errors import SSHError from dstack._internal.server.db import get_db, get_session_ctx from dstack._internal.server.models import ( - GatewayComputeModel, + GatewayReplicaModel, ) from dstack._internal.server.services.gateways import ( GatewayConnection, @@ -32,7 +32,7 @@ async def process_gateways_connections(): async def _remove_inactive_connections(): async with get_session_ctx() as session: res = await session.execute( - select(GatewayComputeModel.ip_address).where(GatewayComputeModel.active == True) + select(GatewayReplicaModel.ip_address).where(GatewayReplicaModel.active == True) ) active_connection_ips = {ip for ip in res.scalars().all() if ip is not None} for conn in await gateway_connections_pool.all(): diff --git a/src/dstack/_internal/server/models.py b/src/dstack/_internal/server/models.py index 70f28eefdd..a1225fb8eb 100644 --- a/src/dstack/_internal/server/models.py +++ b/src/dstack/_internal/server/models.py @@ -656,25 +656,26 @@ class GatewayModel(PipelineModelMixin, BaseModel): """Backend-specific load balancer resource data in JSON. """ - gateway_compute_id: Mapped[Optional[uuid.UUID]] = mapped_column( - ForeignKey("gateway_computes.id", ondelete="CASCADE") + gateway_replica_id: Mapped[Optional[uuid.UUID]] = mapped_column( + "gateway_compute_id", + ForeignKey("gateway_computes.id", ondelete="CASCADE"), ) - gateway_compute: Mapped[Optional["GatewayComputeModel"]] = relationship( - foreign_keys=[gateway_compute_id], + gateway_replica: Mapped[Optional["GatewayReplicaModel"]] = relationship( + foreign_keys=[gateway_replica_id], back_populates="legacy_gateway", ) """ - Relationship with gateway computes for pre-0.20.25 gateways. - Use `get_gateway_compute_models()` for version-agnostic gateway compute retrieval. + Relationship with the gateway replica for pre-0.20.25 gateways. + Use `get_gateway_replica_models()` for version-agnostic gateway replica retrieval. """ - gateway_computes: Mapped[List["GatewayComputeModel"]] = relationship( + gateway_replicas: Mapped[List["GatewayReplicaModel"]] = relationship( back_populates="gateway", - foreign_keys="GatewayComputeModel.gateway_id", + foreign_keys="GatewayReplicaModel.gateway_id", ) """ - Relationship with gateway computes. - Pre-0.20.25 gateways can have an extra compute model referenced by `GatewayModel.gateway_compute`. - Use `get_gateway_compute_models()` for version-agnostic gateway compute retrieval. + Relationship with gateway replicas. + Pre-0.20.25 gateways can have an extra replica referenced by `GatewayModel.gateway_replica`. + Use `get_gateway_replica_models()` for version-agnostic gateway replica retrieval. """ runs: Mapped[List["RunModel"]] = relationship(back_populates="gateway") @@ -684,11 +685,8 @@ class GatewayModel(PipelineModelMixin, BaseModel): # TODO: Add pipeline index ("ix_gateways_pipeline_fetch_q") if gateways become soft-deleted. -class GatewayComputeModel(PipelineModelMixin, BaseModel): - """A single gateway replica. - **TODO**: consider renaming to `GatewayReplicaModel`. - """ - +class GatewayReplicaModel(PipelineModelMixin, BaseModel): + # "gateway compute" is a legacy term superseded by "gateway replica" __tablename__ = "gateway_computes" id: Mapped[uuid.UUID] = mapped_column( @@ -710,7 +708,7 @@ class GatewayComputeModel(PipelineModelMixin, BaseModel): """Replaced by GatewayModel.hostname since 0.21.0""" configuration: Mapped[Optional[str]] = mapped_column(Text) """`configuration` is optional for compatibility with pre-0.18.2 gateways. - Use `get_gateway_compute_configuration` to construct `configuration` for old gateways. + Use `get_gateway_replica_configuration` to construct `configuration` for old gateways. """ backend_data: Mapped[Optional[str]] = mapped_column(Text) region: Mapped[Optional[str]] = mapped_column(String(100)) @@ -723,20 +721,20 @@ class GatewayComputeModel(PipelineModelMixin, BaseModel): ) ) gateway: Mapped[Optional["GatewayModel"]] = relationship( - back_populates="gateway_computes", + back_populates="gateway_replicas", foreign_keys=[gateway_id], ) """ - Gateway. Can be None for pre-0.20.25 gateways, which use GatewayModel.gateway_compute_id to + Gateway. Can be None for pre-0.20.25 gateways, which use GatewayModel.gateway_replica_id to establish the relationship. """ legacy_gateway: Mapped[Optional["GatewayModel"]] = relationship( - back_populates="gateway_compute", - foreign_keys="GatewayModel.gateway_compute_id", + back_populates="gateway_replica", + foreign_keys="GatewayModel.gateway_replica_id", viewonly=True, ) """ - Gateway for pre-0.20.25 gateways, where GatewayModel.gateway_compute_id points to this replica. + Gateway for pre-0.20.25 gateways, where GatewayModel.gateway_replica_id points to this replica. Use `gateway or legacy_gateway` to get the gateway regardless of version. """ @@ -786,7 +784,7 @@ class ServiceRegistrationModel(BaseModel): gateway_replica_id: Mapped[uuid.UUID] = mapped_column( ForeignKey("gateway_computes.id", ondelete="CASCADE"), index=True ) - gateway_replica: Mapped["GatewayComputeModel"] = relationship( + gateway_replica: Mapped["GatewayReplicaModel"] = relationship( back_populates="service_registrations" ) is_registered: Mapped[bool] = mapped_column(Boolean, default=False) @@ -820,7 +818,7 @@ class ServiceReplicaRegistrationModel(BaseModel): gateway_replica_id: Mapped[uuid.UUID] = mapped_column( ForeignKey("gateway_computes.id", ondelete="CASCADE"), index=True ) - gateway_replica: Mapped["GatewayComputeModel"] = relationship( + gateway_replica: Mapped["GatewayReplicaModel"] = relationship( back_populates="service_replica_registrations" ) is_registered: Mapped[bool] = mapped_column(Boolean, default=False) diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py index 60ff7edf32..4c9a696359 100644 --- a/src/dstack/_internal/server/services/gateways/__init__.py +++ b/src/dstack/_internal/server/services/gateways/__init__.py @@ -38,11 +38,11 @@ GATEWAY_REPLICAS_DEFAULT, ApplyGatewayPlanInput, Gateway, - GatewayComputeConfiguration, GatewayConfiguration, GatewayLoadBalancerConfiguration, GatewayPlan, GatewayReplica, + GatewayReplicaConfiguration, GatewayReplicaStatus, GatewaySpec, GatewayStatus, @@ -58,8 +58,8 @@ from dstack._internal.server.models import ( BackendModel, ExportedGatewayModel, - GatewayComputeModel, GatewayModel, + GatewayReplicaModel, ImportModel, ProjectModel, UserModel, @@ -158,7 +158,7 @@ async def list_project_gateways( session=session, project=project, include_imported=include_imported, - load_gateway_compute=True, + load_gateway_replica=True, load_backend_type=True, ) return [ @@ -174,7 +174,7 @@ async def get_gateway_by_name( session=session, project=project, ref=EntityReference(name=name, project=None), - load_gateway_compute=True, + load_gateway_replica=True, load_backend_type=True, ) if gateway is None: @@ -182,20 +182,20 @@ async def get_gateway_by_name( return gateway_model_to_gateway(gateway, default_gateway_id=project.default_gateway_id) -def create_gateway_compute_model( +def create_gateway_replica_model( project_name: str, configuration: GatewayConfiguration, replica_num: int, gateway_id: uuid.UUID, backend_id: uuid.UUID, -) -> GatewayComputeModel: +) -> GatewayReplicaModel: assert configuration.name is not None private_bytes, public_bytes = crypto.generate_rsa_key_pair_bytes() gateway_ssh_private_key = private_bytes.decode() gateway_ssh_public_key = public_bytes.decode() - compute_configuration = GatewayComputeConfiguration( + replica_configuration = GatewayReplicaConfiguration( project_name=project_name, instance_name=f"{configuration.name}-{replica_num}", backend=configuration.backend, @@ -208,11 +208,11 @@ def create_gateway_compute_model( ) now = get_current_datetime() - return GatewayComputeModel( + return GatewayReplicaModel( gateway_id=gateway_id, backend_id=backend_id, replica_num=replica_num, - configuration=compute_configuration.model_dump_json(), + configuration=replica_configuration.model_dump_json(), ssh_private_key=gateway_ssh_private_key, ssh_public_key=gateway_ssh_public_key, status=GatewayReplicaStatus.SUBMITTED, @@ -299,7 +299,7 @@ async def create_gateway( session=session, project=project, ref=EntityReference(name=configuration.name, project=None), - load_gateway_compute=True, + load_gateway_replica=True, load_backend_type=True, ) assert gateway is not None @@ -308,18 +308,18 @@ async def create_gateway( ) -async def connect_to_gateway_with_retry( - gateway_compute: GatewayComputeModel, +async def connect_to_gateway_replica_with_retry( + gateway_replica: GatewayReplicaModel, ) -> Optional[GatewayConnection]: """ - Create gateway connection and add it to connection pool. - Give gateway sufficient time to become available. In the case of gateway + Create a gateway replica connection and add it to the connection pool. + Give the gateway replica sufficient time to become available. In the case of the replica being accessed via domain (e.g. Kubernetes LB), it may take some time before the domain can be resolved. """ - if gateway_compute.ip_address is None: - logger.warning("Gateway replica %s has no ip_address, cannot connect", gateway_compute.id) + if gateway_replica.ip_address is None: + logger.warning("Gateway replica %s has no ip_address, cannot connect", gateway_replica.id) return None connection = None @@ -327,15 +327,19 @@ async def connect_to_gateway_with_retry( for attempt in range(GATEWAY_CONNECT_ATTEMPTS): try: connection = await gateway_connections_pool.get_or_add( - gateway_compute.ip_address, gateway_compute.ssh_private_key + gateway_replica.ip_address, gateway_replica.ssh_private_key ) break except SSHError as e: if attempt < GATEWAY_CONNECT_ATTEMPTS - 1: - logger.debug("Failed to connect to gateway %s: %s", gateway_compute.ip_address, e) + logger.debug( + "Failed to connect to gateway replica %s: %s", gateway_replica.ip_address, e + ) await asyncio.sleep(GATEWAY_CONNECT_DELAY) else: - logger.error("Failed to connect to gateway %s: %s", gateway_compute.ip_address, e) + logger.error( + "Failed to connect to gateway replica %s: %s", gateway_replica.ip_address, e + ) return connection @@ -498,7 +502,7 @@ async def list_project_gateway_models( session: AsyncSession, project: ProjectModel, include_imported: bool = False, - load_gateway_compute: bool = False, + load_gateway_replica: bool = False, load_backend_type: bool = False, ) -> Sequence[GatewayModel]: stmt = select(GatewayModel) @@ -515,15 +519,15 @@ async def list_project_gateway_models( ).options(joinedload(GatewayModel.project).load_only(ProjectModel.id, ProjectModel.name)) else: stmt = stmt.where(GatewayModel.project_id == project.id) - if load_gateway_compute: + if load_gateway_replica: stmt = stmt.options( - joinedload(GatewayModel.gateway_compute) - .joinedload(GatewayComputeModel.backend) + joinedload(GatewayModel.gateway_replica) + .joinedload(GatewayReplicaModel.backend) .load_only(BackendModel.type) ) stmt = stmt.options( - selectinload(GatewayModel.gateway_computes) - .joinedload(GatewayComputeModel.backend) + selectinload(GatewayModel.gateway_replicas) + .joinedload(GatewayReplicaModel.backend) .load_only(BackendModel.type) ) if load_backend_type: @@ -536,7 +540,7 @@ async def get_project_gateway_model_by_reference( session: AsyncSession, project: ProjectModel, ref: EntityReference, - load_gateway_compute: bool = False, + load_gateway_replica: bool = False, load_backend_type: bool = False, ) -> Optional[GatewayModel]: stmt = select(GatewayModel).where(GatewayModel.name == ref.name) @@ -552,15 +556,15 @@ async def get_project_gateway_model_by_reference( ProjectModel.name == ref.project, ) ) - if load_gateway_compute: + if load_gateway_replica: stmt = stmt.options( - joinedload(GatewayModel.gateway_compute) - .joinedload(GatewayComputeModel.backend) + joinedload(GatewayModel.gateway_replica) + .joinedload(GatewayReplicaModel.backend) .load_only(BackendModel.type) ) stmt = stmt.options( - selectinload(GatewayModel.gateway_computes) - .joinedload(GatewayComputeModel.backend) + selectinload(GatewayModel.gateway_replicas) + .joinedload(GatewayReplicaModel.backend) .load_only(BackendModel.type) ) if load_backend_type: @@ -597,13 +601,13 @@ async def get_project_gateway_model_by_name_for_update( select(GatewayModel) .where(GatewayModel.id.in_([gateway_id]), *filters) .options( - joinedload(GatewayModel.gateway_compute) - .joinedload(GatewayComputeModel.backend) + joinedload(GatewayModel.gateway_replica) + .joinedload(GatewayReplicaModel.backend) .load_only(BackendModel.type) ) .options( - selectinload(GatewayModel.gateway_computes) - .joinedload(GatewayComputeModel.backend) + selectinload(GatewayModel.gateway_replicas) + .joinedload(GatewayReplicaModel.backend) .load_only(BackendModel.type) ) .options(joinedload(GatewayModel.backend).load_only(BackendModel.type)) @@ -615,7 +619,7 @@ async def get_project_gateway_model_by_name_for_update( async def get_project_default_gateway_model( session: AsyncSession, project: ProjectModel, - load_gateway_compute: bool = False, + load_gateway_replica: bool = False, load_backend_type: bool = False, ) -> Optional[GatewayModel]: stmt = select(GatewayModel).where( @@ -630,15 +634,15 @@ async def get_project_default_gateway_model( ), ), ) - if load_gateway_compute: + if load_gateway_replica: stmt = stmt.options( - joinedload(GatewayModel.gateway_compute) - .joinedload(GatewayComputeModel.backend) + joinedload(GatewayModel.gateway_replica) + .joinedload(GatewayReplicaModel.backend) .load_only(BackendModel.type) ) stmt = stmt.options( - selectinload(GatewayModel.gateway_computes) - .joinedload(GatewayComputeModel.backend) + selectinload(GatewayModel.gateway_replicas) + .joinedload(GatewayReplicaModel.backend) .load_only(BackendModel.type) ) if load_backend_type: @@ -658,16 +662,16 @@ async def generate_gateway_name(session: AsyncSession, project: ProjectModel) -> # TODO: Connect to gateway outside session async def get_or_add_gateway_connections( - gateway_replicas: Sequence[GatewayComputeModel], + gateway_replicas: Sequence[GatewayReplicaModel], ) -> List[GatewayConnection]: running_replicas = [r for r in gateway_replicas if r.status == GatewayReplicaStatus.RUNNING] if not running_replicas: - raise GatewayError("Gateway compute not found") + raise GatewayError("Gateway replica not found") connections: List[GatewayConnection] = [] for replica in running_replicas: if replica.ip_address is None: logger.warning("Gateway replica %s has no ip_address", replica.id) - raise GatewayError("Failed to connect to gateway") + raise GatewayError("Failed to connect to gateway replica") try: conn = await gateway_connections_pool.get_or_add( hostname=replica.ip_address, @@ -675,13 +679,13 @@ async def get_or_add_gateway_connections( ) connections.append(conn) except Exception as e: - logger.warning("Failed to connect to gateway %s: %s", replica.ip_address, e) - raise GatewayError("Failed to connect to gateway") + logger.warning("Failed to connect to gateway replica %s: %s", replica.ip_address, e) + raise GatewayError("Failed to connect to gateway replica") return connections async def get_combined_gateway_stats( - gateway_replicas: Sequence[GatewayComputeModel], + gateway_replicas: Sequence[GatewayReplicaModel], project_name: str, run_name: str, ) -> Optional[PerWindowStats]: @@ -723,78 +727,86 @@ def _merge_per_window_stats(stats_per_gateway_replica: list[PerWindowStats]) -> async def init_gateways(session: AsyncSession): res = await session.execute( - select(GatewayComputeModel).where( - GatewayComputeModel.status == GatewayReplicaStatus.RUNNING, - GatewayComputeModel.active == True, - GatewayComputeModel.deleted == False, + select(GatewayReplicaModel).where( + GatewayReplicaModel.status == GatewayReplicaStatus.RUNNING, + GatewayReplicaModel.active == True, + GatewayReplicaModel.deleted == False, ) ) - gateway_computes = res.scalars().all() + gateway_replicas = res.scalars().all() - if len(gateway_computes) > 0: - logger.info(f"Connecting to {len(gateway_computes)} gateways...", {"show_path": False}) + if len(gateway_replicas) > 0: + logger.info( + f"Connecting to {len(gateway_replicas)} gateway replicas...", {"show_path": False} + ) async with advisory_lock_ctx( bind=session, dialect_name=get_db().dialect_name, resource="gateway_tunnels", ): - for gateway, error in await gather_map_async( - [g for g in gateway_computes if g.ip_address], - lambda g: gateway_connections_pool.get_or_add( - get_or_error(g.ip_address), g.ssh_private_key, True + for gateway_replica, error in await gather_map_async( + [r for r in gateway_replicas if r.ip_address], + lambda r: gateway_connections_pool.get_or_add( + get_or_error(r.ip_address), r.ssh_private_key, True ), return_exceptions=True, ): if isinstance(error, Exception): - logger.warning("Failed to connect to gateway %s: %s", gateway.ip_address, error) + logger.warning( + "Failed to connect to gateway replica %s: %s", + gateway_replica.ip_address, + error, + ) if settings.SKIP_GATEWAY_UPDATE: - logger.debug("Skipping gateways update due to DSTACK_SKIP_GATEWAY_UPDATE env variable") + logger.debug( + "Skipping gateway replicas update due to DSTACK_SKIP_GATEWAY_UPDATE env variable" + ) else: build = get_dstack_runner_version() or "latest" - for gateway_compute, res in await gather_map_async( - gateway_computes, - lambda c: _update_gateway(c, build), + for gateway_replica, res in await gather_map_async( + gateway_replicas, + lambda r: _update_gateway_replica(r, build), return_exceptions=True, ): if isinstance(res, Exception): logger.warning( - "Failed to update gateway %s: %s", gateway_compute.ip_address, res + "Failed to update gateway replica %s: %s", gateway_replica.ip_address, res ) elif isinstance(res, bool) and res: - gateway_compute.app_updated_at = get_current_datetime() + gateway_replica.app_updated_at = get_current_datetime() - for gateway_compute, error in await gather_map_async( + for connection, error in await gather_map_async( await gateway_connections_pool.all(), - # Need several attempts to handle short gateway downtime after update - partial(configure_gateway, attempts=7), + # Need several attempts to handle short gateway replica downtime after update + partial(configure_gateway_replica, attempts=7), return_exceptions=True, ): if isinstance(error, Exception): logger.warning( - "Failed to configure gateway %s: %r", gateway_compute.ip_address, error + "Failed to configure gateway replica %s: %r", connection.ip_address, error ) -async def _update_gateway(gateway_compute_model: GatewayComputeModel, build: str) -> bool: - if gateway_compute_model.ip_address is None: +async def _update_gateway_replica(gateway_replica_model: GatewayReplicaModel, build: str) -> bool: + if gateway_replica_model.ip_address is None: logger.warning( - "Gateway replica %s has no ip_address, cannot update", gateway_compute_model.id + "Gateway replica %s has no ip_address, cannot update", gateway_replica_model.id ) return False - if _recently_updated(gateway_compute_model): + if _recently_updated(gateway_replica_model): logger.debug( - "Skipping gateway %s update. Gateway was recently updated.", - gateway_compute_model.ip_address, + "Skipping gateway replica %s update. Gateway replica was recently updated.", + gateway_replica_model.ip_address, ) return False connection = await gateway_connections_pool.get_or_add( - gateway_compute_model.ip_address, - gateway_compute_model.ssh_private_key, + gateway_replica_model.ip_address, + gateway_replica_model.ssh_private_key, ) - logger.debug("Updating gateway %s", connection.ip_address) + logger.debug("Updating gateway replica %s", connection.ip_address) # Build package spec with extras and wheel URL gateway_package = get_dstack_gateway_wheel(build) @@ -806,27 +818,27 @@ async def _update_gateway(gateway_compute_model: GatewayComputeModel, build: str ] stdout = await connection.tunnel.aexec("/bin/sh -c '" + " && ".join(commands) + "'") if "Update successfully completed" in stdout: - logger.info("Gateway %s updated", connection.ip_address) + logger.info("Gateway replica %s updated", connection.ip_address) return True return False -def _recently_updated(gateway_compute_model: GatewayComputeModel) -> bool: - return gateway_compute_model.app_updated_at.replace( +def _recently_updated(gateway_replica_model: GatewayReplicaModel) -> bool: + return gateway_replica_model.app_updated_at.replace( tzinfo=datetime.timezone.utc ) > get_current_datetime() - timedelta(seconds=60) -async def configure_gateway( +async def configure_gateway_replica( connection: GatewayConnection, attempts: int = GATEWAY_CONFIGURE_ATTEMPTS, ) -> None: """ - Try submitting gateway config several times in case gateway's HTTP server is not + Try submitting gateway config to the replica several times in case its HTTP server is not running yet """ - logger.debug("Configuring gateway %s", connection.ip_address) + logger.debug("Configuring gateway replica %s", connection.ip_address) for attempt in range(attempts - 1): try: @@ -835,7 +847,7 @@ async def configure_gateway( break except httpx.RequestError as e: logger.debug( - "Failed attempt %s/%s at configuring gateway %s: %r", + "Failed attempt %s/%s at configuring gateway replica %s: %r", attempt + 1, attempts, connection.ip_address, @@ -846,14 +858,14 @@ async def configure_gateway( async with connection.client() as client: await client.submit_gateway_config() - logger.info("Gateway %s configured", connection.ip_address) + logger.info("Gateway replica %s configured", connection.ip_address) -def get_gateway_compute_models(gateway_model: GatewayModel) -> List[GatewayComputeModel]: - computes = list(gateway_model.gateway_computes) - if gateway_model.gateway_compute is not None: # pre-0.20.25 gateway - computes.append(gateway_model.gateway_compute) - return computes +def get_gateway_replica_models(gateway_model: GatewayModel) -> List[GatewayReplicaModel]: + replicas = list(gateway_model.gateway_replicas) + if gateway_model.gateway_replica is not None: # pre-0.20.25 gateway + replicas.append(gateway_model.gateway_replica) + return replicas def get_gateway_configuration(gateway_model: GatewayModel) -> GatewayConfiguration: @@ -868,23 +880,23 @@ def get_gateway_configuration(gateway_model: GatewayModel) -> GatewayConfigurati ) -def get_gateway_compute_configuration( - gateway_compute: GatewayComputeModel, +def get_gateway_replica_configuration( + gateway_replica: GatewayReplicaModel, gateway_model: GatewayModel, -) -> GatewayComputeConfiguration: - if gateway_compute.configuration is not None: +) -> GatewayReplicaConfiguration: + if gateway_replica.configuration is not None: return validate_json_extra_ignore( - GatewayComputeConfiguration, gateway_compute.configuration + GatewayReplicaConfiguration, gateway_replica.configuration ) - # Handle gateways created before GatewayComputeConfiguration was introduced + # Handle gateways created before GatewayReplicaConfiguration was introduced gateway_configuration = get_gateway_configuration(gateway_model) - return GatewayComputeConfiguration( + return GatewayReplicaConfiguration( project_name=gateway_model.project.name, - instance_name=f"{gateway_model.name}-{gateway_compute.replica_num}", + instance_name=f"{gateway_model.name}-{gateway_replica.replica_num}", backend=gateway_configuration.backend, region=gateway_configuration.region, public_ip=True, - ssh_key_pub=gateway_compute.ssh_public_key, + ssh_key_pub=gateway_replica.ssh_public_key, certificate=LetsEncryptGatewayCertificate(), ) @@ -916,25 +928,25 @@ def gateway_model_to_gateway( configuration = get_gateway_configuration(gateway_model) configuration.default = is_default - all_compute_models = sorted( - get_gateway_compute_models(gateway_model), key=lambda c: c.replica_num + all_replica_models = sorted( + get_gateway_replica_models(gateway_model), key=lambda r: r.replica_num ) - relevant_compute_models = [] - for replica_num, compute_models_for_num in itertools.groupby( - all_compute_models, key=lambda c: c.replica_num + relevant_replica_models = [] + for replica_num, replica_models_for_num in itertools.groupby( + all_replica_models, key=lambda r: r.replica_num ): - relevant_compute_models.append(max(compute_models_for_num, key=lambda c: c.created_at)) + relevant_replica_models.append(max(replica_models_for_num, key=lambda r: r.created_at)) replicas = [] - for compute in relevant_compute_models: + for replica_model in relevant_replica_models: replicas.append( GatewayReplica( - hostname=compute.ip_address, - replica_num=compute.replica_num, - backend=compute.backend.type if compute.backend else None, - region=compute.region, - created_at=compute.created_at, - status=compute.status, - status_message=compute.status_message, + hostname=replica_model.ip_address, + replica_num=replica_model.replica_num, + backend=replica_model.backend.type if replica_model.backend else None, + region=replica_model.region, + created_at=replica_model.created_at, + status=replica_model.status, + status_message=replica_model.status_message, ) ) @@ -974,7 +986,7 @@ async def get_plan( session=session, project=project, ref=EntityReference(name=effective_spec.configuration.name, project=None), - load_gateway_compute=True, + load_gateway_replica=True, load_backend_type=True, ) if current_gateway_model is not None: diff --git a/src/dstack/_internal/server/services/gateways/connection.py b/src/dstack/_internal/server/services/gateways/connection.py index fe8187188f..555f9aaa09 100644 --- a/src/dstack/_internal/server/services/gateways/connection.py +++ b/src/dstack/_internal/server/services/gateways/connection.py @@ -32,10 +32,10 @@ def _get_connections_dir() -> Path: class GatewayConnection: """ - `GatewayConnection` instances persist for the lifetime of the gateway. + `GatewayConnection` instances persist for the lifetime of the gateway replica. - The `GatewayConnection.tunnel` is responsible for establishing a bidirectional tunnel with the gateway. - The local tunnel is used for the gateway management. + The `GatewayConnection.tunnel` is responsible for establishing a bidirectional tunnel with the + gateway replica. The local tunnel is used for the gateway replica management. The reverse tunnel is used for authorizing dstack tokens. """ diff --git a/src/dstack/_internal/server/services/runs/__init__.py b/src/dstack/_internal/server/services/runs/__init__.py index 6431373612..47903412e5 100644 --- a/src/dstack/_internal/server/services/runs/__init__.py +++ b/src/dstack/_internal/server/services/runs/__init__.py @@ -56,7 +56,7 @@ from dstack._internal.server.services import events, services from dstack._internal.server.services import projects as projects_services from dstack._internal.server.services import repos as repos_services -from dstack._internal.server.services.gateways import get_gateway_compute_models +from dstack._internal.server.services.gateways import get_gateway_replica_models from dstack._internal.server.services.jobs import ( check_can_attach_job_volumes, get_job_configured_volumes, @@ -160,7 +160,7 @@ def gateway_registration_failed(run_model: RunModel) -> bool: return False running_gateway_replica_ids = { replica.id - for replica in get_gateway_compute_models(run_model.gateway) + for replica in get_gateway_replica_models(run_model.gateway) if replica.status == GatewayReplicaStatus.RUNNING } if not running_gateway_replica_ids: diff --git a/src/dstack/_internal/server/services/runs/replicas.py b/src/dstack/_internal/server/services/runs/replicas.py index b830d18051..e579335e5c 100644 --- a/src/dstack/_internal/server/services/runs/replicas.py +++ b/src/dstack/_internal/server/services/runs/replicas.py @@ -7,7 +7,7 @@ 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.gateways import get_gateway_replica_models from dstack._internal.server.services.jobs import ( get_job_provisioning_data, get_job_spec, @@ -170,7 +170,7 @@ def is_replica_receiving_traffic(run_model: RunModel, jobs: list[JobModel]) -> b return True running_gateway_replica_ids = { replica.id - for replica in get_gateway_compute_models(run_model.gateway) + for replica in get_gateway_replica_models(run_model.gateway) if replica.status == GatewayReplicaStatus.RUNNING } if not running_gateway_replica_ids: diff --git a/src/dstack/_internal/server/services/services/__init__.py b/src/dstack/_internal/server/services/services/__init__.py index b5e03a8be4..a08ed1a8e7 100644 --- a/src/dstack/_internal/server/services/services/__init__.py +++ b/src/dstack/_internal/server/services/services/__init__.py @@ -21,8 +21,8 @@ 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_gateway_replica_models, get_project_default_gateway_model, get_project_gateway_model_by_reference, ) @@ -44,7 +44,7 @@ async def register_service(session: AsyncSession, run_model: RunModel, run_spec: session=session, project=run_model.project, ref=gateway_reference, - load_gateway_compute=True, + load_gateway_replica=True, load_backend_type=True, ) if gateway is None: @@ -62,7 +62,7 @@ async def register_service(session: AsyncSession, run_model: RunModel, run_spec: gateway = await get_project_default_gateway_model( session=session, project=run_model.project, - load_gateway_compute=True, + load_gateway_replica=True, load_backend_type=True, ) if gateway is None and run_spec.configuration.gateway == True: @@ -88,7 +88,7 @@ async def _register_service_in_gateway( ) -> ServiceSpec: assert run_spec.configuration.type == "service" - if not get_gateway_compute_models(gateway): + if not get_gateway_replica_models(gateway): raise ServerClientError("Gateway has no instance associated with it") if gateway.status != GatewayStatus.RUNNING: diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py index 6f0ba64a62..0fa5b35568 100644 --- a/src/dstack/_internal/server/testing/common.py +++ b/src/dstack/_internal/server/testing/common.py @@ -49,8 +49,8 @@ from dstack._internal.core.models.gateways import ( GATEWAY_REPLICAS_DEFAULT, AnyGatewayCertificate, - GatewayComputeConfiguration, GatewayConfiguration, + GatewayReplicaConfiguration, GatewayReplicaStatus, GatewayStatus, LetsEncryptGatewayCertificate, @@ -114,8 +114,8 @@ ExportModel, FileArchiveModel, FleetModel, - GatewayComputeModel, GatewayModel, + GatewayReplicaModel, ImportModel, InstanceHealthCheckModel, InstanceModel, @@ -709,7 +709,7 @@ async def create_gateway( return gateway -async def create_gateway_compute( +async def create_gateway_replica( session: AsyncSession, gateway_id: Optional[UUID] = None, backend_id: Optional[UUID] = None, @@ -726,10 +726,10 @@ async def create_gateway_compute( populate_configuration: bool = True, hostname_deprecated_readonly: Optional[str] = None, backend_data: Optional[str] = None, -) -> GatewayComputeModel: +) -> GatewayReplicaModel: """ Args: - populate_configuration: whether to populate GatewayComputeModel.configuration. + populate_configuration: whether to populate GatewayReplicaModel.configuration. True - 0.18.2+ gateways, False - legacy pre-0.18.2 gateways. Prefer testing against both in major test cases. """ @@ -740,7 +740,7 @@ async def create_gateway_compute( assert backend is not None backend_type = backend.type assert region is not None - configuration = GatewayComputeConfiguration( + configuration = GatewayReplicaConfiguration( project_name="test-project", instance_name=instance_id or "test-instance", backend=backend_type, @@ -749,7 +749,7 @@ async def create_gateway_compute( ssh_key_pub=ssh_public_key, certificate=None, ).model_dump_json() - gateway_compute = GatewayComputeModel( + gateway_replica = GatewayReplicaModel( gateway_id=gateway_id, backend_id=backend_id, ip_address=ip_address, @@ -765,19 +765,19 @@ async def create_gateway_compute( hostname_deprecated_readonly=hostname_deprecated_readonly, backend_data=backend_data, ) - session.add(gateway_compute) + session.add(gateway_replica) await session.commit() - return gateway_compute + return gateway_replica -def get_gateway_compute_configuration( +def get_gateway_replica_configuration( project_name: str = "test-project", instance_name: str = "test-instance", backend: BackendType = BackendType.AWS, region: str = "us", public_ip: bool = True, -) -> GatewayComputeConfiguration: - return GatewayComputeConfiguration( +) -> GatewayReplicaConfiguration: + return GatewayReplicaConfiguration( project_name=project_name, instance_name=instance_name, backend=backend, diff --git a/src/tests/_internal/core/backends/base/test_compute.py b/src/tests/_internal/core/backends/base/test_compute.py index 7892a3f0f5..4fac843447 100644 --- a/src/tests/_internal/core/backends/base/test_compute.py +++ b/src/tests/_internal/core/backends/base/test_compute.py @@ -13,7 +13,7 @@ normalize_arch, ) from dstack._internal.server.testing.common import ( - get_gateway_compute_configuration, + get_gateway_replica_configuration, get_instance_configuration, get_volume, ) @@ -30,7 +30,7 @@ def test_generates_name(self): class TestGenerateUniqueGatewayInstanceName: def test_generates_name(self): - configuration = get_gateway_compute_configuration( + configuration = get_gateway_replica_configuration( project_name="project1", instance_name="my-gateway" ) name = generate_unique_gateway_instance_name(configuration, 60) 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 e685e6cef3..c9c3f66d99 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 @@ -13,7 +13,7 @@ from dstack._internal.core.models.configurations import ServiceConfiguration from dstack._internal.core.models.gateways import ( ACMGatewayCertificate, - GatewayProvisioningData, + GatewayReplicaProvisioningData, GatewayReplicaStatus, GatewayStatus, ) @@ -27,7 +27,7 @@ GatewayReplicaWorker, ) from dstack._internal.server.models import ( - GatewayComputeModel, + GatewayReplicaModel, ServiceRegistrationModel, ServiceReplicaRegistrationModel, ) @@ -37,14 +37,14 @@ create_backend, create_fleet, create_gateway, - create_gateway_compute, + create_gateway_replica, create_instance, create_job, create_project, create_repo, create_run, create_user, - get_gateway_compute_configuration, + get_gateway_replica_configuration, get_job_provisioning_data, get_run_spec, list_events, @@ -68,24 +68,24 @@ def fetcher() -> GatewayReplicaFetcher: ) -def _compute_to_pipeline_item( - compute: GatewayComputeModel, +def _replica_to_pipeline_item( + replica: GatewayReplicaModel, ) -> GatewayReplicaPipelineItem: - assert compute.lock_token is not None - assert compute.lock_expires_at is not None + assert replica.lock_token is not None + assert replica.lock_expires_at is not None return GatewayReplicaPipelineItem( - __tablename__=compute.__tablename__, - id=compute.id, - lock_token=compute.lock_token, - lock_expires_at=compute.lock_expires_at, + __tablename__=replica.__tablename__, + id=replica.id, + lock_token=replica.lock_token, + lock_expires_at=replica.lock_expires_at, prev_lock_expired=False, - status=compute.status, + status=replica.status, ) -def _lock_compute(compute: GatewayComputeModel) -> None: - compute.lock_token = uuid.uuid4() - compute.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) +def _lock_replica(replica: GatewayReplicaModel) -> None: + replica.lock_token = uuid.uuid4() + replica.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) @pytest.mark.asyncio @@ -106,7 +106,7 @@ async def test_fetch_selects_eligible_replicas_and_sets_lock_fields( now = get_current_datetime() stale = now - timedelta(minutes=1) - submitted = await create_gateway_compute( + submitted = await create_gateway_replica( session=session, gateway_id=gateway.id, ip_address=None, @@ -114,35 +114,35 @@ async def test_fetch_selects_eligible_replicas_and_sets_lock_fields( region=None, status=GatewayReplicaStatus.SUBMITTED, last_processed_at=stale - timedelta(seconds=3), - configuration=get_gateway_compute_configuration().model_dump_json(), + configuration=get_gateway_replica_configuration().model_dump_json(), ) - provisioning = await create_gateway_compute( + provisioning = await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.PROVISIONING, last_processed_at=stale - timedelta(seconds=2), ) - terminating = await create_gateway_compute( + terminating = await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.TERMINATING, active=False, last_processed_at=stale - timedelta(seconds=1), ) - running = await create_gateway_compute( + running = await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, last_processed_at=stale, ) - terminated = await create_gateway_compute( + terminated = await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.TERMINATED, active=False, last_processed_at=stale, ) - recent = await create_gateway_compute( + recent = await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.SUBMITTED, @@ -150,11 +150,11 @@ async def test_fetch_selects_eligible_replicas_and_sets_lock_fields( instance_id=None, region=None, last_processed_at=now, - configuration=get_gateway_compute_configuration().model_dump_json(), + configuration=get_gateway_replica_configuration().model_dump_json(), ) recent.created_at = now - timedelta(minutes=2) recent.last_processed_at = now - locked = await create_gateway_compute( + locked = await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.SUBMITTED, @@ -162,7 +162,7 @@ async def test_fetch_selects_eligible_replicas_and_sets_lock_fields( instance_id=None, region=None, last_processed_at=stale + timedelta(seconds=1), - configuration=get_gateway_compute_configuration().model_dump_json(), + configuration=get_gateway_replica_configuration().model_dump_json(), ) locked.lock_expires_at = now + timedelta(minutes=1) locked.lock_token = uuid.uuid4() @@ -184,8 +184,8 @@ async def test_fetch_selects_eligible_replicas_and_sets_lock_fields( (running.id, GatewayReplicaStatus.RUNNING), } - for compute in [submitted, provisioning, terminating, running, terminated, recent, locked]: - await session.refresh(compute) + for replica in [submitted, provisioning, terminating, running, terminated, recent, locked]: + await session.refresh(replica) fetched = [submitted, provisioning, terminating, running] assert all(c.lock_owner == GatewayReplicaPipeline.__name__ for c in fetched) @@ -204,7 +204,7 @@ async def test_fetch_selects_eligible_replicas_and_sets_lock_fields( (GatewayStatus.RUNNING, True), ], ) - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) async def test_fetch_includes_running_replica_needing_cleanup( self, test_db, @@ -212,7 +212,7 @@ async def test_fetch_includes_running_replica_needing_cleanup( fetcher: GatewayReplicaFetcher, gateway_status: GatewayStatus, to_be_deleted: bool, - legacy_compute: bool, + legacy_replica: bool, ): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) @@ -224,15 +224,15 @@ async def test_fetch_includes_running_replica_needing_cleanup( ) gateway.to_be_deleted = to_be_deleted stale = get_current_datetime() - timedelta(minutes=1) - if legacy_compute: - compute = await create_gateway_compute( + if legacy_replica: + replica = await create_gateway_replica( session=session, status=GatewayReplicaStatus.RUNNING, last_processed_at=stale, ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id else: - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, @@ -243,7 +243,7 @@ async def test_fetch_includes_running_replica_needing_cleanup( items = await fetcher.fetch(limit=10) assert len(items) == 1 - assert items[0].id == compute.id + assert items[0].id == replica.id assert items[0].status == GatewayReplicaStatus.RUNNING async def test_fetch_includes_running_replica_with_hard_deleted_gateway( @@ -252,10 +252,10 @@ async def test_fetch_includes_running_replica_with_hard_deleted_gateway( session: AsyncSession, fetcher: GatewayReplicaFetcher, ): - # A compute whose gateway was hard-deleted (orphaned). The fetcher should + # A replica whose gateway was hard-deleted (orphaned). The fetcher should # pick it up so the worker can log the error. stale = get_current_datetime() - timedelta(minutes=1) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=None, status=GatewayReplicaStatus.RUNNING, @@ -266,16 +266,16 @@ async def test_fetch_includes_running_replica_with_hard_deleted_gateway( items = await fetcher.fetch(limit=10) assert len(items) == 1 - assert items[0].id == compute.id + assert items[0].id == replica.id assert items[0].status == GatewayReplicaStatus.RUNNING - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) async def test_fetch_includes_running_replica_with_healthy_gateway( self, test_db, session: AsyncSession, fetcher: GatewayReplicaFetcher, - legacy_compute: bool, + legacy_replica: bool, ): # Healthy running replicas are still fetched periodically so the worker # can run gateway state sync (see _process_running_item). @@ -288,15 +288,15 @@ async def test_fetch_includes_running_replica_with_healthy_gateway( status=GatewayStatus.RUNNING, ) stale = get_current_datetime() - timedelta(minutes=1) - if legacy_compute: - compute = await create_gateway_compute( + if legacy_replica: + replica = await create_gateway_replica( session=session, status=GatewayReplicaStatus.RUNNING, last_processed_at=stale, ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id else: - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, @@ -307,7 +307,7 @@ async def test_fetch_includes_running_replica_with_healthy_gateway( items = await fetcher.fetch(limit=10) assert len(items) == 1 - assert items[0].id == compute.id + assert items[0].id == replica.id assert items[0].status == GatewayReplicaStatus.RUNNING async def test_fetch_includes_running_replica_marked_for_scale_in( @@ -325,19 +325,19 @@ async def test_fetch_includes_running_replica_marked_for_scale_in( status=GatewayStatus.RUNNING, ) stale = get_current_datetime() - timedelta(minutes=1) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, last_processed_at=stale, ) - compute.scale_in = True + replica.scale_in = True await session.commit() items = await fetcher.fetch(limit=10) assert len(items) == 1 - assert items[0].id == compute.id + assert items[0].id == replica.id assert items[0].status == GatewayReplicaStatus.RUNNING @@ -355,7 +355,7 @@ async def test_submitted_to_provisioning( backend_id=backend.id, status=GatewayStatus.PROVISIONING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -363,9 +363,9 @@ async def test_submitted_to_provisioning( instance_id=None, region=None, status=GatewayReplicaStatus.SUBMITTED, - configuration=get_gateway_compute_configuration().model_dump_json(), + configuration=get_gateway_replica_configuration().model_dump_json(), ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with patch( @@ -374,19 +374,21 @@ async def test_submitted_to_provisioning( aws = Mock() m.return_value = [(backend, aws)] aws.compute.return_value = Mock(spec=ComputeMockSpec) - aws.compute.return_value.create_gateway.return_value = GatewayProvisioningData( - instance_id="i-1234567890", - ip_address="2.2.2.2", - region="us", + aws.compute.return_value.create_gateway_replica.return_value = ( + GatewayReplicaProvisioningData( + instance_id="i-1234567890", + ip_address="2.2.2.2", + region="us", + ) ) - await worker.process(_compute_to_pipeline_item(compute)) - aws.compute.return_value.create_gateway.assert_called_once() + await worker.process(_replica_to_pipeline_item(replica)) + aws.compute.return_value.create_gateway_replica.assert_called_once() - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.PROVISIONING - assert compute.ip_address == "2.2.2.2" - assert compute.instance_id == "i-1234567890" - assert compute.region == "us" + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.PROVISIONING + assert replica.ip_address == "2.2.2.2" + assert replica.instance_id == "i-1234567890" + assert replica.region == "us" async def test_submitted_backend_error_marks_terminated( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker @@ -399,7 +401,7 @@ async def test_submitted_backend_error_marks_terminated( backend_id=backend.id, status=GatewayStatus.PROVISIONING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -407,9 +409,9 @@ async def test_submitted_backend_error_marks_terminated( instance_id=None, region=None, status=GatewayReplicaStatus.SUBMITTED, - configuration=get_gateway_compute_configuration().model_dump_json(), + configuration=get_gateway_replica_configuration().model_dump_json(), ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with patch( @@ -418,13 +420,15 @@ async def test_submitted_backend_error_marks_terminated( aws = Mock() m.return_value = [(backend, aws)] aws.compute.return_value = Mock(spec=ComputeMockSpec) - aws.compute.return_value.create_gateway.side_effect = BackendError("Some error") - await worker.process(_compute_to_pipeline_item(compute)) + aws.compute.return_value.create_gateway_replica.side_effect = BackendError( + "Some error" + ) + await worker.process(_replica_to_pipeline_item(replica)) - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATED - assert compute.active is False - assert compute.deleted is True + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATED + assert replica.active is False + assert replica.deleted is True async def test_submitted_backend_not_available_marks_terminated( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker @@ -437,7 +441,7 @@ async def test_submitted_backend_not_available_marks_terminated( backend_id=backend.id, status=GatewayStatus.PROVISIONING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -445,21 +449,21 @@ async def test_submitted_backend_not_available_marks_terminated( instance_id=None, region=None, status=GatewayReplicaStatus.SUBMITTED, - configuration=get_gateway_compute_configuration().model_dump_json(), + configuration=get_gateway_replica_configuration().model_dump_json(), ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with patch( "dstack._internal.server.services.backends.get_project_backends_with_models" ) as m: m.return_value = [] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATED - assert compute.active is False - assert compute.deleted is True + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATED + assert replica.active is False + assert replica.deleted is True async def test_submitted_skips_provisioning_if_gateway_to_be_deleted( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker @@ -473,7 +477,7 @@ async def test_submitted_skips_provisioning_if_gateway_to_be_deleted( status=GatewayStatus.RUNNING, ) gateway.to_be_deleted = True - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -481,21 +485,21 @@ async def test_submitted_skips_provisioning_if_gateway_to_be_deleted( instance_id=None, region=None, status=GatewayReplicaStatus.SUBMITTED, - configuration=get_gateway_compute_configuration().model_dump_json(), + configuration=get_gateway_replica_configuration().model_dump_json(), ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with patch( "dstack._internal.server.services.backends.get_project_backends_with_models" ) as m: - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) m.assert_not_called() - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATED - assert compute.active is False - assert compute.deleted is True + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATED + assert replica.active is False + assert replica.deleted is True async def test_submitted_skips_provisioning_if_gateway_failed( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker @@ -508,7 +512,7 @@ async def test_submitted_skips_provisioning_if_gateway_failed( backend_id=backend.id, status=GatewayStatus.FAILED, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -516,21 +520,21 @@ async def test_submitted_skips_provisioning_if_gateway_failed( instance_id=None, region=None, status=GatewayReplicaStatus.SUBMITTED, - configuration=get_gateway_compute_configuration().model_dump_json(), + configuration=get_gateway_replica_configuration().model_dump_json(), ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with patch( "dstack._internal.server.services.backends.get_project_backends_with_models" ) as m: - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) m.assert_not_called() - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATED - assert compute.active is False - assert compute.deleted is True + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATED + assert replica.active is False + assert replica.deleted is True async def test_submitted_unexpected_error_marks_terminated( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker @@ -543,7 +547,7 @@ async def test_submitted_unexpected_error_marks_terminated( backend_id=backend.id, status=GatewayStatus.PROVISIONING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -551,9 +555,9 @@ async def test_submitted_unexpected_error_marks_terminated( instance_id=None, region=None, status=GatewayReplicaStatus.SUBMITTED, - configuration=get_gateway_compute_configuration().model_dump_json(), + configuration=get_gateway_replica_configuration().model_dump_json(), ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with patch( @@ -562,14 +566,16 @@ async def test_submitted_unexpected_error_marks_terminated( aws = Mock() m.return_value = [(backend, aws)] aws.compute.return_value = Mock(spec=ComputeMockSpec) - aws.compute.return_value.create_gateway.side_effect = RuntimeError("Unexpected!") - await worker.process(_compute_to_pipeline_item(compute)) + aws.compute.return_value.create_gateway_replica.side_effect = RuntimeError( + "Unexpected!" + ) + await worker.process(_replica_to_pipeline_item(replica)) - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATED - assert compute.status_message == "Unexpected error" - assert compute.active is False - assert compute.deleted is True + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATED + assert replica.status_message == "Unexpected error" + assert replica.active is False + assert replica.deleted is True async def test_submitted_to_terminated_when_scaled_in( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker @@ -582,7 +588,7 @@ async def test_submitted_to_terminated_when_scaled_in( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -590,23 +596,23 @@ async def test_submitted_to_terminated_when_scaled_in( instance_id=None, region=None, status=GatewayReplicaStatus.SUBMITTED, - configuration=get_gateway_compute_configuration().model_dump_json(), + configuration=get_gateway_replica_configuration().model_dump_json(), ) - compute.scale_in = True - _lock_compute(compute) + replica.scale_in = True + _lock_replica(replica) await session.commit() with patch( "dstack._internal.server.services.backends.get_project_backends_with_models" ) as m: - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) m.assert_not_called() - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATED - assert compute.active is False - assert compute.deleted is True - assert compute.status_message == "Scaled in" + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATED + assert replica.active is False + assert replica.deleted is True + assert replica.status_message == "Scaled in" @pytest.mark.asyncio @@ -619,7 +625,7 @@ class TestGatewayReplicaWorkerRunning: (GatewayStatus.RUNNING, True), ], ) - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) @pytest.mark.parametrize("populate_configuration", [True, False]) async def test_running_to_terminating( self, @@ -628,7 +634,7 @@ async def test_running_to_terminating( worker: GatewayReplicaWorker, gateway_status: GatewayStatus, to_be_deleted: bool, - legacy_compute: bool, + legacy_replica: bool, populate_configuration: bool, ): project = await create_project(session=session) @@ -641,30 +647,30 @@ async def test_running_to_terminating( populate_configuration=populate_configuration, ) gateway.to_be_deleted = to_be_deleted - if legacy_compute: - compute = await create_gateway_compute( + if legacy_replica: + replica = await create_gateway_replica( session=session, status=GatewayReplicaStatus.RUNNING, active=True, populate_configuration=populate_configuration, ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id else: - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, active=True, populate_configuration=populate_configuration, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATING - assert compute.active is False + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATING + assert replica.active is False async def test_running_to_terminating_when_scaled_in( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker @@ -677,22 +683,22 @@ async def test_running_to_terminating_when_scaled_in( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, active=True, ) - compute.scale_in = True - _lock_compute(compute) + replica.scale_in = True + _lock_replica(replica) await session.commit() - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATING - assert compute.active is False - assert compute.status_message == "Scaled in" + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATING + assert replica.active is False + assert replica.status_message == "Scaled in" def _get_client_mock(mock_gateway_connection: AsyncMock) -> AsyncMock: @@ -776,7 +782,7 @@ async def test_registers_new_service_and_replica( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -786,16 +792,16 @@ async def test_registers_new_service_and_replica( run, job = await self._create_service_run_and_job( session, project, repo, user, gateway, run_name="test-service" ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() client_mock = _get_client_mock(mock_gateway_connection) client_mock.list_services.return_value = [] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) mock_gateway_connection.assert_called_once_with( - hostname=compute.ip_address, id_rsa="replica-private-key" + hostname=replica.ip_address, id_rsa="replica-private-key" ) client_mock.register_service.assert_called_once_with( project=project.name, @@ -829,7 +835,7 @@ async def test_registers_new_service_and_replica( await session.execute( select(ServiceRegistrationModel).where( ServiceRegistrationModel.run_id == run.id, - ServiceRegistrationModel.gateway_replica_id == compute.id, + ServiceRegistrationModel.gateway_replica_id == replica.id, ) ) ).scalar_one() @@ -841,7 +847,7 @@ async def test_registers_new_service_and_replica( await session.execute( select(ServiceReplicaRegistrationModel).where( ServiceReplicaRegistrationModel.job_id == job.id, - ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ServiceReplicaRegistrationModel.gateway_replica_id == replica.id, ) ) ).scalar_one() @@ -849,8 +855,8 @@ async def test_registers_new_service_and_replica( 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}", + f"Service registered on gateway replica {replica.replica_num}", + f"Service replica registered on gateway replica {replica.replica_num}", } async def test_unregisters_dangling_service_and_stale_replica( @@ -870,7 +876,7 @@ async def test_unregisters_dangling_service_and_stale_replica( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -900,7 +906,7 @@ async def test_unregisters_dangling_service_and_stale_replica( job_status=JobStatus.TERMINATED, job_registered=False, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() client_mock = _get_client_mock(mock_gateway_connection) @@ -922,7 +928,7 @@ async def test_unregisters_dangling_service_and_stale_replica( ), ] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) client_mock.register_service.assert_not_called() client_mock.register_replica.assert_not_called() @@ -937,7 +943,7 @@ async def test_unregisters_dangling_service_and_stale_replica( ( await session.execute( select(ServiceRegistrationModel).where( - ServiceRegistrationModel.gateway_replica_id == compute.id, + ServiceRegistrationModel.gateway_replica_id == replica.id, ) ) ) @@ -951,7 +957,7 @@ async def test_unregisters_dangling_service_and_stale_replica( ( await session.execute( select(ServiceReplicaRegistrationModel).where( - ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ServiceReplicaRegistrationModel.gateway_replica_id == replica.id, ) ) ) @@ -963,8 +969,8 @@ async def test_unregisters_dangling_service_and_stale_replica( 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}", + f"Service unregistered from gateway replica {replica.replica_num}", + f"Service replica unregistered from gateway replica {replica.replica_num}", } async def test_deletes_registration_models_for_unregistered_service_and_replica( @@ -984,7 +990,7 @@ async def test_deletes_registration_models_for_unregistered_service_and_replica( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -1017,16 +1023,16 @@ async def test_deletes_registration_models_for_unregistered_service_and_replica( # 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 + run_id=run1.id, gateway_replica_id=replica.id, is_registered=True ) live_replica_registration = ServiceReplicaRegistrationModel( - job_id=job1.id, gateway_replica_id=compute.id, is_registered=True + job_id=job1.id, gateway_replica_id=replica.id, is_registered=True ) stale_replica_registration = ServiceReplicaRegistrationModel( - job_id=stale_job.id, gateway_replica_id=compute.id, is_registered=True + job_id=stale_job.id, gateway_replica_id=replica.id, is_registered=True ) dangling_service_registration = ServiceRegistrationModel( - run_id=run2.id, gateway_replica_id=compute.id, is_registered=True + run_id=run2.id, gateway_replica_id=replica.id, is_registered=True ) session.add_all( [ @@ -1036,7 +1042,7 @@ async def test_deletes_registration_models_for_unregistered_service_and_replica( dangling_service_registration, ] ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() client_mock = _get_client_mock(mock_gateway_connection) @@ -1058,7 +1064,7 @@ async def test_deletes_registration_models_for_unregistered_service_and_replica( ), ] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) client_mock.unregister_service.assert_called_once_with( project=project.name, run_name="dangling-service" @@ -1071,7 +1077,7 @@ async def test_deletes_registration_models_for_unregistered_service_and_replica( ( await session.execute( select(ServiceRegistrationModel).where( - ServiceRegistrationModel.gateway_replica_id == compute.id, + ServiceRegistrationModel.gateway_replica_id == replica.id, ) ) ) @@ -1085,7 +1091,7 @@ async def test_deletes_registration_models_for_unregistered_service_and_replica( ( await session.execute( select(ServiceReplicaRegistrationModel).where( - ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ServiceReplicaRegistrationModel.gateway_replica_id == replica.id, ) ) ) @@ -1129,7 +1135,7 @@ async def test_unregisters_replicas_of_dangling_service_without_extra_gateway_ca backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -1150,13 +1156,13 @@ async def test_unregisters_replicas_of_dangling_service_without_extra_gateway_ca job_registered=False, ) stale_service_registration = ServiceRegistrationModel( - run_id=run.id, gateway_replica_id=compute.id, is_registered=True + run_id=run.id, gateway_replica_id=replica.id, is_registered=True ) stale_replica_registration = ServiceReplicaRegistrationModel( - job_id=job.id, gateway_replica_id=compute.id, is_registered=True + job_id=job.id, gateway_replica_id=replica.id, is_registered=True ) session.add_all([stale_service_registration, stale_replica_registration]) - _lock_compute(compute) + _lock_replica(replica) await session.commit() client_mock = _get_client_mock(mock_gateway_connection) @@ -1169,7 +1175,7 @@ async def test_unregisters_replicas_of_dangling_service_without_extra_gateway_ca ), ] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) client_mock.unregister_service.assert_called_once_with( project=project.name, run_name="dangling-service" @@ -1198,8 +1204,8 @@ async def test_unregisters_replicas_of_dangling_service_without_extra_gateway_ca 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}", + f"Service unregistered from gateway replica {replica.replica_num}", + f"Service replica unregistered from gateway replica {replica.replica_num}", } async def test_no_gateway_calls_when_state_already_in_sync( @@ -1219,7 +1225,7 @@ async def test_no_gateway_calls_when_state_already_in_sync( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -1228,7 +1234,7 @@ async def test_no_gateway_calls_when_state_already_in_sync( run, job = await self._create_service_run_and_job( session, project, repo, user, gateway, run_name="synced-service" ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() client_mock = _get_client_mock(mock_gateway_connection) @@ -1241,7 +1247,7 @@ async def test_no_gateway_calls_when_state_already_in_sync( ), ] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) client_mock.register_service.assert_not_called() client_mock.register_replica.assert_not_called() @@ -1255,7 +1261,7 @@ async def test_no_gateway_calls_when_state_already_in_sync( await session.execute( select(ServiceRegistrationModel).where( ServiceRegistrationModel.run_id == run.id, - ServiceRegistrationModel.gateway_replica_id == compute.id, + ServiceRegistrationModel.gateway_replica_id == replica.id, ) ) ).scalar_one() @@ -1264,7 +1270,7 @@ async def test_no_gateway_calls_when_state_already_in_sync( await session.execute( select(ServiceReplicaRegistrationModel).where( ServiceReplicaRegistrationModel.job_id == job.id, - ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ServiceReplicaRegistrationModel.gateway_replica_id == replica.id, ) ) ).scalar_one() @@ -1290,7 +1296,7 @@ async def test_recovers_legacy_service_id_by_matching_replica( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -1299,7 +1305,7 @@ async def test_recovers_legacy_service_id_by_matching_replica( run, job = await self._create_service_run_and_job( session, project, repo, user, gateway, run_name="legacy-service" ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() client_mock = _get_client_mock(mock_gateway_connection) @@ -1313,7 +1319,7 @@ async def test_recovers_legacy_service_id_by_matching_replica( ), ] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) client_mock.set_service_id.assert_called_once_with( project=project.name, run_name="legacy-service", run_id=run.id @@ -1340,7 +1346,7 @@ async def test_unregisters_and_reregisters_legacy_service_without_id_and_replica backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -1349,7 +1355,7 @@ async def test_unregisters_and_reregisters_legacy_service_without_id_and_replica run, job = await self._create_service_run_and_job( session, project, repo, user, gateway, run_name="legacy-service" ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() client_mock = _get_client_mock(mock_gateway_connection) @@ -1365,7 +1371,7 @@ async def test_unregisters_and_reregisters_legacy_service_without_id_and_replica ), ] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) client_mock.set_service_id.assert_not_called() client_mock.unregister_service.assert_called_once_with( @@ -1400,7 +1406,7 @@ async def test_unregisters_and_reregisters_legacy_service_without_id_and_replica await session.execute( select(ServiceRegistrationModel).where( ServiceRegistrationModel.run_id == run.id, - ServiceRegistrationModel.gateway_replica_id == compute.id, + ServiceRegistrationModel.gateway_replica_id == replica.id, ) ) ).scalar_one() @@ -1409,7 +1415,7 @@ async def test_unregisters_and_reregisters_legacy_service_without_id_and_replica await session.execute( select(ServiceReplicaRegistrationModel).where( ServiceReplicaRegistrationModel.job_id == job.id, - ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ServiceReplicaRegistrationModel.gateway_replica_id == replica.id, ) ) ).scalar_one() @@ -1419,8 +1425,8 @@ async def test_unregisters_and_reregisters_legacy_service_without_id_and_replica # 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}", + f"Service registered on gateway replica {replica.replica_num}", + f"Service replica registered on gateway replica {replica.replica_num}", } async def test_does_nothing_when_in_sync_and_registrations_already_exist( @@ -1440,7 +1446,7 @@ async def test_does_nothing_when_in_sync_and_registrations_already_exist( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -1451,18 +1457,18 @@ async def test_does_nothing_when_in_sync_and_registrations_already_exist( ) existing_service_registration = ServiceRegistrationModel( run_id=run.id, - gateway_replica_id=compute.id, + gateway_replica_id=replica.id, is_registered=True, register_attempt=0, ) existing_replica_registration = ServiceReplicaRegistrationModel( job_id=job.id, - gateway_replica_id=compute.id, + gateway_replica_id=replica.id, is_registered=True, register_attempt=0, ) session.add_all([existing_service_registration, existing_replica_registration]) - _lock_compute(compute) + _lock_replica(replica) await session.commit() client_mock = _get_client_mock(mock_gateway_connection) @@ -1475,7 +1481,7 @@ async def test_does_nothing_when_in_sync_and_registrations_already_exist( ), ] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) client_mock.register_service.assert_not_called() client_mock.register_replica.assert_not_called() @@ -1487,7 +1493,7 @@ async def test_does_nothing_when_in_sync_and_registrations_already_exist( ( await session.execute( select(ServiceRegistrationModel).where( - ServiceRegistrationModel.gateway_replica_id == compute.id, + ServiceRegistrationModel.gateway_replica_id == replica.id, ) ) ) @@ -1503,7 +1509,7 @@ async def test_does_nothing_when_in_sync_and_registrations_already_exist( ( await session.execute( select(ServiceReplicaRegistrationModel).where( - ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ServiceReplicaRegistrationModel.gateway_replica_id == replica.id, ) ) ) @@ -1534,7 +1540,7 @@ async def test_reconciles_out_of_sync_registration_models( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -1548,20 +1554,20 @@ async def test_reconciles_out_of_sync_registration_models( # right after a successful registration, before it could record that). stale_service_registration = ServiceRegistrationModel( run_id=run.id, - gateway_replica_id=compute.id, + gateway_replica_id=replica.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, + gateway_replica_id=replica.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) + _lock_replica(replica) await session.commit() client_mock = _get_client_mock(mock_gateway_connection) @@ -1574,7 +1580,7 @@ async def test_reconciles_out_of_sync_registration_models( ), ] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) # The gateway already reports it registered, so nothing needs to be # (re)registered - only the local bookkeeping needs correcting. @@ -1637,7 +1643,7 @@ async def test_propagates_registration_errors_and_increments_register_attempt( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -1652,13 +1658,13 @@ async def test_propagates_registration_errors_and_increments_register_attempt( # Simulate two earlier failed attempts to register this service. existing_registration = ServiceRegistrationModel( run_id=run_service_fails.id, - gateway_replica_id=compute.id, + gateway_replica_id=replica.id, is_registered=False, register_attempt=2, register_status_message="earlier error", ) session.add(existing_registration) - _lock_compute(compute) + _lock_replica(replica) await session.commit() client_mock = _get_client_mock(mock_gateway_connection) @@ -1675,7 +1681,7 @@ async def register_replica_side_effect(**kwargs): 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)) + await worker.process(_replica_to_pipeline_item(replica)) client_mock.unregister_service.assert_not_called() client_mock.unregister_replica.assert_not_called() @@ -1697,7 +1703,7 @@ async def register_replica_side_effect(**kwargs): await session.execute( select(ServiceRegistrationModel).where( ServiceRegistrationModel.run_id == run_replica_fails.id, - ServiceRegistrationModel.gateway_replica_id == compute.id, + ServiceRegistrationModel.gateway_replica_id == replica.id, ) ) ).scalar_one() @@ -1709,7 +1715,7 @@ async def register_replica_side_effect(**kwargs): await session.execute( select(ServiceReplicaRegistrationModel).where( ServiceReplicaRegistrationModel.job_id == job_replica_fails.id, - ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ServiceReplicaRegistrationModel.gateway_replica_id == replica.id, ) ) ).scalar_one() @@ -1733,7 +1739,7 @@ async def register_replica_side_effect(**kwargs): await session.execute( select(ServiceReplicaRegistrationModel).where( ServiceReplicaRegistrationModel.job_id == job_service_fails.id, - ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ServiceReplicaRegistrationModel.gateway_replica_id == replica.id, ) ) ).scalar_one_or_none() @@ -1741,14 +1747,14 @@ async def register_replica_side_effect(**kwargs): events = await list_events(session) assert {e.message for e in events} == { - f"Service registered on gateway replica {compute.replica_num}", + f"Service registered on gateway replica {replica.replica_num}", ( f"Encountered service registration error on gateway replica " - f"{compute.replica_num}: {expected_service_register_status_message}" + f"{replica.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}" + f"{replica.replica_num}: {expected_replica_register_status_message}" ), } @@ -1769,7 +1775,7 @@ async def test_does_not_emit_duplicate_registration_error_event_for_unchanged_er backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -1785,15 +1791,15 @@ async def test_does_not_emit_duplicate_registration_error_event_for_unchanged_er # First tick: the replica registration fails, recording the error and # emitting one error event. - _lock_compute(compute) + _lock_replica(replica) await session.commit() - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) replica_registration = ( await session.execute( select(ServiceReplicaRegistrationModel).where( ServiceReplicaRegistrationModel.job_id == job.id, - ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ServiceReplicaRegistrationModel.gateway_replica_id == replica.id, ) ) ).scalar_one() @@ -1802,7 +1808,7 @@ async def test_does_not_emit_duplicate_registration_error_event_for_unchanged_er error_message = ( f"Encountered service replica registration error on gateway replica " - f"{compute.replica_num}: boom replica" + f"{replica.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 @@ -1810,10 +1816,10 @@ async def test_does_not_emit_duplicate_registration_error_event_for_unchanged_er # 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.refresh(replica) + _lock_replica(replica) await session.commit() - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) await session.refresh(replica_registration) assert replica_registration.register_attempt == 2 @@ -1859,7 +1865,7 @@ async def test_propagates_unregistration_errors_and_increments_unregister_attemp backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -1889,7 +1895,7 @@ async def test_propagates_unregistration_errors_and_increments_unregister_attemp job_status=JobStatus.TERMINATED, job_registered=False, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() client_mock = _get_client_mock(mock_gateway_connection) @@ -1913,7 +1919,7 @@ async def test_propagates_unregistration_errors_and_increments_unregister_attemp 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)) + await worker.process(_replica_to_pipeline_item(replica)) # Both remain registered as far as the gateway is concerned, since we # failed to remove them - only the unregister bookkeeping changes. @@ -1921,7 +1927,7 @@ async def test_propagates_unregistration_errors_and_increments_unregister_attemp await session.execute( select(ServiceRegistrationModel).where( ServiceRegistrationModel.run_id == dangling_run.id, - ServiceRegistrationModel.gateway_replica_id == compute.id, + ServiceRegistrationModel.gateway_replica_id == replica.id, ) ) ).scalar_one() @@ -1933,7 +1939,7 @@ async def test_propagates_unregistration_errors_and_increments_unregister_attemp await session.execute( select(ServiceReplicaRegistrationModel).where( ServiceReplicaRegistrationModel.job_id == stale_job.id, - ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ServiceReplicaRegistrationModel.gateway_replica_id == replica.id, ) ) ).scalar_one() @@ -1947,11 +1953,11 @@ async def test_propagates_unregistration_errors_and_increments_unregister_attemp assert {e.message for e in events} == { ( f"Encountered service unregistration error on gateway replica " - f"{compute.replica_num}: {expected_service_status_message}" + f"{replica.replica_num}: {expected_service_status_message}" ), ( f"Encountered service replica unregistration error on gateway replica " - f"{compute.replica_num}: {expected_replica_status_message}" + f"{replica.replica_num}: {expected_replica_status_message}" ), } @@ -1972,7 +1978,7 @@ async def test_does_not_emit_duplicate_unregistration_error_event_for_unchanged_ backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -1988,7 +1994,7 @@ async def test_does_not_emit_duplicate_unregistration_error_event_for_unchanged_ registered=False, replica_num=1, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() client_mock = _get_client_mock(mock_gateway_connection) @@ -2007,13 +2013,13 @@ async def test_does_not_emit_duplicate_unregistration_error_event_for_unchanged_ # First tick: unregistering the stale replica fails, recording the error # and emitting one error event. - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) replica_registration = ( await session.execute( select(ServiceReplicaRegistrationModel).where( ServiceReplicaRegistrationModel.job_id == stale_job.id, - ServiceReplicaRegistrationModel.gateway_replica_id == compute.id, + ServiceReplicaRegistrationModel.gateway_replica_id == replica.id, ) ) ).scalar_one() @@ -2022,7 +2028,7 @@ async def test_does_not_emit_duplicate_unregistration_error_event_for_unchanged_ error_message = ( f"Encountered service replica unregistration error on gateway replica " - f"{compute.replica_num}: boom replica" + f"{replica.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 @@ -2030,10 +2036,10 @@ async def test_does_not_emit_duplicate_unregistration_error_event_for_unchanged_ # 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.refresh(replica) + _lock_replica(replica) await session.commit() - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) await session.refresh(replica_registration) assert replica_registration.unregister_attempt == 2 @@ -2046,14 +2052,14 @@ async def test_does_not_emit_duplicate_unregistration_error_event_for_unchanged_ @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) class TestGatewayReplicaWorkerProvisioning: - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [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, + legacy_replica: bool, populate_configuration: bool, ): project = await create_project(session=session) @@ -2065,21 +2071,21 @@ async def test_provisioning_to_running( status=GatewayStatus.PROVISIONING, populate_configuration=populate_configuration, ) - if legacy_compute: - compute = await create_gateway_compute( + if legacy_replica: + replica = await create_gateway_replica( session=session, status=GatewayReplicaStatus.PROVISIONING, populate_configuration=populate_configuration, ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id else: - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.PROVISIONING, populate_configuration=populate_configuration, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with patch( @@ -2087,12 +2093,12 @@ async def test_provisioning_to_running( ) 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)) + await worker.process(_replica_to_pipeline_item(replica)) pool_add.assert_called_once() - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.RUNNING - assert compute.active is True + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.RUNNING + assert replica.active is True async def test_provisioning_to_running_registers_with_load_balancer( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker @@ -2108,13 +2114,13 @@ async def test_provisioning_to_running_registers_with_load_balancer( hostname="gateway-lb.example.com", backend_data="lb-backend-data", ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, status=GatewayReplicaStatus.PROVISIONING, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with ( @@ -2131,20 +2137,20 @@ async def test_provisioning_to_running_registers_with_load_balancer( backend_mock.compute.return_value = Mock(spec=ComputeMockSpec) get_backends_mock.return_value = [(backend, backend_mock)] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) 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[0] == replica.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 + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.RUNNING + assert replica.active is True async def test_provisioning_skips_load_balancer_registration_without_hostname( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker @@ -2157,13 +2163,13 @@ async def test_provisioning_skips_load_balancer_registration_without_hostname( backend_id=backend.id, status=GatewayStatus.PROVISIONING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, status=GatewayReplicaStatus.PROVISIONING, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with ( @@ -2177,13 +2183,13 @@ async def test_provisioning_skips_load_balancer_registration_without_hostname( pool_add.return_value = MagicMock() pool_add.return_value.client.return_value = MagicMock(AsyncContextManager()) - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) get_backends_mock.assert_not_called() - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.RUNNING - assert compute.active is True + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.RUNNING + assert replica.active is True async def test_provisioning_to_terminating_when_load_balancer_registration_fails( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker @@ -2199,13 +2205,13 @@ async def test_provisioning_to_terminating_when_load_balancer_registration_fails hostname="gateway-lb.example.com", backend_data="lb-backend-data", ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, status=GatewayReplicaStatus.PROVISIONING, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with ( @@ -2225,12 +2231,12 @@ async def test_provisioning_to_terminating_when_load_balancer_registration_fails ) get_backends_mock.return_value = [(backend, backend_mock)] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATING - assert compute.active is False - assert compute.status_message == "Error registering with load balancer" + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATING + assert replica.active is False + assert replica.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 @@ -2246,13 +2252,13 @@ async def test_provisioning_to_terminating_when_backend_does_not_support_load_ba hostname="gateway-lb.example.com", backend_data="lb-backend-data", ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, status=GatewayReplicaStatus.PROVISIONING, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with ( @@ -2269,12 +2275,12 @@ async def test_provisioning_to_terminating_when_backend_does_not_support_load_ba 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 worker.process(_replica_to_pipeline_item(replica)) - 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" + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATING + assert replica.active is False + assert replica.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 @@ -2289,31 +2295,31 @@ async def test_provisioning_waits_for_pending_acm_gateway_migration( certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), hostname=None, # migration not yet performed ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( 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 + _lock_replica(replica) + original_last_processed_at = replica.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)) + await worker.process(_replica_to_pipeline_item(replica)) pool_add.assert_not_called() - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.PROVISIONING - assert compute.last_processed_at > original_last_processed_at - assert compute.lock_token is None + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.PROVISIONING + assert replica.last_processed_at > original_last_processed_at + assert replica.lock_token is None - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) async def test_provisioning_to_terminating_if_connect_fails( - self, test_db, session: AsyncSession, worker: GatewayReplicaWorker, legacy_compute: bool + self, test_db, session: AsyncSession, worker: GatewayReplicaWorker, legacy_replica: bool ): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) @@ -2323,36 +2329,36 @@ async def test_provisioning_to_terminating_if_connect_fails( backend_id=backend.id, status=GatewayStatus.PROVISIONING, ) - if legacy_compute: - compute = await create_gateway_compute( + if legacy_replica: + replica = await create_gateway_replica( session=session, status=GatewayReplicaStatus.PROVISIONING, ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id else: - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.PROVISIONING, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with patch( - "dstack._internal.server.services.gateways.connect_to_gateway_with_retry" + "dstack._internal.server.services.gateways.connect_to_gateway_replica_with_retry" ) as connect_mock: connect_mock.return_value = None - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) connect_mock.assert_called_once() - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATING - assert compute.active is False - assert compute.status_message == "Failed to connect to gateway" + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATING + assert replica.active is False + assert replica.status_message == "Failed to connect to gateway replica" - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) async def test_provisioning_to_terminating_if_configure_fails( - self, test_db, session: AsyncSession, worker: GatewayReplicaWorker, legacy_compute: bool + self, test_db, session: AsyncSession, worker: GatewayReplicaWorker, legacy_replica: bool ): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) @@ -2362,37 +2368,39 @@ async def test_provisioning_to_terminating_if_configure_fails( backend_id=backend.id, status=GatewayStatus.PROVISIONING, ) - if legacy_compute: - compute = await create_gateway_compute( + if legacy_replica: + replica = await create_gateway_replica( session=session, status=GatewayReplicaStatus.PROVISIONING, ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id else: - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.PROVISIONING, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with ( patch( - "dstack._internal.server.services.gateways.connect_to_gateway_with_retry" + "dstack._internal.server.services.gateways.connect_to_gateway_replica_with_retry" ) as connect_mock, - patch("dstack._internal.server.services.gateways.configure_gateway") as configure_mock, + patch( + "dstack._internal.server.services.gateways.configure_gateway_replica" + ) as configure_mock, ): connect_mock.return_value = MagicMock() configure_mock.side_effect = Exception("Configure failed") - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) connect_mock.assert_called_once() configure_mock.assert_called_once() - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATING - assert compute.active is False - assert compute.status_message == "Failed to configure gateway" + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATING + assert replica.active is False + assert replica.status_message == "Failed to configure gateway replica" @pytest.mark.parametrize( "gateway_status,to_be_deleted", @@ -2401,7 +2409,7 @@ async def test_provisioning_to_terminating_if_configure_fails( (GatewayStatus.RUNNING, True), ], ) - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) async def test_provisioning_to_terminating_if_gateway_needs_cleanup( self, test_db, @@ -2409,7 +2417,7 @@ async def test_provisioning_to_terminating_if_gateway_needs_cleanup( worker: GatewayReplicaWorker, gateway_status: GatewayStatus, to_be_deleted: bool, - legacy_compute: bool, + legacy_replica: bool, ): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) @@ -2420,43 +2428,43 @@ async def test_provisioning_to_terminating_if_gateway_needs_cleanup( status=gateway_status, ) gateway.to_be_deleted = to_be_deleted - if legacy_compute: - compute = await create_gateway_compute( + if legacy_replica: + replica = await create_gateway_replica( session=session, status=GatewayReplicaStatus.PROVISIONING, ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id else: - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.PROVISIONING, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with patch( "dstack._internal.server.background.pipeline_tasks.gateway_replicas._connect_and_configure_gateway_replica" ) as connect_mock: - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) connect_mock.assert_not_called() - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATING - assert compute.active is False + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATING + assert replica.active is False @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) class TestGatewayReplicaWorkerTerminating: - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) @pytest.mark.parametrize("populate_configuration", [True, False]) async def test_terminating_to_terminated( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker, - legacy_compute: bool, + legacy_replica: bool, populate_configuration: bool, ): project = await create_project(session=session) @@ -2468,17 +2476,17 @@ async def test_terminating_to_terminated( status=GatewayStatus.FAILED, populate_configuration=populate_configuration, ) - if legacy_compute: - compute = await create_gateway_compute( + if legacy_replica: + replica = await create_gateway_replica( session=session, backend_id=backend.id, status=GatewayReplicaStatus.TERMINATING, active=False, populate_configuration=populate_configuration, ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id else: - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -2486,7 +2494,7 @@ async def test_terminating_to_terminated( active=False, populate_configuration=populate_configuration, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with ( @@ -2501,16 +2509,16 @@ async def test_terminating_to_terminated( backend_mock.compute.return_value = Mock(spec=ComputeMockSpec) get_backends_mock.return_value = [(backend, backend_mock)] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) get_backends_mock.assert_called_once() - backend_mock.compute.return_value.terminate_gateway.assert_called_once() - remove_mock.assert_called_once_with(compute.ip_address) + backend_mock.compute.return_value.terminate_gateway_replica.assert_called_once() + remove_mock.assert_called_once_with(replica.ip_address) - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATED - assert compute.active is False - assert compute.deleted is True + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATED + assert replica.active is False + assert replica.deleted is True async def test_terminating_to_terminated_deletes_only_own_registration_records( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker @@ -2525,7 +2533,7 @@ async def test_terminating_to_terminated_deletes_only_own_registration_records( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute_to_terminate = await create_gateway_compute( + replica_to_terminate = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -2533,7 +2541,7 @@ async def test_terminating_to_terminated_deletes_only_own_registration_records( active=False, replica_num=0, ) - other_compute = await create_gateway_compute( + other_replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -2545,16 +2553,16 @@ async def test_terminating_to_terminated_deletes_only_own_registration_records( 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 + run_id=run.id, gateway_replica_id=replica_to_terminate.id, is_registered=True ) terminated_replica_registration = ServiceReplicaRegistrationModel( - job_id=job.id, gateway_replica_id=compute_to_terminate.id, is_registered=True + job_id=job.id, gateway_replica_id=replica_to_terminate.id, is_registered=True ) other_service_registration = ServiceRegistrationModel( - run_id=run.id, gateway_replica_id=other_compute.id, is_registered=True + run_id=run.id, gateway_replica_id=other_replica.id, is_registered=True ) other_replica_registration = ServiceReplicaRegistrationModel( - job_id=job.id, gateway_replica_id=other_compute.id, is_registered=True + job_id=job.id, gateway_replica_id=other_replica.id, is_registered=True ) session.add_all( [ @@ -2564,7 +2572,7 @@ async def test_terminating_to_terminated_deletes_only_own_registration_records( other_replica_registration, ] ) - _lock_compute(compute_to_terminate) + _lock_replica(replica_to_terminate) await session.commit() with ( @@ -2579,20 +2587,20 @@ async def test_terminating_to_terminated_deletes_only_own_registration_records( 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 worker.process(_replica_to_pipeline_item(replica_to_terminate)) - await session.refresh(compute_to_terminate) - assert compute_to_terminate.status == GatewayReplicaStatus.TERMINATED - assert compute_to_terminate.deleted is True + await session.refresh(replica_to_terminate) + assert replica_to_terminate.status == GatewayReplicaStatus.TERMINATED + assert replica_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 + assert remaining_service_registration.gateway_replica_id == other_replica.id remaining_replica_registration = ( (await session.execute(select(ServiceReplicaRegistrationModel))).scalars().one() ) - assert remaining_replica_registration.gateway_replica_id == other_compute.id + assert remaining_replica_registration.gateway_replica_id == other_replica.id async def test_terminating_deregisters_from_load_balancer_before_terminating( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker @@ -2608,14 +2616,14 @@ async def test_terminating_deregisters_from_load_balancer_before_terminating( hostname="gateway-lb.example.com", backend_data="lb-backend-data", ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, status=GatewayReplicaStatus.TERMINATING, active=False, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with ( @@ -2630,22 +2638,22 @@ async def test_terminating_deregisters_from_load_balancer_before_terminating( backend_mock.compute.return_value = Mock(spec=ComputeMockSpec) get_backends_mock.return_value = [(backend, backend_mock)] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) deregister_mock = ( backend_mock.compute.return_value.deregister_gateway_replica_from_load_balancer ) deregister_mock.assert_called_once() call_args = deregister_mock.call_args.args - assert call_args[0] == compute.instance_id + assert call_args[0] == replica.instance_id assert call_args[1].gateway_name == gateway.name assert call_args[2] == "lb-backend-data" - backend_mock.compute.return_value.terminate_gateway.assert_called_once() + backend_mock.compute.return_value.terminate_gateway_replica.assert_called_once() - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATED - assert compute.active is False - assert compute.deleted is True + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATED + assert replica.active is False + assert replica.deleted is True async def test_terminating_proceeds_when_load_balancer_deregistration_raises( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker @@ -2661,14 +2669,14 @@ async def test_terminating_proceeds_when_load_balancer_deregistration_raises( hostname="gateway-lb.example.com", backend_data="lb-backend-data", ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, status=GatewayReplicaStatus.TERMINATING, active=False, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with ( @@ -2686,20 +2694,20 @@ async def test_terminating_proceeds_when_load_balancer_deregistration_raises( ) get_backends_mock.return_value = [(backend, backend_mock)] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) - backend_mock.compute.return_value.terminate_gateway.assert_called_once() + backend_mock.compute.return_value.terminate_gateway_replica.assert_called_once() deregister_mock = ( backend_mock.compute.return_value.deregister_gateway_replica_from_load_balancer ) deregister_mock.assert_called_once() - await session.refresh(compute) + await session.refresh(replica) # Deregistration failures do not block termination: the load balancer is expected # to eventually deregister the (now-terminated) target automatically. - assert compute.status == GatewayReplicaStatus.TERMINATED - assert compute.active is False - assert compute.deleted is True + assert replica.status == GatewayReplicaStatus.TERMINATED + assert replica.active is False + assert replica.deleted is True async def test_terminating_skips_deregistration_when_gateway_has_no_hostname( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker @@ -2712,14 +2720,14 @@ async def test_terminating_skips_deregistration_when_gateway_has_no_hostname( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, status=GatewayReplicaStatus.TERMINATING, active=False, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with ( @@ -2734,13 +2742,13 @@ async def test_terminating_skips_deregistration_when_gateway_has_no_hostname( backend_mock.compute.return_value = Mock(spec=ComputeMockSpec) get_backends_mock.return_value = [(backend, backend_mock)] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) backend_mock.compute.return_value.deregister_gateway_replica_from_load_balancer.assert_not_called() - backend_mock.compute.return_value.terminate_gateway.assert_called_once() + backend_mock.compute.return_value.terminate_gateway_replica.assert_called_once() - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATED + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATED async def test_terminating_waits_for_pending_acm_gateway_migration( self, test_db, session: AsyncSession, worker: GatewayReplicaWorker @@ -2755,7 +2763,7 @@ async def test_terminating_waits_for_pending_acm_gateway_migration( certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), hostname=None, # migration not yet performed by the gateway pipeline ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -2763,8 +2771,8 @@ async def test_terminating_waits_for_pending_acm_gateway_migration( active=False, hostname_deprecated_readonly="legacy-lb.example.com", ) - _lock_compute(compute) - original_last_processed_at = compute.last_processed_at + _lock_replica(replica) + original_last_processed_at = replica.last_processed_at await session.commit() with patch( @@ -2774,19 +2782,19 @@ async def test_terminating_waits_for_pending_acm_gateway_migration( backend_mock.compute.return_value = Mock(spec=ComputeMockSpec) get_backends_mock.return_value = [(backend, backend_mock)] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) - backend_mock.compute.return_value.terminate_gateway.assert_not_called() + backend_mock.compute.return_value.terminate_gateway_replica.assert_not_called() backend_mock.compute.return_value.deregister_gateway_replica_from_load_balancer.assert_not_called() - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATING - assert compute.last_processed_at > original_last_processed_at - assert compute.lock_token is None + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATING + assert replica.last_processed_at > original_last_processed_at + assert replica.lock_token is None - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) async def test_terminating_to_terminated_if_backend_not_available( - self, test_db, session: AsyncSession, worker: GatewayReplicaWorker, legacy_compute: bool + self, test_db, session: AsyncSession, worker: GatewayReplicaWorker, legacy_replica: bool ): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) @@ -2796,39 +2804,39 @@ async def test_terminating_to_terminated_if_backend_not_available( backend_id=backend.id, status=GatewayStatus.FAILED, ) - if legacy_compute: - compute = await create_gateway_compute( + if legacy_replica: + replica = await create_gateway_replica( session=session, backend_id=backend.id, status=GatewayReplicaStatus.TERMINATING, active=False, ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id else: - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, status=GatewayReplicaStatus.TERMINATING, active=False, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with patch( "dstack._internal.server.services.backends.get_project_backends_with_models" ) as get_backends_mock: get_backends_mock.return_value = [] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATED - assert compute.active is False - assert compute.deleted is True + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATED + assert replica.active is False + assert replica.deleted is True - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) async def test_terminating_to_terminated_with_no_instance_id( - self, test_db, session: AsyncSession, worker: GatewayReplicaWorker, legacy_compute: bool + self, test_db, session: AsyncSession, worker: GatewayReplicaWorker, legacy_replica: bool ): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) @@ -2838,17 +2846,17 @@ async def test_terminating_to_terminated_with_no_instance_id( backend_id=backend.id, status=GatewayStatus.FAILED, ) - if legacy_compute: - compute = await create_gateway_compute( + if legacy_replica: + replica = await create_gateway_replica( session=session, backend_id=backend.id, instance_id=None, status=GatewayReplicaStatus.TERMINATING, active=False, ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id else: - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, @@ -2856,7 +2864,7 @@ async def test_terminating_to_terminated_with_no_instance_id( status=GatewayReplicaStatus.TERMINATING, active=False, ) - _lock_compute(compute) + _lock_replica(replica) await session.commit() with ( @@ -2871,19 +2879,19 @@ async def test_terminating_to_terminated_with_no_instance_id( backend_mock.compute.return_value = Mock(spec=ComputeMockSpec) get_backends_mock.return_value = [(backend, backend_mock)] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) - backend_mock.compute.return_value.terminate_gateway.assert_not_called() + backend_mock.compute.return_value.terminate_gateway_replica.assert_not_called() remove_mock.assert_not_called() - await session.refresh(compute) - assert compute.status == GatewayReplicaStatus.TERMINATED - assert compute.active is False - assert compute.deleted is True + await session.refresh(replica) + assert replica.status == GatewayReplicaStatus.TERMINATED + assert replica.active is False + assert replica.deleted is True - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) async def test_terminating_retries_if_terminate_fails( - self, test_db, session: AsyncSession, worker: GatewayReplicaWorker, legacy_compute: bool + self, test_db, session: AsyncSession, worker: GatewayReplicaWorker, legacy_replica: bool ): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) @@ -2893,24 +2901,24 @@ async def test_terminating_retries_if_terminate_fails( backend_id=backend.id, status=GatewayStatus.FAILED, ) - if legacy_compute: - compute = await create_gateway_compute( + if legacy_replica: + replica = await create_gateway_replica( session=session, backend_id=backend.id, status=GatewayReplicaStatus.TERMINATING, active=False, ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id else: - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id, status=GatewayReplicaStatus.TERMINATING, active=False, ) - _lock_compute(compute) - original_last_processed_at = compute.last_processed_at + _lock_replica(replica) + original_last_processed_at = replica.last_processed_at await session.commit() with ( @@ -2923,21 +2931,21 @@ async def test_terminating_retries_if_terminate_fails( ): backend_mock = Mock() backend_mock.compute.return_value = Mock(spec=ComputeMockSpec) - backend_mock.compute.return_value.terminate_gateway.side_effect = Exception( + backend_mock.compute.return_value.terminate_gateway_replica.side_effect = Exception( "Terminate failed" ) get_backends_mock.return_value = [(backend, backend_mock)] - await worker.process(_compute_to_pipeline_item(compute)) + await worker.process(_replica_to_pipeline_item(replica)) get_backends_mock.assert_called_once() - backend_mock.compute.return_value.terminate_gateway.assert_called_once() + backend_mock.compute.return_value.terminate_gateway_replica.assert_called_once() remove_mock.assert_not_called() - await session.refresh(compute) + await session.refresh(replica) # Not TERMINATED, should retry termination - assert compute.status == GatewayReplicaStatus.TERMINATING - assert compute.last_processed_at > original_last_processed_at - assert compute.lock_token is None - assert compute.lock_expires_at is None - assert compute.lock_owner is None + assert replica.status == GatewayReplicaStatus.TERMINATING + assert replica.last_processed_at > original_last_processed_at + assert replica.lock_token is None + assert replica.lock_expires_at is None + assert replica.lock_owner is None diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_gateways.py b/src/tests/_internal/server/background/pipeline_tasks/test_gateways.py index e8c8b45a54..e2d38f6516 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_gateways.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_gateways.py @@ -24,15 +24,15 @@ GatewayPipelineItem, GatewayWorker, ) -from dstack._internal.server.models import GatewayComputeModel, GatewayModel -from dstack._internal.server.services.gateways import get_gateway_compute_models +from dstack._internal.server.models import GatewayModel, GatewayReplicaModel +from dstack._internal.server.services.gateways import get_gateway_replica_models from dstack._internal.server.testing.common import ( ComputeMockSpec, create_backend, create_gateway, - create_gateway_compute, + create_gateway_replica, create_project, - get_gateway_compute_configuration, + get_gateway_replica_configuration, list_events, ) from dstack._internal.utils.common import get_current_datetime @@ -68,17 +68,17 @@ def _gateway_to_pipeline_item(gateway_model: GatewayModel) -> GatewayPipelineIte ) -async def _fetch_all_gateway_computes( +async def _fetch_all_gateway_replicas( session: AsyncSession, gateway_id: uuid.UUID -) -> list[GatewayComputeModel]: +) -> list[GatewayReplicaModel]: res = await session.execute( select(GatewayModel) .where(GatewayModel.id == gateway_id) - .options(selectinload(GatewayModel.gateway_computes)) - .options(selectinload(GatewayModel.gateway_compute)) + .options(selectinload(GatewayModel.gateway_replicas)) + .options(selectinload(GatewayModel.gateway_replica)) ) gateway = res.unique().scalar_one() - return get_gateway_compute_models(gateway) + return get_gateway_replica_models(gateway) @pytest.mark.asyncio @@ -259,9 +259,9 @@ async def test_fetch_returns_oldest_gateways_first_up_to_limit( assert middle.lock_owner == GatewayPipeline.__name__ assert newest.lock_owner is None - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) async def test_fetch_excludes_running_gateway_when_replica_count_matches( - self, test_db, session: AsyncSession, fetcher: GatewayFetcher, legacy_compute: bool + self, test_db, session: AsyncSession, fetcher: GatewayFetcher, legacy_replica: bool ): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) @@ -274,15 +274,15 @@ async def test_fetch_excludes_running_gateway_when_replica_count_matches( replicas=1, last_processed_at=stale, ) - if legacy_compute: - compute = await create_gateway_compute( + if legacy_replica: + replica = await create_gateway_replica( session=session, backend_id=backend.id, status=GatewayReplicaStatus.RUNNING, ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id else: - await create_gateway_compute( + await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, @@ -307,14 +307,14 @@ async def test_fetch_includes_running_gateway_with_pending_scale_attempt_even_if replicas=1, last_processed_at=stale, ) - await create_gateway_compute( + await create_gateway_replica( session=session, gateway_id=gateway.id, ip_address=None, instance_id=None, region=None, status=GatewayReplicaStatus.SUBMITTED, - configuration=get_gateway_compute_configuration().model_dump_json(), + configuration=get_gateway_replica_configuration().model_dump_json(), ) gateway.replica_scale_attempt = 1 await session.commit() @@ -322,9 +322,9 @@ async def test_fetch_includes_running_gateway_with_pending_scale_attempt_even_if items = await fetcher.fetch(limit=10) assert {item.id for item in items} == {gateway.id} - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) async def test_fetch_includes_running_gateway_when_replica_count_not_matches( - self, test_db, session: AsyncSession, fetcher: GatewayFetcher, legacy_compute: bool + self, test_db, session: AsyncSession, fetcher: GatewayFetcher, legacy_replica: bool ): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) @@ -337,15 +337,15 @@ async def test_fetch_includes_running_gateway_when_replica_count_not_matches( replicas=2, last_processed_at=stale, ) - if legacy_compute: - compute = await create_gateway_compute( + if legacy_replica: + replica = await create_gateway_replica( session=session, backend_id=backend.id, status=GatewayReplicaStatus.RUNNING, ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id else: - await create_gateway_compute( + await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, @@ -355,9 +355,9 @@ async def test_fetch_includes_running_gateway_when_replica_count_not_matches( items = await fetcher.fetch(limit=10) assert {item.id for item in items} == {gateway.id} - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) async def test_fetch_includes_running_gateway_with_unmigrated_legacy_hostname( - self, test_db, session: AsyncSession, fetcher: GatewayFetcher, legacy_compute: bool + self, test_db, session: AsyncSession, fetcher: GatewayFetcher, legacy_replica: bool ): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) @@ -372,16 +372,16 @@ async def test_fetch_includes_running_gateway_with_unmigrated_legacy_hostname( certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), hostname=None, # not yet migrated ) - if legacy_compute: - compute = await create_gateway_compute( + if legacy_replica: + replica = await create_gateway_replica( session=session, backend_id=backend.id, status=GatewayReplicaStatus.RUNNING, hostname_deprecated_readonly="legacy-lb.example.com", ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id else: - await create_gateway_compute( + await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, @@ -410,7 +410,7 @@ async def test_fetch_excludes_running_gateway_without_legacy_hostname_to_migrate certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), hostname=None, ) - await create_gateway_compute( + await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, @@ -437,7 +437,7 @@ async def test_fetch_excludes_already_migrated_gateway( certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), hostname="gateway-lb.example.com", # already migrated ) - await create_gateway_compute( + await create_gateway_replica( session=session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, @@ -478,15 +478,15 @@ async def test_submitted_to_provisioning( await session.refresh(gateway) assert gateway.status == GatewayStatus.PROVISIONING - computes = sorted( - await _fetch_all_gateway_computes(session, gateway.id), key=lambda c: c.replica_num + replicas = sorted( + await _fetch_all_gateway_replicas(session, gateway.id), key=lambda r: r.replica_num ) - assert len(computes) == 2 - assert computes[0].status == GatewayReplicaStatus.SUBMITTED - assert computes[0].replica_num == 0 - assert computes[1].status == GatewayReplicaStatus.SUBMITTED - assert computes[1].replica_num == 1 - assert all(c.ip_address is None for c in computes) + assert len(replicas) == 2 + assert replicas[0].status == GatewayReplicaStatus.SUBMITTED + assert replicas[0].replica_num == 0 + assert replicas[1].status == GatewayReplicaStatus.SUBMITTED + assert replicas[1].replica_num == 1 + assert all(r.ip_address is None for r in replicas) async def test_submitted_to_provisioning_creates_load_balancer_for_acm_gateway( self, test_db, session: AsyncSession, worker: GatewayWorker @@ -630,14 +630,14 @@ async def test_submitted_to_failed_when_load_balancer_creation_raises( @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) class TestGatewayWorkerProvisioning: - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) @pytest.mark.parametrize("populate_configuration", [True, False]) async def test_provisioning_to_running( self, test_db, session: AsyncSession, worker: GatewayWorker, - legacy_compute: bool, + legacy_replica: bool, populate_configuration: bool, ): project = await create_project(session=session) @@ -649,16 +649,16 @@ async def test_provisioning_to_running( status=GatewayStatus.PROVISIONING, populate_configuration=populate_configuration, ) - if legacy_compute: - gateway_compute = await create_gateway_compute( + if legacy_replica: + gateway_replica = await create_gateway_replica( session=session, backend_id=backend.id, status=GatewayReplicaStatus.RUNNING, populate_configuration=populate_configuration, ) - gateway.gateway_compute_id = gateway_compute.id + gateway.gateway_replica_id = gateway_replica.id else: - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, @@ -689,7 +689,7 @@ async def test_provisioning_migrates_hostname_and_backend_data_from_legacy_repli certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), hostname=None, # migration not yet performed ) - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, @@ -721,7 +721,7 @@ async def test_provisioning_does_not_overwrite_already_migrated_hostname( hostname="already-migrated.example.com", backend_data="current-backend-data", ) - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, @@ -751,14 +751,14 @@ async def test_provisioning_to_running_with_multiple_replicas( status=GatewayStatus.PROVISIONING, replicas=2, ) - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, ip_address="1.1.1.1", status=GatewayReplicaStatus.RUNNING, replica_num=0, ) - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, ip_address="2.2.2.2", @@ -789,14 +789,14 @@ async def test_still_provisioning_if_not_all_replicas_running( status=GatewayStatus.PROVISIONING, replicas=2, ) - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, ip_address="1.1.1.1", status=GatewayReplicaStatus.RUNNING, replica_num=0, ) - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, ip_address="2.2.2.2", @@ -816,7 +816,7 @@ async def test_still_provisioning_if_not_all_replicas_running( events = await list_events(session) assert len(events) == 0 - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) @pytest.mark.parametrize( "replica_status", [GatewayReplicaStatus.TERMINATING, GatewayReplicaStatus.TERMINATED] ) @@ -825,7 +825,7 @@ async def test_marks_gateway_as_failed_if_replica_failed( test_db, session: AsyncSession, worker: GatewayWorker, - legacy_compute: bool, + legacy_replica: bool, replica_status: GatewayReplicaStatus, ): project = await create_project(session=session) @@ -836,16 +836,16 @@ async def test_marks_gateway_as_failed_if_replica_failed( backend_id=backend.id, status=GatewayStatus.PROVISIONING, ) - if legacy_compute: - gateway_compute = await create_gateway_compute( + if legacy_replica: + gateway_replica = await create_gateway_replica( session=session, backend_id=backend.id, status=replica_status, active=False, ) - gateway.gateway_compute_id = gateway_compute.id + gateway.gateway_replica_id = gateway_replica.id else: - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, status=replica_status, @@ -878,14 +878,14 @@ async def test_still_provisioning_with_submitted_replica( backend_id=backend.id, status=GatewayStatus.PROVISIONING, ) - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, ip_address=None, instance_id=None, region=None, status=GatewayReplicaStatus.SUBMITTED, - configuration=get_gateway_compute_configuration().model_dump_json(), + configuration=get_gateway_replica_configuration().model_dump_json(), ) gateway.lock_token = uuid.uuid4() gateway.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) @@ -900,14 +900,14 @@ async def test_still_provisioning_with_submitted_replica( events = await list_events(session) assert len(events) == 0 - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) @pytest.mark.parametrize("populate_configuration", [True, False]) async def test_still_provisioning_when_scale_out_adds_new_replicas( self, test_db, session: AsyncSession, worker: GatewayWorker, - legacy_compute: bool, + legacy_replica: bool, populate_configuration: bool, ): project = await create_project(session=session) @@ -920,17 +920,17 @@ async def test_still_provisioning_when_scale_out_adds_new_replicas( replicas=2, populate_configuration=populate_configuration, ) - if legacy_compute: - gateway_compute = await create_gateway_compute( + if legacy_replica: + gateway_replica = await create_gateway_replica( session=session, backend_id=backend.id, ip_address="1.1.1.1", status=GatewayReplicaStatus.RUNNING, populate_configuration=populate_configuration, ) - gateway.gateway_compute_id = gateway_compute.id + gateway.gateway_replica_id = gateway_replica.id else: - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, ip_address="1.1.1.1", @@ -946,11 +946,11 @@ async def test_still_provisioning_when_scale_out_adds_new_replicas( await session.refresh(gateway) assert gateway.status == GatewayStatus.PROVISIONING - computes = sorted( - await _fetch_all_gateway_computes(session, gateway.id), key=lambda c: c.replica_num + replicas = sorted( + await _fetch_all_gateway_replicas(session, gateway.id), key=lambda r: r.replica_num ) - assert [c.replica_num for c in computes] == [0, 1] - assert computes[1].status == GatewayReplicaStatus.SUBMITTED + assert [r.replica_num for r in replicas] == [0, 1] + assert replicas[1].status == GatewayReplicaStatus.SUBMITTED assert gateway.replica_scale_attempt == 1 assert gateway.last_replica_scale_attempt_at is not None events = await list_events(session) @@ -968,7 +968,7 @@ async def test_provisioning_to_running_when_scale_in_removes_surplus_replicas( status=GatewayStatus.PROVISIONING, replicas=1, ) - older = await create_gateway_compute( + older = await create_gateway_replica( session, gateway_id=gateway.id, ip_address="1.1.1.1", @@ -976,7 +976,7 @@ async def test_provisioning_to_running_when_scale_in_removes_surplus_replicas( replica_num=0, ) older.created_at = datetime(2025, 1, 1) - newer = await create_gateway_compute( + newer = await create_gateway_replica( session, gateway_id=gateway.id, ip_address="2.2.2.2", @@ -1012,14 +1012,14 @@ async def test_ignores_previously_scaled_in_replica_when_determining_status( status=GatewayStatus.PROVISIONING, replicas=1, ) - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, ip_address="1.1.1.1", status=GatewayReplicaStatus.RUNNING, replica_num=0, ) - scaled_in_compute = await create_gateway_compute( + scaled_in_replica = await create_gateway_replica( session, gateway_id=gateway.id, ip_address="2.2.2.2", @@ -1027,7 +1027,7 @@ async def test_ignores_previously_scaled_in_replica_when_determining_status( active=False, replica_num=1, ) - scaled_in_compute.scale_in = True + scaled_in_replica.scale_in = True gateway.lock_token = uuid.uuid4() gateway.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) await session.commit() @@ -1044,9 +1044,9 @@ async def test_ignores_previously_scaled_in_replica_when_determining_status( @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) class TestGatewayWorkerRunning: - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) async def test_no_scaling_when_replica_count_matches( - self, test_db, session: AsyncSession, worker: GatewayWorker, legacy_compute: bool + self, test_db, session: AsyncSession, worker: GatewayWorker, legacy_replica: bool ): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) @@ -1057,15 +1057,15 @@ async def test_no_scaling_when_replica_count_matches( status=GatewayStatus.RUNNING, replicas=1, ) - if legacy_compute: - compute = await create_gateway_compute( + if legacy_replica: + replica = await create_gateway_replica( session=session, backend_id=backend.id, status=GatewayReplicaStatus.RUNNING, ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id else: - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, @@ -1082,9 +1082,9 @@ async def test_no_scaling_when_replica_count_matches( await session.refresh(gateway) assert gateway.status == GatewayStatus.RUNNING assert gateway.last_processed_at > original_last_processed_at - computes = await _fetch_all_gateway_computes(session, gateway.id) - assert len(computes) == 1 - assert computes[0].scale_in is False + replicas = await _fetch_all_gateway_replicas(session, gateway.id) + assert len(replicas) == 1 + assert replicas[0].scale_in is False assert gateway.replica_scale_attempt == 0 # The desired count is met, reset counter async def test_running_migrates_hostname_and_backend_data_from_legacy_replica( @@ -1101,7 +1101,7 @@ async def test_running_migrates_hostname_and_backend_data_from_legacy_replica( certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), hostname=None, # migration not yet performed ) - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, @@ -1118,14 +1118,14 @@ async def test_running_migrates_hostname_and_backend_data_from_legacy_replica( assert gateway.hostname == "legacy-lb.example.com" assert gateway.backend_data == "legacy-backend-data" - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) @pytest.mark.parametrize("populate_configuration", [True, False]) async def test_scales_out_when_desired_replica_count_increased( self, test_db, session: AsyncSession, worker: GatewayWorker, - legacy_compute: bool, + legacy_replica: bool, populate_configuration: bool, ): project = await create_project(session=session) @@ -1138,16 +1138,16 @@ async def test_scales_out_when_desired_replica_count_increased( replicas=3, populate_configuration=populate_configuration, ) - if legacy_compute: - gateway_compute = await create_gateway_compute( + if legacy_replica: + gateway_replica = await create_gateway_replica( session=session, backend_id=backend.id, status=GatewayReplicaStatus.RUNNING, populate_configuration=populate_configuration, ) - gateway.gateway_compute_id = gateway_compute.id + gateway.gateway_replica_id = gateway_replica.id else: - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, @@ -1162,11 +1162,11 @@ async def test_scales_out_when_desired_replica_count_increased( await session.refresh(gateway) assert gateway.status == GatewayStatus.RUNNING - computes = sorted( - await _fetch_all_gateway_computes(session, gateway.id), key=lambda c: c.replica_num + replicas = sorted( + await _fetch_all_gateway_replicas(session, gateway.id), key=lambda r: r.replica_num ) - assert [c.replica_num for c in computes] == [0, 1, 2] - assert [c.status for c in computes] == [ + assert [r.replica_num for r in replicas] == [0, 1, 2] + assert [r.status for r in replicas] == [ GatewayReplicaStatus.RUNNING, GatewayReplicaStatus.SUBMITTED, GatewayReplicaStatus.SUBMITTED, @@ -1186,18 +1186,18 @@ async def test_scales_in_oldest_replicas_when_desired_replica_count_decreased( status=GatewayStatus.RUNNING, replicas=1, ) - compute0 = await create_gateway_compute( + replica0 = await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 ) - compute0.created_at = datetime(2025, 1, 1) - compute1 = await create_gateway_compute( + replica0.created_at = datetime(2025, 1, 1) + replica1 = await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=1 ) - compute1.created_at = datetime(2025, 1, 2) - compute2 = await create_gateway_compute( + replica1.created_at = datetime(2025, 1, 2) + replica2 = await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=2 ) - compute2.created_at = datetime(2025, 1, 3) + replica2.created_at = datetime(2025, 1, 3) gateway.lock_token = uuid.uuid4() gateway.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) await session.commit() @@ -1206,12 +1206,12 @@ async def test_scales_in_oldest_replicas_when_desired_replica_count_decreased( await session.refresh(gateway) assert gateway.status == GatewayStatus.RUNNING - await session.refresh(compute0) - await session.refresh(compute1) - await session.refresh(compute2) - assert compute0.scale_in is True - assert compute1.scale_in is True - assert compute2.scale_in is False + await session.refresh(replica0) + await session.refresh(replica1) + await session.refresh(replica2) + assert replica0.scale_in is True + assert replica1.scale_in is True + assert replica2.scale_in is False async def test_scale_in_prefers_less_advanced_replicas_over_older_running_ones( self, test_db, session: AsyncSession, worker: GatewayWorker @@ -1225,11 +1225,11 @@ async def test_scale_in_prefers_less_advanced_replicas_over_older_running_ones( status=GatewayStatus.RUNNING, replicas=1, ) - running = await create_gateway_compute( + running = await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 ) running.created_at = datetime(2025, 1, 1) - submitted = await create_gateway_compute( + submitted = await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.SUBMITTED, @@ -1237,7 +1237,7 @@ async def test_scale_in_prefers_less_advanced_replicas_over_older_running_ones( ip_address=None, instance_id=None, region=None, - configuration=get_gateway_compute_configuration().model_dump_json(), + configuration=get_gateway_replica_configuration().model_dump_json(), ) submitted.created_at = datetime(2025, 1, 2) gateway.lock_token = uuid.uuid4() @@ -1251,9 +1251,9 @@ async def test_scale_in_prefers_less_advanced_replicas_over_older_running_ones( assert running.scale_in is False assert submitted.scale_in is True - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) async def test_no_scaling_for_legacy_gateway_without_desired_replica_count( - self, test_db, session: AsyncSession, worker: GatewayWorker, legacy_compute: bool + self, test_db, session: AsyncSession, worker: GatewayWorker, legacy_replica: bool ): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) @@ -1263,15 +1263,15 @@ async def test_no_scaling_for_legacy_gateway_without_desired_replica_count( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - if legacy_compute: - compute = await create_gateway_compute( + if legacy_replica: + replica = await create_gateway_replica( session=session, backend_id=backend.id, status=GatewayReplicaStatus.RUNNING, ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id else: - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 ) gateway.desired_replica_count = None @@ -1283,9 +1283,9 @@ async def test_no_scaling_for_legacy_gateway_without_desired_replica_count( await session.refresh(gateway) assert gateway.status == GatewayStatus.RUNNING - computes = await _fetch_all_gateway_computes(session, gateway.id) - assert len(computes) == 1 - assert computes[0].scale_in is False + replicas = await _fetch_all_gateway_replicas(session, gateway.id) + assert len(replicas) == 1 + assert replicas[0].scale_in is False async def test_scale_out_skipped_before_retry_delay_elapses( self, test_db, session: AsyncSession, worker: GatewayWorker @@ -1299,7 +1299,7 @@ async def test_scale_out_skipped_before_retry_delay_elapses( status=GatewayStatus.RUNNING, replicas=2, ) - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 ) gateway.replica_scale_attempt = 1 @@ -1311,8 +1311,8 @@ async def test_scale_out_skipped_before_retry_delay_elapses( await worker.process(_gateway_to_pipeline_item(gateway)) await session.refresh(gateway) - computes = await _fetch_all_gateway_computes(session, gateway.id) - assert len(computes) == 1 + replicas = await _fetch_all_gateway_replicas(session, gateway.id) + assert len(replicas) == 1 assert gateway.replica_scale_attempt == 1 async def test_scale_out_retries_after_retry_delay_elapses( @@ -1327,7 +1327,7 @@ async def test_scale_out_retries_after_retry_delay_elapses( status=GatewayStatus.RUNNING, replicas=2, ) - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 ) gateway.replica_scale_attempt = 1 @@ -1339,8 +1339,8 @@ async def test_scale_out_retries_after_retry_delay_elapses( await worker.process(_gateway_to_pipeline_item(gateway)) await session.refresh(gateway) - computes = await _fetch_all_gateway_computes(session, gateway.id) - assert len(computes) == 2 + replicas = await _fetch_all_gateway_replicas(session, gateway.id) + assert len(replicas) == 2 assert gateway.replica_scale_attempt == 2 async def test_scale_out_stops_after_reaching_attempt_limit( @@ -1355,7 +1355,7 @@ async def test_scale_out_stops_after_reaching_attempt_limit( status=GatewayStatus.RUNNING, replicas=2, ) - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 ) gateway.replica_scale_attempt = _MAX_REPLICA_SCALE_ATTEMPTS @@ -1367,8 +1367,8 @@ async def test_scale_out_stops_after_reaching_attempt_limit( await worker.process(_gateway_to_pipeline_item(gateway)) await session.refresh(gateway) - computes = await _fetch_all_gateway_computes(session, gateway.id) - assert len(computes) == 1 + replicas = await _fetch_all_gateway_replicas(session, gateway.id) + assert len(replicas) == 1 assert gateway.replica_scale_attempt == _MAX_REPLICA_SCALE_ATTEMPTS events = await list_events(session) assert len(events) == 0 @@ -1385,7 +1385,7 @@ async def test_scale_out_emits_event_on_reaching_attempt_limit( status=GatewayStatus.RUNNING, replicas=2, ) - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 ) gateway.replica_scale_attempt = _MAX_REPLICA_SCALE_ATTEMPTS - 1 @@ -1397,9 +1397,9 @@ async def test_scale_out_emits_event_on_reaching_attempt_limit( await worker.process(_gateway_to_pipeline_item(gateway)) await session.refresh(gateway) - computes = await _fetch_all_gateway_computes(session, gateway.id) + replicas = await _fetch_all_gateway_replicas(session, gateway.id) # Last allowed attempt still creates the missing replica - assert len(computes) == 2 + assert len(replicas) == 2 assert gateway.replica_scale_attempt == _MAX_REPLICA_SCALE_ATTEMPTS events = await list_events(session) assert len(events) == 1 @@ -1417,7 +1417,7 @@ async def test_attempt_counter_not_reset_while_replacement_replica_still_provisi status=GatewayStatus.RUNNING, replicas=1, ) - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, ip_address=None, @@ -1425,7 +1425,7 @@ async def test_attempt_counter_not_reset_while_replacement_replica_still_provisi region=None, status=GatewayReplicaStatus.PROVISIONING, replica_num=0, - configuration=get_gateway_compute_configuration().model_dump_json(), + configuration=get_gateway_replica_configuration().model_dump_json(), ) gateway.replica_scale_attempt = 2 gateway.lock_token = uuid.uuid4() @@ -1449,7 +1449,7 @@ async def test_attempt_counter_resets_and_scales_out_immediately_after_in_place_ status=GatewayStatus.RUNNING, replicas=2, ) - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 ) gateway.replica_scale_attempt = _MAX_REPLICA_SCALE_ATTEMPTS @@ -1463,8 +1463,8 @@ async def test_attempt_counter_resets_and_scales_out_immediately_after_in_place_ await worker.process(_gateway_to_pipeline_item(gateway)) await session.refresh(gateway) - computes = await _fetch_all_gateway_computes(session, gateway.id) - assert len(computes) == 2 + replicas = await _fetch_all_gateway_replicas(session, gateway.id) + assert len(replicas) == 2 assert gateway.replica_scale_attempt == 1 async def test_attempt_counter_not_reset_when_update_precedes_last_scale_attempt( @@ -1479,7 +1479,7 @@ async def test_attempt_counter_not_reset_when_update_precedes_last_scale_attempt status=GatewayStatus.RUNNING, replicas=2, ) - await create_gateway_compute( + await create_gateway_replica( session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 ) gateway.replica_scale_attempt = _MAX_REPLICA_SCALE_ATTEMPTS @@ -1492,15 +1492,15 @@ async def test_attempt_counter_not_reset_when_update_precedes_last_scale_attempt await worker.process(_gateway_to_pipeline_item(gateway)) await session.refresh(gateway) - computes = await _fetch_all_gateway_computes(session, gateway.id) - assert len(computes) == 1 + replicas = await _fetch_all_gateway_replicas(session, gateway.id) + assert len(replicas) == 1 assert gateway.replica_scale_attempt == _MAX_REPLICA_SCALE_ATTEMPTS @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) class TestGatewayWorkerDeleted: - async def test_deletes_gateway_with_no_computes( + async def test_deletes_gateway_with_no_replicas( self, test_db, session: AsyncSession, worker: GatewayWorker ): project = await create_project(session=session) @@ -1524,14 +1524,14 @@ async def test_deletes_gateway_with_no_computes( assert len(events) == 1 assert events[0].message == "Gateway deleted" - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) @pytest.mark.parametrize("populate_configuration", [True, False]) async def test_deletes_gateway_when_all_replicas_terminated( self, test_db, session: AsyncSession, worker: GatewayWorker, - legacy_compute: bool, + legacy_replica: bool, populate_configuration: bool, ): project = await create_project(session=session) @@ -1543,17 +1543,17 @@ async def test_deletes_gateway_when_all_replicas_terminated( status=GatewayStatus.RUNNING, populate_configuration=populate_configuration, ) - if legacy_compute: - gateway_compute = await create_gateway_compute( + if legacy_replica: + gateway_replica = await create_gateway_replica( session=session, backend_id=backend.id, status=GatewayReplicaStatus.TERMINATED, active=False, populate_configuration=populate_configuration, ) - gateway.gateway_compute_id = gateway_compute.id + gateway.gateway_replica_id = gateway_replica.id else: - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -1588,7 +1588,7 @@ async def test_deletes_gateway_and_terminates_load_balancer_when_hostname_set( hostname="gateway-lb.example.com", backend_data="lb-backend-data", ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -1632,7 +1632,7 @@ async def test_delete_skips_load_balancer_termination_when_hostname_not_set( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -1667,7 +1667,7 @@ async def test_delete_deferred_when_load_balancer_termination_fails( hostname="gateway-lb.example.com", backend_data="lb-backend-data", ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -1717,7 +1717,7 @@ async def test_delete_deferred_when_backend_does_not_support_load_balancer( hostname="gateway-lb.example.com", backend_data="lb-backend-data", ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -1755,7 +1755,7 @@ async def test_delete_migrates_hostname_before_evaluating_termination( certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), hostname=None, # migration not yet performed ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -1808,7 +1808,7 @@ async def test_waits_when_replicas_not_yet_terminated( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -1847,7 +1847,7 @@ async def test_deletes_gateway_with_multiple_replicas_all_terminated( status=GatewayStatus.RUNNING, replicas=2, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -1856,7 +1856,7 @@ async def test_deletes_gateway_with_multiple_replicas_all_terminated( active=False, replica_num=0, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, 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 4d6cf5b9a0..1e854148a0 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 @@ -81,7 +81,7 @@ create_export, create_fleet, create_gateway, - create_gateway_compute, + create_gateway_replica, create_instance, create_job, create_job_metrics_point, @@ -1888,7 +1888,7 @@ async def test_terminates_job_on_gateway_registration_failure( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - gateway_compute = await create_gateway_compute(session=session, gateway_id=gateway.id) + gateway_replica = await create_gateway_replica(session=session, gateway_id=gateway.id) run = await create_run( session=session, project=project, @@ -1920,7 +1920,7 @@ async def test_terminates_job_on_gateway_registration_failure( session.add( ServiceReplicaRegistrationModel( job_id=job.id, - gateway_replica_id=gateway_compute.id, + gateway_replica_id=gateway_replica.id, is_registered=False, register_attempt=2, register_status_message="Connection refused", @@ -1966,10 +1966,10 @@ async def test_does_not_terminate_job_when_registered_with_at_least_one_gateway_ backend_id=backend.id, status=GatewayStatus.RUNNING, ) - gateway_compute_1 = await create_gateway_compute( + gateway_replica_1 = await create_gateway_replica( session=session, gateway_id=gateway.id, replica_num=0 ) - gateway_compute_2 = await create_gateway_compute( + gateway_replica_2 = await create_gateway_replica( session=session, gateway_id=gateway.id, replica_num=1 ) run = await create_run( @@ -2003,7 +2003,7 @@ async def test_does_not_terminate_job_when_registered_with_at_least_one_gateway_ session.add( ServiceReplicaRegistrationModel( job_id=job.id, - gateway_replica_id=gateway_compute_1.id, + gateway_replica_id=gateway_replica_1.id, is_registered=False, register_attempt=2, register_status_message="Connection refused", @@ -2012,7 +2012,7 @@ async def test_does_not_terminate_job_when_registered_with_at_least_one_gateway_ session.add( ServiceReplicaRegistrationModel( job_id=job.id, - gateway_replica_id=gateway_compute_2.id, + gateway_replica_id=gateway_replica_2.id, is_registered=True, register_attempt=0, ) @@ -2057,11 +2057,11 @@ async def test_does_not_terminate_job_when_gateway_replica_has_not_attempted_reg backend_id=backend.id, status=GatewayStatus.RUNNING, ) - gateway_compute_1 = await create_gateway_compute( + gateway_replica_1 = await create_gateway_replica( 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) + await create_gateway_replica(session=session, gateway_id=gateway.id, replica_num=1) run = await create_run( session=session, project=project, @@ -2093,7 +2093,7 @@ async def test_does_not_terminate_job_when_gateway_replica_has_not_attempted_reg session.add( ServiceReplicaRegistrationModel( job_id=job.id, - gateway_replica_id=gateway_compute_1.id, + gateway_replica_id=gateway_replica_1.id, is_registered=False, register_attempt=2, register_status_message="Connection refused", @@ -2139,12 +2139,12 @@ async def test_terminates_job_ignoring_registration_on_non_running_replica( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - gateway_compute_running = await create_gateway_compute( + gateway_replica_running = await create_gateway_replica( 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( + gateway_replica_terminating = await create_gateway_replica( session=session, gateway_id=gateway.id, replica_num=1, @@ -2181,7 +2181,7 @@ async def test_terminates_job_ignoring_registration_on_non_running_replica( session.add( ServiceReplicaRegistrationModel( job_id=job.id, - gateway_replica_id=gateway_compute_running.id, + gateway_replica_id=gateway_replica_running.id, is_registered=False, register_attempt=2, register_status_message="Connection refused", @@ -2190,7 +2190,7 @@ async def test_terminates_job_ignoring_registration_on_non_running_replica( session.add( ServiceReplicaRegistrationModel( job_id=job.id, - gateway_replica_id=gateway_compute_terminating.id, + gateway_replica_id=gateway_replica_terminating.id, is_registered=True, register_attempt=0, ) @@ -2442,7 +2442,7 @@ async def test_registers_service_replica_in_gateway( name="test-gateway", wildcard_domain="example.com", ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -2526,7 +2526,7 @@ async def test_registers_service_replica_in_gateway_when_running_on_imported_ins name="test-gateway", wildcard_domain="example.com", ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -2815,7 +2815,7 @@ async def test_registers_router_replica_but_not_worker_replica_in_gateway( name="test-gateway", wildcard_domain="example.com", ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, 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 9ea9962c6e..fe24a8cdef 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 @@ -41,7 +41,7 @@ create_backend, create_fleet, create_gateway, - create_gateway_compute, + create_gateway_replica, create_instance, create_job, create_project, @@ -196,7 +196,7 @@ async def test_terminates_run_on_gateway_registration_failure( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - gateway_compute = await create_gateway_compute(session=session, gateway_id=gateway.id) + gateway_replica = await create_gateway_replica(session=session, gateway_id=gateway.id) run = await create_run( session=session, project=project, @@ -220,7 +220,7 @@ async def test_terminates_run_on_gateway_registration_failure( session.add( ServiceRegistrationModel( run_id=run.id, - gateway_replica_id=gateway_compute.id, + gateway_replica_id=gateway_replica.id, is_registered=False, register_attempt=3, register_status_message="Connection refused", @@ -249,10 +249,10 @@ async def test_does_not_terminate_run_when_service_registered_despite_failed_att backend_id=backend.id, status=GatewayStatus.RUNNING, ) - gateway_compute_1 = await create_gateway_compute( + gateway_replica_1 = await create_gateway_replica( session=session, gateway_id=gateway.id, replica_num=0 ) - gateway_compute_2 = await create_gateway_compute( + gateway_replica_2 = await create_gateway_replica( session=session, gateway_id=gateway.id, replica_num=1 ) run = await create_run( @@ -278,7 +278,7 @@ async def test_does_not_terminate_run_when_service_registered_despite_failed_att session.add( ServiceRegistrationModel( run_id=run.id, - gateway_replica_id=gateway_compute_1.id, + gateway_replica_id=gateway_replica_1.id, is_registered=False, register_attempt=3, register_status_message="Connection refused", @@ -287,7 +287,7 @@ async def test_does_not_terminate_run_when_service_registered_despite_failed_att session.add( ServiceRegistrationModel( run_id=run.id, - gateway_replica_id=gateway_compute_2.id, + gateway_replica_id=gateway_replica_2.id, is_registered=True, register_attempt=0, ) @@ -315,11 +315,11 @@ async def test_does_not_terminate_run_when_one_running_replica_has_not_attempted backend_id=backend.id, status=GatewayStatus.RUNNING, ) - gateway_compute_1 = await create_gateway_compute( + gateway_replica_1 = await create_gateway_replica( 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) + await create_gateway_replica(session=session, gateway_id=gateway.id, replica_num=1) run = await create_run( session=session, project=project, @@ -343,7 +343,7 @@ async def test_does_not_terminate_run_when_one_running_replica_has_not_attempted session.add( ServiceRegistrationModel( run_id=run.id, - gateway_replica_id=gateway_compute_1.id, + gateway_replica_id=gateway_replica_1.id, is_registered=False, register_attempt=3, register_status_message="Connection refused", @@ -372,12 +372,12 @@ async def test_terminates_run_ignoring_registration_on_non_running_replica( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - gateway_compute_running = await create_gateway_compute( + gateway_replica_running = await create_gateway_replica( 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( + gateway_replica_terminating = await create_gateway_replica( session=session, gateway_id=gateway.id, replica_num=1, @@ -406,7 +406,7 @@ async def test_terminates_run_ignoring_registration_on_non_running_replica( session.add( ServiceRegistrationModel( run_id=run.id, - gateway_replica_id=gateway_compute_running.id, + gateway_replica_id=gateway_replica_running.id, is_registered=False, register_attempt=3, register_status_message="Connection refused", @@ -415,7 +415,7 @@ async def test_terminates_run_ignoring_registration_on_non_running_replica( session.add( ServiceRegistrationModel( run_id=run.id, - gateway_replica_id=gateway_compute_terminating.id, + gateway_replica_id=gateway_replica_terminating.id, is_registered=True, register_attempt=0, ) @@ -1407,10 +1407,10 @@ async def test_service_rolling_deployment_keeps_old_replica_until_new_replica_re backend_id=backend.id, status=GatewayStatus.RUNNING, ) - gateway_compute_1 = await create_gateway_compute( + gateway_replica_1 = await create_gateway_replica( session=session, gateway_id=gateway.id, replica_num=0 ) - gateway_compute_2 = await create_gateway_compute( + gateway_replica_2 = await create_gateway_replica( session=session, gateway_id=gateway.id, replica_num=1 ) run_spec = get_run_spec( @@ -1457,7 +1457,7 @@ async def test_service_rolling_deployment_keeps_old_replica_until_new_replica_re session.add( ServiceReplicaRegistrationModel( job_id=old_job.id, - gateway_replica_id=gateway_compute_1.id, + gateway_replica_id=gateway_replica_1.id, is_registered=True, register_attempt=0, ) @@ -1465,7 +1465,7 @@ async def test_service_rolling_deployment_keeps_old_replica_until_new_replica_re session.add( ServiceReplicaRegistrationModel( job_id=old_job.id, - gateway_replica_id=gateway_compute_2.id, + gateway_replica_id=gateway_replica_2.id, is_registered=True, register_attempt=0, ) @@ -1474,7 +1474,7 @@ async def test_service_rolling_deployment_keeps_old_replica_until_new_replica_re session.add( ServiceReplicaRegistrationModel( job_id=new_job.id, - gateway_replica_id=gateway_compute_1.id, + gateway_replica_id=gateway_replica_1.id, is_registered=True, register_attempt=0, ) @@ -1484,7 +1484,7 @@ async def test_service_rolling_deployment_keeps_old_replica_until_new_replica_re session.add( ServiceReplicaRegistrationModel( job_id=new_job.id, - gateway_replica_id=gateway_compute_2.id, + gateway_replica_id=gateway_replica_2.id, is_registered=False, register_attempt=2, ) @@ -1517,10 +1517,10 @@ async def test_service_rolling_deployment_scales_down_old_replica_once_new_repli backend_id=backend.id, status=GatewayStatus.RUNNING, ) - gateway_compute_1 = await create_gateway_compute( + gateway_replica_1 = await create_gateway_replica( session=session, gateway_id=gateway.id, replica_num=0 ) - gateway_compute_2 = await create_gateway_compute( + gateway_replica_2 = await create_gateway_replica( session=session, gateway_id=gateway.id, replica_num=1 ) run_spec = get_run_spec( @@ -1563,11 +1563,11 @@ async def test_service_rolling_deployment_scales_down_old_replica_once_new_repli ready=True, replica_num=1, ) - for gateway_compute in (gateway_compute_1, gateway_compute_2): + for gateway_replica in (gateway_replica_1, gateway_replica_2): session.add( ServiceReplicaRegistrationModel( job_id=old_job.id, - gateway_replica_id=gateway_compute.id, + gateway_replica_id=gateway_replica.id, is_registered=True, register_attempt=0, ) @@ -1575,7 +1575,7 @@ async def test_service_rolling_deployment_scales_down_old_replica_once_new_repli session.add( ServiceReplicaRegistrationModel( job_id=new_job.id, - gateway_replica_id=gateway_compute.id, + gateway_replica_id=gateway_replica.id, is_registered=True, register_attempt=0, ) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pending.py b/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pending.py index fb3842a357..af41ba8393 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pending.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pending.py @@ -21,7 +21,7 @@ from dstack._internal.server.testing.common import ( create_backend, create_gateway, - create_gateway_compute, + create_gateway_replica, create_job, create_project, create_repo, @@ -381,7 +381,7 @@ async def test_terminates_run_on_gateway_registration_failure( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - gateway_compute = await create_gateway_compute(session=session, gateway_id=gateway.id) + gateway_replica = await create_gateway_replica(session=session, gateway_id=gateway.id) run_spec = get_run_spec( run_name="test-run", repo_id=repo.name, @@ -409,7 +409,7 @@ async def test_terminates_run_on_gateway_registration_failure( session.add( ServiceRegistrationModel( run_id=run.id, - gateway_replica_id=gateway_compute.id, + gateway_replica_id=gateway_replica.id, is_registered=False, register_attempt=3, register_status_message="Connection refused", diff --git a/src/tests/_internal/server/routers/test_gateways.py b/src/tests/_internal/server/routers/test_gateways.py index 7753ccc6a4..0ff8e2ca5a 100644 --- a/src/tests/_internal/server/routers/test_gateways.py +++ b/src/tests/_internal/server/routers/test_gateways.py @@ -15,7 +15,7 @@ create_backend, create_export, create_gateway, - create_gateway_compute, + create_gateway_replica, create_project, create_user, get_auth_headers, @@ -32,14 +32,14 @@ async def test_returns_40x_if_not_authenticated(self, client: AsyncClient): @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) @pytest.mark.parametrize("populate_configuration", [True, False]) async def test_list( self, test_db, session: AsyncSession, client: AsyncClient, - legacy_compute: bool, + legacy_replica: bool, populate_configuration: bool, ): user = await create_user(session, global_role=GlobalRole.USER) @@ -54,15 +54,15 @@ async def test_list( backend_id=backend.id, populate_configuration=populate_configuration, ) - if legacy_compute: - gateway_compute = await create_gateway_compute( + if legacy_replica: + gateway_replica = await create_gateway_replica( session=session, backend_id=backend.id, populate_configuration=populate_configuration, ) - gateway.gateway_compute_id = gateway_compute.id # pre-0.20.25 relationship style + gateway.gateway_replica_id = gateway_replica.id # pre-0.20.25 relationship style else: - gateway_compute = await create_gateway_compute( + gateway_replica = await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -85,7 +85,7 @@ async def test_list( "status_message": None, "replicas": [ { - "hostname": gateway_compute.ip_address, + "hostname": gateway_replica.ip_address, "replica_num": 0, "backend": backend.type.value, "region": "us", @@ -118,14 +118,14 @@ async def test_list( @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) - @pytest.mark.parametrize("legacy_compute", [False, True]) + @pytest.mark.parametrize("legacy_replica", [False, True]) @pytest.mark.parametrize("populate_configuration", [True, False]) async def test_get( self, test_db, session: AsyncSession, client: AsyncClient, - legacy_compute: bool, + legacy_replica: bool, populate_configuration: bool, ): user = await create_user(session, global_role=GlobalRole.USER) @@ -140,15 +140,15 @@ async def test_get( backend_id=backend.id, populate_configuration=populate_configuration, ) - if legacy_compute: - gateway_compute = await create_gateway_compute( + if legacy_replica: + gateway_replica = await create_gateway_replica( session=session, backend_id=backend.id, populate_configuration=populate_configuration, ) - gateway.gateway_compute_id = gateway_compute.id # pre-0.20.25 relationship style + gateway.gateway_replica_id = gateway_replica.id # pre-0.20.25 relationship style else: - gateway_compute = await create_gateway_compute( + gateway_replica = await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -171,7 +171,7 @@ async def test_get( "status_message": None, "replicas": [ { - "hostname": gateway_compute.ip_address, + "hostname": gateway_replica.ip_address, "replica_num": 0, "backend": backend.type.value, "region": "us", @@ -218,7 +218,7 @@ async def test_list_legacy_client_populates_compat_fields( project_id=project.id, backend_id=backend.id, ) - gateway_compute = await create_gateway_compute( + gateway_replica = await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -230,9 +230,9 @@ async def test_list_legacy_client_populates_compat_fields( assert response.status_code == 200 assert len(response.json()) == 1 gw = response.json()[0] - assert gw["ip_address"] == gateway_compute.ip_address + assert gw["ip_address"] == gateway_replica.ip_address assert gw["instance_id"] == "" - assert gw["hostname"] == gateway_compute.ip_address + assert gw["hostname"] == gateway_replica.ip_address @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) @@ -247,7 +247,7 @@ async def test_list_non_member_public_project( project_id=project.id, backend_id=backend.id, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -273,7 +273,7 @@ async def test_get_non_member_public_project( project_id=project.id, backend_id=backend.id, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -326,7 +326,7 @@ async def test_list_returns_imported_gateway_with_include_imported( backend_id=backend.id, name="exported-gateway", ) - await create_gateway_compute(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) await create_export( session=session, exporter_project=exporter_project, @@ -369,7 +369,7 @@ async def test_list_not_returns_imported_gateway_without_include_imported( backend_id=backend.id, name="exported-gateway", ) - await create_gateway_compute(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) await create_export( session=session, exporter_project=exporter_project, @@ -410,7 +410,7 @@ async def test_get_returns_imported_gateway( backend_id=backend.id, name="exported-gateway", ) - await create_gateway_compute(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) await create_export( session=session, exporter_project=exporter_project, @@ -457,7 +457,7 @@ async def test_get_returns_403_on_foreign_gateway_if_not_imported( backend_id=backend.id, name="exported-gateway", ) - await create_gateway_compute(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) await create_export( session=session, exporter_project=exporter_project, @@ -818,7 +818,7 @@ async def test_set_default_gateway( name="first_gateway", populate_configuration=populate_configuration, ) - gateway_compute = await create_gateway_compute( + gateway_replica = await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -847,7 +847,7 @@ async def test_set_default_gateway( "status_message": None, "replicas": [ { - "hostname": gateway_compute.ip_address, + "hostname": gateway_replica.ip_address, "replica_num": 0, "backend": backend.type.value, "region": "us", @@ -887,7 +887,7 @@ async def test_set_default_gateway( name="second_gateway", populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=second_gateway.id, @@ -994,7 +994,7 @@ async def test_set_imported_gateway_as_default( backend_id=backend.id, name="exported-gateway", ) - await create_gateway_compute(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) await create_export( session=session, exporter_project=exporter_project, @@ -1038,7 +1038,7 @@ async def test_cannot_set_non_imported_foreign_gateway_as_default( backend_id=backend.id, name="exported-gateway", ) - await create_gateway_compute(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) await create_export( session=session, exporter_project=exporter_project, @@ -1095,7 +1095,7 @@ async def test_marks_gateways_to_be_deleted( name="gateway-aws", populate_configuration=populate_configuration, ) - gateway_compute_aws = await create_gateway_compute( + gateway_replica_aws = await create_gateway_replica( session=session, backend_id=backend_aws.id, gateway_id=gateway_aws.id, @@ -1108,7 +1108,7 @@ async def test_marks_gateways_to_be_deleted( name="gateway-gcp", populate_configuration=populate_configuration, ) - gateway_compute_gcp = await create_gateway_compute( + gateway_replica_gcp = await create_gateway_replica( session=session, backend_id=backend_gcp.id, gateway_id=gateway_gcp.id, @@ -1123,14 +1123,14 @@ async def test_marks_gateways_to_be_deleted( await session.refresh(gateway_aws) await session.refresh(gateway_gcp) - await session.refresh(gateway_compute_aws) - await session.refresh(gateway_compute_gcp) + await session.refresh(gateway_replica_aws) + await session.refresh(gateway_replica_gcp) assert gateway_aws.to_be_deleted is True assert gateway_gcp.to_be_deleted is True - assert gateway_compute_aws.active is True - assert gateway_compute_aws.deleted is False - assert gateway_compute_gcp.active is True - assert gateway_compute_gcp.deleted is False + assert gateway_replica_aws.active is True + assert gateway_replica_aws.deleted is False + assert gateway_replica_gcp.active is True + assert gateway_replica_gcp.deleted is False response = await client.post( f"/api/project/{project.name}/gateways/list", @@ -1221,7 +1221,7 @@ async def test_set_wildcard_domain( wildcard_domain="old.example", populate_configuration=populate_configuration, ) - gateway_compute = await create_gateway_compute( + gateway_replica = await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -1243,7 +1243,7 @@ async def test_set_wildcard_domain( "default": False, "replicas": [ { - "hostname": gateway_compute.ip_address, + "hostname": gateway_replica.ip_address, "replica_num": 0, "backend": backend.type.value, "region": "us", @@ -1466,7 +1466,7 @@ async def test_get_plan_with_existing_gateway_no_changes( region="us-east-1", populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -1512,7 +1512,7 @@ async def test_get_plan_with_domain_change_is_update( wildcard_domain="old.example.com", populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -1560,7 +1560,7 @@ async def test_get_plan_rejects_failed_gateway( status=GatewayStatus.FAILED, populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -1604,7 +1604,7 @@ async def test_get_plan_with_region_change_is_create( region="us-east-1", populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -1679,7 +1679,7 @@ async def test_get_plan_rejects_to_be_deleted_gateway( wildcard_domain="old.example.com", populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -1963,7 +1963,7 @@ async def test_updates_in_place( wildcard_domain="old.example.com", populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -2025,7 +2025,7 @@ async def test_updates_in_place_with_force_apply( wildcard_domain="old.example.com", populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -2080,7 +2080,7 @@ async def test_force_apply_no_changes_succeeds( region="us-east-1", populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -2128,7 +2128,7 @@ async def test_rejects_update( region="us-east-1", populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -2182,7 +2182,7 @@ async def test_rejects_update_with_force_apply( region="us-east-1", populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -2229,7 +2229,7 @@ async def test_returns_error_on_missing_current_resource( region="us-east-1", populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -2277,7 +2277,7 @@ async def test_returns_error_on_current_resource_mismatch( region="us-east-1", populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -2334,7 +2334,7 @@ async def test_rejects_apply_on_to_be_deleted_gateway( wildcard_domain="old.example.com", populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -2399,7 +2399,7 @@ async def test_rejects_apply_on_failed_gateway( status=GatewayStatus.FAILED, populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -2456,7 +2456,7 @@ async def test_sets_default_in_place( region="us-east-1", populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=first_gateway.id ) second_gateway = await create_gateway( @@ -2467,7 +2467,7 @@ async def test_sets_default_in_place( region="us-east-1", populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=second_gateway.id ) response = await client.post( @@ -2539,7 +2539,7 @@ async def test_unsets_default_in_place( region="us-east-1", populate_configuration=populate_configuration, ) - await create_gateway_compute(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) response = await client.post( f"/api/project/{project.name}/gateways/set_default", json={"name": gateway.name}, @@ -2604,7 +2604,7 @@ async def test_omitted_default_leaves_current_status_unchanged( name="my-gateway", region="us-east-1", ) - await create_gateway_compute(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) if initial_default: response = await client.post( f"/api/project/{project.name}/gateways/set_default", @@ -2669,7 +2669,7 @@ async def test_legacy_client_default_false_is_treated_as_omitted( name="my-gateway", region="us-east-1", ) - await create_gateway_compute(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) if initial_default: response = await client.post( f"/api/project/{project.name}/gateways/set_default", diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py index 60d958f6ad..6010d7c7d5 100644 --- a/src/tests/_internal/server/routers/test_runs.py +++ b/src/tests/_internal/server/routers/test_runs.py @@ -59,7 +59,7 @@ create_export, create_fleet, create_gateway, - create_gateway_compute, + create_gateway_replica, create_instance, create_job, create_project, @@ -3974,7 +3974,7 @@ async def test_submit_to_correct_proxy( name=gateway_name, wildcard_domain=f"{gateway_name}.example", ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -4038,7 +4038,7 @@ async def test_submit_to_gateway_by_name( wildcard_domain="my-gateway.example", populate_configuration=populate_configuration, ) - await create_gateway_compute( + await create_gateway_replica( session=session, backend_id=backend.id, gateway_id=gateway.id, @@ -4149,7 +4149,7 @@ async def test_submit_to_foreign_gateway_only_if_imported( name="exported-gateway", wildcard_domain="exported-gateway.example", ) - await create_gateway_compute(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) importer_user = await create_user( session=session, global_role=GlobalRole.USER, name="importer_user" @@ -4243,7 +4243,7 @@ async def test_not_submits_to_default_gateway_if_not_imported( backend_id=backend.id, status=GatewayStatus.RUNNING, ) - await create_gateway_compute(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) service_project = await create_project(session=session, owner=user, name="service-project") # The project's default_gateway_id may point to the gateway (e.g., if the gateway was @@ -4303,7 +4303,7 @@ async def test_interpolates_project_name_in_imported_gateway_domain( name="exported-gateway", wildcard_domain="${{ run.project_name }}.example.com", ) - await create_gateway_compute(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) importer_user = await create_user( session=session, global_role=GlobalRole.USER, name="importer_user" @@ -4367,7 +4367,7 @@ async def test_returns_error_if_imported_gateway_domain_has_unknown_variable( name="exported-gateway", wildcard_domain="${{ run.unknown_variable }}.example.com", ) - await create_gateway_compute(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) importer_user = await create_user( session=session, global_role=GlobalRole.USER, name="importer_user" @@ -4438,7 +4438,7 @@ async def test_return_error_if_default_gateway_forbids_new_services( wildcard_domain="example.com", forbid_new_services=True, ) - await create_gateway_compute(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) project.default_gateway_id = gateway.id await session.commit() @@ -4482,7 +4482,7 @@ async def test_return_error_if_explicitly_specified_gateway_forbids_new_services wildcard_domain="example.com", forbid_new_services=True, ) - await create_gateway_compute(session=session, backend_id=backend.id, gateway_id=gateway.id) + await create_gateway_replica(session=session, backend_id=backend.id, gateway_id=gateway.id) response = await client.post( f"/api/project/{project.name}/runs/apply", diff --git a/src/tests/_internal/server/services/gateways/test_gateways.py b/src/tests/_internal/server/services/gateways/test_gateways.py index aaf8fe6d52..0e4fbe9386 100644 --- a/src/tests/_internal/server/services/gateways/test_gateways.py +++ b/src/tests/_internal/server/services/gateways/test_gateways.py @@ -5,12 +5,12 @@ from dstack._internal.proxy.gateway.schemas.stats import Stat from dstack._internal.server.services.gateways import ( _merge_per_window_stats, - get_gateway_compute_models, + get_gateway_replica_models, ) from dstack._internal.server.testing.common import ( create_backend, create_gateway, - create_gateway_compute, + create_gateway_replica, create_project, ) @@ -48,41 +48,41 @@ def test_zero_requests_across_all_replicas_returns_zero_time(self): @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) -class TestGetGatewayComputeModels: - async def test_new_style_returns_gateway_computes(self, test_db, session: AsyncSession): +class TestGetGatewayReplicaModels: + async def test_new_style_returns_gateway_replicas(self, test_db, session: AsyncSession): 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 ) - compute = await create_gateway_compute( + replica = await create_gateway_replica( session=session, gateway_id=gateway.id, backend_id=backend.id ) - await session.refresh(gateway, ["gateway_computes", "gateway_compute"]) - result = get_gateway_compute_models(gateway) + await session.refresh(gateway, ["gateway_replicas", "gateway_replica"]) + result = get_gateway_replica_models(gateway) assert len(result) == 1 - assert result[0].id == compute.id + assert result[0].id == replica.id - async def test_old_style_returns_single_compute(self, test_db, session: AsyncSession): + async def test_old_style_returns_single_replica(self, test_db, session: AsyncSession): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) - compute = await create_gateway_compute(session=session, backend_id=backend.id) + replica = await create_gateway_replica(session=session, backend_id=backend.id) gateway = await create_gateway( session=session, project_id=project.id, backend_id=backend.id ) - gateway.gateway_compute_id = compute.id + gateway.gateway_replica_id = replica.id await session.commit() - await session.refresh(gateway, ["gateway_computes", "gateway_compute"]) - result = get_gateway_compute_models(gateway) + await session.refresh(gateway, ["gateway_replicas", "gateway_replica"]) + result = get_gateway_replica_models(gateway) assert len(result) == 1 - assert result[0].id == compute.id + assert result[0].id == replica.id - async def test_no_computes_returns_empty(self, test_db, session: AsyncSession): + async def test_no_replicas_returns_empty(self, test_db, session: AsyncSession): 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 ) - await session.refresh(gateway, ["gateway_computes", "gateway_compute"]) - result = get_gateway_compute_models(gateway) + await session.refresh(gateway, ["gateway_replicas", "gateway_replica"]) + result = get_gateway_replica_models(gateway) assert result == []