Optimize ServiceRouterWorkerSyncWorkerProcess select query

Author: Bihan Rana

src/dstack/_internal/core/models/configurations.py

@@ -1046,6 +1046,10 @@ def validate_at_most_one_router_replica_group(cls, values):
         router_groups = [g for g in replicas if g.router is not None]
         if len(router_groups) > 1:
             raise ValueError("At most one replica group may specify `router`.")
+        if router_groups:
+            router_group = router_groups[0]
+            if router_group.count.min != 1 or router_group.count.max != 1:
+                raise ValueError("For now replica group with `router` must have `count: 1`.")
         return values
 
 

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

@@ -4,10 +4,10 @@
 from datetime import timedelta
 from typing import Sequence
 
-from sqlalchemy import delete, or_, select, update
-from sqlalchemy.orm import load_only, selectinload
+from sqlalchemy import delete, or_, select, true, update
+from sqlalchemy.orm import joinedload, load_only, selectinload
 
-from dstack._internal.core.models.runs import RunStatus
+from dstack._internal.core.models.runs import JobStatus, RunStatus
 from dstack._internal.server.background.pipeline_tasks.base import (
     Fetcher,
     Heartbeater,
@@ -25,6 +25,7 @@
 from dstack._internal.server.models import (
     InstanceModel,
     JobModel,
+    ProjectModel,
     RunModel,
     ServiceRouterWorkerSyncModel,
 )
@@ -105,7 +106,7 @@ def _workers(self) -> Sequence["ServiceRouterWorkerSyncWorker"]:
 
 
 class ServiceRouterWorkerSyncFetcher(Fetcher[ServiceRouterWorkerSyncPipelineItem]):
-    @sentry_utils.instrument_named_task("pipeline_tasks.ServiceRouterWorkerSyncFetcher.fetch")
+    @sentry_utils.instrument_pipeline_task("ServiceRouterWorkerSyncFetcher.fetch")
     async def fetch(self, limit: int) -> list[ServiceRouterWorkerSyncPipelineItem]:
         sync_lock, _ = get_locker(get_db().dialect_name).get_lockset(
             ServiceRouterWorkerSyncModel.__tablename__
@@ -183,7 +184,7 @@ def __init__(
             pipeline_hinter=pipeline_hinter,
         )
 
-    @sentry_utils.instrument_named_task("pipeline_tasks.ServiceRouterWorkerSyncWorker.process")
+    @sentry_utils.instrument_pipeline_task("ServiceRouterWorkerSyncWorker.process")
     async def process(self, item: ServiceRouterWorkerSyncPipelineItem) -> None:
         async with get_session_ctx() as session:
             res = await session.execute(
@@ -199,10 +200,6 @@ async def process(self, item: ServiceRouterWorkerSyncPipelineItem) -> None:
                 log_lock_token_mismatch(logger, item)
                 return
             run_model = sync_row.run
-            if run_model is None:
-                await session.delete(sync_row)
-                await session.commit()
-                return
             if (
                 run_model.deleted
                 or run_model.status.is_finished()
@@ -218,11 +215,30 @@ async def process(self, item: ServiceRouterWorkerSyncPipelineItem) -> None:
                 select(RunModel)
                 .where(RunModel.id == item.run_id)
                 .options(
-                    selectinload(RunModel.project),
-                    selectinload(RunModel.jobs).selectinload(JobModel.project),
-                    selectinload(RunModel.jobs)
-                    .selectinload(JobModel.instance)
-                    .selectinload(InstanceModel.project),
+                    load_only(RunModel.id, RunModel.run_spec),
+                    selectinload(
+                        RunModel.jobs.and_(
+                            JobModel.status == JobStatus.RUNNING,
+                            JobModel.registered == true(),
+                        )
+                    )
+                    .load_only(
+                        JobModel.id,
+                        JobModel.status,
+                        JobModel.registered,
+                        JobModel.job_spec_data,
+                        JobModel.job_provisioning_data,
+                        JobModel.job_runtime_data,
+                    )
+                    .options(
+                        joinedload(JobModel.project).load_only(
+                            ProjectModel.id, ProjectModel.ssh_private_key
+                        ),
+                        joinedload(JobModel.instance)
+                        .load_only(InstanceModel.id, InstanceModel.remote_connection_info)
+                        .joinedload(InstanceModel.project)
+                        .load_only(ProjectModel.id, ProjectModel.ssh_private_key),
+                    ),
                 )
             )
             run_for_sync = res.unique().scalar_one_or_none()

src/dstack/_internal/server/migrations/versions/2026/03_29_1200_e7f4a91b2c3d_add_service_router_worker_sync.py

@@ -1,7 +1,7 @@
 """Add service_router_worker_sync for router-worker reconcile pipeline.
 
 Revision ID: e7f4a91b2c3d
-Revises: e9d81c97c042
+Revises: ad8c50120507
 Create Date: 2026-03-29 12:00:00.000000+00:00
 
 """
@@ -13,7 +13,7 @@
 import dstack._internal.server.models
 
 revision = "e7f4a91b2c3d"
-down_revision = "e9d81c97c042"
+down_revision = "ad8c50120507"
 branch_labels = None
 depends_on = None
 

src/dstack/_internal/server/services/router_worker_sync.py

@@ -109,6 +109,9 @@ def _get_router_job(run_model: RunModel, router_group) -> Optional[JobModel]:
     ]
     if not router_jobs or not is_replica_registered(router_jobs):
         return None
+    # Router replica group is currently validated to have count=1, so we assume a single active
+    # router job here. When we support multiple router replicas for HA, this should be updated
+    # to handle syncing across all active router jobs.
     return router_jobs[0]
 
 

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

@@ -521,6 +521,7 @@ async def apply_plan(
             deployment_num=new_deployment_num,
         )
     )
+    await ensure_service_router_worker_sync_row(session, current_resource_model, run_spec)
     events.emit(
         session,
         (