From 62c4c19a9485be205ee63fe5de84f686ea0391ba Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Wed, 23 Sep 2026 15:37:02 +0000 Subject: [PATCH] Fix errors in service router worker sync * Fixed two `DetachedInstanceError` due to missing `RunModel.run_name` and `JobModel.job_name`. * Fixed `Could not fetch server_info for worker ...: ReadError('')` logged as an error with a traceback, repeating until the worker is ready. An SSH tunnel to a port nobody listens on yet fails with `httpx.ReadError`, so this is the normal startup path. Both `_probe_http_worker()` and `_get_router_workers()` now catch `httpx.RequestError` and log it at the debug level. Also: * `_get_router_workers()` returns `None` instead of `[]` when the response cannot be read. `[]` means "the router has no workers", so an unexpected status, an unparsable body, or a missing `workers` key made the sync re-register every worker. * Bare `except Exception` clauses are replaced with the exceptions actually expected, so a programming error is no longer reported as an unready replica. `sync_router_workers_for_run_model()` logs such errors itself instead of letting them escape to the pipeline worker, which leaves the sync row locked until the lock expires. * `_get_worker()` takes the replica address instead of two pre-built URLs, and each probe builds the URL it registers the worker under. `_get_http_worker()`, `_get_grpc_worker()`, and `_discover_grpc_server_info()` are folded into loops over the connection modes and the runtime types to probe. Fixes: https://github.com/dstackai/dstack/issues/4300 --- .../service_router_worker_sync.py | 16 +- .../services/runs/router_worker_sync.py | 350 +++++------- .../test_service_router_worker_sync.py | 122 ++++- .../services/runs/test_router_worker_sync.py | 506 ++++++++++++------ 4 files changed, 616 insertions(+), 378 deletions(-) diff --git a/src/dstack/_internal/server/background/pipeline_tasks/service_router_worker_sync.py b/src/dstack/_internal/server/background/pipeline_tasks/service_router_worker_sync.py index 48fe522108..ca42017571 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/service_router_worker_sync.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/service_router_worker_sync.py @@ -201,7 +201,14 @@ async def process(self, item: ServiceRouterWorkerSyncPipelineItem) -> None: ServiceRouterWorkerSyncModel.id == item.id, ServiceRouterWorkerSyncModel.lock_token == item.lock_token, ) - .options(selectinload(ServiceRouterWorkerSyncModel.run)) + .options( + joinedload(ServiceRouterWorkerSyncModel.run).load_only( + RunModel.id, + RunModel.deleted, + RunModel.status, + RunModel.run_spec, + ) + ) ) sync_row = res.unique().scalar_one_or_none() if sync_row is None: @@ -229,7 +236,11 @@ async def process(self, item: ServiceRouterWorkerSyncPipelineItem) -> None: select(RunModel) .where(RunModel.id == item.run_id) .options( - load_only(RunModel.id, RunModel.run_spec), + load_only( + RunModel.id, + RunModel.run_name, + RunModel.run_spec, + ), selectinload( RunModel.jobs.and_( JobModel.status == JobStatus.RUNNING, @@ -238,6 +249,7 @@ async def process(self, item: ServiceRouterWorkerSyncPipelineItem) -> None: ) .load_only( JobModel.id, + JobModel.job_name, JobModel.status, JobModel.job_spec_data, JobModel.job_provisioning_data, 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 e8d13fa4c6..918711be1e 100644 --- a/src/dstack/_internal/server/services/runs/router_worker_sync.py +++ b/src/dstack/_internal/server/services/runs/router_worker_sync.py @@ -6,14 +6,7 @@ import grpc from google.protobuf.json_format import MessageToDict -from httpx import ( - AsyncClient, - ConnectError, - ConnectTimeout, - ReadTimeout, - RemoteProtocolError, - Response, -) +from httpx import AsyncClient, RequestError, Response from smg_grpc_proto import ( sglang_scheduler_pb2, sglang_scheduler_pb2_grpc, @@ -30,7 +23,6 @@ from dstack._internal.server.services.jobs import get_job_provisioning_data, get_job_spec from dstack._internal.server.services.jobs.job_replica_grpc_client import ( get_service_replica_grpc_channel_over_uds, - get_service_replica_grpc_client, ) from dstack._internal.server.services.jobs.job_replica_http_client import ( get_service_replica_client, @@ -45,13 +37,14 @@ logger = get_logger(__name__) -_ROUTER_HTTP = "http://dstack" -_ROUTER_HTTP_TIMEOUT = 10.0 +# Requests are made over a UDS tunnel to the replica, so the authority is a placeholder. +_HTTP_BASE_URL = "http://dstack" +_HTTP_TIMEOUT = 10.0 _MAX_SERVER_INFO_RESPONSE_BYTES = 256 * 1024 _MAX_WORKERS_RESPONSE_BYTES = 2 * 1024 * 1024 _MAX_WORKERS_COMMAND_ACK_BYTES = 64 * 1024 _MAX_WORKERS_LIST_ITEMS = 8192 -_GRPC_DISCOVERY_TIMEOUT = 30.0 +_GRPC_TIMEOUT = 30.0 class _ResponseTooLargeError(Exception): @@ -75,7 +68,7 @@ async def _request_json_limited( max_response_bytes: int, ok_statuses: set[int], json_body: Optional[dict] = None, - timeout: float = _ROUTER_HTTP_TIMEOUT, + timeout: float = _HTTP_TIMEOUT, ) -> Any: kwargs: dict[str, Any] = {"timeout": timeout} if json_body is not None: @@ -105,7 +98,7 @@ async def _request_json_limited( return None -class _TargetWorker(TypedDict): +class _Worker(TypedDict): url: str worker_type: str bootstrap_port: NotRequired[Optional[int]] @@ -115,14 +108,13 @@ class _TargetWorker(TypedDict): kv_role: NotRequired[str] -class _WorkerPayloadResult(TypedDict): - status: Literal["ready", "not_ready"] - worker: Optional[_TargetWorker] - - _ConnectionMode = Literal["grpc", "http"] +# The order does matter -- we discover connection modes in the specified order +_CONNECTION_MODES: tuple[_ConnectionMode, ...] = ("http", "grpc") + _RuntimeType = Literal["sglang", "vllm"] -_GRPC_RUNTIME_TYPES: tuple[_RuntimeType, ...] = ("sglang", "vllm") +# The order does matter -- we discover runtime types in the specified order +_RUNTIME_TYPES: tuple[_RuntimeType, ...] = ("sglang", "vllm") def run_model_has_sglang_router_replica_group(run_model: RunModel) -> bool: @@ -183,7 +175,7 @@ def _get_runtime_type_from_workers( if worker.get("connection_mode") != "grpc": continue runtime_type = worker.get("runtime_type") - if isinstance(runtime_type, str) and runtime_type in _GRPC_RUNTIME_TYPES: + if isinstance(runtime_type, str) and runtime_type in _RUNTIME_TYPES: runtimes.add(runtime_type) if runtimes == {"sglang"}: return "sglang" @@ -192,45 +184,25 @@ def _get_runtime_type_from_workers( return None -def _is_expected_router_workers_fetch_error(error: Exception) -> bool: - """SMG router may not accept HTTP yet during startup.""" - if isinstance( - error, - ( - RemoteProtocolError, - ConnectError, - ConnectTimeout, - ReadTimeout, - TimeoutError, - ), - ): - return True - if isinstance(error, OSError) and error.errno in {61, 111}: - return True - return False - - -def _log_router_workers_fetch_failure(error: Exception) -> None: - if _is_expected_router_workers_fetch_error(error): - logger.debug("Router /workers not ready yet: %r", error) - return - logger.exception("Error getting router /workers") - - -async def _get_router_workers(client: AsyncClient) -> List[dict]: +async def _get_router_workers(client: AsyncClient) -> Optional[List[dict]]: try: data = await _request_json_limited( client, "GET", - f"{_ROUTER_HTTP}/workers", + f"{_HTTP_BASE_URL}/workers", max_response_bytes=_MAX_WORKERS_RESPONSE_BYTES, ok_statuses={200}, ) if not isinstance(data, dict): - return [] - workers = data.get("workers", []) + # Non-200 status or unparsable response or response is not a JSON object + return None + workers = data.get("workers") if not isinstance(workers, list): - return [] + # Unexpected response structure -- `workers` is missing or is not an array + return None + # TODO: Truncating a long list and/or dropping unexpectedly shaped items doesn't seem + # right. We should add validation (an item must be a dict with some required fields, see + # _update_workers_in_router_replica) and decide what to do with a partially valid list if len(workers) > _MAX_WORKERS_LIST_ITEMS: logger.warning( "Router /workers list exceeds %s items, truncating", @@ -240,9 +212,9 @@ async def _get_router_workers(client: AsyncClient) -> List[dict]: return [w for w in workers if isinstance(w, dict)] except _ResponseTooLargeError: logger.warning("Router /workers response exceeded size limit") - except Exception as e: - _log_router_workers_fetch_failure(e) - return [] + except RequestError as e: + logger.debug("Router /workers not ready yet: %r", e) + return None async def _add_worker_to_router( @@ -271,18 +243,20 @@ async def _add_worker_to_router( body = await _request_json_limited( client, "POST", - f"{_ROUTER_HTTP}/workers", + f"{_HTTP_BASE_URL}/workers", max_response_bytes=_MAX_WORKERS_COMMAND_ACK_BYTES, ok_statuses={202}, json_body=payload, ) - return isinstance(body, dict) and body.get("status") == "accepted" + added = isinstance(body, dict) and body.get("status") == "accepted" + if not added: + logger.warning("Unexpected add-worker response for %s: %s", url, body) + return added except _ResponseTooLargeError: logger.warning("Router add-worker response exceeded size limit for %s", url) - return False - except Exception: - logger.exception("Error adding worker %s", url) - return False + except RequestError as e: + logger.warning("Error adding worker %s: %r", url, e) + return False async def _remove_worker_from_router_by_id( @@ -292,22 +266,24 @@ async def _remove_worker_from_router_by_id( body = await _request_json_limited( client, "DELETE", - f"{_ROUTER_HTTP}/workers/{worker_id}", + f"{_HTTP_BASE_URL}/workers/{worker_id}", max_response_bytes=_MAX_WORKERS_COMMAND_ACK_BYTES, ok_statuses={202}, ) - return isinstance(body, dict) and body.get("status") == "accepted" + removed = isinstance(body, dict) and body.get("status") == "accepted" + if not removed: + logger.warning("Unexpected remove-worker response for %s: %s", worker_url, body) + return removed except _ResponseTooLargeError: logger.warning("Router remove-worker response exceeded size limit for %s", worker_url) - return False - except Exception: - logger.exception("Error removing worker %s", worker_url) - return False + except RequestError as e: + logger.warning("Error removing worker %s: %r", worker_url, e) + return False async def _update_workers_in_router_replica( client: AsyncClient, - target_workers: List[_TargetWorker], + target_workers: List[_Worker], *, current_workers: List[dict], ) -> None: @@ -339,7 +315,7 @@ async def _update_workers_in_router_replica( kv_role=tw.get("kv_role"), ) if not ok: - logger.warning("Failed to add worker %s, continuing with others", tw["url"]) + logger.debug("Failed to add worker %s, continuing with others", tw["url"]) for url in to_remove: wid = current_ids_by_norm_url.get(url) if not wid: @@ -348,7 +324,7 @@ async def _update_workers_in_router_replica( 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) + logger.debug("Failed to remove worker %s, continuing with others", url) def _vllm_kv_role_to_worker_type(kv_role: str) -> str: @@ -359,33 +335,33 @@ def _vllm_kv_role_to_worker_type(kv_role: str) -> str: return "regular" -def _is_expected_grpc_discovery_error(error: Exception) -> bool: +def _is_expected_grpc_error(error: grpc.aio.AioRpcError) -> bool: """Expected while a gRPC worker is still starting or the wrong stub is probed.""" - if isinstance(error, grpc.aio.AioRpcError): - return error.code() in ( - grpc.StatusCode.UNAVAILABLE, - grpc.StatusCode.DEADLINE_EXCEEDED, - grpc.StatusCode.UNIMPLEMENTED, - ) - return False + return error.code() in ( + grpc.StatusCode.UNAVAILABLE, + grpc.StatusCode.DEADLINE_EXCEEDED, + grpc.StatusCode.UNIMPLEMENTED, + ) -async def _probe_http_worker(client: AsyncClient, *, worker_url: str) -> _WorkerPayloadResult: +async def _probe_http_worker(client: AsyncClient, *, address: str) -> Optional[_Worker]: + # The request goes over the tunnel, `worker_url` is the address the router itself dials. + worker_url = f"http://{address}" try: data = await _request_json_limited( client, "GET", - f"{_ROUTER_HTTP}/server_info", + f"{_HTTP_BASE_URL}/server_info", max_response_bytes=_MAX_SERVER_INFO_RESPONSE_BYTES, ok_statuses={200}, ) if isinstance(data, dict): if data.get("status") != "ready": - return {"status": "not_ready", "worker": None} + return None mode = data.get("disaggregation_mode", "") if mode == "prefill": bootstrap_port = data.get("disaggregation_bootstrap_port") - worker: _TargetWorker = { + worker: _Worker = { "url": worker_url, "worker_type": "prefill", "connection_mode": "http", @@ -393,38 +369,25 @@ async def _probe_http_worker(client: AsyncClient, *, worker_url: str) -> _Worker } if bootstrap_port is not None: worker["bootstrap_port"] = bootstrap_port - return {"status": "ready", "worker": worker} + return worker if mode == "decode": return { - "status": "ready", - "worker": { - "url": worker_url, - "worker_type": "decode", - "connection_mode": "http", - "runtime_type": "sglang", - }, - } - return { - "status": "ready", - "worker": { "url": worker_url, - "worker_type": "regular", + "worker_type": "decode", "connection_mode": "http", "runtime_type": "sglang", - }, + } + return { + "url": worker_url, + "worker_type": "regular", + "connection_mode": "http", + "runtime_type": "sglang", } except _ResponseTooLargeError: logger.warning("server_info response too large for worker %s", worker_url) - except RemoteProtocolError as e: - logger.debug("HTTP server_info not available for worker %s: %r", worker_url, e) - except Exception as e: - logger.exception("Could not fetch server_info for worker %s: %r", worker_url, e) - return {"status": "not_ready", "worker": None} - - -async def _get_http_worker(job_model: JobModel, *, worker_url: str) -> _WorkerPayloadResult: - async with get_service_replica_client(job_model) as client: - return await _probe_http_worker(client, worker_url=worker_url) + except RequestError as e: + logger.debug("Could not fetch server_info for worker %s: %r", worker_url, e) + return None async def _get_grpc_server_info( @@ -437,33 +400,18 @@ async def _get_grpc_server_info( else: stub = vllm_engine_pb2_grpc.VllmEngineStub(channel) request = vllm_engine_pb2.GetServerInfoRequest() - return await stub.GetServerInfo(request, timeout=_GRPC_DISCOVERY_TIMEOUT) - - -async def _discover_grpc_server_info( - channel: grpc.aio.Channel, -) -> tuple[Optional[_RuntimeType], Optional[Any]]: - # Bootstrap only: router workers list has no runtime_type yet. - for runtime_type in _GRPC_RUNTIME_TYPES: - try: - response = await _get_grpc_server_info(channel, runtime_type) - except Exception as e: - if _is_expected_grpc_discovery_error(e): - continue - raise - return runtime_type, response - return None, None + return await stub.GetServerInfo(request, timeout=_GRPC_TIMEOUT) def _grpc_server_info_to_worker( worker_url: str, runtime_type: _RuntimeType, response: Any, -) -> _TargetWorker: +) -> _Worker: if runtime_type == "vllm": kv_role = response.kv_role or "" kv_connector = response.kv_connector or "" - worker: _TargetWorker = { + worker: _Worker = { "url": worker_url, "connection_mode": "grpc", "runtime_type": runtime_type, @@ -498,81 +446,67 @@ def _grpc_server_info_to_worker( async def _probe_grpc_worker( channel: grpc.aio.Channel, *, - worker_url: str, + address: str, runtime_type: Optional[_RuntimeType] = None, -) -> _WorkerPayloadResult: - if runtime_type is not None: +) -> Optional[_Worker]: + # The RPC goes over the tunnel, `worker_url` is the address the router itself dials. + worker_url = f"grpc://{address}" + runtime_types: tuple[_RuntimeType, ...] + if runtime_type is None: + # Bootstrap only: router workers list has no runtime_type yet, should try all + runtime_types = _RUNTIME_TYPES + else: + runtime_types = (runtime_type,) + for runtime_type in runtime_types: try: response = await _get_grpc_server_info(channel, runtime_type) - except Exception as e: - if _is_expected_grpc_discovery_error(e): - logger.debug("gRPC worker %s not ready (GetServerInfo)", worker_url) - return {"status": "not_ready", "worker": None} + break + except grpc.aio.AioRpcError as e: + if _is_expected_grpc_error(e): + continue raise else: - runtime_type, response = await _discover_grpc_server_info(channel) - if runtime_type is None or response is None: - logger.debug("gRPC worker %s not ready (GetServerInfo)", worker_url) - return {"status": "not_ready", "worker": None} - - worker = _grpc_server_info_to_worker(worker_url, runtime_type, response) - return {"status": "ready", "worker": worker} - - -async def _get_grpc_worker( - job_model: JobModel, - *, - worker_url: str, - runtime_type: Optional[_RuntimeType] = None, -) -> _WorkerPayloadResult: - try: - async with get_service_replica_grpc_client(job_model) as channel: - return await _probe_grpc_worker( - channel, worker_url=worker_url, runtime_type=runtime_type - ) - except Exception as e: - logger.exception( - "Could not fetch gRPC GetServerInfo for worker %s: %r", - worker_url, - e, - ) - return {"status": "not_ready", "worker": None} + logger.debug("gRPC worker %s not ready (GetServerInfo)", worker_url) + return None + return _grpc_server_info_to_worker(worker_url, runtime_type, response) async def _get_worker( job_model: JobModel, *, - http_worker_url: str, - grpc_worker_url: str, + address: str, connection_mode: Optional[_ConnectionMode] = None, runtime_type: Optional[_RuntimeType] = None, -) -> _WorkerPayloadResult: - if connection_mode == "grpc": - return await _get_grpc_worker( - job_model, worker_url=grpc_worker_url, runtime_type=runtime_type - ) - if connection_mode == "http": - return await _get_http_worker(job_model, worker_url=http_worker_url) - # Router workers list is empty and no connection_mode discovered. - async with get_service_replica_tunnel(job_model) as uds_path: - async with get_service_replica_http_client_over_uds(uds_path) as client: - result = await _probe_http_worker(client, worker_url=http_worker_url) - if result["status"] == "ready": - return result - async with get_service_replica_grpc_channel_over_uds(uds_path) as channel: - try: - return await _probe_grpc_worker( - channel, - worker_url=grpc_worker_url, - runtime_type=runtime_type, - ) - except Exception as e: - logger.exception( - "Could not fetch gRPC GetServerInfo for worker %s: %r", - grpc_worker_url, - e, - ) - return {"status": "not_ready", "worker": None} +) -> Optional[_Worker]: + connection_modes: tuple[_ConnectionMode, ...] + if connection_mode is None: + # No connection_mode discovered -- should probe all + connection_modes = _CONNECTION_MODES + else: + connection_modes = (connection_mode,) + try: + async with get_service_replica_tunnel(job_model) as uds_path: + for connection_mode in connection_modes: + if connection_mode == "grpc": + async with get_service_replica_grpc_channel_over_uds(uds_path) as channel: + worker = await _probe_grpc_worker( + channel, address=address, runtime_type=runtime_type + ) + elif connection_mode == "http": + async with get_service_replica_http_client_over_uds(uds_path) as client: + worker = await _probe_http_worker(client, address=address) + if worker is not None: + return worker + except SSHError as e: + # An unreachable worker is reported as not ready rather than aborting the sync, so that + # one dead replica cannot hold back registration of the healthy ones. The cost is that a + # transient failure deregisters a healthy worker until the next sync re-adds it. + # TODO: `_update_workers_in_router_replica` cannot tell "not serving" from "could not be + # reached" -- both mean "absent from the target list", hence "remove". A third, unknown + # outcome should be excluded from both `to_add` and `to_remove`, leaving an unreachable + # worker as the router last saw it. + logger.warning("%s: failed to connect to worker replica: %r", fmt(job_model), e) + return None async def _build_target_workers( @@ -582,8 +516,8 @@ async def _build_target_workers( *, connection_mode: Optional[_ConnectionMode] = None, runtime_type: Optional[_RuntimeType] = None, -) -> List[_TargetWorker]: - workers: List[_TargetWorker] = [] +) -> List[_Worker]: + workers: List[_Worker] = [] config = run_spec.configuration if not isinstance(config, ServiceConfiguration): return workers @@ -601,28 +535,21 @@ async def _build_target_workers( jpd = get_job_provisioning_data(job) if jpd is None: continue - ip = jpd.internal_ip or jpd.hostname - if not ip: + hostname = jpd.internal_ip or jpd.hostname + if not hostname: continue job_spec = get_job_spec(job) port = get_service_port(job_spec, config) - http_worker_url = f"http://{ip}:{port}" - grpc_worker_url = f"grpc://{ip}:{port}" - result = await _get_worker( + worker = await _get_worker( job, - http_worker_url=http_worker_url, - grpc_worker_url=grpc_worker_url, + address=f"{hostname}:{port}", connection_mode=connection_mode, runtime_type=runtime_type, ) - if result["status"] == "ready" and result["worker"]: - workers.append(result["worker"]) - elif result["status"] == "not_ready": - logger.debug( - "Worker not ready http=%s grpc=%s", - http_worker_url, - grpc_worker_url, - ) + if worker is not None: + workers.append(worker) + else: + logger.debug("%s: worker replica not ready", fmt(job)) return workers @@ -644,9 +571,20 @@ async def sync_router_workers_for_run_model(run_model: RunModel) -> None: router_group.name, ) return + # A tunnel is opened here for the router, and inside `_build_target_workers` for every + # worker. Only the router being unreachable aborts the sync -- without a client there is + # nothing to reconcile against. An unreachable worker is skipped instead, see `_get_worker`. try: async with get_service_replica_client(router_job) as client: current_workers = await _get_router_workers(client) + if current_workers is None: + logger.debug( + "%s: failed to get current workers from the router in group %s," + " skipping worker sync", + fmt(run_model), + router_group.name, + ) + return # connection_mode can be grpc or http, runtime_type can be sglang or vllm. connection_mode = _get_connection_mode_from_workers(current_workers) runtime_type = _get_runtime_type_from_workers(current_workers) @@ -664,8 +602,14 @@ async def sync_router_workers_for_run_model(run_model: RunModel) -> None: client, target_workers, current_workers=current_workers ) except SSHError as e: + # Only the router's own tunnel reaches here, worker tunnels are handled in + # `_get_worker`. Warning is the right level: a job only reaches `RUNNING` after the + # server has talked to its runner over SSH, so an unreachable replica is always a + # regression, never a replica that has not started yet. logger.warning( "%s: failed to sync workers with router: %r", fmt(router_job), e, ) + except Exception: + logger.exception("%s: unexpected error when syncing workers with router", fmt(run_model)) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_service_router_worker_sync.py b/src/tests/_internal/server/background/pipeline_tasks/test_service_router_worker_sync.py index 5827b14966..26633a8baf 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_service_router_worker_sync.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_service_router_worker_sync.py @@ -1,13 +1,15 @@ import asyncio +import logging import uuid from datetime import timedelta -from unittest.mock import AsyncMock, Mock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest from sqlalchemy.ext.asyncio import AsyncSession +from dstack._internal.core.errors import SSHError from dstack._internal.core.models.configurations import parse_run_configuration -from dstack._internal.core.models.runs import RunStatus +from dstack._internal.core.models.runs import JobStatus, RunStatus from dstack._internal.server.background.pipeline_tasks.service_router_worker_sync import ( ServiceRouterWorkerSyncFetcher, ServiceRouterWorkerSyncPipeline, @@ -15,13 +17,19 @@ ServiceRouterWorkerSyncWorker, ) from dstack._internal.server.models import RunModel, ServiceRouterWorkerSyncModel +from dstack._internal.server.services.runs import router_worker_sync from dstack._internal.server.testing.common import ( + create_instance, + create_job, create_project, create_repo, create_run, create_user, get_run_spec, ) +from dstack._internal.server.testing.common import ( + get_job_provisioning_data as make_job_provisioning_data, +) from dstack._internal.utils.common import get_current_datetime @@ -451,3 +459,113 @@ async def test_process_calls_sync_and_unlocks_on_success( assert sync_row.lock_expires_at is None assert sync_row.lock_owner is None assert sync_row.last_processed_at is not None + + async def test_process_skips_sync_when_router_replica_not_ready( + self, + test_db, + session: AsyncSession, + worker: ServiceRouterWorkerSyncWorker, + caplog: pytest.LogCaptureFixture, + ): + caplog.set_level(level=logging.DEBUG, logger=router_worker_sync.__name__) + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + status=RunStatus.RUNNING, + run_spec=_router_service_run_spec(repo.name), + ) + instance = await create_instance(session=session, project=project) + # The router replica is still starting up, the worker replica is already serving. + await create_job( + session=session, + run=run, + instance=instance, + status=JobStatus.RUNNING, + ready=False, + replica_group_name="router", + job_provisioning_data=make_job_provisioning_data(), + ) + await create_job( + session=session, + run=run, + instance=instance, + status=JobStatus.RUNNING, + ready=True, + replica_num=1, + replica_group_name="worker", + job_provisioning_data=make_job_provisioning_data(), + ) + sync_row = await _add_service_router_worker_sync_row(session, run.id) + sync_row.lock_token = uuid.uuid4() + sync_row.lock_expires_at = get_current_datetime() + timedelta(seconds=30) + sync_row.lock_owner = ServiceRouterWorkerSyncPipeline.__name__ + await session.commit() + item = _sync_row_to_pipeline_item(sync_row) + + await worker.process(item) + + assert "no ready router job in group router, skipping worker sync" in caplog.text + # The run stays eligible for the next sync attempt. + await session.refresh(sync_row) + assert sync_row.deleted is False + assert sync_row.lock_token is None + + async def test_process_logs_router_job_when_router_connection_fails( + self, + test_db, + session: AsyncSession, + worker: ServiceRouterWorkerSyncWorker, + caplog: pytest.LogCaptureFixture, + ): + caplog.set_level(level=logging.WARNING, logger=router_worker_sync.__name__) + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + status=RunStatus.RUNNING, + run_spec=_router_service_run_spec(repo.name), + ) + instance = await create_instance(session=session, project=project) + router_job = await create_job( + session=session, + run=run, + instance=instance, + status=JobStatus.RUNNING, + ready=True, + replica_group_name="router", + # `dockerized` makes `get_container_ssh_credentials()` also read + # `instance.project` and `job_runtime_data`. + job_provisioning_data=make_job_provisioning_data(dockerized=True), + ) + sync_row = await _add_service_router_worker_sync_row(session, run.id) + sync_row.lock_token = uuid.uuid4() + sync_row.lock_expires_at = get_current_datetime() + timedelta(seconds=30) + sync_row.lock_owner = ServiceRouterWorkerSyncPipeline.__name__ + await session.commit() + item = _sync_row_to_pipeline_item(sync_row) + + # Patching the tunnel itself keeps `get_container_ssh_credentials()` real, so the test + # covers the job/instance/project attributes it reads. + tunnel_mock = MagicMock() + ssh_error = SSHError("connection refused") + tunnel_mock.return_value.__aenter__ = AsyncMock(side_effect=ssh_error) + tunnel_mock.return_value.__aexit__ = AsyncMock(return_value=False) + with patch("dstack._internal.server.services.ssh.SSHTunnel", tunnel_mock): + await worker.process(item) + + assert ( + f"job({router_job.id.hex[:6]}){router_job.job_name}:" + f" failed to sync workers with router: {ssh_error!r}" in caplog.text + ) + await session.refresh(sync_row) + assert sync_row.deleted is False + assert sync_row.lock_token is None diff --git a/src/tests/_internal/server/services/runs/test_router_worker_sync.py b/src/tests/_internal/server/services/runs/test_router_worker_sync.py index 0932939942..80a9137e0e 100644 --- a/src/tests/_internal/server/services/runs/test_router_worker_sync.py +++ b/src/tests/_internal/server/services/runs/test_router_worker_sync.py @@ -1,15 +1,25 @@ -from contextlib import asynccontextmanager, contextmanager +import json +import logging +from contextlib import contextmanager from pathlib import Path +from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch +import grpc +import httpx import pytest +from httpx import AsyncClient +from dstack._internal.core.errors import SSHError +from dstack._internal.server.services.runs import router_worker_sync from dstack._internal.server.services.runs.router_worker_sync import ( _get_connection_mode_from_workers, - _get_grpc_worker, + _get_router_workers, _get_runtime_type_from_workers, _get_worker, _grpc_server_info_to_worker, + _probe_grpc_worker, + _probe_http_worker, ) @@ -76,10 +86,127 @@ def test_sglang_prefill(self): } +def _router_client(handler) -> AsyncClient: + return AsyncClient(transport=httpx.MockTransport(handler)) + + +def _json_response(status_code: int, payload) -> httpx.Response: + return httpx.Response(status_code, content=json.dumps(payload).encode()) + + +@pytest.mark.asyncio +class TestGetRouterWorkers: + """ + `[]` must mean "the router really has no workers", so every response that does not carry a + usable worker list has to come back as `None`. Otherwise the caller reconciles against a + fabricated empty list and re-adds every worker, see + https://github.com/dstackai/dstack/issues/4300. + """ + + async def test_returns_workers(self): + payload = {"workers": [{"id": "1", "url": "http://10.0.0.1:8000"}]} + async with _router_client(lambda _: _json_response(200, payload)) as client: + assert await _get_router_workers(client) == payload["workers"] + + async def test_returns_empty_list_when_router_has_no_workers(self): + async with _router_client(lambda _: _json_response(200, {"workers": []})) as client: + assert await _get_router_workers(client) == [] + + async def test_returns_none_on_unexpected_status(self): + async with _router_client(lambda _: _json_response(500, {"workers": []})) as client: + assert await _get_router_workers(client) is None + + async def test_returns_none_on_unparsable_body(self): + async with _router_client(lambda _: httpx.Response(200, content=b"not json")) as client: + assert await _get_router_workers(client) is None + + async def test_returns_none_when_body_is_not_an_object(self): + async with _router_client(lambda _: _json_response(200, ["a"])) as client: + assert await _get_router_workers(client) is None + + async def test_returns_none_when_workers_key_is_missing(self): + async with _router_client(lambda _: _json_response(200, {})) as client: + assert await _get_router_workers(client) is None + + async def test_returns_none_when_workers_is_not_an_array(self): + async with _router_client(lambda _: _json_response(200, {"workers": "oops"})) as client: + assert await _get_router_workers(client) is None + + async def test_returns_none_on_request_error(self, caplog: pytest.LogCaptureFixture): + def handler(_): + # What an SSH tunnel to a port nobody listens on yet produces. + raise httpx.ReadError("") + + caplog.set_level(level=logging.DEBUG, logger=router_worker_sync.__name__) + async with _router_client(handler) as client: + assert await _get_router_workers(client) is None + # A router that has not bound its port yet is expected, not an error. + assert "Router /workers not ready yet" in caplog.text + assert not [r for r in caplog.records if r.levelno > logging.DEBUG] + + +@pytest.mark.asyncio +class TestProbeHttpWorker: + """ + The probe talks to the replica over the tunnel, but the `url` it reports is the address the + router dials. It must stay byte-identical to what the router echoes back in `/workers`, + otherwise `_update_workers_in_router_replica` re-registers every worker on each sync. + """ + + async def test_regular_worker(self): + async with _router_client(lambda _: _json_response(200, {"status": "ready"})) as client: + assert await _probe_http_worker(client, address="10.0.0.1:8000") == { + "url": "http://10.0.0.1:8000", + "worker_type": "regular", + "connection_mode": "http", + "runtime_type": "sglang", + } + + async def test_prefill_worker(self): + payload = { + "status": "ready", + "disaggregation_mode": "prefill", + "disaggregation_bootstrap_port": 8998, + } + async with _router_client(lambda _: _json_response(200, payload)) as client: + assert await _probe_http_worker(client, address="10.0.0.1:8000") == { + "url": "http://10.0.0.1:8000", + "worker_type": "prefill", + "connection_mode": "http", + "runtime_type": "sglang", + "bootstrap_port": 8998, + } + + async def test_decode_worker(self): + payload = {"status": "ready", "disaggregation_mode": "decode"} + async with _router_client(lambda _: _json_response(200, payload)) as client: + worker = await _probe_http_worker(client, address="10.0.0.1:8000") + assert worker is not None + assert worker["worker_type"] == "decode" + + async def test_returns_none_when_not_ready(self): + payload = {"status": "starting"} + async with _router_client(lambda _: _json_response(200, payload)) as client: + assert await _probe_http_worker(client, address="10.0.0.1:8000") is None + + async def test_returns_none_on_request_error(self, caplog: pytest.LogCaptureFixture): + def handler(_): + # A gRPC-only worker, or one that has not bound its port yet. + raise httpx.ReadError("") + + caplog.set_level(level=logging.DEBUG, logger=router_worker_sync.__name__) + async with _router_client(handler) as client: + assert await _probe_http_worker(client, address="10.0.0.1:8000") is None + # Repeats every sync until the worker is up, so it must not be logged as an error, + # see https://github.com/dstackai/dstack/issues/4300. + assert "Could not fetch server_info for worker http://10.0.0.1:8000" in caplog.text + assert not [r for r in caplog.records if r.levelno > logging.DEBUG] + + @contextmanager -def _fake_vllm_grpc_proto(*, server_info: MagicMock): +def _fake_vllm_grpc_proto(*, server_info=None, error: Optional[Exception] = None): stub = MagicMock() - stub.GetServerInfo = AsyncMock(return_value=server_info) + stub.GetServerInfo = AsyncMock(return_value=server_info, side_effect=error) pb2 = MagicMock(GetServerInfoRequest=MagicMock(return_value="req")) pb2_grpc = MagicMock(VllmEngineStub=MagicMock(return_value=stub)) with ( @@ -96,9 +223,9 @@ def _fake_vllm_grpc_proto(*, server_info: MagicMock): @contextmanager -def _fake_sglang_grpc_proto(*, server_info: MagicMock): +def _fake_sglang_grpc_proto(*, server_info=None, error: Optional[Exception] = None): stub = MagicMock() - stub.GetServerInfo = AsyncMock(return_value=server_info) + stub.GetServerInfo = AsyncMock(return_value=server_info, side_effect=error) pb2 = MagicMock(GetServerInfoRequest=MagicMock(return_value="req")) pb2_grpc = MagicMock(SglangSchedulerStub=MagicMock(return_value=stub)) with ( @@ -114,195 +241,232 @@ def _fake_sglang_grpc_proto(*, server_info: MagicMock): yield -@pytest.mark.asyncio -async def test_get_grpc_worker_ready(): - job = MagicMock() - channel = MagicMock() - - @asynccontextmanager - async def _fake_grpc_client(_job): - yield channel - - server_info = MagicMock(kv_role="kv_producer", kv_connector="NixlConnector") - - with ( - _fake_vllm_grpc_proto(server_info=server_info), - patch( - "dstack._internal.server.services.runs.router_worker_sync.get_service_replica_grpc_client", - _fake_grpc_client, - ), - ): - result = await _get_grpc_worker( - job, - worker_url="grpc://10.0.0.1:50051", - runtime_type="vllm", - ) - - assert result["status"] == "ready" - assert result["worker"] == { - "url": "grpc://10.0.0.1:50051", - "worker_type": "prefill", - "connection_mode": "grpc", - "runtime_type": "vllm", - "kv_connector": "NixlConnector", - "kv_role": "kv_producer", - } - - -@pytest.mark.asyncio -async def test_get_grpc_worker_not_ready_on_error(): - job = MagicMock() - - @asynccontextmanager - async def _failing_client(_job): - raise OSError("ssh failed") - yield # pragma: no cover - - with patch( - "dstack._internal.server.services.runs.router_worker_sync.get_service_replica_grpc_client", - _failing_client, - ): - result = await _get_grpc_worker(job, worker_url="grpc://10.0.0.1:50051") - - assert result == {"status": "not_ready", "worker": None} +def _rpc_error(code: grpc.StatusCode) -> grpc.aio.AioRpcError: + return grpc.aio.AioRpcError(code, grpc.aio.Metadata(), grpc.aio.Metadata(), details=code.name) @pytest.mark.asyncio -async def test_get_grpc_worker_sglang_bootstrap(): - job = MagicMock() - channel = MagicMock() - sglang_server_info = MagicMock(server_args=MagicMock()) - - @asynccontextmanager - async def _fake_grpc_client(_job): - yield channel - - with ( - _fake_sglang_grpc_proto(server_info=sglang_server_info), - patch( - "dstack._internal.server.services.runs.router_worker_sync.MessageToDict", - return_value={ - "disaggregation_mode": "prefill", - "disaggregation_bootstrap_port": 8998, - }, - ), - patch( - "dstack._internal.server.services.runs.router_worker_sync" - ".get_service_replica_grpc_client", - _fake_grpc_client, - ), - ): - result = await _get_grpc_worker(job, worker_url="grpc://10.0.0.1:8000") - - assert result["status"] == "ready" - assert result["worker"] == { - "url": "grpc://10.0.0.1:8000", - "worker_type": "prefill", - "connection_mode": "grpc", - "runtime_type": "sglang", - "bootstrap_port": 8998, - } - - -@pytest.mark.asyncio -async def test_get_worker_grpc_preference_skips_http(): - job = MagicMock() - grpc_not_ready = {"status": "not_ready", "worker": None} - - with ( - patch( - "dstack._internal.server.services.runs.router_worker_sync._get_grpc_worker", - new_callable=AsyncMock, - return_value=grpc_not_ready, - ) as grpc_mock, - patch( - "dstack._internal.server.services.runs.router_worker_sync._get_http_worker", - new_callable=AsyncMock, - ) as http_mock, - ): - result = await _get_worker( - job, - http_worker_url="http://10.0.0.1:8000", - grpc_worker_url="grpc://10.0.0.1:8000", - connection_mode="grpc", - ) - - assert result == grpc_not_ready - grpc_mock.assert_awaited_once() - http_mock.assert_not_awaited() - +class TestProbeGrpcWorker: + async def test_known_runtime_type(self): + server_info = MagicMock(kv_role="kv_producer", kv_connector="NixlConnector") + with _fake_vllm_grpc_proto(server_info=server_info): + worker = await _probe_grpc_worker( + MagicMock(), address="10.0.0.1:50051", runtime_type="vllm" + ) + assert worker == { + "url": "grpc://10.0.0.1:50051", + "worker_type": "prefill", + "connection_mode": "grpc", + "runtime_type": "vllm", + "kv_connector": "NixlConnector", + "kv_role": "kv_producer", + } -@pytest.mark.asyncio -async def test_get_worker_bootstrap_uses_single_tunnel(): - job = MagicMock() - uds_path = Path("/tmp/replica.sock") - grpc_ready: dict = { - "status": "ready", - "worker": { + async def test_bootstrap_tries_sglang_first(self): + with ( + _fake_sglang_grpc_proto(server_info=MagicMock(server_args=MagicMock())), + patch( + "dstack._internal.server.services.runs.router_worker_sync.MessageToDict", + return_value={ + "disaggregation_mode": "prefill", + "disaggregation_bootstrap_port": 8998, + }, + ), + ): + worker = await _probe_grpc_worker(MagicMock(), address="10.0.0.1:8000") + assert worker == { "url": "grpc://10.0.0.1:8000", "worker_type": "prefill", "connection_mode": "grpc", - "runtime_type": "vllm", - }, - } + "runtime_type": "sglang", + "bootstrap_port": 8998, + } + + async def test_bootstrap_falls_back_to_vllm(self): + # A vLLM worker does not implement the SGLang scheduler service. + with ( + _fake_sglang_grpc_proto(error=_rpc_error(grpc.StatusCode.UNIMPLEMENTED)), + _fake_vllm_grpc_proto( + server_info=MagicMock(kv_role="kv_consumer", kv_connector="NixlConnector") + ), + ): + worker = await _probe_grpc_worker(MagicMock(), address="10.0.0.1:8000") + assert worker is not None + assert worker["runtime_type"] == "vllm" + assert worker["worker_type"] == "decode" + + @pytest.mark.parametrize( + "code", + [ + # No listener yet, or an SSH tunnel to a dead remote port, or an HTTP-only worker. + grpc.StatusCode.UNAVAILABLE, + grpc.StatusCode.DEADLINE_EXCEEDED, + grpc.StatusCode.UNIMPLEMENTED, + ], + ) + async def test_returns_none_on_expected_error(self, code: grpc.StatusCode): + with _fake_vllm_grpc_proto(error=_rpc_error(code)): + worker = await _probe_grpc_worker( + MagicMock(), address="10.0.0.1:8000", runtime_type="vllm" + ) + assert worker is None + + async def test_reraises_unexpected_error(self): + error = _rpc_error(grpc.StatusCode.PERMISSION_DENIED) + with _fake_vllm_grpc_proto(error=error), pytest.raises(grpc.aio.AioRpcError) as exc_info: + await _probe_grpc_worker(MagicMock(), address="10.0.0.1:8000", runtime_type="vllm") + assert exc_info.value is error + + async def test_returns_none_when_no_runtime_type_matches(self): + with ( + _fake_sglang_grpc_proto(error=_rpc_error(grpc.StatusCode.UNAVAILABLE)), + _fake_vllm_grpc_proto(error=_rpc_error(grpc.StatusCode.UNAVAILABLE)), + ): + worker = await _probe_grpc_worker(MagicMock(), address="10.0.0.1:8000") + assert worker is None + + +_HTTP_WORKER = { + "url": "http://10.0.0.1:8000", + "worker_type": "regular", + "connection_mode": "http", + "runtime_type": "sglang", +} +_GRPC_WORKER = { + "url": "grpc://10.0.0.1:8000", + "worker_type": "prefill", + "connection_mode": "grpc", + "runtime_type": "vllm", +} - tunnel_cm = AsyncMock() - tunnel_cm.__aenter__.return_value = uds_path - tunnel_cm.__aexit__.return_value = None - http_cm = AsyncMock() - http_cm.__aenter__.return_value = MagicMock() - http_cm.__aexit__.return_value = None +def _async_cm(enter_value=None, enter_error: Optional[Exception] = None) -> AsyncMock: + cm = AsyncMock() + cm.__aenter__ = AsyncMock(return_value=enter_value, side_effect=enter_error) + cm.__aexit__ = AsyncMock(return_value=False) + return cm - grpc_cm = AsyncMock() - grpc_cm.__aenter__.return_value = MagicMock() - grpc_cm.__aexit__.return_value = None +@contextmanager +def _fake_replica_transports( + *, + uds_path: Path = Path("/tmp/replica.sock"), + tunnel_error: Optional[Exception] = None, + http_worker=None, + grpc_worker=None, +): + """Patch the tunnel and both probes, leaving `_get_worker`'s own logic intact.""" + mocks = MagicMock() with ( patch( "dstack._internal.server.services.runs.router_worker_sync.get_service_replica_tunnel", - return_value=tunnel_cm, - ) as tunnel_mock, + return_value=_async_cm(uds_path, tunnel_error), + ) as mocks.tunnel, patch( "dstack._internal.server.services.runs.router_worker_sync" ".get_service_replica_http_client_over_uds", - return_value=http_cm, - ) as http_over_uds_mock, + return_value=_async_cm(MagicMock()), + ) as mocks.http_client, patch( "dstack._internal.server.services.runs.router_worker_sync" ".get_service_replica_grpc_channel_over_uds", - return_value=grpc_cm, - ) as grpc_over_uds_mock, + return_value=_async_cm(MagicMock()), + ) as mocks.grpc_channel, patch( "dstack._internal.server.services.runs.router_worker_sync._probe_http_worker", new_callable=AsyncMock, - return_value={"status": "not_ready", "worker": None}, - ) as http_probe_mock, + return_value=http_worker, + ) as mocks.http_probe, patch( "dstack._internal.server.services.runs.router_worker_sync._probe_grpc_worker", new_callable=AsyncMock, - return_value=grpc_ready, - ) as grpc_probe_mock, - patch( - "dstack._internal.server.services.runs.router_worker_sync._get_http_worker", - new_callable=AsyncMock, - ) as get_http_mock, - patch( - "dstack._internal.server.services.runs.router_worker_sync._get_grpc_worker", - new_callable=AsyncMock, - ) as get_grpc_mock, + return_value=grpc_worker, + ) as mocks.grpc_probe, + ): + yield mocks + + +@pytest.mark.asyncio +class TestGetWorker: + async def test_connection_mode_grpc_skips_http(self): + with _fake_replica_transports(grpc_worker=_GRPC_WORKER) as mocks: + worker = await _get_worker( + MagicMock(), + address="10.0.0.1:8000", + connection_mode="grpc", + ) + assert worker == _GRPC_WORKER + mocks.grpc_probe.assert_awaited_once() + mocks.http_probe.assert_not_awaited() + + async def test_connection_mode_http_skips_grpc(self): + with _fake_replica_transports(http_worker=_HTTP_WORKER) as mocks: + worker = await _get_worker( + MagicMock(), + address="10.0.0.1:8000", + connection_mode="http", + ) + assert worker == _HTTP_WORKER + mocks.http_probe.assert_awaited_once() + mocks.grpc_probe.assert_not_awaited() + + async def test_bootstrap_probes_http_first(self): + # An HTTP worker must not pay for two gRPC `GetServerInfo` timeouts first. + with _fake_replica_transports(http_worker=_HTTP_WORKER, grpc_worker=_GRPC_WORKER) as mocks: + worker = await _get_worker( + MagicMock(), + address="10.0.0.1:8000", + ) + assert worker == _HTTP_WORKER + mocks.http_probe.assert_awaited_once() + mocks.grpc_probe.assert_not_awaited() + + async def test_bootstrap_falls_back_to_grpc_over_one_tunnel(self): + job = MagicMock() + uds_path = Path("/tmp/replica.sock") + with _fake_replica_transports(uds_path=uds_path, grpc_worker=_GRPC_WORKER) as mocks: + worker = await _get_worker( + job, + address="10.0.0.1:8000", + ) + assert worker == _GRPC_WORKER + mocks.tunnel.assert_called_once_with(job) + mocks.http_client.assert_called_once_with(uds_path) + mocks.grpc_channel.assert_called_once_with(uds_path) + mocks.http_probe.assert_awaited_once() + mocks.grpc_probe.assert_awaited_once() + + async def test_returns_none_when_no_mode_reports_ready(self): + with _fake_replica_transports() as mocks: + worker = await _get_worker( + MagicMock(), + address="10.0.0.1:8000", + ) + assert worker is None + mocks.http_probe.assert_awaited_once() + mocks.grpc_probe.assert_awaited_once() + + async def test_unreachable_worker_is_skipped_not_raised( + self, caplog: pytest.LogCaptureFixture ): - result = await _get_worker( - job, - http_worker_url="http://10.0.0.1:8000", - grpc_worker_url="grpc://10.0.0.1:8000", - ) - - assert result == grpc_ready - tunnel_mock.assert_called_once_with(job) - http_over_uds_mock.assert_called_once_with(uds_path) - grpc_over_uds_mock.assert_called_once_with(uds_path) - http_probe_mock.assert_awaited_once() - grpc_probe_mock.assert_awaited_once() - get_http_mock.assert_not_awaited() - get_grpc_mock.assert_not_awaited() + # One dead replica must not abort the sync for the healthy ones. + caplog.set_level(level=logging.WARNING, logger=router_worker_sync.__name__) + ssh_error = SSHError("connection refused") + with _fake_replica_transports(tunnel_error=ssh_error) as mocks: + worker = await _get_worker( + MagicMock(), + address="10.0.0.1:8000", + ) + assert worker is None + mocks.http_probe.assert_not_awaited() + mocks.grpc_probe.assert_not_awaited() + assert f"failed to connect to worker replica: {ssh_error!r}" in caplog.text + + async def test_unexpected_tunnel_error_propagates(self): + # Only `SSHError` means "unreachable"; anything else is a bug and must not be swallowed. + with _fake_replica_transports(tunnel_error=RuntimeError("boom")): + with pytest.raises(RuntimeError, match="boom"): + await _get_worker( + MagicMock(), + address="10.0.0.1:8000", + )