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,