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 29370a07b..b1b96b364 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/gateway_replicas.py @@ -157,6 +157,7 @@ async def fetch(self, limit: int) -> list[GatewayReplicaPipelineItem]: and_( GatewayComputeModel.status == GatewayReplicaStatus.RUNNING, or_( + GatewayComputeModel.scale_in == True, GatewayModel.to_be_deleted == True, GatewayModel.status == GatewayStatus.FAILED, # Gateway was hard-deleted (unexpected, fetch to log an error) @@ -317,24 +318,32 @@ def _get_loaded_gateway_model(replica_model: GatewayComputeModel) -> Optional[Ga return gateway_model -def _mark_terminating_if_gateway_terminating( +def _mark_terminating_if_needed( gateway_model: GatewayModel, replica_model: GatewayComputeModel ) -> Optional[_GatewayReplicaUpdateMap]: if gateway_model.to_be_deleted or gateway_model.status == GatewayStatus.FAILED: - if replica_model.status == GatewayReplicaStatus.SUBMITTED: - new_status = GatewayReplicaStatus.TERMINATED - deleted = True - else: - new_status = GatewayReplicaStatus.TERMINATING - deleted = False - logger.info( - "%s replica %d: marked %s, gateway is being deleted or failed", - fmt(gateway_model), - replica_model.replica_num, - new_status.value, - ) - return _GatewayReplicaUpdateMap(status=new_status, active=False, deleted=deleted) - return None + status_message = None + elif replica_model.scale_in: + status_message = "Scaled in" + else: + return + if replica_model.status == GatewayReplicaStatus.SUBMITTED: + new_status = GatewayReplicaStatus.TERMINATED + deleted = True + else: + new_status = GatewayReplicaStatus.TERMINATING + deleted = False + logger.info( + "%s replica %d: marked %s (%s)", + fmt(gateway_model), + replica_model.replica_num, + new_status.value, + status_message or "-", + ) + update_map = _GatewayReplicaUpdateMap(status=new_status, active=False, deleted=deleted) + if status_message: + update_map["status_message"] = status_message + return update_map async def _commit_update( @@ -369,6 +378,7 @@ async def _process_submitted_item(item: GatewayReplicaPipelineItem): GatewayComputeModel.backend_id, GatewayComputeModel.configuration, GatewayComputeModel.ssh_public_key, + GatewayComputeModel.scale_in, ], gateway_fields=_GATEWAY_FIELDS_MIN + [ @@ -385,7 +395,7 @@ async def _process_submitted_item(item: GatewayReplicaPipelineItem): if gateway_model is None: await _commit_update(item, replica_model, update_map={}) return - if update_map := _mark_terminating_if_gateway_terminating(gateway_model, replica_model): + if update_map := _mark_terminating_if_needed(gateway_model, replica_model): await _commit_update(item, replica_model, update_map=update_map) return update_map = await _provision_gateway_replica(gateway_model, replica_model) @@ -477,6 +487,7 @@ async def _process_provisioning_item(item: GatewayReplicaPipelineItem): + [ GatewayComputeModel.ip_address, GatewayComputeModel.ssh_private_key, + GatewayComputeModel.scale_in, ], gateway_fields=_GATEWAY_FIELDS_MIN, ) @@ -486,7 +497,7 @@ async def _process_provisioning_item(item: GatewayReplicaPipelineItem): if gateway_model is None: await _commit_update(item, replica_model, update_map={}) return - if update_map := _mark_terminating_if_gateway_terminating(gateway_model, replica_model): + if update_map := _mark_terminating_if_needed(gateway_model, replica_model): await _commit_update(item, replica_model, update_map=update_map) return error = await _connect_and_configure_gateway_replica(gateway_model, replica_model) @@ -553,7 +564,10 @@ async def _connect_and_configure_gateway_replica( async def _process_running_item(item: GatewayReplicaPipelineItem): replica_model = await _load_gateway_replica( item, - replica_fields=_REPLICA_FIELDS_MIN, + replica_fields=_REPLICA_FIELDS_MIN + + [ + GatewayComputeModel.scale_in, + ], gateway_fields=_GATEWAY_FIELDS_MIN, ) if replica_model is None: @@ -562,7 +576,7 @@ async def _process_running_item(item: GatewayReplicaPipelineItem): if gateway_model is None: await _commit_update(item, replica_model, update_map={}) return - if update_map := _mark_terminating_if_gateway_terminating(gateway_model, replica_model): + if update_map := _mark_terminating_if_needed(gateway_model, replica_model): await _commit_update(item, replica_model, update_map=update_map) return logger.warning( diff --git a/src/dstack/_internal/server/background/pipeline_tasks/gateways.py b/src/dstack/_internal/server/background/pipeline_tasks/gateways.py index 8e71568cc..7d143e430 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/gateways.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/gateways.py @@ -1,13 +1,14 @@ import asyncio +import itertools import uuid from dataclasses import dataclass, field from datetime import timedelta from typing import Sequence -from sqlalchemy import delete, or_, select, update +from sqlalchemy import ColumnElement, and_, delete, func, or_, select, update +from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import joinedload, load_only, selectinload -from dstack._internal.core.errors import BackendNotAvailable from dstack._internal.core.models.gateways import ( GATEWAY_REPLICAS_DEFAULT, GatewayReplicaStatus, @@ -33,7 +34,6 @@ GatewayModel, ProjectModel, ) -from dstack._internal.server.services import backends as backends_services from dstack._internal.server.services import events from dstack._internal.server.services import gateways as gateways_services from dstack._internal.server.services.gateways import ( @@ -44,7 +44,7 @@ from dstack._internal.server.services.logging import fmt from dstack._internal.server.services.pipelines import PipelineHinterProtocol from dstack._internal.server.utils import sentry_utils -from dstack._internal.utils.common import get_current_datetime +from dstack._internal.utils.common import get_current_datetime, get_lowest_unused_nums from dstack._internal.utils.logging import get_logger logger = get_logger(__name__) @@ -139,6 +139,18 @@ async def fetch(self, limit: int) -> list[GatewayPipelineItem]: async with gateway_lock: async with get_session_ctx() as session: now = get_current_datetime() + active_replica_count_subquery = ( + select(func.count(GatewayComputeModel.id)) + .where( + or_( + GatewayComputeModel.gateway_id == GatewayModel.id, + GatewayComputeModel.id == GatewayModel.gateway_compute_id, + ), + *_get_active_replica_filters(), + ) + .correlate(GatewayModel) + .scalar_subquery() + ) res = await session.execute( select(GatewayModel) .where( @@ -146,6 +158,12 @@ async def fetch(self, limit: int) -> list[GatewayPipelineItem]: GatewayModel.status.in_( [GatewayStatus.SUBMITTED, GatewayStatus.PROVISIONING] ), + and_( + GatewayModel.status == GatewayStatus.RUNNING, + GatewayModel.desired_replica_count.is_not(None), + GatewayModel.desired_replica_count + != active_replica_count_subquery, + ), GatewayModel.to_be_deleted == True, ), or_( @@ -219,6 +237,14 @@ async def process(self, item: GatewayPipelineItem): await _process_submitted_item(item) elif item.status == GatewayStatus.PROVISIONING: await _process_provisioning_item(item) + elif item.status == GatewayStatus.RUNNING: + await _process_running_item(item) + + +@dataclass +class _ReplicaScalingResult: + new_gateway_compute_models: list[GatewayComputeModel] = field(default_factory=list) + scale_in_replica_ids: list[uuid.UUID] = field(default_factory=list) async def _process_submitted_item(item: GatewayPipelineItem): @@ -229,7 +255,7 @@ async def _process_submitted_item(item: GatewayPipelineItem): GatewayModel.id == item.id, GatewayModel.lock_token == item.lock_token, ) - .options(joinedload(GatewayModel.project).joinedload(ProjectModel.backends)) + .options(joinedload(GatewayModel.project).load_only(ProjectModel.name)) .options(joinedload(GatewayModel.backend).load_only(BackendModel.type)) ) gateway_model = res.unique().scalar_one_or_none() @@ -243,8 +269,7 @@ async def _process_submitted_item(item: GatewayPipelineItem): set_processed_update_map_fields(update_map) set_unlock_update_map_fields(update_map) async with get_session_ctx() as session: - for gateway_compute_model in result.gateway_compute_models: - session.add(gateway_compute_model) + await _apply_replica_scaling(session, result.scale_result) now = get_current_datetime() resolve_now_placeholders(update_map, now=now) res = await session.execute( @@ -277,48 +302,16 @@ class _GatewayUpdateMap(ItemUpdateMap, total=False): @dataclass class _SubmittedResult: update_map: _GatewayUpdateMap = field(default_factory=_GatewayUpdateMap) - gateway_compute_models: list[GatewayComputeModel] = field(default_factory=list) + scale_result: _ReplicaScalingResult = field(default_factory=_ReplicaScalingResult) async def _process_submitted_gateway(gateway_model: GatewayModel) -> _SubmittedResult: - configuration = gateways_services.get_gateway_configuration(gateway_model) - try: - ( - backend_model, - _, - ) = await backends_services.get_project_backend_with_model_by_type_or_error( - project=gateway_model.project, backend_type=configuration.backend - ) - except BackendNotAvailable: - return _SubmittedResult( - update_map={ - "status": GatewayStatus.FAILED, - "status_message": "Backend not available", - } - ) # NOTE: On a later stage of #3959, the SUBMITTED status may also be responsible for # setting up the load balancer (e.g., AWS ALB) before replicas are created. - replicas = ( - configuration.replicas if configuration.replicas is not None else GATEWAY_REPLICAS_DEFAULT - ) - gateway_compute_models = [] - for replica_num in range(replicas): - gateway_compute_model = gateways_services.create_gateway_compute_model( - project_name=gateway_model.project.name, - configuration=configuration, - replica_num=replica_num, - gateway_id=gateway_model.id, - backend_id=backend_model.id, - ) - gateway_compute_models.append(gateway_compute_model) - logger.info( - "%s: created %d replica record(s) in submitted state", - fmt(gateway_model), - len(gateway_compute_models), - ) + scale_result = _reconcile_gateway_replica_count(gateway_model, gateway_replicas=[]) return _SubmittedResult( update_map={"status": GatewayStatus.PROVISIONING}, - gateway_compute_models=gateway_compute_models, + scale_result=scale_result, ) @@ -330,10 +323,16 @@ async def _process_provisioning_item(item: GatewayPipelineItem): GatewayModel.id == item.id, GatewayModel.lock_token == item.lock_token, ) + .options(joinedload(GatewayModel.project).load_only(ProjectModel.name)) + .options(joinedload(GatewayModel.backend).load_only(BackendModel.type)) .options(joinedload(GatewayModel.gateway_compute)) .options( selectinload(GatewayModel.gateway_computes).load_only( - GatewayComputeModel.id, GatewayComputeModel.status + GatewayComputeModel.id, + GatewayComputeModel.status, + GatewayComputeModel.replica_num, + GatewayComputeModel.created_at, + GatewayComputeModel.scale_in, ) ) ) @@ -348,6 +347,7 @@ async def _process_provisioning_item(item: GatewayPipelineItem): set_unlock_update_map_fields(gateway_update_map) async with get_session_ctx() as session: + await _apply_replica_scaling(session, result.scale_result) now = get_current_datetime() resolve_now_placeholders(gateway_update_map, now=now) res = await session.execute( @@ -375,6 +375,7 @@ async def _process_provisioning_item(item: GatewayPipelineItem): @dataclass class _ProvisioningResult: gateway_update_map: _GatewayUpdateMap = field(default_factory=_GatewayUpdateMap) + scale_result: _ReplicaScalingResult = field(default_factory=_ReplicaScalingResult) def _process_provisioning_gateway(gateway_model: GatewayModel) -> _ProvisioningResult: @@ -382,7 +383,12 @@ def _process_provisioning_gateway(gateway_model: GatewayModel) -> _ProvisioningR # Provisioning gateways must have compute. assert len(gateway_computes) > 0 - statuses = {gc.status for gc in gateway_computes} + scale_result = _reconcile_gateway_replica_count(gateway_model, gateway_computes) + statuses = { + gc.status + for gc in gateway_computes + if not gc.scale_in and gc.id not in scale_result.scale_in_replica_ids + } if statuses & {GatewayReplicaStatus.TERMINATING, GatewayReplicaStatus.TERMINATED}: return _ProvisioningResult( @@ -390,15 +396,68 @@ def _process_provisioning_gateway(gateway_model: GatewayModel) -> _ProvisioningR "status": GatewayStatus.FAILED, "status_message": "Failed to provision gateway replica", }, + scale_result=_ReplicaScalingResult(), # do not scale, gateway failed ) - if statuses == {GatewayReplicaStatus.RUNNING}: + if statuses == {GatewayReplicaStatus.RUNNING} and not scale_result.new_gateway_compute_models: return _ProvisioningResult( gateway_update_map={"status": GatewayStatus.RUNNING}, + scale_result=scale_result, ) # Replicas are still being provisioned - return _ProvisioningResult() + return _ProvisioningResult(scale_result=scale_result) + + +async def _process_running_item(item: GatewayPipelineItem): + async with get_session_ctx() as session: + res = await session.execute( + select(GatewayModel) + .where( + GatewayModel.id == item.id, + GatewayModel.lock_token == item.lock_token, + ) + .options(joinedload(GatewayModel.project).load_only(ProjectModel.name)) + .options(joinedload(GatewayModel.backend).load_only(BackendModel.type)) + .options(joinedload(GatewayModel.gateway_compute)) + .options( + selectinload(GatewayModel.gateway_computes).load_only( + GatewayComputeModel.id, + GatewayComputeModel.status, + GatewayComputeModel.replica_num, + GatewayComputeModel.created_at, + GatewayComputeModel.scale_in, + ) + ) + ) + gateway_model = res.unique().scalar_one_or_none() + if gateway_model is None: + 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) + + update_map = _GatewayUpdateMap() + set_processed_update_map_fields(update_map) + set_unlock_update_map_fields(update_map) + async with get_session_ctx() as session: + await _apply_replica_scaling(session, scale_result) + now = get_current_datetime() + resolve_now_placeholders(update_map, now=now) + res = await session.execute( + update(GatewayModel) + .where( + GatewayModel.id == gateway_model.id, + GatewayModel.lock_token == gateway_model.lock_token, + ) + .values(**update_map) + .returning(GatewayModel.id) + ) + updated_ids = list(res.scalars().all()) + if len(updated_ids) == 0: + log_lock_token_changed_after_processing(logger, item) + return async def _process_to_be_deleted_item(item: GatewayPipelineItem): @@ -476,3 +535,102 @@ def _process_to_be_deleted_gateway(gateway_model: GatewayModel) -> _ProcessToBeD gateway_computes = get_gateway_compute_models(gateway_model) all_terminated = all(gc.status == GatewayReplicaStatus.TERMINATED for gc in gateway_computes) return _ProcessToBeDeletedResult(delete_gateway=all_terminated) + + +REPLICA_SCALE_IN_PRIORITY: dict[GatewayReplicaStatus, int] = { + GatewayReplicaStatus.SUBMITTED: 0, + GatewayReplicaStatus.PROVISIONING: 1, + GatewayReplicaStatus.RUNNING: 2, +} + + +def _is_replica_active(replica: GatewayComputeModel) -> bool: + # should match _get_active_replica_filters + return not replica.scale_in and replica.status not in ( + GatewayReplicaStatus.TERMINATING, + GatewayReplicaStatus.TERMINATED, + ) + + +def _get_active_replica_filters() -> list[ColumnElement[bool]]: + # should match _is_replica_active + return [ + GatewayComputeModel.scale_in == False, + GatewayComputeModel.status.not_in( + [GatewayReplicaStatus.TERMINATING, GatewayReplicaStatus.TERMINATED] + ), + ] + + +def _reconcile_gateway_replica_count( + gateway_model: GatewayModel, + gateway_replicas: list[GatewayComputeModel], +) -> _ReplicaScalingResult: + desired_replica_count = gateway_model.desired_replica_count + if desired_replica_count is None: # pre-0.20.28 gateway + if gateway_model.status != GatewayStatus.SUBMITTED: + return _ReplicaScalingResult() + desired_replica_count = gateways_services.get_gateway_configuration(gateway_model).replicas + if desired_replica_count is None: + desired_replica_count = GATEWAY_REPLICAS_DEFAULT + + active_replicas = [r for r in gateway_replicas if _is_replica_active(r)] + diff = desired_replica_count - len(active_replicas) + if diff == 0: + return _ReplicaScalingResult() + + if diff > 0: + configuration = gateways_services.get_gateway_configuration(gateway_model) + used_nums = { + 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( + project_name=gateway_model.project.name, + configuration=configuration, + replica_num=replica_num, + gateway_id=gateway_model.id, + backend_id=gateway_model.backend_id, + ) + for replica_num in new_nums + ] + logger.info( + "%s: scaling out, adding %d replica(s)", + fmt(gateway_model), + diff, + ) + return _ReplicaScalingResult(new_gateway_compute_models=new_gateway_compute_models) + + replicas_redundant = -diff + active_replicas.sort( + key=lambda r: ( + # Stop replicas with lower priority statuses first. + REPLICA_SCALE_IN_PRIORITY.get(r.status, max(REPLICA_SCALE_IN_PRIORITY.values()) + 1), + # Stop older replicas first. This allows to migrate off old instances + # (e.g., to update the gateway to a new OS image). + r.created_at, + ) + ) + scale_in_replica_ids = [r.id for r in active_replicas[:replicas_redundant]] + logger.info( + "%s: scaling in, marking %d replica(s) for scale-in", + fmt(gateway_model), + len(scale_in_replica_ids), + ) + return _ReplicaScalingResult(scale_in_replica_ids=scale_in_replica_ids) + + +async def _apply_replica_scaling( + session: AsyncSession, scale_result: _ReplicaScalingResult +) -> None: + for gateway_compute_model in scale_result.new_gateway_compute_models: + session.add(gateway_compute_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)) + .values(scale_in=True) + ) diff --git a/src/dstack/_internal/server/migrations/versions/2026/07_03_1544_f8336ffdac34_gateway_scaling.py b/src/dstack/_internal/server/migrations/versions/2026/07_03_1544_f8336ffdac34_gateway_scaling.py new file mode 100644 index 000000000..2eb1130f1 --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/07_03_1544_f8336ffdac34_gateway_scaling.py @@ -0,0 +1,40 @@ +"""Gateway scaling + +Revision ID: f8336ffdac34 +Revises: e9c5e7e26c78 +Create Date: 2026-07-03 15:44:09.684229+00:00 + +""" + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "f8336ffdac34" +down_revision = "e9c5e7e26c78" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("gateway_computes", schema=None) as batch_op: + batch_op.add_column( + sa.Column("scale_in", sa.Boolean(), server_default=sa.false(), nullable=False) + ) + + with op.batch_alter_table("gateways", schema=None) as batch_op: + batch_op.add_column(sa.Column("desired_replica_count", sa.Integer(), nullable=True)) + + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("gateways", schema=None) as batch_op: + batch_op.drop_column("desired_replica_count") + + with op.batch_alter_table("gateway_computes", schema=None) as batch_op: + batch_op.drop_column("scale_in") + + # ### end Alembic commands ### diff --git a/src/dstack/_internal/server/models.py b/src/dstack/_internal/server/models.py index fc7301026..7fa0afde1 100644 --- a/src/dstack/_internal/server/models.py +++ b/src/dstack/_internal/server/models.py @@ -613,6 +613,8 @@ class GatewayModel(PipelineModelMixin, BaseModel): created_at: Mapped[datetime] = mapped_column(NaiveDateTime, default=get_current_datetime) status: Mapped[GatewayStatus] = mapped_column(EnumAsString(GatewayStatus, 100)) status_message: Mapped[Optional[str]] = mapped_column(Text) + desired_replica_count: Mapped[Optional[int]] = mapped_column(Integer) + """Only `None` for pre-0.20.28 gateways that were never scaled""" last_processed_at: Mapped[datetime] = mapped_column(NaiveDateTime) to_be_deleted: Mapped[bool] = mapped_column(Boolean, server_default=false()) forbid_new_services: Mapped[bool] = mapped_column(Boolean, server_default=false()) @@ -642,7 +644,8 @@ class GatewayModel(PipelineModelMixin, BaseModel): foreign_keys="GatewayComputeModel.gateway_id", ) """ - Relationship with gateway computes for 0.20.25+ gateways. + 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. """ @@ -667,6 +670,8 @@ class GatewayComputeModel(PipelineModelMixin, BaseModel): last_processed_at: Mapped[datetime] = mapped_column(NaiveDateTime) status: Mapped[GatewayReplicaStatus] = mapped_column(EnumAsString(GatewayReplicaStatus, 100)) status_message: Mapped[Optional[str]] = mapped_column(Text) + scale_in: Mapped[bool] = mapped_column(Boolean, server_default=false()) + """Indicates that replica termination is requested due to the gateway scaling in""" replica_num: Mapped[int] = mapped_column(Integer, server_default="0") instance_id: Mapped[Optional[str]] = mapped_column(String(100)) ip_address: Mapped[Optional[str]] = mapped_column(String(100)) diff --git a/src/dstack/_internal/server/services/fleets.py b/src/dstack/_internal/server/services/fleets.py index 547f91d52..ea0adc87c 100644 --- a/src/dstack/_internal/server/services/fleets.py +++ b/src/dstack/_internal/server/services/fleets.py @@ -94,6 +94,7 @@ EntityName, EntityNameOrID, get_current_datetime, + get_lowest_unused_nums, ) from dstack._internal.utils.logging import get_logger from dstack._internal.utils.ssh import pkey_from_str @@ -964,16 +965,7 @@ def get_fleet_requirements(fleet_spec: FleetSpec) -> Requirements: def get_next_instance_num(taken_instance_nums: set[int]) -> int: - if not taken_instance_nums: - return 0 - min_instance_num = min(taken_instance_nums) - if min_instance_num > 0: - return 0 - instance_num = min_instance_num + 1 - while True: - if instance_num not in taken_instance_nums: - return instance_num - instance_num += 1 + return next(get_lowest_unused_nums(used_nums=taken_instance_nums)) def get_fleet_master_instance_provisioning_data( diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py index 1601800c7..c7f0cb1e6 100644 --- a/src/dstack/_internal/server/services/gateways/__init__.py +++ b/src/dstack/_internal/server/services/gateways/__init__.py @@ -1,5 +1,6 @@ import asyncio import datetime +import itertools import uuid from collections.abc import AsyncGenerator from contextlib import asynccontextmanager @@ -77,6 +78,7 @@ from dstack._internal.server.services.pipelines import PipelineHinterProtocol from dstack._internal.server.services.plugins import apply_plugin_policies from dstack._internal.server.utils.common import gather_map_async +from dstack._internal.settings import FeatureFlags from dstack._internal.utils.common import ( get_current_datetime, get_or_error, @@ -87,6 +89,8 @@ logger = get_logger(__name__) _CONF_UPDATABLE_FIELDS = frozenset({"domain"}) +if FeatureFlags.GATEWAY_SCALING: + _CONF_UPDATABLE_FIELDS |= {"replicas"} def switch_gateway_status( @@ -267,6 +271,11 @@ async def create_gateway( wildcard_domain=configuration.domain, configuration=configuration.json(), status=GatewayStatus.SUBMITTED, + desired_replica_count=( + configuration.replicas + if configuration.replicas is not None + else GATEWAY_REPLICAS_DEFAULT + ), created_at=now, last_processed_at=now, ) @@ -837,11 +846,10 @@ async def configure_gateway( def get_gateway_compute_models(gateway_model: GatewayModel) -> List[GatewayComputeModel]: - if gateway_model.gateway_computes: # 0.20.25+ gateway - return list(gateway_model.gateway_computes) + computes = list(gateway_model.gateway_computes) if gateway_model.gateway_compute is not None: # pre-0.20.25 gateway - return [gateway_model.gateway_compute] - return [] + computes.append(gateway_model.gateway_compute) + return computes def get_gateway_configuration(gateway_model: GatewayModel) -> GatewayConfiguration: @@ -889,10 +897,17 @@ def gateway_model_to_gateway( configuration = get_gateway_configuration(gateway_model) configuration.default = is_default - compute_models = sorted(get_gateway_compute_models(gateway_model), key=lambda c: c.replica_num) + all_compute_models = sorted( + get_gateway_compute_models(gateway_model), key=lambda c: c.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_compute_models.append(max(compute_models_for_num, key=lambda c: c.created_at)) gateway_hostname = None replicas = [] - for compute in compute_models: + for compute in relevant_compute_models: replicas.append( GatewayReplica( hostname=compute.ip_address, @@ -1035,6 +1050,12 @@ async def apply_plan( ) gateway_model.wildcard_domain = new_configuration.domain + if new_configuration.replicas != current_configuration.replicas: + gateway_model.desired_replica_count = ( + new_configuration.replicas + if new_configuration.replicas is not None + else GATEWAY_REPLICAS_DEFAULT + ) gateway_model.configuration = new_configuration.json() events.emit( session, diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py index 3b24ef96e..4371175e5 100644 --- a/src/dstack/_internal/server/testing/common.py +++ b/src/dstack/_internal/server/testing/common.py @@ -45,6 +45,7 @@ SSHParams, ) from dstack._internal.core.models.gateways import ( + GATEWAY_REPLICAS_DEFAULT, GatewayComputeConfiguration, GatewayConfiguration, GatewayReplicaStatus, @@ -644,6 +645,7 @@ async def create_gateway( region: str = "us", wildcard_domain: Optional[str] = None, status: Optional[GatewayStatus] = GatewayStatus.SUBMITTED, + replicas: Optional[int] = None, last_processed_at: datetime = datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), forbid_new_services: bool = False, populate_configuration: bool = True, @@ -663,6 +665,7 @@ async def create_gateway( backend=backend.type, region=region, domain=wildcard_domain, + replicas=replicas, ).json() gateway = GatewayModel( project_id=project_id, @@ -672,6 +675,7 @@ async def create_gateway( wildcard_domain=wildcard_domain, configuration=configuration, status=status, + desired_replica_count=replicas if replicas is not None else GATEWAY_REPLICAS_DEFAULT, last_processed_at=last_processed_at, forbid_new_services=forbid_new_services, ) diff --git a/src/dstack/_internal/settings.py b/src/dstack/_internal/settings.py index afb8bb5f0..3dd7acf65 100644 --- a/src/dstack/_internal/settings.py +++ b/src/dstack/_internal/settings.py @@ -50,3 +50,8 @@ class FeatureFlags: """If DSTACK_FF_CLI_PRINT_JOB_CONNECTION_INFO enabled, `dstack apply` command prints server-provided IDE URL(s) and SSH command(s) before job logs (for dev-environments only). """ + GATEWAY_SCALING = os.getenv("DSTACK_FF_GATEWAY_SCALING") is not None + """Allows in-place update of gateway `replicas`. + Has limited applicability at the current stage — without a state sync mechanism, newly started + gateway replicas are not aware of existing services. + """ diff --git a/src/dstack/_internal/utils/common.py b/src/dstack/_internal/utils/common.py index 2ee1f337a..217e03f7c 100644 --- a/src/dstack/_internal/utils/common.py +++ b/src/dstack/_internal/utils/common.py @@ -3,7 +3,7 @@ import itertools import re import time -from collections.abc import Callable +from collections.abc import Callable, Iterator from dataclasses import dataclass from datetime import datetime, timedelta, timezone from functools import partial @@ -309,6 +309,13 @@ def batched(seq: Iterable[T], n: int) -> Iterable[List[T]]: return iter(lambda: list(itertools.islice(it, n)), []) +def get_lowest_unused_nums(used_nums: set[int]) -> Iterator[int]: + for num in itertools.count(): + if num not in used_nums: + yield num + assert False, "unreachable" + + StrT = TypeVar("StrT", str, bytes) 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 950af1fac..638c5ff8d 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 @@ -80,6 +80,7 @@ async def test_fetch_selects_eligible_replicas_and_sets_lock_fields( project_id=project.id, backend_id=backend.id, status=GatewayStatus.PROVISIONING, + replicas=7, ) now = get_current_datetime() stale = now - timedelta(minutes=1) @@ -279,6 +280,36 @@ async def test_fetch_excludes_running_replica_with_healthy_gateway( assert len(items) == 0 + async def test_fetch_includes_running_replica_marked_for_scale_in( + self, + test_db, + session: AsyncSession, + fetcher: GatewayReplicaFetcher, + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + stale = get_current_datetime() - timedelta(minutes=1) + compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + last_processed_at=stale, + ) + compute.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].status == GatewayReplicaStatus.RUNNING + @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) @@ -510,6 +541,43 @@ async def test_submitted_unexpected_error_marks_terminated( assert compute.active is False assert compute.deleted is True + async def test_submitted_to_terminated_when_scaled_in( + self, test_db, session: AsyncSession, worker: GatewayReplicaWorker + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + backend_id=backend.id, + ip_address=None, + instance_id=None, + region=None, + status=GatewayReplicaStatus.SUBMITTED, + configuration=get_gateway_compute_configuration().json(), + ) + compute.scale_in = True + _lock_compute(compute) + 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)) + 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" + @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) @@ -568,6 +636,34 @@ async def test_running_to_terminating( assert compute.status == GatewayReplicaStatus.TERMINATING assert compute.active is False + async def test_running_to_terminating_when_scaled_in( + self, test_db, session: AsyncSession, worker: GatewayReplicaWorker + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + compute = await create_gateway_compute( + session=session, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + active=True, + ) + compute.scale_in = True + _lock_compute(compute) + await session.commit() + + await worker.process(_compute_to_pipeline_item(compute)) + + await session.refresh(compute) + assert compute.status == GatewayReplicaStatus.TERMINATING + assert compute.active is False + assert compute.status_message == "Scaled in" + @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) 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 f44d89b1f..f378edad7 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_gateways.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_gateways.py @@ -1,17 +1,14 @@ import asyncio import uuid from datetime import datetime, timedelta, timezone -from unittest.mock import Mock, patch +from unittest.mock import Mock import pytest from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import selectinload -from dstack._internal.core.errors import BackendNotAvailable -from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.gateways import ( - GatewayConfiguration, GatewayReplicaStatus, GatewayStatus, ) @@ -21,7 +18,8 @@ GatewayPipelineItem, GatewayWorker, ) -from dstack._internal.server.models import GatewayModel +from dstack._internal.server.models import GatewayComputeModel, GatewayModel +from dstack._internal.server.services.gateways import get_gateway_compute_models from dstack._internal.server.testing.common import ( create_backend, create_gateway, @@ -63,6 +61,19 @@ def _gateway_to_pipeline_item(gateway_model: GatewayModel) -> GatewayPipelineIte ) +async def _fetch_all_gateway_computes( + session: AsyncSession, gateway_id: uuid.UUID +) -> list[GatewayComputeModel]: + res = await session.execute( + select(GatewayModel) + .where(GatewayModel.id == gateway_id) + .options(selectinload(GatewayModel.gateway_computes)) + .options(selectinload(GatewayModel.gateway_compute)) + ) + gateway = res.unique().scalar_one() + return get_gateway_compute_models(gateway) + + @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) class TestGatewayFetcher: @@ -100,6 +111,15 @@ async def test_fetch_selects_eligible_gateways_and_sets_lock_fields( ) to_be_deleted.to_be_deleted = True + running_with_missing_replica = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + name="running-with-missing-replica", + status=GatewayStatus.RUNNING, + last_processed_at=stale - timedelta(seconds=4), + ) + just_created = await create_gateway( session=session, project_id=project.id, @@ -111,12 +131,12 @@ async def test_fetch_selects_eligible_gateways_and_sets_lock_fields( just_created.created_at = now just_created.last_processed_at = now - ineligible_status = await create_gateway( + failed = await create_gateway( session=session, project_id=project.id, backend_id=backend.id, - name="ineligible-status", - status=GatewayStatus.RUNNING, + name="failed", + status=GatewayStatus.FAILED, last_processed_at=stale, ) recent = await create_gateway( @@ -149,12 +169,14 @@ async def test_fetch_selects_eligible_gateways_and_sets_lock_fields( submitted.id, provisioning.id, to_be_deleted.id, + running_with_missing_replica.id, just_created.id, } assert {(item.id, item.status, item.to_be_deleted) for item in items} == { (submitted.id, GatewayStatus.SUBMITTED, False), (provisioning.id, GatewayStatus.PROVISIONING, False), (to_be_deleted.id, GatewayStatus.RUNNING, True), + (running_with_missing_replica.id, GatewayStatus.RUNNING, False), (just_created.id, GatewayStatus.SUBMITTED, False), } @@ -162,20 +184,27 @@ async def test_fetch_selects_eligible_gateways_and_sets_lock_fields( submitted, provisioning, to_be_deleted, + running_with_missing_replica, just_created, - ineligible_status, + failed, recent, locked, ]: await session.refresh(gateway) - fetched_gateways = [submitted, provisioning, to_be_deleted, just_created] + fetched_gateways = [ + submitted, + provisioning, + to_be_deleted, + running_with_missing_replica, + just_created, + ] assert all(gateway.lock_owner == GatewayPipeline.__name__ for gateway in fetched_gateways) assert all(gateway.lock_expires_at is not None for gateway in fetched_gateways) assert all(gateway.lock_token is not None for gateway in fetched_gateways) assert len({gateway.lock_token for gateway in fetched_gateways}) == 1 - assert ineligible_status.lock_owner is None + assert failed.lock_owner is None assert recent.lock_owner is None assert locked.lock_owner == "OtherPipeline" @@ -223,56 +252,84 @@ 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.asyncio -@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) -class TestGatewayWorkerSubmitted: - async def test_submitted_to_provisioning( - self, test_db, session: AsyncSession, worker: GatewayWorker + @pytest.mark.parametrize("legacy_compute", [False, True]) + async def test_fetch_excludes_running_gateway_when_replica_count_matches( + self, test_db, session: AsyncSession, fetcher: GatewayFetcher, legacy_compute: bool ): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) + stale = get_current_datetime() - timedelta(minutes=1) gateway = await create_gateway( session=session, project_id=project.id, backend_id=backend.id, - status=GatewayStatus.SUBMITTED, + status=GatewayStatus.RUNNING, + replicas=1, + last_processed_at=stale, ) - config = GatewayConfiguration( - name=gateway.name, - backend=BackendType.AWS, - region=gateway.region, + if legacy_compute: + compute = await create_gateway_compute( + session=session, + backend_id=backend.id, + status=GatewayReplicaStatus.RUNNING, + ) + gateway.gateway_compute_id = compute.id + else: + await create_gateway_compute( + session=session, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + ) + await session.commit() + + items = await fetcher.fetch(limit=10) + + assert items == [] + + @pytest.mark.parametrize("legacy_compute", [False, True]) + async def test_fetch_includes_running_gateway_when_replica_count_not_matches( + self, test_db, session: AsyncSession, fetcher: GatewayFetcher, legacy_compute: bool + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + stale = get_current_datetime() - timedelta(minutes=1) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, replicas=2, + last_processed_at=stale, ) - gateway.configuration = config.json() - gateway.lock_token = uuid.uuid4() - gateway.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) + if legacy_compute: + compute = await create_gateway_compute( + session=session, + backend_id=backend.id, + status=GatewayReplicaStatus.RUNNING, + ) + gateway.gateway_compute_id = compute.id + else: + await create_gateway_compute( + session=session, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + ) await session.commit() - with patch( - "dstack._internal.server.services.backends.get_project_backend_with_model_by_type_or_error" - ) as m: - m.return_value = (backend, Mock()) - await worker.process(_gateway_to_pipeline_item(gateway)) + items = await fetcher.fetch(limit=10) + assert {item.id for item in items} == {gateway.id} - await session.refresh(gateway) - res = await session.execute( - select(GatewayModel) - .where(GatewayModel.id == gateway.id) - .options(selectinload(GatewayModel.gateway_computes)) - ) - gateway = res.unique().scalar_one() - assert gateway.status == GatewayStatus.PROVISIONING - computes = sorted(gateway.gateway_computes, key=lambda c: c.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) - async def test_marks_gateway_as_failed_if_backend_not_available( - self, test_db, session: AsyncSession, worker: GatewayWorker +@pytest.mark.asyncio +@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) +class TestGatewayWorkerSubmitted: + @pytest.mark.parametrize("populate_configuration", [True, False]) + async def test_submitted_to_provisioning( + self, + test_db, + session: AsyncSession, + worker: GatewayWorker, + populate_configuration: bool, ): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) @@ -281,26 +338,26 @@ async def test_marks_gateway_as_failed_if_backend_not_available( project_id=project.id, backend_id=backend.id, status=GatewayStatus.SUBMITTED, + replicas=2, + populate_configuration=populate_configuration, ) gateway.lock_token = uuid.uuid4() gateway.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) await session.commit() - with patch( - "dstack._internal.server.services.backends.get_project_backend_with_model_by_type_or_error" - ) as m: - m.side_effect = BackendNotAvailable() - await worker.process(_gateway_to_pipeline_item(gateway)) + await worker.process(_gateway_to_pipeline_item(gateway)) await session.refresh(gateway) - assert gateway.status == GatewayStatus.FAILED - assert gateway.status_message == "Backend not available" - events = await list_events(session) - assert len(events) == 1 - assert ( - events[0].message - == "Gateway status changed SUBMITTED -> FAILED (Backend not available)" + assert gateway.status == GatewayStatus.PROVISIONING + computes = sorted( + await _fetch_all_gateway_computes(session, gateway.id), key=lambda c: c.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) @pytest.mark.asyncio @@ -362,6 +419,7 @@ async def test_provisioning_to_running_with_multiple_replicas( project_id=project.id, backend_id=backend.id, status=GatewayStatus.PROVISIONING, + replicas=2, ) await create_gateway_compute( session, @@ -399,6 +457,7 @@ async def test_still_provisioning_if_not_all_replicas_running( project_id=project.id, backend_id=backend.id, status=GatewayStatus.PROVISIONING, + replicas=2, ) await create_gateway_compute( session, @@ -511,6 +570,356 @@ 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("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, + populate_configuration: bool, + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.PROVISIONING, + replicas=2, + populate_configuration=populate_configuration, + ) + if legacy_compute: + gateway_compute = await create_gateway_compute( + 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 + else: + await create_gateway_compute( + session, + gateway_id=gateway.id, + ip_address="1.1.1.1", + status=GatewayReplicaStatus.RUNNING, + replica_num=0, + populate_configuration=populate_configuration, + ) + gateway.lock_token = uuid.uuid4() + gateway.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) + await session.commit() + + await worker.process(_gateway_to_pipeline_item(gateway)) + + 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 + ) + assert [c.replica_num for c in computes] == [0, 1] + assert computes[1].status == GatewayReplicaStatus.SUBMITTED + events = await list_events(session) + assert len(events) == 0 + + async def test_provisioning_to_running_when_scale_in_removes_surplus_replicas( + self, test_db, session: AsyncSession, worker: GatewayWorker + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.PROVISIONING, + replicas=1, + ) + older = await create_gateway_compute( + session, + gateway_id=gateway.id, + ip_address="1.1.1.1", + status=GatewayReplicaStatus.PROVISIONING, + replica_num=0, + ) + older.created_at = datetime(2025, 1, 1) + newer = await create_gateway_compute( + session, + gateway_id=gateway.id, + ip_address="2.2.2.2", + status=GatewayReplicaStatus.RUNNING, + replica_num=1, + ) + newer.created_at = datetime(2025, 1, 2) + gateway.lock_token = uuid.uuid4() + gateway.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) + await session.commit() + + await worker.process(_gateway_to_pipeline_item(gateway)) + + await session.refresh(gateway) + assert gateway.status == GatewayStatus.RUNNING + await session.refresh(older) + await session.refresh(newer) + assert older.scale_in is True + assert newer.scale_in is False + events = await list_events(session) + assert len(events) == 1 + assert events[0].message == "Gateway status changed PROVISIONING -> RUNNING" + + async def test_ignores_previously_scaled_in_replica_when_determining_status( + self, test_db, session: AsyncSession, worker: GatewayWorker + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.PROVISIONING, + replicas=1, + ) + await create_gateway_compute( + session, + gateway_id=gateway.id, + ip_address="1.1.1.1", + status=GatewayReplicaStatus.RUNNING, + replica_num=0, + ) + scaled_in_compute = await create_gateway_compute( + session, + gateway_id=gateway.id, + ip_address="2.2.2.2", + status=GatewayReplicaStatus.TERMINATING, + active=False, + replica_num=1, + ) + scaled_in_compute.scale_in = True + gateway.lock_token = uuid.uuid4() + gateway.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) + await session.commit() + + await worker.process(_gateway_to_pipeline_item(gateway)) + + await session.refresh(gateway) + assert gateway.status == GatewayStatus.RUNNING + events = await list_events(session) + assert len(events) == 1 + assert events[0].message == "Gateway status changed PROVISIONING -> RUNNING" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) +class TestGatewayWorkerRunning: + @pytest.mark.parametrize("legacy_compute", [False, True]) + async def test_no_scaling_when_replica_count_matches( + self, test_db, session: AsyncSession, worker: GatewayWorker, legacy_compute: bool + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + replicas=1, + ) + if legacy_compute: + compute = await create_gateway_compute( + session=session, + backend_id=backend.id, + status=GatewayReplicaStatus.RUNNING, + ) + gateway.gateway_compute_id = compute.id + else: + await create_gateway_compute( + session, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + replica_num=0, + ) + gateway.lock_token = uuid.uuid4() + gateway.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) + original_last_processed_at = gateway.last_processed_at + await session.commit() + + await worker.process(_gateway_to_pipeline_item(gateway)) + + 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 + + @pytest.mark.parametrize("legacy_compute", [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, + populate_configuration: bool, + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + replicas=3, + populate_configuration=populate_configuration, + ) + if legacy_compute: + gateway_compute = await create_gateway_compute( + session=session, + backend_id=backend.id, + status=GatewayReplicaStatus.RUNNING, + populate_configuration=populate_configuration, + ) + gateway.gateway_compute_id = gateway_compute.id + else: + await create_gateway_compute( + session, + gateway_id=gateway.id, + status=GatewayReplicaStatus.RUNNING, + replica_num=0, + populate_configuration=populate_configuration, + ) + gateway.lock_token = uuid.uuid4() + gateway.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) + await session.commit() + + await worker.process(_gateway_to_pipeline_item(gateway)) + + 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 + ) + assert [c.replica_num for c in computes] == [0, 1, 2] + assert [c.status for c in computes] == [ + GatewayReplicaStatus.RUNNING, + GatewayReplicaStatus.SUBMITTED, + GatewayReplicaStatus.SUBMITTED, + ] + + async def test_scales_in_oldest_replicas_when_desired_replica_count_decreased( + self, test_db, session: AsyncSession, worker: GatewayWorker + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + replicas=1, + ) + compute0 = await create_gateway_compute( + session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 + ) + compute0.created_at = datetime(2025, 1, 1) + compute1 = await create_gateway_compute( + session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=1 + ) + compute1.created_at = datetime(2025, 1, 2) + compute2 = await create_gateway_compute( + session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=2 + ) + compute2.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() + + await worker.process(_gateway_to_pipeline_item(gateway)) + + 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 + + async def test_scale_in_prefers_less_advanced_replicas_over_older_running_ones( + self, test_db, session: AsyncSession, worker: GatewayWorker + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + replicas=1, + ) + running = await create_gateway_compute( + session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 + ) + running.created_at = datetime(2025, 1, 1) + submitted = await create_gateway_compute( + session, + gateway_id=gateway.id, + status=GatewayReplicaStatus.SUBMITTED, + replica_num=1, + ip_address=None, + instance_id=None, + region=None, + configuration=get_gateway_compute_configuration().json(), + ) + submitted.created_at = datetime(2025, 1, 2) + gateway.lock_token = uuid.uuid4() + gateway.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) + await session.commit() + + await worker.process(_gateway_to_pipeline_item(gateway)) + + await session.refresh(running) + await session.refresh(submitted) + assert running.scale_in is False + assert submitted.scale_in is True + + @pytest.mark.parametrize("legacy_compute", [False, True]) + async def test_no_scaling_for_legacy_gateway_without_desired_replica_count( + self, test_db, session: AsyncSession, worker: GatewayWorker, legacy_compute: bool + ): + project = await create_project(session=session) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + ) + if legacy_compute: + compute = await create_gateway_compute( + session=session, + backend_id=backend.id, + status=GatewayReplicaStatus.RUNNING, + ) + gateway.gateway_compute_id = compute.id + else: + await create_gateway_compute( + session, gateway_id=gateway.id, status=GatewayReplicaStatus.RUNNING, replica_num=0 + ) + gateway.desired_replica_count = None + gateway.lock_token = uuid.uuid4() + gateway.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) + await session.commit() + + await worker.process(_gateway_to_pipeline_item(gateway)) + + 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 + @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) @@ -650,6 +1059,7 @@ async def test_deletes_gateway_with_multiple_replicas_all_terminated( project_id=project.id, backend_id=backend.id, status=GatewayStatus.RUNNING, + replicas=2, ) await create_gateway_compute( session=session, diff --git a/src/tests/_internal/utils/test_common.py b/src/tests/_internal/utils/test_common.py index 70d12c8f3..6f26151bd 100644 --- a/src/tests/_internal/utils/test_common.py +++ b/src/tests/_internal/utils/test_common.py @@ -1,3 +1,4 @@ +import itertools from datetime import datetime, timedelta, timezone from typing import Any, Iterable @@ -8,6 +9,7 @@ batched, concat_url_path, format_duration_multiunit, + get_lowest_unused_nums, has_duplicates, local_time, make_proxy_url, @@ -314,6 +316,28 @@ def test_gpu_zero_count_range_with_vendor(self): ) +class TestGetLowestUnusedNums: + @pytest.mark.parametrize( + ("used_nums", "count", "expected"), + [ + (set(), 1, [0]), + (set(), 3, [0, 1, 2]), + ({0, 1, 2}, 3, [3, 4, 5]), + ({0, 2, 4}, 3, [1, 3, 5]), + ({1, 2, 3}, 2, [0, 4]), + ], + ) + def test_get_lowest_unused_nums( + self, used_nums: set[int], count: int, expected: list[int] + ) -> None: + assert list(itertools.islice(get_lowest_unused_nums(used_nums), count)) == expected + + def test_does_not_mutate_used_nums(self) -> None: + used_nums = {0, 1} + next(get_lowest_unused_nums(used_nums)) + assert used_nums == {0, 1} + + class TestSizeofFmt: @pytest.mark.parametrize( ("num", "suffix", "expected"),