Skip to content

Commit 765fb31

Browse files
authored
Add GatewayReplicaModel.name (#4184)
This is a prerequisite to adding gateway replica events, since our event framework requires event targets to be named. The naming pattern is `<gateway-name>-<replica-num>`, 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-<replica-num>`. No user-facing changes in this commit.
1 parent d667f12 commit 765fb31

5 files changed

Lines changed: 85 additions & 15 deletions

File tree

src/dstack/_internal/server/background/pipeline_tasks/gateways.py

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -764,18 +764,14 @@ def _reconcile_gateway_replica_count(
764764
gateway_model
765765
):
766766
return _ReplicaScalingResult(needs_more_replicas=True)
767-
configuration = gateways_services.get_gateway_configuration(gateway_model)
768767
used_nums = {
769768
r.replica_num for r in gateway_replicas if r.status != GatewayReplicaStatus.TERMINATED
770769
}
771770
new_nums = itertools.islice(get_lowest_unused_nums(used_nums), diff)
772771
new_gateway_replica_models = [
773772
gateways_services.create_gateway_replica_model(
774-
project_name=gateway_model.project.name,
775-
configuration=configuration,
773+
gateway_model=gateway_model,
776774
replica_num=replica_num,
777-
gateway_id=gateway_model.id,
778-
backend_id=gateway_model.backend_id,
779775
)
780776
for replica_num in new_nums
781777
]
Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
"""Add GatewayReplicaModel.name
2+
3+
Revision ID: dbbe9f32ec66
4+
Revises: eee3e79f29e9
5+
Create Date: 2026-08-20 19:40:13.516516+00:00
6+
7+
"""
8+
9+
import uuid
10+
11+
import sqlalchemy as sa
12+
from alembic import op
13+
from sqlalchemy_utils import UUIDType
14+
15+
# revision identifiers, used by Alembic.
16+
revision = "dbbe9f32ec66"
17+
down_revision = "eee3e79f29e9"
18+
branch_labels = None
19+
depends_on = None
20+
21+
# Partial table descriptions - only columns needed for the data migration below.
22+
gateway_computes_table = sa.Table(
23+
"gateway_computes",
24+
sa.MetaData(),
25+
sa.Column("id", UUIDType(binary=False), primary_key=True, default=uuid.uuid4),
26+
sa.Column("name", sa.String(100)),
27+
sa.Column("gateway_id", UUIDType(binary=False), nullable=True),
28+
sa.Column("replica_num", sa.Integer()),
29+
)
30+
gateways_table = sa.Table(
31+
"gateways",
32+
sa.MetaData(),
33+
sa.Column("id", UUIDType(binary=False), primary_key=True, default=uuid.uuid4),
34+
sa.Column("name", sa.String(100)),
35+
sa.Column("gateway_compute_id", UUIDType(binary=False), nullable=True),
36+
)
37+
38+
39+
def upgrade() -> None:
40+
with op.batch_alter_table("gateway_computes", schema=None) as batch_op:
41+
batch_op.add_column(sa.Column("name", sa.String(length=100), nullable=True))
42+
43+
bind = op.get_bind()
44+
gateway_name = sa.func.coalesce(
45+
sa.select(gateways_table.c.name)
46+
.where(
47+
sa.or_(
48+
gateways_table.c.id == gateway_computes_table.c.gateway_id,
49+
gateways_table.c.gateway_compute_id == gateway_computes_table.c.id,
50+
)
51+
)
52+
.scalar_subquery(),
53+
"unknown", # fallback for deleted gateways
54+
)
55+
bind.execute(
56+
gateway_computes_table.update().values(
57+
name=gateway_name.concat("-").concat(
58+
sa.cast(gateway_computes_table.c.replica_num, sa.String)
59+
)
60+
)
61+
)
62+
63+
with op.batch_alter_table("gateway_computes", schema=None) as batch_op:
64+
batch_op.alter_column("name", existing_type=sa.String(length=100), nullable=False)
65+
66+
67+
def downgrade() -> None:
68+
# ### commands auto generated by Alembic - please adjust! ###
69+
with op.batch_alter_table("gateway_computes", schema=None) as batch_op:
70+
batch_op.drop_column("name")
71+
72+
# ### end Alembic commands ###

src/dstack/_internal/server/models.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -692,6 +692,7 @@ class GatewayReplicaModel(PipelineModelMixin, BaseModel):
692692
id: Mapped[uuid.UUID] = mapped_column(
693693
UUIDType(binary=False), primary_key=True, default=uuid.uuid4
694694
)
695+
name: Mapped[str] = mapped_column(String(100))
695696
created_at: Mapped[datetime] = mapped_column(NaiveDateTime, default=get_current_datetime)
696697
last_processed_at: Mapped[datetime] = mapped_column(NaiveDateTime)
697698
skip_min_processing_interval: Mapped[bool] = mapped_column(

src/dstack/_internal/server/services/gateways/__init__.py

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -183,21 +183,19 @@ async def get_gateway_by_name(
183183

184184

185185
def create_gateway_replica_model(
186-
project_name: str,
187-
configuration: GatewayConfiguration,
186+
gateway_model: GatewayModel,
188187
replica_num: int,
189-
gateway_id: uuid.UUID,
190-
backend_id: uuid.UUID,
191188
) -> GatewayReplicaModel:
192-
assert configuration.name is not None
189+
configuration = get_gateway_configuration(gateway_model)
190+
replica_name = f"{gateway_model.name}-{replica_num}"
193191

194192
private_bytes, public_bytes = crypto.generate_rsa_key_pair_bytes()
195193
gateway_ssh_private_key = private_bytes.decode()
196194
gateway_ssh_public_key = public_bytes.decode()
197195

198196
replica_configuration = GatewayReplicaConfiguration(
199-
project_name=project_name,
200-
instance_name=f"{configuration.name}-{replica_num}",
197+
project_name=gateway_model.project.name,
198+
instance_name=replica_name,
201199
backend=configuration.backend,
202200
region=configuration.region,
203201
instance_type=configuration.instance_type,
@@ -209,8 +207,9 @@ def create_gateway_replica_model(
209207

210208
now = get_current_datetime()
211209
return GatewayReplicaModel(
212-
gateway_id=gateway_id,
213-
backend_id=backend_id,
210+
name=replica_name,
211+
gateway_id=gateway_model.id,
212+
backend_id=gateway_model.backend_id,
214213
replica_num=replica_num,
215214
configuration=replica_configuration.model_dump_json(),
216215
ssh_private_key=gateway_ssh_private_key,

src/dstack/_internal/server/testing/common.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -733,18 +733,20 @@ async def create_gateway_replica(
733733
True - 0.18.2+ gateways, False - legacy pre-0.18.2 gateways. Prefer
734734
testing against both in major test cases.
735735
"""
736+
name = f"test-gateway-{replica_num}"
736737
if configuration is None and populate_configuration:
737738
assert region is not None
738739
configuration = GatewayReplicaConfiguration(
739740
project_name="test-project",
740-
instance_name=instance_id or "test-instance",
741+
instance_name=name,
741742
backend=backend.type,
742743
region=region,
743744
public_ip=True,
744745
ssh_key_pub=ssh_public_key,
745746
certificate=None,
746747
).model_dump_json()
747748
gateway_replica = GatewayReplicaModel(
749+
name=name,
748750
gateway_id=gateway_id,
749751
backend_id=backend.id,
750752
ip_address=ip_address,

0 commit comments

Comments
 (0)