From f3293bc19c9929f43fe3c240b599ffe4b972200d Mon Sep 17 00:00:00 2001 From: Jvst Me Date: Fri, 21 Aug 2026 00:16:40 +0200 Subject: [PATCH] Add `GatewayReplicaModel.name` This is a prerequisite to adding gateway replica events, since our event framework requires event targets to be named. The naming pattern is `-`, consistent with other dstack entities. The new field is backfilled for existing gateways. Restoring the gateway name for deleted gateways is not feasible, so replicas of deleted gateways are backfilled with `unknown-`. No user-facing changes in this commit. --- .../background/pipeline_tasks/gateways.py | 6 +- ...be9f32ec66_add_gatewayreplicamodel_name.py | 72 +++++++++++++++++++ src/dstack/_internal/server/models.py | 1 + .../server/services/gateways/__init__.py | 17 +++-- src/dstack/_internal/server/testing/common.py | 4 +- 5 files changed, 85 insertions(+), 15 deletions(-) create mode 100644 src/dstack/_internal/server/migrations/versions/2026/08_20_1940_dbbe9f32ec66_add_gatewayreplicamodel_name.py diff --git a/src/dstack/_internal/server/background/pipeline_tasks/gateways.py b/src/dstack/_internal/server/background/pipeline_tasks/gateways.py index 0a879fd09..18d950618 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/gateways.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/gateways.py @@ -764,18 +764,14 @@ def _reconcile_gateway_replica_count( gateway_model ): return _ReplicaScalingResult(needs_more_replicas=True) - 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_replica_models = [ gateways_services.create_gateway_replica_model( - project_name=gateway_model.project.name, - configuration=configuration, + gateway_model=gateway_model, replica_num=replica_num, - gateway_id=gateway_model.id, - backend_id=gateway_model.backend_id, ) for replica_num in new_nums ] diff --git a/src/dstack/_internal/server/migrations/versions/2026/08_20_1940_dbbe9f32ec66_add_gatewayreplicamodel_name.py b/src/dstack/_internal/server/migrations/versions/2026/08_20_1940_dbbe9f32ec66_add_gatewayreplicamodel_name.py new file mode 100644 index 000000000..bb5729205 --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/08_20_1940_dbbe9f32ec66_add_gatewayreplicamodel_name.py @@ -0,0 +1,72 @@ +"""Add GatewayReplicaModel.name + +Revision ID: dbbe9f32ec66 +Revises: eee3e79f29e9 +Create Date: 2026-08-20 19:40:13.516516+00:00 + +""" + +import uuid + +import sqlalchemy as sa +from alembic import op +from sqlalchemy_utils import UUIDType + +# revision identifiers, used by Alembic. +revision = "dbbe9f32ec66" +down_revision = "eee3e79f29e9" +branch_labels = None +depends_on = None + +# Partial table descriptions - only columns needed for the data migration below. +gateway_computes_table = sa.Table( + "gateway_computes", + sa.MetaData(), + sa.Column("id", UUIDType(binary=False), primary_key=True, default=uuid.uuid4), + sa.Column("name", sa.String(100)), + sa.Column("gateway_id", UUIDType(binary=False), nullable=True), + sa.Column("replica_num", sa.Integer()), +) +gateways_table = sa.Table( + "gateways", + sa.MetaData(), + sa.Column("id", UUIDType(binary=False), primary_key=True, default=uuid.uuid4), + sa.Column("name", sa.String(100)), + sa.Column("gateway_compute_id", UUIDType(binary=False), nullable=True), +) + + +def upgrade() -> None: + with op.batch_alter_table("gateway_computes", schema=None) as batch_op: + batch_op.add_column(sa.Column("name", sa.String(length=100), nullable=True)) + + bind = op.get_bind() + gateway_name = sa.func.coalesce( + sa.select(gateways_table.c.name) + .where( + sa.or_( + gateways_table.c.id == gateway_computes_table.c.gateway_id, + gateways_table.c.gateway_compute_id == gateway_computes_table.c.id, + ) + ) + .scalar_subquery(), + "unknown", # fallback for deleted gateways + ) + bind.execute( + gateway_computes_table.update().values( + name=gateway_name.concat("-").concat( + sa.cast(gateway_computes_table.c.replica_num, sa.String) + ) + ) + ) + + with op.batch_alter_table("gateway_computes", schema=None) as batch_op: + batch_op.alter_column("name", existing_type=sa.String(length=100), nullable=False) + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("gateway_computes", schema=None) as batch_op: + batch_op.drop_column("name") + + # ### end Alembic commands ### diff --git a/src/dstack/_internal/server/models.py b/src/dstack/_internal/server/models.py index 24822d7c8..53d971075 100644 --- a/src/dstack/_internal/server/models.py +++ b/src/dstack/_internal/server/models.py @@ -692,6 +692,7 @@ class GatewayReplicaModel(PipelineModelMixin, BaseModel): id: Mapped[uuid.UUID] = mapped_column( UUIDType(binary=False), primary_key=True, default=uuid.uuid4 ) + name: Mapped[str] = mapped_column(String(100)) created_at: Mapped[datetime] = mapped_column(NaiveDateTime, default=get_current_datetime) last_processed_at: Mapped[datetime] = mapped_column(NaiveDateTime) skip_min_processing_interval: Mapped[bool] = mapped_column( diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py index acaa890cb..4b6c9fcb4 100644 --- a/src/dstack/_internal/server/services/gateways/__init__.py +++ b/src/dstack/_internal/server/services/gateways/__init__.py @@ -183,21 +183,19 @@ async def get_gateway_by_name( def create_gateway_replica_model( - project_name: str, - configuration: GatewayConfiguration, + gateway_model: GatewayModel, replica_num: int, - gateway_id: uuid.UUID, - backend_id: uuid.UUID, ) -> GatewayReplicaModel: - assert configuration.name is not None + configuration = get_gateway_configuration(gateway_model) + replica_name = f"{gateway_model.name}-{replica_num}" private_bytes, public_bytes = crypto.generate_rsa_key_pair_bytes() gateway_ssh_private_key = private_bytes.decode() gateway_ssh_public_key = public_bytes.decode() replica_configuration = GatewayReplicaConfiguration( - project_name=project_name, - instance_name=f"{configuration.name}-{replica_num}", + project_name=gateway_model.project.name, + instance_name=replica_name, backend=configuration.backend, region=configuration.region, instance_type=configuration.instance_type, @@ -209,8 +207,9 @@ def create_gateway_replica_model( now = get_current_datetime() return GatewayReplicaModel( - gateway_id=gateway_id, - backend_id=backend_id, + name=replica_name, + gateway_id=gateway_model.id, + backend_id=gateway_model.backend_id, replica_num=replica_num, configuration=replica_configuration.model_dump_json(), ssh_private_key=gateway_ssh_private_key, diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py index 0f11da4a2..bf4e4f027 100644 --- a/src/dstack/_internal/server/testing/common.py +++ b/src/dstack/_internal/server/testing/common.py @@ -733,11 +733,12 @@ async def create_gateway_replica( True - 0.18.2+ gateways, False - legacy pre-0.18.2 gateways. Prefer testing against both in major test cases. """ + name = f"test-gateway-{replica_num}" if configuration is None and populate_configuration: assert region is not None configuration = GatewayReplicaConfiguration( project_name="test-project", - instance_name=instance_id or "test-instance", + instance_name=name, backend=backend.type, region=region, public_ip=True, @@ -745,6 +746,7 @@ async def create_gateway_replica( certificate=None, ).model_dump_json() gateway_replica = GatewayReplicaModel( + name=name, gateway_id=gateway_id, backend_id=backend.id, ip_address=ip_address,