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 @@ -20,7 +20,6 @@
)
from dstack._internal.core.models.configurations import (
DevEnvironmentConfiguration,
ServiceConfiguration,
)
from dstack._internal.core.models.files import FileArchiveMapping
from dstack._internal.core.models.instances import InstanceStatus, SSHConnectionParams
Expand All @@ -34,7 +33,6 @@
Job,
JobProvisioningData,
JobRuntimeData,
JobSpec,
JobStatus,
JobSubmission,
JobTerminationReason,
Expand Down Expand Up @@ -366,6 +364,7 @@ class _JobUpdateMap(ItemUpdateMap, total=False):
disconnected_at: Optional[datetime]
inactivity_secs: Optional[int]
exit_status: Optional[int]
ready: bool
registered: bool
image_pull_progress: Optional[str]
skip_min_processing_interval: bool
Expand Down Expand Up @@ -1083,6 +1082,12 @@ def _emit_result_events(
job_model.disconnected_at,
),
)
_emit_readiness_change_event(
session=session,
job_model=job_model,
old_ready=job_model.ready,
new_ready=result.job_update_map.get("ready", job_model.ready),
)
if result.replica_registration is not None:
targets = [events.Target.from_model(job_model)]
if result.replica_registration.gateway_target is not None:
Expand Down Expand Up @@ -1175,13 +1180,32 @@ async def _maybe_register_replica(
) -> None:
if (
context.run.run_spec.configuration.type != "service"
or _get_result_registered(context.job_model, result)
or context.job_model.job_num != 0
or result.new_probe_models
or not is_job_ready(context.job_model.probes, context.job.job_spec.probes)
):
return

is_ready = is_job_ready(context.job_model.probes, context.job.job_spec.probes)
if is_ready and not context.job_model.ready:
result.job_update_map["ready"] = True

router_group = next(
(g for g in context.run.run_spec.configuration.replica_groups if g.router is not None),
None,
)
is_router_replica = (
router_group is not None and context.job.job_spec.replica_group == router_group.name
)
# non-router replicas aren't registered if the service has a router
if router_group is not None and not is_router_replica:
if context.job_model.registered:
# migration edge case: a pre-0.21.1 server replica incorrectly set registered=True
result.job_update_map["registered"] = False
return

if not is_ready or _get_result_registered(context.job_model, result):
return

ssh_head_proxy: Optional[SSHConnectionParams] = None
ssh_head_proxy_private_key: Optional[str] = None
instance = get_or_error(context.job_model.instance)
Expand Down Expand Up @@ -1220,23 +1244,6 @@ async def _register_service_replica(
) -> Optional[events.Target]:
if context.run_model.gateway_id is None:
return None

job_spec = validate_json_extra_ignore(JobSpec, context.job_model.job_spec_data)

# For router-based services (e.g. PD disaggregation), only router replicas should be
# registered with the gateway. Worker replicas are discovered by the router-worker
# sync pipeline and should not be routed to directly by the gateway.
config = context.run.run_spec.configuration
assert isinstance(config, ServiceConfiguration)
router_group = next((g for g in config.replica_groups if g.router is not None), None)
is_router_replica = router_group is not None and job_spec.replica_group == router_group.name
if router_group is not None and not is_router_replica:
logger.debug(
"%s: skipping gateway replica registration (non-router replica)",
fmt(context.job_model),
)
return None

async with get_session_ctx() as session:
gateway_model, connections = await get_or_add_gateway_connections(
session, context.run_model.gateway_id
Expand All @@ -1261,7 +1268,7 @@ async def _register_service_replica(
async with conn.client() as gateway_client:
await gateway_client.register_replica(
run=context.run,
job_spec=job_spec,
job_spec=context.job.job_spec,
job_submission=job_submission,
instance_project_ssh_private_key=instance_project_ssh_private_key,
ssh_head_proxy=ssh_head_proxy,
Expand Down Expand Up @@ -1877,6 +1884,23 @@ def _emit_reachability_change_event(
)


def _emit_readiness_change_event(
session: AsyncSession,
job_model: JobModel,
old_ready: bool,
new_ready: bool,
) -> None:
# ready: False -> True
if not old_ready and new_ready:
events.emit(
session,
"Service replica ready to receive requests",
actor=events.SystemActor(),
targets=[events.Target.from_model(job_model)],
)
# ready: True -> False is not possible as of this writing


def _terminate_job(
job_model: JobModel,
job_update_map: _JobUpdateMap,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -686,11 +686,11 @@ async def _build_rolling_deployment_maps(
max_new = max(j.replica_num for j in new_jobs)
next_replica_num = max(next_replica_num, max_new + 1)

# Scale down: terminate unregistered out-of-date + excess registered replicas
replicas_to_stop = state.unregistered_out_of_date_replica_count
# Scale down: terminate unready out-of-date + excess ready replicas
replicas_to_stop = state.unready_out_of_date_replica_count
replicas_to_stop += max(
0,
state.registered_non_terminating_replica_count - group_desired,
state.ready_non_terminating_replica_count - group_desired,
)
if replicas_to_stop > 0:
scale_down_maps = _build_scale_down_job_update_maps(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -233,13 +233,12 @@ async def process(self, item: ServiceRouterWorkerSyncPipelineItem) -> None:
selectinload(
RunModel.jobs.and_(
JobModel.status == JobStatus.RUNNING,
JobModel.registered == True,
JobModel.ready == True,
)
)
.load_only(
JobModel.id,
JobModel.status,
JobModel.registered,
JobModel.job_spec_data,
JobModel.job_provisioning_data,
JobModel.job_runtime_data,
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,120 @@
"""Add JobModel.ready

Revision ID: 72cfa56364ad
Revises: a1c3f5e7b209
Create Date: 2026-08-04 07:02:02.041000+00:00

"""

import json
import uuid
from typing import Optional

import sqlalchemy as sa
from alembic import op
from sqlalchemy_utils import UUIDType

# revision identifiers, used by Alembic.
revision = "72cfa56364ad"
down_revision = "a1c3f5e7b209"
branch_labels = None
depends_on = None

# Partial table descriptions - only columns needed for the data migration below.
jobs_table = sa.Table(
"jobs",
sa.MetaData(),
sa.Column("id", UUIDType(binary=False), primary_key=True, default=uuid.uuid4),
sa.Column("run_id", UUIDType(binary=False)),
sa.Column("status", sa.String(100)),
sa.Column("registered", sa.Boolean()),
sa.Column("ready", sa.Boolean()),
sa.Column("job_spec_data", sa.Text()),
)
runs_table = sa.Table(
"runs",
sa.MetaData(),
sa.Column("id", UUIDType(binary=False), primary_key=True, default=uuid.uuid4),
sa.Column("run_spec", sa.Text()),
sa.Column("service_spec", sa.Text(), nullable=True),
)


def _get_router_group_name(run_spec_data: str) -> Optional[str]:
configuration = json.loads(run_spec_data).get("configuration") or {}
if configuration.get("type") != "service":
return None
replica_groups = configuration.get("replicas")
if not isinstance(replica_groups, list):
return None
for group in replica_groups:
if isinstance(group, dict) and group.get("router") is not None:
return group.get("name")
return None


def _get_job_replica_group(job_spec_data: str) -> str:
return json.loads(job_spec_data).get("replica_group", "0")


def upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table("jobs", schema=None) as batch_op:
batch_op.add_column(
sa.Column("ready", sa.Boolean(), server_default=sa.false(), nullable=False)
)

# ### end Alembic commands ###

bind = op.get_bind()

# backfill ready=True for existing registered replicas

bind.execute(jobs_table.update().where(jobs_table.c.registered == True).values(ready=True))

# set registered=False for non-router replicas in services with a router

candidate_jobs = bind.execute(
sa.select(jobs_table.c.id, jobs_table.c.run_id, jobs_table.c.job_spec_data)
.select_from(jobs_table.join(runs_table, jobs_table.c.run_id == runs_table.c.id))
.where(
jobs_table.c.registered == True,
jobs_table.c.status == "RUNNING",
runs_table.c.service_spec.is_not(None),
sa.or_(
runs_table.c.run_spec.like('%"sglang"%'),
runs_table.c.run_spec.like('%"dynamo"%'),
),
)
).all()

run_ids = {row.run_id for row in candidate_jobs}
router_group_name_by_run_id = {}
if run_ids:
for row in bind.execute(
sa.select(runs_table.c.id, runs_table.c.run_spec).where(runs_table.c.id.in_(run_ids))
).all():
router_group_name_by_run_id[row.id] = _get_router_group_name(row.run_spec)

non_router_job_ids = []
for row in candidate_jobs:
router_group_name = router_group_name_by_run_id.get(row.run_id)
if router_group_name is None:
continue
if _get_job_replica_group(row.job_spec_data) != router_group_name:
non_router_job_ids.append(row.id)

if non_router_job_ids:
bind.execute(
jobs_table.update()
.where(jobs_table.c.id.in_(non_router_job_ids))
.values(registered=False)
)


def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table("jobs", schema=None) as batch_op:
batch_op.drop_column("ready")

# ### end Alembic commands ###
9 changes: 7 additions & 2 deletions src/dstack/_internal/server/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -578,9 +578,14 @@ class JobModel(PipelineModelMixin, BaseModel):
probes: Mapped[list["ProbeModel"]] = relationship(
back_populates="job", order_by="ProbeModel.probe_num"
)
ready: Mapped[bool] = mapped_column(Boolean, server_default=false())
"""Whether the replica is ready to receive service requests based on probe statuses.
Always `False` for non-service runs.
"""
registered: Mapped[bool] = mapped_column(Boolean, server_default=false())
"""`registered` shows whether the replica is registered to receive service requests.
It is always `False` for non-service runs.
"""Whether the replica is registered to receive service requests from dstack-proxy.
Always `False` for non-service runs or jobs that shouldn't be registered
(e.g., non-router replicas for services with routers).
"""
waiting_master_job: Mapped[Optional[bool]] = mapped_column(Boolean)
"""`waiting_master_job` is `True` for non-master jobs that have to wait for master processing before
Expand Down
4 changes: 0 additions & 4 deletions src/dstack/_internal/server/services/proxy/repo.py
Original file line number Diff line number Diff line change
Expand Up @@ -113,10 +113,6 @@ async def get_service(self, project_name: str, run_name: str) -> Optional[Servic
ssh_head_proxy = rci.ssh_proxy
ssh_head_proxy_private_key = get_or_error(rci.ssh_proxy_keys)[0].private
job_spec = get_job_spec(job)
if router_group is not None and job_spec.replica_group != router_group.name:
# Strict router-only: when a router is configured, the proxy should only be aware
# of router replicas.
continue
Comment on lines -116 to -119

@r4victor r4victor Aug 10, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What about Dynamo workers? I suppose registered remains True for dynamo workers, which means you cannot rely on registered=True to filter router jobs.

#3868

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

registered doesn't depend on the router type, so Dynamo workers also have registered=False and are filtered out here

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was looking at the migration that sets registered=False only for jobs that are in the "service_router_worker_sync" table

replica = Replica(
id=job.id.hex,
app_port=get_service_port(job_spec, run_spec.configuration),
Expand Down
26 changes: 13 additions & 13 deletions src/dstack/_internal/server/services/runs/replicas.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,8 @@ class GroupRolloutState:
inactive_replicas: List[Tuple[int, bool, int, List[JobModel]]]
has_out_of_date_replicas: bool
non_terminated_replica_count: int
unregistered_out_of_date_replica_count: int
registered_non_terminating_replica_count: int
unready_out_of_date_replica_count: int
ready_non_terminating_replica_count: int


class RouterEnvStatus(str, Enum):
Expand Down Expand Up @@ -77,7 +77,7 @@ def build_replica_lists(
elif {JobStatus.PROVISIONING, JobStatus.PULLING} & statuses:
# if there are any provisioning or pulling jobs, the replica is active and has the importance of 1
active_replicas.append((1, is_out_of_date, replica_num, replica_jobs))
elif not is_replica_registered(replica_jobs):
elif not is_replica_ready(replica_jobs):
# all jobs are running, but not receiving traffic, the replica is active and has the importance of 2
active_replicas.append((2, is_out_of_date, replica_num, replica_jobs))
else:
Expand All @@ -98,8 +98,8 @@ def get_group_rollout_state(run_model: RunModel, group: ReplicaGroup) -> GroupRo
)

non_terminated_replica_nums = set()
unregistered_out_of_date_replica_count = 0
registered_non_terminating_replica_count = 0
unready_out_of_date_replica_count = 0
ready_non_terminating_replica_count = 0

for _, jobs in group_jobs_by_replica_latest(run_model.jobs):
if not job_belongs_to_group(jobs[0], group.name):
Expand All @@ -114,20 +114,20 @@ def get_group_rollout_state(run_model: RunModel, group: ReplicaGroup) -> GroupRo
j.status not in [JobStatus.TERMINATING] + JobStatus.finished_statuses()
for j in jobs
)
and not is_replica_registered(jobs)
and not is_replica_ready(jobs)
):
unregistered_out_of_date_replica_count += 1
unready_out_of_date_replica_count += 1

if is_replica_registered(jobs) and all(j.status != JobStatus.TERMINATING for j in jobs):
registered_non_terminating_replica_count += 1
if is_replica_ready(jobs) and all(j.status != JobStatus.TERMINATING for j in jobs):
ready_non_terminating_replica_count += 1

return GroupRolloutState(
active_replicas=active_replicas,
inactive_replicas=inactive_replicas,
has_out_of_date_replicas=has_out_of_date_replicas(run_model, group_filter=group.name),
non_terminated_replica_count=len(non_terminated_replica_nums),
unregistered_out_of_date_replica_count=unregistered_out_of_date_replica_count,
registered_non_terminating_replica_count=registered_non_terminating_replica_count,
unready_out_of_date_replica_count=unready_out_of_date_replica_count,
ready_non_terminating_replica_count=ready_non_terminating_replica_count,
)


Expand All @@ -149,9 +149,9 @@ def has_out_of_date_replicas(run: RunModel, group_filter: Optional[str] = None)
return False


def is_replica_registered(jobs: list[JobModel]) -> bool:
def is_replica_ready(jobs: list[JobModel]) -> bool:
# Only job_num=0 is supposed to receive service requests
return jobs[0].registered
return jobs[0].ready


def get_router_replica_group(run_spec: RunSpec) -> Optional[ReplicaGroup]:
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
"""Reconcile SGLang router /workers with dstack's registered worker replicas (async, SSH-tunneled)."""
"""Reconcile SGLang router /workers with dstack's ready worker replicas (async, SSH-tunneled)."""

import json
from typing import Any, List, Literal, Optional, TypedDict
Expand Down
Loading
Loading