diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py index 02860daf3..cc55a4267 100644 --- a/src/dstack/_internal/core/models/configurations.py +++ b/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 diff --git a/src/dstack/_internal/server/services/jobs/configurators/base.py b/src/dstack/_internal/server/services/jobs/configurators/base.py index 811066348..c2da60797 100644 --- a/src/dstack/_internal/server/services/jobs/configurators/base.py +++ b/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() diff --git a/src/dstack/_internal/server/services/runs/router_worker_sync.py b/src/dstack/_internal/server/services/runs/router_worker_sync.py index 9876d2415..e8d13fa4c 100644 --- a/src/dstack/_internal/server/services/runs/router_worker_sync.py +++ b/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: diff --git a/src/tests/_internal/server/services/jobs/configurators/test_service.py b/src/tests/_internal/server/services/jobs/configurators/test_service.py index a4d24ddc5..8a962542a 100644 --- a/src/tests/_internal/server/services/jobs/configurators/test_service.py +++ b/src/tests/_internal/server/services/jobs/configurators/test_service.py @@ -6,6 +6,7 @@ from dstack._internal import settings from dstack._internal.core.models.configurations import ( OPENAI_MODEL_PROBE_TIMEOUT, + ROUTER_HEALTH_PROBE_URL, ProbeConfig, PythonVersion, ReplicaGroup, @@ -13,6 +14,7 @@ ) 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(