Resolve review comments

Author: Bihan Rana

src/dstack/_internal/proxy/gateway/routers/registry.py

@@ -82,7 +82,6 @@ async def register_replica(
         ssh_head_proxy=body.ssh_head_proxy,
         ssh_head_proxy_private_key=body.ssh_head_proxy_private_key,
         internal_ip=body.internal_ip,
-        is_router_replica=body.is_router_replica,
         repo=repo,
         nginx=nginx,
         service_conn_pool=service_conn_pool,

src/dstack/_internal/proxy/gateway/schemas/registry.py

@@ -58,7 +58,6 @@ class RegisterReplicaRequest(BaseModel):
     ssh_head_proxy: Optional[SSHConnectionParams]
     ssh_head_proxy_private_key: Optional[str]
     internal_ip: Optional[str] = None
-    is_router_replica: bool = False
 
 
 class RegisterEntrypointRequest(BaseModel):

src/dstack/_internal/proxy/gateway/services/registry.py

@@ -141,7 +141,6 @@ async def register_replica(
     nginx: Nginx,
     service_conn_pool: ServiceConnectionPool,
     internal_ip: Optional[str] = None,
-    is_router_replica: bool = False,
 ) -> None:
     replica = models.Replica(
         id=replica_id,
@@ -153,7 +152,6 @@ async def register_replica(
         ssh_head_proxy=ssh_head_proxy,
         ssh_head_proxy_private_key=ssh_head_proxy_private_key,
         internal_ip=internal_ip,
-        is_router_replica=is_router_replica,
     )
 
     async with lock:

src/dstack/_internal/proxy/lib/models.py

@@ -30,7 +30,6 @@ class Replica(ImmutableModel):
     ssh_head_proxy: Optional[SSHConnectionParams] = None
     ssh_head_proxy_private_key: Optional[str] = None
     internal_ip: Optional[str] = None
-    is_router_replica: bool = False
 
 
 class IPAddressPartitioningKey(ImmutableModel):

src/dstack/_internal/proxy/lib/services/service_connection.py

@@ -151,9 +151,7 @@ async def get_service_replica_client(
             timeout=HTTP_TIMEOUT,
         )
     # Nginx not available, forward directly to the tunnel
-    router_replicas = [r for r in service.replicas if r.is_router_replica]
-    replicas_to_use = router_replicas if router_replicas else service.replicas
-    replica = random.choice(replicas_to_use)
+    replica = random.choice(service.replicas)
     connection = await service_conn_pool.get(replica.id)
     if connection is None:
         project = await repo.get_project(service.project_name)

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

@@ -6,7 +6,7 @@
 
 from sqlalchemy import or_, select, update
 from sqlalchemy.orm import joinedload, load_only, selectinload
-from sqlalchemy.sql import false
+from sqlalchemy.sql import false, true
 
 from dstack._internal.core.models.runs import JobStatus, RunStatus
 from dstack._internal.server.background.pipeline_tasks.base import (
@@ -120,7 +120,13 @@ async def fetch(self, limit: int) -> list[ServiceRouterWorkerSyncPipelineItem]:
                     .join(RunModel, RunModel.id == ServiceRouterWorkerSyncModel.run_id)
                     .where(
                         ServiceRouterWorkerSyncModel.deleted == false(),
-                        RunModel.status == RunStatus.RUNNING,
+                        # Fetch RUNNING runs for normal processing, and finished/deleted runs so
+                        # the worker can mark their sync rows deleted.
+                        or_(
+                            RunModel.status == RunStatus.RUNNING,
+                            RunModel.status.in_(RunStatus.finished_statuses()),
+                            RunModel.deleted == true(),
+                        ),
                         or_(
                             ServiceRouterWorkerSyncModel.last_processed_at
                             <= now - self._min_processing_interval,
@@ -170,7 +176,7 @@ async def fetch(self, limit: int) -> list[ServiceRouterWorkerSyncPipelineItem]:
 
 
 class _SyncRowUpdateMap(ItemUpdateMap, total=False):
-    pass
+    deleted: bool
 
 
 class ServiceRouterWorkerSyncWorker(Worker[ServiceRouterWorkerSyncPipelineItem]):
@@ -224,6 +230,7 @@ async def process(self, item: ServiceRouterWorkerSyncPipelineItem) -> None:
                     selectinload(
                         RunModel.jobs.and_(
                             JobModel.status == JobStatus.RUNNING,
+                            JobModel.registered == True,
                         )
                     )
                     .load_only(
@@ -248,20 +255,23 @@ async def process(self, item: ServiceRouterWorkerSyncPipelineItem) -> None:
             run_for_sync = res.unique().scalar_one_or_none()
 
         if run_for_sync is None:
+            cleanup_update_map: _SyncRowUpdateMap = {"deleted": True}
+            set_processed_update_map_fields(cleanup_update_map)
+            set_unlock_update_map_fields(cleanup_update_map)
             async with get_session_ctx() as session:
-                await session.execute(
+                now = get_current_datetime()
+                resolve_now_placeholders(cleanup_update_map, now=now)
+                res2 = await session.execute(
                     update(ServiceRouterWorkerSyncModel)
                     .where(
                         ServiceRouterWorkerSyncModel.id == item.id,
                         ServiceRouterWorkerSyncModel.lock_token == item.lock_token,
                     )
-                    .values(
-                        deleted=True,
-                        lock_expires_at=None,
-                        lock_token=None,
-                        lock_owner=None,
-                    )
+                    .values(**cleanup_update_map)
+                    .returning(ServiceRouterWorkerSyncModel.id)
                 )
+                if not list(res2.scalars().all()):
+                    log_lock_token_changed_after_processing(logger, item)
                 await session.commit()
             return
 

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

@@ -90,17 +90,11 @@ async def register_replica(
         ssh_head_proxy_private_key: Optional[str],
     ):
         assert run.run_spec.configuration.type == "service"
-        config = run.run_spec.configuration
-        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
-        )
         payload = {
             "job_id": job_submission.id.hex,
             "app_port": get_service_port(job_spec, run.run_spec.configuration),
             "ssh_head_proxy": ssh_head_proxy.dict() if ssh_head_proxy is not None else None,
             "ssh_head_proxy_private_key": ssh_head_proxy_private_key,
-            "is_router_replica": is_router_replica,
         }
         jpd = job_submission.job_provisioning_data
         assert jpd is not None

src/dstack/_internal/server/services/proxy/repo.py

@@ -71,6 +71,11 @@ async def get_service(self, project_name: str, run_name: str) -> Optional[Servic
         run_spec = get_run_spec(run)
         if not isinstance(run_spec.configuration, ServiceConfiguration):
             return None
+        router_group = next(
+            (g for g in run_spec.configuration.replica_groups if g.router is not None),
+            None,
+        )
+        router = run_spec.configuration.router
         replicas = []
         for job in jobs:
             jpd: JobProvisioningData = JobProvisioningData.__response__.parse_raw(
@@ -109,13 +114,10 @@ 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)
-            router_group = next(
-                (g for g in run_spec.configuration.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 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
             replica = Replica(
                 id=job.id.hex,
                 app_port=get_service_port(job_spec, run_spec.configuration),
@@ -126,7 +128,6 @@ async def get_service(self, project_name: str, run_name: str) -> Optional[Servic
                 ssh_head_proxy=ssh_head_proxy,
                 ssh_head_proxy_private_key=ssh_head_proxy_private_key,
                 internal_ip=jpd.internal_ip,
-                is_router_replica=is_router_replica,
             )
             replicas.append(replica)
         return Service(
@@ -138,6 +139,7 @@ async def get_service(self, project_name: str, run_name: str) -> Optional[Servic
             client_max_body_size=DEFAULT_SERVICE_CLIENT_MAX_BODY_SIZE,
             strip_prefix=run_spec.configuration.strip_prefix,
             replicas=tuple(replicas),
+            router=router,
         )
 
     async def list_models(self, project_name: str) -> List[ChatModel]:

src/dstack/_internal/server/services/proxy/services/service_proxy.py

@@ -5,6 +5,7 @@
 from fastapi import status
 from starlette.requests import ClientDisconnect
 
+from dstack._internal.core.models.routers import RouterType
 from dstack._internal.proxy.lib.deps import ProxyAuthContext
 from dstack._internal.proxy.lib.errors import ProxyError
 from dstack._internal.proxy.lib.repo import BaseProxyRepo
@@ -36,11 +37,16 @@ async def proxy(
     if service.auth:
         await auth.enforce()
 
-    client = await get_service_replica_client(service, repo, service_conn_pool)
-
     if not service.strip_prefix:
         path = concat_url_path(request.scope.get("root_path", "/"), request.url.path)
 
+    if service.router is not None and service.router.type == RouterType.SGLANG:
+        path_for_match = path if path.startswith("/") else f"/{path}"
+        if not _is_whitelisted_path(path_for_match, _SGLANG_WHITELISTED_PATHS):
+            raise ProxyError("Path is not allowed for this service", status.HTTP_404_NOT_FOUND)
+
+    client = await get_service_replica_client(service, repo, service_conn_pool)
+
     try:
         upstream_request = await build_upstream_request(request, path, client)
     except ClientDisconnect:
@@ -68,6 +74,23 @@ async def proxy(
     )
 
 
+_SGLANG_WHITELISTED_PATHS = (
+    "/generate",
+    "/v1/",
+    "/chat/completions",
+)
+
+
+def _is_whitelisted_path(path: str, whitelisted_paths: tuple[str, ...]) -> bool:
+    for allowed in whitelisted_paths:
+        if allowed.endswith("/"):
+            if path.startswith(allowed):
+                return True
+        elif path == allowed:
+            return True
+    return False
+
+
 async def stream_response(response: httpx.Response) -> AsyncGenerator[bytes, None]:
     try:
         async for chunk in response.aiter_raw():

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

@@ -176,20 +176,10 @@ async def _add_worker_to_router(
         return False
 
 
-async def _remove_worker_from_router(client: AsyncClient, worker_url: str) -> bool:
+async def _remove_worker_from_router_by_id(
+    client: AsyncClient, worker_id: str, *, worker_url: str
+) -> bool:
     try:
-        current = await _get_router_workers(client)
-        worker_id = None
-        for w in current:
-            u = w.get("url")
-            if u and isinstance(u, str) and u.rstrip("/") == worker_url.rstrip("/"):
-                wid = w.get("id")
-                if wid and isinstance(wid, str):
-                    worker_id = wid
-                    break
-        if not worker_id:
-            logger.error("No worker id found for url %s", worker_url)
-            return False
         body = await _request_json_limited(
             client,
             "DELETE",
@@ -212,10 +202,16 @@ async def _update_workers_in_router_replica(
 ) -> None:
     current = await _get_router_workers(client)
     current_urls: set[str] = set()
+    current_ids_by_norm_url: dict[str, str] = {}
     for w in current:
         u = w.get("url")
-        if isinstance(u, str) and u:
-            current_urls.add(_normalize_worker_url(u))
+        if not isinstance(u, str) or not u:
+            continue
+        norm_u = _normalize_worker_url(u)
+        current_urls.add(norm_u)
+        wid = w.get("id")
+        if isinstance(wid, str) and wid:
+            current_ids_by_norm_url[norm_u] = wid
     target_by_norm = {_normalize_worker_url(t["url"]): t for t in target_workers}
     target_urls = set(target_by_norm.keys())
     to_add = sorted(target_urls - current_urls)
@@ -231,7 +227,12 @@ async def _update_workers_in_router_replica(
         if not ok:
             logger.warning("Failed to add worker %s, continuing with others", tw["url"])
     for url in to_remove:
-        ok = await _remove_worker_from_router(client, url)
+        wid = current_ids_by_norm_url.get(url)
+        if not wid:
+            logger.error("No worker id found for url %s", url)
+            ok = False
+        else:
+            ok = await _remove_worker_from_router_by_id(client, wid, worker_url=url)
         if not ok:
             logger.warning("Failed to remove worker %s, continuing with others", url)
 
@@ -270,9 +271,9 @@ async def _get_worker_payload(job_model: JobModel, worker_url: str) -> _WorkerPa
                     "payload": {"url": worker_url, "worker_type": "regular"},
                 }
     except _ResponseTooLargeError:
-        logger.debug("server_info response too large for worker %s", worker_url)
+        logger.warning("server_info response too large for worker %s", worker_url)
     except Exception as e:
-        logger.debug("Could not fetch server_info for worker %s: %r", worker_url, e)
+        logger.exception("Could not fetch server_info for worker %s: %r", worker_url, e)
     return {"status": "not_ready", "payload": None}