Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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
]
Expand Down
Original file line number Diff line number Diff line change
@@ -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 ###
1 change: 1 addition & 0 deletions src/dstack/_internal/server/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
17 changes: 8 additions & 9 deletions src/dstack/_internal/server/services/gateways/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand Down
4 changes: 3 additions & 1 deletion src/dstack/_internal/server/testing/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -733,18 +733,20 @@ 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,
ssh_key_pub=ssh_public_key,
certificate=None,
).model_dump_json()
gateway_replica = GatewayReplicaModel(
name=name,
gateway_id=gateway_id,
backend_id=backend.id,
ip_address=ip_address,
Expand Down
Loading