diff --git a/pyproject.toml b/pyproject.toml index ae7440a..ab9e369 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -153,7 +153,7 @@ dev = [ "ruff>=0.14.10", "types-colorama>=0.4.15.20250801", "types-docker>=7.1.0.20251202", - "types-grpcio>=1.0.0.20251009", + "types-grpcio>=1.0.0.20260614", "types-pexpect>=4.9.0.20250916", "types-protobuf>=6.32.1.20251210", "types-psutil>=7.1.3.20251211", diff --git a/src/server/clients/redis.py b/src/server/clients/redis.py index 86136b9..408b02f 100644 --- a/src/server/clients/redis.py +++ b/src/server/clients/redis.py @@ -3,7 +3,7 @@ import socket import ssl from collections.abc import Awaitable, Iterable -from typing import Any +from typing import Any, cast from urllib.parse import urlparse, urlunparse import redis @@ -18,26 +18,17 @@ def _keepalive_kwargs() -> dict[str, Any]: """Connection kwargs that keep an idle Redis socket alive. - A pub/sub SUBSCRIBE that receives no messages is an idle TCP connection. - Intermediaries (load balancers, NAT gateways) cull idle server-side - connections on their own timeout, silently dropping a node's dispatch - subscription so the publisher sees zero receivers. TCP keepalive probes - keep the connection observably alive well below such timeouts, and - ``health_check_interval`` makes the client re-establish a connection it - finds dead rather than blocking on a stale socket. + An idle pub/sub SUBSCRIBE gets culled by intermediaries (load balancers, NAT + gateways), dropping a node's dispatch subscription. Keepalive probes hold it open + and ``health_check_interval`` reconnects a dead socket instead of blocking on it. """ - options: dict[int, int] = {} - for name, value in ( - ("TCP_KEEPIDLE", 60), - ("TCP_KEEPINTVL", 15), - ("TCP_KEEPCNT", 4), - ): - opt = getattr(socket, name, None) - if opt is not None: - options[opt] = value return { "socket_keepalive": True, - "socket_keepalive_options": options, + "socket_keepalive_options": { + socket.TCP_KEEPIDLE: 60, + socket.TCP_KEEPINTVL: 15, + socket.TCP_KEEPCNT: 4, + }, "health_check_interval": 30, } @@ -148,6 +139,25 @@ def ssh_connection_key(connection_id: str) -> str: SSH_CONNECTION_IDS_KEY = "ssh:connections:active" +def parse_pubsub_message(msg: dict[str, Any] | None) -> Any | None: + """Decode a redis-py pub/sub frame into its JSON payload. + + Returns ``None`` for control frames, empty payloads, and malformed JSON. Payloads + are always JSON objects, so a ``None`` return means "nothing to deliver". + """ + if msg is None or msg.get("type") != "message": + return None + raw = msg.get("data") + if raw is None: + return None + if isinstance(raw, bytes): + raw = raw.decode("utf-8", errors="ignore") + try: + return json.loads(raw) + except json.JSONDecodeError: + return None + + def iter_pubsub_messages(pubsub: PubSub) -> Iterable[Any]: """Iterate over messages from a Redis PubSub instance. @@ -157,23 +167,16 @@ def iter_pubsub_messages(pubsub: PubSub) -> Iterable[Any]: """ try: for msg in pubsub.listen(): - if msg.get("type") != "message": - continue - raw = msg.get("data") - if raw is None: - continue - if isinstance(raw, bytes): - raw = raw.decode("utf-8", errors="ignore") - try: - yield json.loads(raw) - except json.JSONDecodeError: + parsed = parse_pubsub_message(msg) + if parsed is None: continue + yield parsed except (ConnectionError, OSError): return def _sync[T](value: Awaitable[T] | T) -> T: - return value # type: ignore + return cast(T, value) def _with_redis_auth(url: str, acl_enabled: bool, username: str, password: str) -> str: @@ -280,6 +283,9 @@ def ttl(self, key: str) -> float | None: def incr(self, key: str) -> int: return int(_sync(self._control.incr(key))) + def eval(self, script: str, numkeys: int, *keys_and_args: str) -> str: + return _sync(self._control.eval(script, numkeys, *keys_and_args)) + def set_value(self, key: str, value: str) -> None: self._control.set(key, value) @@ -402,7 +408,7 @@ def flush_all(self) -> None: def _awaitable[T](value: Awaitable[T] | T) -> Awaitable[T]: - return value # type: ignore + return cast(Awaitable[T], value) class AsyncRedisClient: @@ -483,6 +489,9 @@ async def ttl(self, key: str) -> float | None: async def incr(self, key: str) -> int: return int(await _awaitable(self._control.incr(key))) + async def eval(self, script: str, numkeys: int, *keys_and_args: str) -> str: + return await _awaitable(self._control.eval(script, numkeys, *keys_and_args)) + async def set_value(self, key: str, value: str) -> None: await _awaitable(self._control.set(key, value)) diff --git a/src/server/main.py b/src/server/main.py index 9a9608f..ca9bbed 100644 --- a/src/server/main.py +++ b/src/server/main.py @@ -355,11 +355,18 @@ async def _lifespan(_: FastAPI): # --- Supervisor (all nodes with worker management) --- if SUPERVISOR is not None: await SUPERVISOR.start(system_principal) - app.state.node_id = SUPERVISOR.node_id - # Tell EventMonitor which node this server belongs to so that it can wait - # for the supervisor's SV_UNREGISTER event on shutdown. - if EVENT_MONITOR is not None: - EVENT_MONITOR.set_own_node(SUPERVISOR.node_id) + + def _on_node_id_change(new_node_id: str) -> None: + app.state.node_id = new_node_id + # Tell EventMonitor which node this server belongs to so that it can + # wait for the supervisor's SV_UNREGISTER event on shutdown. + if EVENT_MONITOR is not None: + EVENT_MONITOR.set_own_node(new_node_id) + + # Keep the node_id updated for request auth scope + shutdown + # self-identification. + _on_node_id_change(SUPERVISOR.node_id) + SUPERVISOR.add_node_id_listener(_on_node_id_change) # --- Startup reconcile --- # Runs after the supervisor handshake so this node is in NODE_REGISTRY diff --git a/src/server/supervisor/resource_manager.py b/src/server/supervisor/resource_manager.py index 95d4596..607d72b 100644 --- a/src/server/supervisor/resource_manager.py +++ b/src/server/supervisor/resource_manager.py @@ -49,11 +49,9 @@ def get_instance(cls) -> "ResourceManager": cls._instance = cls() return cls._instance - @property def total_gpu_count(self) -> int: return self._env.gpu_count - @property def available_gpu_count(self) -> int: return len(self._env.available_gpus) diff --git a/src/server/supervisor/services/command_listener.py b/src/server/supervisor/services/command_listener.py index 736c9d7..2b3761f 100644 --- a/src/server/supervisor/services/command_listener.py +++ b/src/server/supervisor/services/command_listener.py @@ -5,7 +5,7 @@ import secrets from collections.abc import Callable, Iterable from concurrent import futures -from threading import Thread +from threading import Event, Lock, Thread from pydantic import ValidationError from redis.client import PubSub @@ -19,8 +19,8 @@ from ...clients.redis import ( NODE_RESPONSE_CHANNEL, SyncRedisClient, - iter_pubsub_messages, node_cmd_channel, + parse_pubsub_message, ) from ...utils.concurrent import Sentinel, TaskReceiver from ..adapters.docker import DockerWorkerConfig @@ -36,6 +36,7 @@ _CREATE_WORKER_TIMEOUT = 600.0 _DESTROY_WORKER_TIMEOUT = 60.0 _DESTROY_WORKERS_TIMEOUT = 120.0 +_POLL_TIMEOUT_SEC = 0.25 def _cmd_receiver_loop( @@ -53,25 +54,6 @@ def send_response(resp: CommandResponse) -> None: q.put((cmd, make_handler(cmd_id))) -def _pubsub_loop( - redis: SyncRedisClient, - pubsub: PubSub, - q: queue.Queue[tuple[CommandMessage, ResponseHandler]], - logger: logging.Logger, -) -> None: - for data in iter_pubsub_messages(pubsub): - try: - cmd = CommandMessage.model_validate(data) - except ValidationError as e: - logger.error("Invalid command message: %s", e) - continue - - def send_response(resp: CommandResponse) -> None: - redis.publish_control(NODE_RESPONSE_CHANNEL, resp.model_dump_json()) - - q.put((cmd, send_response)) - - class _CommandStream: _SENTINEL = Sentinel() @@ -82,6 +64,7 @@ def __init__( cmd_receiver: TaskReceiver[CommandMessage, CommandResponse] | None, logger: logging.Logger, ): + # node_id is mutated only by the pubsub reader thread once running. self.node_id = node_id self.redis = redis self.cmd_receiver = cmd_receiver @@ -89,7 +72,11 @@ def __init__( self._cmd_queue: ( queue.Queue[tuple[CommandMessage, ResponseHandler] | Sentinel] | None ) = None - self._pubsub: PubSub | None = None + self._pubsub_running = False + + self._rebind_lock = Lock() + self._pending_node_id: str | None = None + self._rebind_applied = Event() def close(self) -> None: if self._cmd_queue is not None: @@ -97,19 +84,64 @@ def close(self) -> None: self._cmd_queue = None def rebind(self, node_id: str) -> None: - """Move the command subscription to a new node id on the live pubsub.""" + """Request moving the command subscription to a new node id. + + Records the target under a lock; the reader thread applies the actual + ``subscribe``/``unsubscribe`` between polls (mutating redis-py ``PubSub`` is not + thread-safe). ``wait_rebound`` blocks until the switch has taken effect. + """ if node_id == self.node_id: + self._rebind_applied.set() return - old_node_id = self.node_id - self.node_id = node_id - pubsub = self._pubsub - if pubsub is None: - return - pubsub.subscribe(node_cmd_channel(node_id)) - pubsub.unsubscribe(node_cmd_channel(old_node_id)) + self._rebind_applied.clear() + with self._rebind_lock: + self._pending_node_id = node_id + + def wait_rebound(self, timeout: float) -> bool: + return self._rebind_applied.wait(timeout) + + def _apply_pending_rebind(self, pubsub: PubSub, current_id: str) -> str: + with self._rebind_lock: + pending = self._pending_node_id + self._pending_node_id = None + if pending is None or pending == current_id: + return current_id + pubsub.subscribe(node_cmd_channel(pending)) + pubsub.unsubscribe(node_cmd_channel(current_id)) + self.node_id = pending + self._rebind_applied.set() self.logger.info( - "Command stream rebound from node %s to %s", old_node_id, node_id + "Command stream rebound from node %s to %s", current_id, pending ) + return pending + + def _run_pubsub( + self, + pubsub: PubSub, + cmd_queue: queue.Queue[tuple[CommandMessage, ResponseHandler] | Sentinel], + ) -> None: + def send_response(resp: CommandResponse) -> None: + self.redis.publish_control(NODE_RESPONSE_CHANNEL, resp.model_dump_json()) + + current_id = self.node_id + try: + while self._pubsub_running: + current_id = self._apply_pending_rebind(pubsub, current_id) + msg = pubsub.get_message(timeout=_POLL_TIMEOUT_SEC) + data = parse_pubsub_message(msg) + if data is None: + continue + try: + cmd = CommandMessage.model_validate(data) + except ValidationError as e: + self.logger.error("Invalid command message: %s", e) + continue + cmd_queue.put((cmd, send_response)) + except (ConnectionError, OSError): + return + except Exception as exc: + if self._pubsub_running: + self.logger.exception("Command pubsub loop error: %s", exc) def iter_stream(self) -> Iterable[tuple[CommandMessage, ResponseHandler]]: if self._cmd_queue is not None: @@ -127,10 +159,10 @@ def iter_stream(self) -> Iterable[tuple[CommandMessage, ResponseHandler]]: task_thread.start() pubsub = self.redis.subscribe_control(node_cmd_channel(self.node_id)) - self._pubsub = pubsub + self._pubsub_running = True pubsub_thread = Thread( - target=_pubsub_loop, - args=(self.redis, pubsub, cmd_queue, self.logger), + target=self._run_pubsub, + args=(pubsub, cmd_queue), name="CommandPubSubThread", daemon=True, ) @@ -144,7 +176,8 @@ def iter_stream(self) -> Iterable[tuple[CommandMessage, ResponseHandler]]: break yield item finally: - self._pubsub = None + self._pubsub_running = False + pubsub_thread.join() pubsub.close() @@ -218,12 +251,18 @@ def start(self) -> None: ) def rebind(self, node_id: str) -> None: - """Move the command subscription to a new node id without a restart.""" + """Request moving the command subscription to a new node id. + + The reader thread applies the switch; ``wait_rebound`` blocks until it has taken + effect. + """ self._node_id = node_id - stream = self._cmd_stream - if stream is not None: + if stream := self._cmd_stream: stream.rebind(node_id) + def wait_rebound(self, timeout: float) -> bool: + return stream.wait_rebound(timeout) if (stream := self._cmd_stream) else True + async def stop(self) -> None: if self._thread is None or self._cmd_stream is None: self.logger.warning("Command listener not started") diff --git a/src/server/supervisor/services/grpc_server.py b/src/server/supervisor/services/grpc_server.py index da050f4..cbcd585 100644 --- a/src/server/supervisor/services/grpc_server.py +++ b/src/server/supervisor/services/grpc_server.py @@ -1,7 +1,8 @@ import asyncio import logging -from collections.abc import AsyncIterator, Iterable +from collections.abc import AsyncIterator from pathlib import Path +from threading import Lock import grpc import grpc.aio @@ -28,14 +29,32 @@ from ..services.relay_service import RelayService from ..services.task_listener import TaskListener - -def _token_from_metadata(metadata: Iterable[tuple[str, str]]) -> WorkerTokenType | None: - for key, value in metadata: - key_l = key.lower() - if key_l == "authorization" and value.startswith("Bearer "): - return value[7:] # type: ignore - if key_l == "x-worker-token": - return value # type: ignore +# Rewrite node_id for each worker key that still exists, atomically. KEYS are +# worker keys; ARGV[1] is the new node id. Returns the count actually rewritten. +_REHOME_LUA = """ +local rehomed = 0 +for _, key in ipairs(KEYS) do + if redis.call('EXISTS', key) == 1 then + redis.call('HSET', key, 'node_id', ARGV[1]) + rehomed = rehomed + 1 + end +end +return rehomed +""" + + +def _token_from_context(context: grpc.aio.ServicerContext) -> WorkerTokenType | None: + if metadata := context.invocation_metadata(): + for key, value in metadata: + key_l = key.lower() + if isinstance(value, (bytes, bytearray)): + value = value.decode("utf-8", errors="ignore") + elif isinstance(value, memoryview): + value = value.tobytes().decode("utf-8", errors="ignore") + if key_l == "authorization" and value.startswith("Bearer "): + return WorkerTokenType(value[7:]) + if key_l == "x-worker-token": + return WorkerTokenType(value) return None @@ -86,31 +105,47 @@ def __init__( self._node_id = node_id self._node_alias = node_alias self._logger = logger + # Guards _node_id and the registry-vs-rehome window against concurrent + # RegisterWorker (grpc loop thread) and rebind_node (heartbeat thread). + self._lock = Lock() def rebind_node(self, node_id: str) -> None: """Re-home this node's workers under a new node id. Future registrations stamp the new id, and every already-registered worker's ``node_id`` field is rewritten in Redis so the dispatcher - routes tasks to them on the node's new dispatch channel. + routes tasks to them on the node's new dispatch channel. A worker whose + record no longer exists (e.g. Redis was wiped) is skipped rather than + resurrected as a partial record. """ - if node_id == self._node_id: - return - old_node_id = self._node_id - self._node_id = node_id - rehomed = 0 - for worker in self._registry.all_workers(): - worker_id = self._registry.get_worker_id(worker.token) - if worker_id is None: - continue - self._redis.hash_set(worker_key(worker_id), {"node_id": node_id}) - rehomed += 1 - self._logger.info( - "Re-homed %d worker(s) from node %s to %s", - rehomed, - old_node_id, - node_id, - ) + with self._lock: + if node_id == self._node_id: + return + old_node_id = self._node_id + self._node_id = node_id + worker_ids = [ + worker_id + for worker in self._registry.all_workers() + if (worker_id := self._registry.get_worker_id(worker.token)) is not None + ] + rehomed, skipped = self._rehome_workers(worker_ids, node_id) + self._logger.info( + "Re-homed %d worker(s) (%d skipped: no record) from node %s to %s", + rehomed, + skipped, + old_node_id, + node_id, + ) + + def _rehome_workers(self, worker_ids: list[str], node_id: str) -> tuple[int, int]: + """Rewrite node_id only for workers whose record still exists, atomically + so a worker deleted mid-rebind is skipped rather than resurrected as a + partial record. Returns (rehomed, skipped).""" + if not worker_ids: + return 0, 0 + keys = [worker_key(worker_id) for worker_id in worker_ids] + rehomed = int(self._redis.eval(_REHOME_LUA, len(keys), *keys, node_id)) + return rehomed, len(worker_ids) - rehomed async def RegisterWorker( self, @@ -123,11 +158,15 @@ async def RegisterWorker( worker_meta = _payload_from_struct(request.meta) worker_id = new_worker_id(self._redis.incr(WORKER_ID_SEQ_KEY)) worker_meta["id"] = worker_id - worker_meta["node_id"] = self._node_id worker_meta["node_alias"] = self._node_alias - self._redis.sadd(WORKERS_SET_KEY, worker_id) - self._redis.hash_set(worker_key(worker_id), worker_meta) - self._registry.set_worker_id(worker.token, worker_id) + # Stamp node_id, persist the record, and insert into the registry as one unit so + # a concurrent rebind_node either sees this worker in its snapshot or stamps it + # with the new id. + with self._lock: + worker_meta["node_id"] = self._node_id + self._redis.sadd(WORKERS_SET_KEY, worker_id) + self._redis.hash_set(worker_key(worker_id), worker_meta) + self._registry.set_worker_id(worker.token, worker_id) self._task_listener.add_worker(worker_id) try: worker.set_worker_id(worker_id) @@ -227,18 +266,16 @@ async def PushLogs( def _get_worker_from_context( self, context: grpc.aio.ServicerContext ) -> WorkerAdapter | None: - token = _token_from_metadata(context.invocation_metadata()) # type: ignore - if token is None: - return None - return self._registry.try_get(token) + if token := _token_from_context(context): + return self._registry.try_get(token) + return None def _get_worker_id_from_context( self, context: grpc.aio.ServicerContext ) -> str | None: - token = _token_from_metadata(context.invocation_metadata()) # type: ignore - if token is None: - return None - return self._registry.get_worker_id(token) + if token := _token_from_context(context): + return self._registry.get_worker_id(token) + return None _GRPC_MAX_MSG_BYTES = 1024 * 1024 * 1024 # 1 GB @@ -264,10 +301,6 @@ def __init__( ) self._listen_addr = f"{host}:{port}" - def rebind_node(self, node_id: str) -> None: - """Re-home registered workers under a new node id.""" - self._servicer.rebind_node(node_id) - async def start(self) -> None: if self._server is not None: self._logger.warning("Server gRPC server already started") @@ -314,3 +347,7 @@ async def stop(self, grace: float = 5.0) -> None: await self._server.stop(grace) self._server = None self._logger.info("Server gRPC server stopped") + + def rebind_node(self, node_id: str) -> None: + """Re-home registered workers under a new node id.""" + self._servicer.rebind_node(node_id) diff --git a/src/server/supervisor/services/lifecycle.py b/src/server/supervisor/services/lifecycle.py index 9fa8479..88efcdf 100644 --- a/src/server/supervisor/services/lifecycle.py +++ b/src/server/supervisor/services/lifecycle.py @@ -60,12 +60,7 @@ def node_id(self) -> str: return self._node_id def set_reregister_callback(self, callback: Callable[[str], None]) -> None: - """Register a hook invoked with the new node id after a re-register. - - Lets components that captured the original node id (dispatch/command - subscriptions, worker homing) rebind once the node re-registers under a - new id, so dispatch is restored — not just registration. - """ + """Register a hook invoked with the new node id after a re-register.""" self._on_reregister = callback # ------------------------------------------------------------------ # @@ -167,12 +162,9 @@ def start(self) -> str: def _reregister_if_lost(self) -> None: """Re-register when the root registry no longer holds this node. - The node record lives in the shared control-plane Redis. When the root - is redeployed its registry is rebuilt from scratch, dropping records for - nodes that registered before the restart. Those nodes keep heartbeating, - but ``update_node_hb`` early-returns for an unknown node, so they never - reappear and are orphaned. Detect the missing record and re-register so - the fleet self-heals within one heartbeat interval, without a restart. + In case the node registry is reachable but the current node_id is no longer + present in the root registry (e.g., when the control-plane Redis is cleared on + root redeploy), the node must re-register to obtain a new node_id. """ if self._node_registry.node_exists(self.node_id): return @@ -188,9 +180,7 @@ def _reregister_if_lost(self) -> None: self._on_reregister(self._node_id) except Exception as exc: self.logger.warning( - "Re-register callback failed for node %s: %s", - self._node_id, - exc, + "Re-register callback failed for node %s: %s", self._node_id, exc ) def _hb_loop(self) -> None: diff --git a/src/server/supervisor/services/task_listener.py b/src/server/supervisor/services/task_listener.py index 06e3482..1ef8316 100644 --- a/src/server/supervisor/services/task_listener.py +++ b/src/server/supervisor/services/task_listener.py @@ -1,6 +1,6 @@ import asyncio import logging -from threading import Thread +from threading import Event, Lock, Thread from typing import Any from redis.client import PubSub @@ -13,11 +13,13 @@ from ...clients.redis import ( SyncRedisClient, - iter_pubsub_messages, node_dispatch_channel, + parse_pubsub_message, ) from ...utils.helpers import TSQueue +_POLL_TIMEOUT_SEC = 0.25 + class TaskListener: def __init__( @@ -25,6 +27,7 @@ def __init__( ) -> None: self.logger = logger self._redis = redis + # node_id is mutated only by the pubsub reader thread once running. self._node_id = node_id # TODO(kaiitunnz): Consider cleaning up old queues @@ -34,6 +37,10 @@ def __init__( self._loop: asyncio.AbstractEventLoop | None = None self._running: bool = False + self._rebind_lock = Lock() + self._pending_node_id: str | None = None + self._rebind_applied = Event() + def start(self) -> None: if self._thread is not None: self.logger.warning("Task listener already started") @@ -67,32 +74,45 @@ def stop(self) -> None: return assert self._running self._running = False - self._pubsub.close() self._thread.join() + self._pubsub.close() self._thread = None self._pubsub = None self._loop = None self.logger.info("Task listener stopped") def rebind(self, node_id: str) -> None: - """Move the dispatch subscription to a new node id. + """Request moving the dispatch subscription to a new node id. - Subscribes the new node's dispatch channel and drops the old one on the - live pubsub connection so an already-running listener keeps receiving - without a restart. Registered worker queues are preserved. + Records the target under a lock; the reader thread applies the actual + ``subscribe``/``unsubscribe`` between polls (mutating redis-py ``PubSub`` is not + thread-safe). ``wait_rebound`` blocks until the switch has taken effect. + Registered worker queues are preserved across the rebind. """ if node_id == self._node_id: + self._rebind_applied.set() return - old_node_id = self._node_id - self._node_id = node_id - pubsub = self._pubsub - if pubsub is None: - return - pubsub.subscribe(node_dispatch_channel(node_id)) - pubsub.unsubscribe(node_dispatch_channel(old_node_id)) + self._rebind_applied.clear() + with self._rebind_lock: + self._pending_node_id = node_id + + def wait_rebound(self, timeout: float) -> bool: + return self._rebind_applied.wait(timeout) + + def _apply_pending_rebind(self, pubsub: PubSub, current_id: str) -> str: + with self._rebind_lock: + pending = self._pending_node_id + self._pending_node_id = None + if pending is None or pending == current_id: + return current_id + pubsub.subscribe(node_dispatch_channel(pending)) + pubsub.unsubscribe(node_dispatch_channel(current_id)) + self._node_id = pending + self._rebind_applied.set() self.logger.info( - "Task listener rebound from node %s to %s", old_node_id, node_id + "Task listener rebound from node %s to %s", current_id, pending ) + return pending def add_worker(self, worker_id: str) -> None: if worker_id not in self._qs: @@ -113,8 +133,14 @@ def _run(self) -> None: if pubsub is None or loop is None: self.logger.error("Task listener not properly initialized") return + current_id = self._node_id try: - for data in iter_pubsub_messages(pubsub): + while self._running: + current_id = self._apply_pending_rebind(pubsub, current_id) + msg = pubsub.get_message(timeout=_POLL_TIMEOUT_SEC) + data = parse_pubsub_message(msg) + if data is None: + continue if "kind" not in data: self.logger.warning( "Received dispatch message without kind: %s", data @@ -153,6 +179,8 @@ def _run(self) -> None: continue queue = self._qs[worker_id] asyncio.run_coroutine_threadsafe(queue.put(payload), loop) + except (ConnectionError, OSError): + return except Exception as exc: if self._running: self.logger.exception("Task listener loop error: %s", exc) diff --git a/src/server/supervisor/supervisor.py b/src/server/supervisor/supervisor.py index c5d2d8a..54674e6 100644 --- a/src/server/supervisor/supervisor.py +++ b/src/server/supervisor/supervisor.py @@ -4,9 +4,12 @@ import logging import os import signal +from collections.abc import Callable from multiprocessing.process import BaseProcess from multiprocessing.queues import Queue as MPQueue from queue import Empty as QueueEmpty +from queue import Full as QueueFull +from threading import Thread from shared.schemas.command import CommandMessage, CommandResponse @@ -27,6 +30,8 @@ _CMD_TIMEOUT = 120.0 _NODE_ID_HANDSHAKE_TIMEOUT = 30.0 +_NODE_ID_WATCH_POLL_SEC = 0.5 +_REBIND_APPLY_TIMEOUT_SEC = 2.0 class WorkerSupervisor: @@ -50,12 +55,15 @@ def __init__( self._process: BaseProcess | None = None self._cmd_sender: TaskSender[CommandMessage, CommandResponse] | None = None self._cmd_receiver: TaskReceiver[CommandMessage, CommandResponse] | None = None - self._node_id_queue: MPQueue[str] = MP_CTX.Queue(maxsize=1) + self._node_id_queue: MPQueue[str] = MP_CTX.Queue(maxsize=8) self._node_id: str | None = None + self._node_id_listeners: list[Callable[[str], None]] = [] + self._node_id_watcher: Thread | None = None + self._node_id_stop = False @property def node_id(self) -> str: - """Return the node_id assigned by the child after `start()`.""" + """Return the current node_id. Reflects re-registrations after `start()`.""" if self._node_id is None: raise RuntimeError("Supervisor not started; node_id not yet assigned") return self._node_id @@ -112,8 +120,19 @@ async def start(self, system_principal: PrincipalContext) -> None: self._node_id = node_id self._logger.info("Supervisor handshake complete: node_id=%s", node_id) + self._node_id_stop = False + self._node_id_watcher = Thread( + target=self._watch_node_id, name="SupervisorNodeIdWatcher", daemon=True + ) + self._node_id_watcher.start() + async def stop(self, timeout: float = 3.0) -> None: """Gracefully stop the supervisor child process.""" + self._node_id_stop = True + if self._node_id_watcher is not None: + await asyncio.to_thread(self._node_id_watcher.join) + self._node_id_watcher = None + proc = self._process if ( proc is None @@ -157,12 +176,57 @@ async def exec_cmd( raise RuntimeError("Supervisor command channel not initialized") return await sender.send(cmd.command_id, cmd, timeout=timeout) + # ------------------------------------------------------------------ # + # Node ID watching + # ------------------------------------------------------------------ # + + def add_node_id_listener(self, listener: Callable[[str], None]) -> None: + """Register a hook invoked with the new node_id whenever the child re-registers + under a fresh id. Lets the parent refresh its cached copies.""" + self._node_id_listeners.append(listener) + + def _watch_node_id(self) -> None: + """Drain the node_id queue for re-registrations after the initial handshake, + updating the cached id and notifying listeners.""" + while not self._node_id_stop: + try: + new_id = self._node_id_queue.get(timeout=_NODE_ID_WATCH_POLL_SEC) + except QueueEmpty: + continue + self._node_id = new_id + self._logger.info("Supervisor node re-registered as %s", new_id) + for listener in self._node_id_listeners: + try: + listener(new_id) + except Exception as exc: + self._logger.warning("node_id listener failed: %s", exc) + # ------------------------------------------------------------------ # # Child-process entry point # ------------------------------------------------------------------ # +def _enqueue_latest_node_id( + queue: MPQueue[str], node_id: str, logger: logging.Logger +) -> None: + """Put node_id on the handshake queue, dropping the stale head when it is + full so the parent always learns the latest id.""" + try: + queue.put_nowait(node_id) + except QueueFull: + try: + queue.get_nowait() + except QueueEmpty: + pass + try: + queue.put_nowait(node_id) + except QueueFull: + logger.warning( + "node_id queue full; parent will sync on the next re-register" + ) + + def _run_supervisor( identity: IdentityConfig, redis_cfg: RedisConfig, @@ -221,8 +285,8 @@ def _run_supervisor( current_gpu_count_getter = None try: rm = ResourceManager.get_instance() - max_gpu_count = rm.total_gpu_count - current_gpu_count_getter = lambda: rm.available_gpu_count # noqa: E731 + max_gpu_count = rm.total_gpu_count() + current_gpu_count_getter = rm.available_gpu_count except Exception as exc: logger.warning("Failed to detect node GPU capacity: %s", exc) @@ -299,11 +363,31 @@ def _run_supervisor( ) def _on_reregister(new_node_id: str) -> None: - grpc_server.rebind_node(new_node_id) + # Rebind the subscriptions first and wait for the reader threads to + # switch channels, THEN re-home workers, so the dispatcher never + # publishes to a dispatch channel nobody is listening on yet. task_listener.rebind(new_node_id) command_listener.rebind(new_node_id) - - lifecycle.set_reregister_callback(_on_reregister) + # The rebind is reader-owned and idempotent; a timeout means the reader + # is unresponsive, not that the switch was lost. Surface it and proceed + # rather than stall the heartbeat thread — the reader applies the pending + # id once it recovers, and leaving the node half-migrated is worse. + if not task_listener.wait_rebound(_REBIND_APPLY_TIMEOUT_SEC): + logger.error( + "Task dispatch did not rebind to %s within %ss", + new_node_id, + _REBIND_APPLY_TIMEOUT_SEC, + ) + if not command_listener.wait_rebound(_REBIND_APPLY_TIMEOUT_SEC): + logger.error( + "Command channel did not rebind to %s within %ss", + new_node_id, + _REBIND_APPLY_TIMEOUT_SEC, + ) + grpc_server.rebind_node(new_node_id) + # Propagate the new id to the parent process (non-blocking: this runs on + # the heartbeat thread, which must not stall on a full queue). + _enqueue_latest_node_id(node_id_queue, new_node_id, logger) # --- Event loop with signal handling --- loop = asyncio.new_event_loop() @@ -324,6 +408,8 @@ async def _run() -> None: await worker_manager.start() command_listener.start() await grpc_server.start() + # Wire the re-register callback only once the reader threads are up + lifecycle.set_reregister_callback(_on_reregister) logger.info("Supervisor ready for node %s", node_id) # Wait for termination signal diff --git a/tests/server/supervisor_helpers.py b/tests/server/supervisor_helpers.py new file mode 100644 index 0000000..e46dee4 --- /dev/null +++ b/tests/server/supervisor_helpers.py @@ -0,0 +1,40 @@ +"""Shared stubs for the supervisor lifecycle / re-register tests.""" + +import logging + +from server.registries.node import NodeRegistry +from server.supervisor.services.lifecycle import Lifecycle + +_LOGGER = logging.getLogger("test.supervisor") + + +class StubRegistry(NodeRegistry): + """NodeRegistry whose ``node_exists`` returns a fixed value and records the + ids it was queried with.""" + + def __init__(self, exists: bool) -> None: + self._exists = exists + self.exists_calls: list[str] = [] + + def node_exists(self, node_id: str) -> bool: + self.exists_calls.append(node_id) + return self._exists + + +class StubLifecycle(Lifecycle): + """Lifecycle with registration stubbed: re-register mints ``nde-2`` and each + published event type is recorded instead of sent.""" + + def __init__(self, node_registry: NodeRegistry, node_id: str) -> None: + self._node_registry = node_registry + self._node_id = node_id + self.logger = _LOGGER + self._unregister_published = True + self._on_reregister = None + self.published_events: list[str] = [] + + def _register(self) -> str: + return "nde-2" + + def _publish_event(self, event_type: str, **extra: object) -> None: + self.published_events.append(event_type) diff --git a/tests/server/test_pubsub_parse.py b/tests/server/test_pubsub_parse.py new file mode 100644 index 0000000..3491635 --- /dev/null +++ b/tests/server/test_pubsub_parse.py @@ -0,0 +1,32 @@ +"""Decoding of redis-py pub/sub frames shared by the blocking and polling readers.""" + +import json + +from server.clients.redis import parse_pubsub_message + + +def test_parses_json_object_message() -> None: + msg = {"type": "message", "data": json.dumps({"kind": "task", "id": 1})} + assert parse_pubsub_message(msg) == {"kind": "task", "id": 1} + + +def test_decodes_bytes_payload() -> None: + msg = {"type": "message", "data": b'{"kind": "stop"}'} + assert parse_pubsub_message(msg) == {"kind": "stop"} + + +def test_none_frame_returns_none() -> None: + # get_message(timeout=...) yields None when nothing arrived + assert parse_pubsub_message(None) is None + + +def test_control_frame_returns_none() -> None: + assert parse_pubsub_message({"type": "subscribe", "data": 1}) is None + + +def test_missing_data_returns_none() -> None: + assert parse_pubsub_message({"type": "message", "data": None}) is None + + +def test_malformed_json_returns_none() -> None: + assert parse_pubsub_message({"type": "message", "data": "{not json"}) is None diff --git a/tests/server/test_supervisor_lifecycle.py b/tests/server/test_supervisor_lifecycle.py index 73d3bfa..1316c75 100644 --- a/tests/server/test_supervisor_lifecycle.py +++ b/tests/server/test_supervisor_lifecycle.py @@ -8,6 +8,7 @@ from server.supervisor.services import lifecycle as lifecycle_module from server.supervisor.services.lifecycle import Lifecycle from shared.schemas.node import NodeInfo +from tests.server.supervisor_helpers import StubLifecycle, StubRegistry def _build_lifecycle(base_url: str = "http://root:8000") -> Lifecycle: @@ -77,42 +78,43 @@ def fake_post(url: str, **kwargs: Any) -> _StubResponse: assert captured["headers"] == {} -class _StubRegistry: - def __init__(self, exists: bool) -> None: - self._exists = exists - self.exists_calls: list[str] = [] - - def node_exists(self, node_id: str) -> bool: - self.exists_calls.append(node_id) - return self._exists - - def test_reregister_if_lost_noops_when_present() -> None: - instance = _build_lifecycle() - registry = _StubRegistry(exists=True) - instance._node_registry = registry # type: ignore[assignment] - instance._node_id = "nde-1" - instance._register = lambda: pytest.fail("must not re-register") # type: ignore[method-assign] + registry = StubRegistry(exists=True) + instance = StubLifecycle(registry, "nde-1") + calls: list[str] = [] + instance.set_reregister_callback(calls.append) instance._reregister_if_lost() assert registry.exists_calls == ["nde-1"] assert instance._node_id == "nde-1" + assert instance.published_events == [] + assert calls == [] # callback not fired when the record is present def test_reregister_if_lost_reregisters_when_missing() -> None: - events: list[str] = [] - instance = _build_lifecycle() - registry = _StubRegistry(exists=False) - instance._node_registry = registry # type: ignore[assignment] - instance._node_id = "nde-1" - instance._unregister_published = True - instance._register = lambda: "nde-2" # type: ignore[method-assign] - instance._publish_event = lambda event_type, **extra: events.append(event_type) # type: ignore[method-assign] + registry = StubRegistry(exists=False) + instance = StubLifecycle(registry, "nde-1") + calls: list[str] = [] + instance.set_reregister_callback(calls.append) instance._reregister_if_lost() assert registry.exists_calls == ["nde-1"] assert instance._node_id == "nde-2" assert instance._unregister_published is False - assert events == ["SV_REGISTER"] + assert instance.published_events == ["SV_REGISTER"] + assert calls == ["nde-2"] # callback fired with the new id + + +def test_reregister_if_lost_swallows_callback_error() -> None: + instance = StubLifecycle(StubRegistry(exists=False), "nde-1") + + def _boom(_new_id: str) -> None: + raise RuntimeError("rebind failed") + + instance.set_reregister_callback(_boom) + + # a failing callback must be logged, not crash the heartbeat loop + instance._reregister_if_lost() + assert instance._node_id == "nde-2" diff --git a/tests/server/test_supervisor_node_id_watch.py b/tests/server/test_supervisor_node_id_watch.py new file mode 100644 index 0000000..864764d --- /dev/null +++ b/tests/server/test_supervisor_node_id_watch.py @@ -0,0 +1,103 @@ +"""WorkerSupervisor keeps the parent-process node_id current after re-register. + +The child re-puts its new node id on the handshake queue; the parent watcher +picks it up, updates the cached id, and notifies listeners (which refresh +``app.state.node_id`` and the EventMonitor's own-node in production). +""" + +import logging +from multiprocessing import Queue +from multiprocessing.queues import Queue as MPQueue +from queue import Queue as ThreadQueue +from threading import Event, Thread +from typing import cast + +from server.supervisor.supervisor import WorkerSupervisor, _enqueue_latest_node_id + +_LOGGER = logging.getLogger("test.node_id_watch") + + +def _bounded(maxsize: int) -> MPQueue[str]: + # A synchronous thread queue stands in for the mp queue: same put_nowait / + # get_nowait contract, without the feeder-thread flush race that makes an + # mp queue's get_nowait unreliable immediately after a put. + return cast(MPQueue[str], ThreadQueue(maxsize=maxsize)) + + +def _build_supervisor() -> WorkerSupervisor: + sup = WorkerSupervisor.__new__(WorkerSupervisor) + sup._logger = _LOGGER + sup._node_id = "nde-1" + sup._node_id_queue = Queue() + sup._node_id_listeners = [] + sup._node_id_stop = False + sup._node_id_watcher = None + return sup + + +def test_watcher_updates_id_and_notifies_listeners() -> None: + sup = _build_supervisor() + seen: list[str] = [] + notified = Event() + + def _listener(new_id: str) -> None: + seen.append(new_id) + notified.set() + + sup.add_node_id_listener(_listener) + + watcher = Thread(target=sup._watch_node_id, daemon=True) + watcher.start() + try: + sup._node_id_queue.put("nde-2") + assert notified.wait(2.0) + finally: + sup._node_id_stop = True + watcher.join(timeout=2.0) + + assert seen == ["nde-2"] + assert sup.node_id == "nde-2" + + +def test_watcher_survives_a_failing_listener() -> None: + sup = _build_supervisor() + good_seen: list[str] = [] + second_ran = Event() + + def _boom(_new_id: str) -> None: + raise RuntimeError("listener failed") + + def _good(new_id: str) -> None: + good_seen.append(new_id) + second_ran.set() + + sup.add_node_id_listener(_boom) + sup.add_node_id_listener(_good) + + watcher = Thread(target=sup._watch_node_id, daemon=True) + watcher.start() + try: + sup._node_id_queue.put("nde-2") + assert second_ran.wait(2.0) + finally: + sup._node_id_stop = True + watcher.join(timeout=2.0) + + # a raising listener must not stop later listeners or the watcher + assert good_seen == ["nde-2"] + assert sup.node_id == "nde-2" + + +def test_enqueue_puts_when_not_full() -> None: + queue = _bounded(8) + _enqueue_latest_node_id(queue, "nde-2", _LOGGER) + assert queue.get_nowait() == "nde-2" + + +def test_enqueue_drops_oldest_when_full() -> None: + queue = _bounded(1) + _enqueue_latest_node_id(queue, "nde-1", _LOGGER) + # queue is full; the stale id is dropped so the newest still lands + _enqueue_latest_node_id(queue, "nde-2", _LOGGER) + assert queue.get_nowait() == "nde-2" + assert queue.empty() diff --git a/tests/server/test_supervisor_reregister.py b/tests/server/test_supervisor_reregister.py index ccf0ef2..4448deb 100644 --- a/tests/server/test_supervisor_reregister.py +++ b/tests/server/test_supervisor_reregister.py @@ -4,23 +4,34 @@ node id. These tests prove the node also (a) moves its dispatch/command subscriptions to the new channel and (b) re-homes its already-registered workers under the new id, so the dispatcher can reach them again. + +The rebind is split in two: the heartbeat thread only records the target id +(``rebind``), and the pubsub reader thread applies the actual subscribe / +unsubscribe (``_apply_pending_rebind``). A redis-py ``PubSub`` must not be +mutated from a second thread while its owner is reading it, so these tests +assert ``rebind`` touches nothing and the reader-side apply does the switch. """ import logging -from typing import Any +from threading import Lock +from typing import Any, cast + +from redis.client import PubSub from server.clients.redis import ( + SyncRedisClient, node_cmd_channel, node_dispatch_channel, worker_key, ) +from server.supervisor.registry import WorkerRegistry from server.supervisor.services.command_listener import CommandListener, _CommandStream from server.supervisor.services.grpc_server import SupervisorServicer -from server.supervisor.services.lifecycle import Lifecycle from server.supervisor.services.task_listener import TaskListener -from shared.schemas.node import NodeInfo +from tests.server.supervisor_helpers import StubLifecycle, StubRegistry _LOGGER = logging.getLogger("test.reregister") +_ANY_OBJECT: Any = None class _FakePubSub: @@ -40,34 +51,47 @@ def unsubscribe(self, channel: str) -> None: self.unsubscribe_log.append(channel) +def _as_pubsub(fake: _FakePubSub) -> PubSub: + return cast(PubSub, fake) + + # --------------------------------------------------------------------------- # -# TaskListener.rebind +# TaskListener.rebind (heartbeat thread) + _apply_pending_rebind (reader thread) # --------------------------------------------------------------------------- # def _build_task_listener(node_id: str) -> tuple[TaskListener, _FakePubSub]: - listener = TaskListener.__new__(TaskListener) - listener.logger = _LOGGER - listener._node_id = node_id - listener._qs = {} + listener = TaskListener(_ANY_OBJECT, node_id, _LOGGER) pubsub = _FakePubSub() pubsub.subscribe(node_dispatch_channel(node_id)) - listener._pubsub = pubsub # type: ignore[assignment] - listener._thread = object() # type: ignore[assignment] - listener._loop = None - listener._running = True return listener, pubsub -def test_task_listener_rebind_moves_subscription() -> None: +def test_task_listener_rebind_records_target_without_touching_pubsub() -> None: + # Regression guard for the cross-thread race: rebind runs on the heartbeat + # thread and must NOT mutate the live pubsub. listener, pubsub = _build_task_listener("nde-1") - listener.add_worker("wkr-1") listener.rebind("nde-2") + assert listener._pending_node_id == "nde-2" + assert pubsub.subscribe_log == [node_dispatch_channel("nde-1")] + assert pubsub.unsubscribe_log == [] + assert not listener._rebind_applied.is_set() + + +def test_task_listener_apply_pending_moves_subscription() -> None: + listener, pubsub = _build_task_listener("nde-1") + listener.add_worker("wkr-1") + listener.rebind("nde-2") + + live = listener._apply_pending_rebind(_as_pubsub(pubsub), "nde-1") + + assert live == "nde-2" assert listener._node_id == "nde-2" assert node_dispatch_channel("nde-2") in pubsub.subscribed assert node_dispatch_channel("nde-1") not in pubsub.subscribed + assert listener._rebind_applied.is_set() # worker queues survive the rebind so in-flight streams keep working assert "wkr-1" in listener._qs @@ -75,53 +99,47 @@ def test_task_listener_rebind_moves_subscription() -> None: def test_task_listener_rebind_same_id_is_noop() -> None: listener, pubsub = _build_task_listener("nde-1") listener.rebind("nde-1") - assert pubsub.subscribe_log == [node_dispatch_channel("nde-1")] + # same id resolves immediately, nothing pending, no pubsub mutation on apply + assert listener._rebind_applied.is_set() + assert listener._apply_pending_rebind(_as_pubsub(pubsub), "nde-1") == "nde-1" assert pubsub.unsubscribe_log == [] -def test_task_listener_rebind_before_start_only_updates_id() -> None: - listener = TaskListener.__new__(TaskListener) - listener.logger = _LOGGER - listener._node_id = "nde-1" - listener._qs = {} - listener._pubsub = None - listener.rebind("nde-2") - assert listener._node_id == "nde-2" - - # --------------------------------------------------------------------------- # -# CommandListener.rebind +# _CommandStream.rebind + _apply_pending_rebind # --------------------------------------------------------------------------- # -def test_command_listener_rebind_moves_subscription() -> None: - stream = _CommandStream.__new__(_CommandStream) - stream.node_id = "nde-1" - stream.logger = _LOGGER +def _build_command_stream(node_id: str) -> tuple[_CommandStream, _FakePubSub]: + stream = _CommandStream(node_id, _ANY_OBJECT, None, _LOGGER) pubsub = _FakePubSub() - pubsub.subscribe(node_cmd_channel("nde-1")) - stream._pubsub = pubsub # type: ignore[assignment] + pubsub.subscribe(node_cmd_channel(node_id)) + return stream, pubsub + - listener = CommandListener.__new__(CommandListener) - listener.logger = _LOGGER - listener._node_id = "nde-1" +def test_command_listener_rebind_records_target_without_touching_pubsub() -> None: + stream, pubsub = _build_command_stream("nde-1") + listener = CommandListener(_ANY_OBJECT, "nde-1", _ANY_OBJECT, _LOGGER) listener._cmd_stream = stream listener.rebind("nde-2") assert listener._node_id == "nde-2" + assert stream._pending_node_id == "nde-2" + assert pubsub.unsubscribe_log == [] + + +def test_command_stream_apply_pending_moves_subscription() -> None: + stream, pubsub = _build_command_stream("nde-1") + stream.rebind("nde-2") + + live = stream._apply_pending_rebind(_as_pubsub(pubsub), "nde-1") + + assert live == "nde-2" assert stream.node_id == "nde-2" assert node_cmd_channel("nde-2") in pubsub.subscribed assert node_cmd_channel("nde-1") not in pubsub.subscribed - - -def test_command_listener_rebind_before_start_only_updates_id() -> None: - listener = CommandListener.__new__(CommandListener) - listener.logger = _LOGGER - listener._node_id = "nde-1" - listener._cmd_stream = None - listener.rebind("nde-2") - assert listener._node_id == "nde-2" + assert stream._rebind_applied.is_set() # --------------------------------------------------------------------------- # @@ -146,23 +164,38 @@ def get_worker_id(self, token: str) -> str | None: class _FakeRedis: - def __init__(self) -> None: + """Simulates the re-home Lua: HSET node_id only for keys that exist.""" + + def __init__(self, existing: set[str] | None = None) -> None: + self.existing = existing if existing is not None else set() self.hash_writes: list[tuple[str, dict[str, Any]]] = [] - def hash_set(self, key: str, mapping: dict[str, Any]) -> None: - self.hash_writes.append((key, mapping)) + def eval(self, script: str, numkeys: int, *args: str) -> int: + keys = args[:numkeys] + node_id = args[numkeys] + rehomed = 0 + for key in keys: + if key in self.existing: + self.hash_writes.append((key, {"node_id": node_id})) + rehomed += 1 + return rehomed def _build_servicer( - node_id: str, token_to_id: dict[str, str | None] + node_id: str, + token_to_id: dict[str, str | None], + existing: set[str] | None = None, ) -> tuple[SupervisorServicer, _FakeRedis]: + if existing is None: + existing = {worker_key(w) for w in token_to_id.values() if w is not None} + redis = _FakeRedis(existing) servicer = SupervisorServicer.__new__(SupervisorServicer) - servicer._registry = _FakeWorkerRegistry(token_to_id) # type: ignore[assignment] - redis = _FakeRedis() - servicer._redis = redis # type: ignore[assignment] + servicer._registry = cast(WorkerRegistry, _FakeWorkerRegistry(token_to_id)) + servicer._redis = cast(SyncRedisClient, redis) servicer._node_id = node_id servicer._node_alias = "worker-box" servicer._logger = _LOGGER + servicer._lock = Lock() return servicer, redis @@ -180,6 +213,18 @@ def test_rebind_node_rehomes_all_registered_workers() -> None: assert len(redis.hash_writes) == 2 +def test_rebind_node_skips_workers_without_a_record() -> None: + # Full-wipe scenario: the worker key was evicted; re-home must NOT resurrect + # a partial record. + servicer, redis = _build_servicer( + "nde-1", {"tok-a": "wkr-1", "tok-b": "wkr-2"}, existing={worker_key("wkr-1")} + ) + + servicer.rebind_node("nde-2") + + assert redis.hash_writes == [(worker_key("wkr-1"), {"node_id": "nde-2"})] + + def test_rebind_node_stamps_future_registrations() -> None: servicer, _ = _build_servicer("nde-1", {}) servicer.rebind_node("nde-9") @@ -193,83 +238,6 @@ def test_rebind_node_same_id_is_noop() -> None: assert redis.hash_writes == [] -# --------------------------------------------------------------------------- # -# Lifecycle wiring — callback fires on re-register, restoring dispatch -# --------------------------------------------------------------------------- # - - -def _build_lifecycle() -> Lifecycle: - instance = Lifecycle.__new__(Lifecycle) - instance._base_url = "http://root:8000" - instance._node_info = NodeInfo( - namespace="ns", - cluster="cl", - alias="worker-1", - version="0.1.0", - started_at="2026-05-13T00:00:00Z", - tags=[], - last_seen="2026-05-13T00:00:00Z", - max_gpu_count=0, - ) - instance.logger = _LOGGER - instance._on_reregister = None - return instance - - -class _StubRegistry: - def __init__(self, exists: bool) -> None: - self._exists = exists - - def node_exists(self, node_id: str) -> bool: - return self._exists - - -def test_reregister_invokes_callback_with_new_id() -> None: - seen: list[str] = [] - instance = _build_lifecycle() - instance._node_registry = _StubRegistry(exists=False) # type: ignore[assignment] - instance._node_id = "nde-1" - instance._unregister_published = True - instance._register = lambda: "nde-2" # type: ignore[method-assign] - instance._publish_event = lambda *a, **k: None # type: ignore[method-assign] - instance.set_reregister_callback(seen.append) - - instance._reregister_if_lost() - - assert instance._node_id == "nde-2" - assert seen == ["nde-2"] - - -def test_reregister_present_does_not_invoke_callback() -> None: - seen: list[str] = [] - instance = _build_lifecycle() - instance._node_registry = _StubRegistry(exists=True) # type: ignore[assignment] - instance._node_id = "nde-1" - instance.set_reregister_callback(seen.append) - - instance._reregister_if_lost() - - assert seen == [] - - -def test_reregister_callback_error_does_not_break_heartbeat() -> None: - instance = _build_lifecycle() - instance._node_registry = _StubRegistry(exists=False) # type: ignore[assignment] - instance._node_id = "nde-1" - instance._unregister_published = True - instance._register = lambda: "nde-2" # type: ignore[method-assign] - instance._publish_event = lambda *a, **k: None # type: ignore[method-assign] - - def _boom(_new_id: str) -> None: - raise RuntimeError("rebind failed") - - instance.set_reregister_callback(_boom) - - # must not raise — a failed rebind should be logged, not crash the hb loop - instance._reregister_if_lost() - assert instance._node_id == "nde-2" - - # --------------------------------------------------------------------------- # # End-to-end: registry loss -> re-register-under-new-id -> dispatchable again # --------------------------------------------------------------------------- # @@ -278,34 +246,28 @@ def _boom(_new_id: str) -> None: def test_full_reregister_rebinds_and_rehomes() -> None: """After a simulated registry loss and re-register under a new id, the node is subscribed on the NEW dispatch/command channels and its workers are - homed under the NEW id (so the host routes tasks to them).""" + homed under the NEW id (so the host routes tasks to them). The reader-side + apply is driven inline to stand in for the pubsub reader thread.""" task_listener, task_pubsub = _build_task_listener("nde-1") task_listener.add_worker("wkr-1") - cmd_stream = _CommandStream.__new__(_CommandStream) - cmd_stream.node_id = "nde-1" - cmd_stream.logger = _LOGGER - cmd_pubsub = _FakePubSub() - cmd_pubsub.subscribe(node_cmd_channel("nde-1")) - cmd_stream._pubsub = cmd_pubsub # type: ignore[assignment] - cmd_listener = CommandListener.__new__(CommandListener) - cmd_listener.logger = _LOGGER - cmd_listener._node_id = "nde-1" + cmd_stream, cmd_pubsub = _build_command_stream("nde-1") + cmd_listener = CommandListener(_ANY_OBJECT, "nde-1", _ANY_OBJECT, _LOGGER) cmd_listener._cmd_stream = cmd_stream servicer, redis = _build_servicer("nde-1", {"tok-a": "wkr-1"}) def _on_reregister(new_node_id: str) -> None: - servicer.rebind_node(new_node_id) task_listener.rebind(new_node_id) cmd_listener.rebind(new_node_id) + # stand in for the reader threads applying the pending rebind + task_listener._apply_pending_rebind(_as_pubsub(task_pubsub), "nde-1") + cmd_stream._apply_pending_rebind(_as_pubsub(cmd_pubsub), "nde-1") + assert task_listener.wait_rebound(1.0) + assert cmd_listener.wait_rebound(1.0) + servicer.rebind_node(new_node_id) - lifecycle = _build_lifecycle() - lifecycle._node_registry = _StubRegistry(exists=False) # type: ignore[assignment] - lifecycle._node_id = "nde-1" - lifecycle._unregister_published = True - lifecycle._register = lambda: "nde-2" # type: ignore[method-assign] - lifecycle._publish_event = lambda *a, **k: None # type: ignore[method-assign] + lifecycle = StubLifecycle(StubRegistry(exists=False), "nde-1") lifecycle.set_reregister_callback(_on_reregister) lifecycle._reregister_if_lost() diff --git a/tests/server/test_worker_manager_gpu.py b/tests/server/test_worker_manager_gpu.py index a7a8417..0b3d1e1 100644 --- a/tests/server/test_worker_manager_gpu.py +++ b/tests/server/test_worker_manager_gpu.py @@ -54,7 +54,7 @@ def test_single_gpu_returns_lowest_index_and_reserves(self) -> None: rm = _resource_manager({0, 1, 2, 3}) devices, _ = rm.reserve_gpus(n=1) assert devices == [0] - assert rm.available_gpu_count == 3 + assert rm.available_gpu_count() == 3 assert 0 not in rm._env.available_gpus def test_single_gpu_picks_minimum_when_indices_nonzero(self) -> None: @@ -66,7 +66,7 @@ def test_two_gpus_returns_two_lowest_sorted(self) -> None: rm = _resource_manager({0, 1, 2, 3}) devices, _ = rm.reserve_gpus(n=2) assert devices == [0, 1] - assert rm.available_gpu_count == 2 + assert rm.available_gpu_count() == 2 def test_two_gpus_sparse_indices(self) -> None: rm = _resource_manager({1, 3}) @@ -77,13 +77,13 @@ def test_four_gpus_returns_all_sorted(self) -> None: rm = _resource_manager({0, 1, 2, 3}) devices, _ = rm.reserve_gpus(n=4) assert devices == [0, 1, 2, 3] - assert rm.available_gpu_count == 0 + assert rm.available_gpu_count() == 0 def test_requesting_more_than_available_raises(self) -> None: rm = _resource_manager({0, 1}) with pytest.raises(ValueError, match="Not enough available GPUs"): rm.reserve_gpus(n=3) - assert rm.available_gpu_count == 2 + assert rm.available_gpu_count() == 2 def test_zero_raises(self) -> None: rm = _resource_manager({0, 1}) @@ -107,7 +107,7 @@ def test_unavailable_device_raises(self) -> None: rm = _resource_manager({0, 1}) with pytest.raises(ValueError, match="Requested GPUs are not available"): rm.reserve_gpus(devices=[5]) - assert rm.available_gpu_count == 2 + assert rm.available_gpu_count() == 2 def test_partially_unavailable_raises_without_partial_reserve(self) -> None: rm = _resource_manager({0, 1}) @@ -140,7 +140,7 @@ def test_repeated_calls_yield_disjoint_devices(self) -> None: d2, _ = rm.reserve_gpus(n=2) assert set(d1).isdisjoint(d2) assert sorted(d1 + d2) == [0, 1, 2, 3] - assert rm.available_gpu_count == 0 + assert rm.available_gpu_count() == 0 def test_mixed_arch_raises(self) -> None: rm = object.__new__(ResourceManager) @@ -157,13 +157,13 @@ def test_mixed_arch_raises(self) -> None: class TestAvailableGpuCount: def test_four_gpus(self) -> None: - assert _resource_manager({0, 1, 2, 3}).available_gpu_count == 4 + assert _resource_manager({0, 1, 2, 3}).available_gpu_count() == 4 def test_two_gpus(self) -> None: - assert _resource_manager({0, 1}).available_gpu_count == 2 + assert _resource_manager({0, 1}).available_gpu_count() == 2 def test_no_gpus(self) -> None: - assert _resource_manager(set()).available_gpu_count == 0 + assert _resource_manager(set()).available_gpu_count() == 0 class TestMachineEnvDetection: diff --git a/uv.lock b/uv.lock index 4d430e8..ad22651 100644 --- a/uv.lock +++ b/uv.lock @@ -2247,7 +2247,7 @@ ci = [ { name = "trl", specifier = ">=1.6.0" }, { name = "types-colorama", specifier = ">=0.4.15.20250801" }, { name = "types-docker", specifier = ">=7.1.0.20251202" }, - { name = "types-grpcio", specifier = ">=1.0.0.20251009" }, + { name = "types-grpcio", specifier = ">=1.0.0.20260614" }, { name = "types-pexpect", specifier = ">=4.9.0.20250916" }, { name = "types-protobuf", specifier = ">=6.32.1.20251210" }, { name = "types-psutil", specifier = ">=7.1.3.20251211" }, @@ -2282,7 +2282,7 @@ dev = [ { name = "ruff", specifier = ">=0.14.10" }, { name = "types-colorama", specifier = ">=0.4.15.20250801" }, { name = "types-docker", specifier = ">=7.1.0.20251202" }, - { name = "types-grpcio", specifier = ">=1.0.0.20251009" }, + { name = "types-grpcio", specifier = ">=1.0.0.20260614" }, { name = "types-pexpect", specifier = ">=4.9.0.20250916" }, { name = "types-protobuf", specifier = ">=6.32.1.20251210" }, { name = "types-psutil", specifier = ">=7.1.3.20251211" }, @@ -7859,11 +7859,11 @@ wheels = [ [[package]] name = "types-grpcio" -version = "1.0.0.20251009" +version = "1.0.0.20260614" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/de/93/78aa083216853c667c9412df4ef8284b2a68c6bcd2aef833f970b311f3c1/types_grpcio-1.0.0.20251009.tar.gz", hash = "sha256:a8f615ea7a47b31f10da028ab5258d4f1611fbd70719ca450fc0ab3fb9c62b63", size = 14479, upload-time = "2025-10-09T02:54:14.539Z" } +sdist = { url = "https://files.pythonhosted.org/packages/4f/1c/e299150827a2224fe9ddb59e9f789f0eb8e72e0befc23fd47195024311fa/types_grpcio-1.0.0.20260614.tar.gz", hash = "sha256:e86d1aabb7e7eedae487c3ac10e010a13d5db1fd350e766562053af557366aab", size = 15063, upload-time = "2026-06-14T06:23:33.591Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d4/93/66d28f41b16bb4e6b611bd608ef28dffc740facec93250b30cf83138da21/types_grpcio-1.0.0.20251009-py3-none-any.whl", hash = "sha256:112ac4312a5b0a273a4c414f7f2c7668f342990d9c6ab0f647391c36331f95ed", size = 15208, upload-time = "2025-10-09T02:54:13.588Z" }, + { url = "https://files.pythonhosted.org/packages/33/f5/a8b8ca4d891bc4e835ef758a8fb50c7539c7f45bc0e42fae2b847cbc59b7/types_grpcio-1.0.0.20260614-py3-none-any.whl", hash = "sha256:050472cb4ef329ae8fe20278c62bc0da5b961aad3dd889cc35f6e0c70d368dae", size = 15629, upload-time = "2026-06-14T06:23:32.616Z" }, ] [[package]]