Fix gRPC router workers never becoming ready with model set

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

@@ -72,6 +72,7 @@
 DEFAULT_PROBE_READY_AFTER = 1
 DEFAULT_PROBE_METHOD = "get"
 DEFAULT_PROBE_UNTIL_READY = False
+ROUTER_HEALTH_PROBE_URL = "/health"
 MAX_PROBE_URL_LEN = 2048
 DEFAULT_REPLICA_GROUP_NAME = "0"
 OPENAI_MODEL_PROBE_TIMEOUT = 30

src/dstack/_internal/server/services/jobs/configurators/base.py

@@ -27,11 +27,13 @@
     DEFAULT_REPLICA_GROUP_NAME,
     LEGACY_REPO_DIR,
     OPENAI_MODEL_PROBE_TIMEOUT,
+    ROUTER_HEALTH_PROBE_URL,
     HTTPHeaderSpec,
     NodeGroup,
     PortMapping,
     ProbeConfig,
     PythonVersion,
+    ReplicaGroup,
     RepoExistsAction,
     RunConfigurationType,
     ServiceConfiguration,
@@ -489,15 +491,40 @@ def _service_port(self) -> Optional[int]:
             return self.run_spec.configuration.port.container_port
         return None
 
+    def _replica_group(self) -> Optional[ReplicaGroup]:
+        conf = self.run_spec.configuration
+        if not isinstance(conf, ServiceConfiguration):
+            return None
+        # `replica_group_name` is unset for services declaring `replicas` instead of
+        # `groups`; `replica_groups` synthesizes a group under the default name for them.
+        name = self.replica_group_name or DEFAULT_REPLICA_GROUP_NAME
+        for group in conf.replica_groups:
+            if group.name == name:
+                return group
+        return None
+
     def _probes(self) -> list[ProbeSpec]:
-        if isinstance(self.run_spec.configuration, ServiceConfiguration):
-            probes = self.run_spec.configuration.probes
-            if probes is not None:
-                return list(map(_probe_config_to_spec, probes))
-            # Generate default probe if model is set
-            model = self.run_spec.configuration.model
-            if isinstance(model, OpenAIChatModel):
-                return [_openai_model_probe_spec(model.name, model.prefix)]
+        conf = self.run_spec.configuration
+        if not isinstance(conf, ServiceConfiguration):
+            return []
+        if conf.probes is not None:
+            return list(map(_probe_config_to_spec, conf.probes))
+        # Generate default probe if model is set
+        model = conf.model
+        if not isinstance(model, OpenAIChatModel):
+            return []
+        if all(group.router is None for group in conf.replica_groups):
+            # No router: every replica serves the model itself, so a chat completions
+            # request is a genuine end-to-end readiness check.
+            return [_openai_model_probe_spec(model.name, model.prefix)]
+        group = self._replica_group()
+        if group is not None and group.router is not None:
+            # A router only answers chat completions once dstack has registered workers
+            # with it, and registration skips routers that are not ready yet. Probing
+            # chat completions here would deadlock: readiness would wait on registration
+            # while registration waits on readiness. Probe the router's own liveness
+            # endpoint instead, which does not depend on any worker.
+            return [_router_health_probe_spec()]
         return []
 
 
@@ -573,6 +600,19 @@ def _openai_model_probe_spec(model_name: str, prefix: str) -> ProbeSpec:
     )
 
 
+def _router_health_probe_spec() -> ProbeSpec:
+    # Both supported routers (SGLang/SMG and Dynamo) serve `/health` independently of
+    # whether any worker is registered.
+    return ProbeSpec(
+        type="http",
+        method=DEFAULT_PROBE_METHOD,
+        url=ROUTER_HEALTH_PROBE_URL,
+        timeout=DEFAULT_PROBE_TIMEOUT,
+        interval=DEFAULT_PROBE_INTERVAL,
+        ready_after=DEFAULT_PROBE_READY_AFTER,
+    )
+
+
 def _join_shell_commands(commands: List[str]) -> str:
     for i, cmd in enumerate(commands):
         cmd = cmd.strip()

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

@@ -638,6 +638,11 @@ async def sync_router_workers_for_run_model(run_model: RunModel) -> None:
 
     router_job = _get_router_job(run_model, router_group)
     if router_job is None:
+        logger.debug(
+            "%s: no ready router job in group %s, skipping worker sync",
+            fmt(run_model),
+            router_group.name,
+        )
         return
     try:
         async with get_service_replica_client(router_job) as client:

src/tests/_internal/server/services/jobs/configurators/test_service.py

@@ -6,13 +6,15 @@
 from dstack._internal import settings
 from dstack._internal.core.models.configurations import (
     OPENAI_MODEL_PROBE_TIMEOUT,
+    ROUTER_HEALTH_PROBE_URL,
     ProbeConfig,
     PythonVersion,
     ReplicaGroup,
     ServiceConfiguration,
 )
 from dstack._internal.core.models.profiles import SpotPolicy
 from dstack._internal.core.models.resources import Range
+from dstack._internal.core.models.routers import ReplicaGroupRouterConfig
 from dstack._internal.core.models.services import OpenAIChatModel
 from dstack._internal.server.services.docker import ImageConfig
 from dstack._internal.server.services.jobs.configurators.base import get_default_image
@@ -92,6 +94,67 @@ async def test_explicit_empty_probes(self):
         assert len(job_specs) == 1
         assert len(job_specs[0].probes) == 0
 
+    @staticmethod
+    def _router_worker_configuration() -> ServiceConfiguration:
+        return ServiceConfiguration(
+            port=8000,
+            image="debian",
+            model=OpenAIChatModel(
+                name="meta-llama/Meta-Llama-3.1-8B-Instruct",
+                format="openai",
+            ),
+            groups=[
+                ReplicaGroup(
+                    name="router",
+                    replicas=Range[int](min=1, max=1),
+                    router=ReplicaGroupRouterConfig(type="sglang"),
+                ),
+                ReplicaGroup(
+                    name="worker",
+                    replicas=Range[int](min=1, max=1),
+                ),
+            ],
+        )
+
+    async def test_router_group_gets_health_probe(self):
+        """The router must not be probed with chat completions: it only answers those once
+        dstack has registered workers, and registration requires the router to be ready."""
+        run_spec = get_run_spec(
+            run_name="run", repo_id="id", configuration=self._router_worker_configuration()
+        )
+        configurator = ServiceJobConfigurator(run_spec, replica_group_name="router")
+
+        job_specs = await configurator.get_job_specs(replica_num=0)
+
+        probes = job_specs[0].probes
+        assert len(probes) == 1
+        assert probes[0].url == ROUTER_HEALTH_PROBE_URL
+        assert probes[0].method == "get"
+        assert probes[0].body is None
+
+    async def test_worker_group_gets_no_derived_probe(self):
+        """Workers behind a router may speak gRPC, so no probe can be derived from `model`."""
+        run_spec = get_run_spec(
+            run_name="run", repo_id="id", configuration=self._router_worker_configuration()
+        )
+        configurator = ServiceJobConfigurator(run_spec, replica_group_name="worker")
+
+        job_specs = await configurator.get_job_specs(replica_num=0)
+
+        assert job_specs[0].probes == []
+
+    async def test_service_wide_explicit_probes_win_in_router_service(self):
+        """`probes` is service-wide, so an explicit list still applies to every group,
+        router and worker alike, taking precedence over the derived probes."""
+        configuration = self._router_worker_configuration()
+        configuration.probes = [ProbeConfig(type="http", url="/custom")]
+        run_spec = get_run_spec(run_name="run", repo_id="id", configuration=configuration)
+
+        for group_name in ("router", "worker"):
+            configurator = ServiceJobConfigurator(run_spec, replica_group_name=group_name)
+            job_specs = await configurator.get_job_specs(replica_num=0)
+            assert [p.url for p in job_specs[0].probes] == ["/custom"]
+
     async def test_no_probe_when_no_model(self):
         """When neither model nor probes are set, no probes should be generated."""
         configuration = ServiceConfiguration(