diff --git a/gigl/distributed/graph_store/shared_dist_sampling_producer.py b/gigl/distributed/graph_store/shared_dist_sampling_producer.py index 608504d8b..f2d94df83 100644 --- a/gigl/distributed/graph_store/shared_dist_sampling_producer.py +++ b/gigl/distributed/graph_store/shared_dist_sampling_producer.py @@ -60,7 +60,8 @@ ├─────────────────────────────────────────────────┤ │ Phase 3: Idle wait │ │ if no commands and no batches submitted: │ - │ task_queue.get(timeout=SCHEDULER_TICK_SECS) │ + │ wake_event.wait(timeout=SCHEDULER_TICK_SECS)│ + │ (a completion's _on_batch_done set()s it) │ └─────────────────────────────────────────────────┘ """ @@ -375,8 +376,11 @@ def _shared_sampling_worker_loop( b. Submitting batches round-robin from ``runnable_channel_ids`` — a FIFO queue of channels that have pending work. Each channel gets one batch submitted per round to prevent starvation. - c. If no commands were processed and no batches submitted, blocking on - ``task_queue`` with a short timeout to avoid busy-waiting. + c. If no commands were processed and no batches submitted, waiting on + a completion wake-event with a short fallback timeout to avoid + busy-waiting. ``_on_batch_done`` sets the event when it frees an + in-flight slot, so a parked channel refills immediately rather than + after the full tick. 3. Completion callbacks from the sampler update per-channel state and emit ``EPOCH_DONE_EVENT`` to ``event_queue`` when all batches for an epoch are finished. @@ -393,8 +397,32 @@ def _shared_sampling_worker_loop( removing_channel_ids: set[int] = set() cleanup_ready_channel_ids: set[int] = set() state_lock = threading.RLock() + # Level-triggered event the completion callback (``_on_batch_done``) sets + # after re-enqueueing a parked channel, so the scheduler's Phase-3 idle wait + # returns at once instead of sleeping the full ``SCHEDULER_TICK_SECS``. + # Loop-local: the single scheduler thread is the only waiter, callback + # threads only ``set()``, so nothing leaks across worker instances. + wake_event = threading.Event() last_state_log_time = 0.0 current_device: Optional[torch.device] = None + # Per-channel in-flight submission cap. The scheduler parks a channel + # (skips its submit) once it has this many batches submitted but not yet + # completed, so the single scheduler thread never issues a submit that could + # block in the sampler's ``BoundedSemaphore(worker_concurrency).acquire()`` + # on a saturated output channel. The cap mirrors that semaphore exactly: + # each in-flight coroutine holds one slot from submit until its + # ``channel.send`` completes, so at ``worker_concurrency`` in flight the next + # submit would block this thread -- parking co-located channels and, at high + # fan-out, closing a cross-rank deadlock cycle. Keeping the scheduler + # wait-free (rotate past saturated channels, keep draining commands) is what + # stops high fan-out from deadlocking. + # + # The ``submitted - completed`` mirror of the semaphore holds WITHIN an + # epoch: ``START_EPOCH`` resets ``submitted_batches`` to 0, so the invariant + # relies on the "previous epoch fully completes before the next + # ``START_EPOCH``" contract that ``BaseDistLoader.__iter__`` enforces today + # (it drains every batch via ``__next__`` before re-entering the loop). + worker_concurrency = worker_options.worker_concurrency # --- Scheduler helper functions --- @@ -411,6 +439,24 @@ def _enqueue_channel_if_runnable_locked(channel_id: int) -> None: runnable_channel_ids.append(channel_id) runnable_channel_id_set.add(channel_id) + def _is_channel_parked_locked(channel_id: int) -> bool: + """Return whether a channel has hit its in-flight submission cap. + + A parked channel has ``worker_concurrency`` batches submitted but not + yet completed, so submitting another could block the scheduler thread in + the sampler's bounded semaphore on a saturated output channel. + + The pump skips a parked channel (and drops it from the runnable queue); + ``_on_batch_done`` re-enqueues it once an in-flight batch completes and + frees a slot. + + Must be called while holding ``state_lock``. + """ + state = active_epoch_by_channel_id.get(channel_id) + if state is None: + return False + return state.submitted_batches - state.completed_batches >= worker_concurrency + def _drain_channel_locked(channel_id: int) -> int: """Lossily drain buffered sampled messages for a removing channel. @@ -555,6 +601,22 @@ def _on_batch_done(channel_id: int, epoch: int) -> None: if state is None or state.epoch != epoch: return state.completed_batches += 1 + # WAKE: an in-flight slot just freed, so the channel may have been + # parked by the in-flight cap. Re-enter it into the runnable queue. + # This is the ONLY path that wakes a parked channel (the pump never + # re-enqueues a channel it parked), so a channel whose consumer had + # paused resumes within one scheduler tick of draining a batch. + # Idempotent via the membership guard and a no-op for cancelled or + # fully-submitted channels, so it is safe to call on every batch. + _enqueue_channel_if_runnable_locked(channel_id) + # Wake the scheduler's Phase-3 idle wait so the just-freed in-flight + # slot is refilled at once rather than after up to + # ``SCHEDULER_TICK_SECS``. ``Event.set()`` is non-blocking and safe + # under ``state_lock``; being level-triggered it can never be a lost + # wakeup even if the scheduler is not yet waiting. Never + # ``task_queue.put`` here: that queue is bounded and would + # backpressure this callback thread. + wake_event.set() if channel_id in removing_channel_ids: drained_messages = _drain_channel_locked(channel_id) if drained_messages > 0: @@ -580,16 +642,53 @@ def _submit_one_batch(channel_id: int) -> bool: """Submit the next batch for a channel to its sampler. Re-enqueues the channel into ``runnable_channel_ids`` if more batches - remain. - Returns True if a batch was submitted, False if the channel had no - pending work. + remain (via the membership-guarded ``_enqueue_channel_if_runnable_locked``). + + Parks the channel — returning without submitting and without + re-enqueueing — when it already has ``worker_concurrency`` batches in + flight, keeping the scheduler thread wait-free. A parked channel is + woken only by ``_on_batch_done`` when an in-flight batch completes. + + TODO(kmonte): failed-task finalization is deferred. GLT's + ``ConcurrentEventLoop.on_done`` fires this channel's callback only after + ``f.result()`` succeeds; a coroutine that truly raises/cancels releases + the semaphore without incrementing ``completed_batches``, so the channel + would park forever and any deferred unregister would stick at + ``completed < submitted``. In practice GiGL's ``_send_adapter`` override + (``base_sampler.py``) swallows sampling/collate errors in channel mode, + sends a one-time poison pill, and returns ``None`` normally — so the + callback still fires and in-flight decrements — which covers the common + failure path. Closing the residual raise/cancel gap needs an + always-fired finalization wrapping GLT's future (``add_task`` returns + ``None`` and does not expose it), so it is left as a follow-up. + + Returns True if a batch was submitted, False if the channel was parked + or had no pending work (so a park-only pump cycle leaves + ``made_progress`` unset). """ - # Hold the lock only to read state and advance the cursor. - # Release before the sampler call to avoid blocking other threads. + # Hold the lock only to read state and advance the cursor; release + # before the sampler call. + # + # INVARIANT (load-bearing): the sampler ``sample_*`` call MUST stay + # OUTSIDE ``state_lock``. Moving it under the lock turns the bounded + # semaphore handoff into a real deadlock, since ``_on_batch_done`` needs + # ``state_lock`` before the sampler's coroutine can release its + # semaphore slot -- the scheduler would then hold the lock while waiting + # on ``add_task``'s ``_sem.acquire()`` for a slot only the callback can + # free. with state_lock: state = active_epoch_by_channel_id.get(channel_id) if state is None: return False + # PARK: skip a channel already at its in-flight cap (the + # ``worker_concurrency`` block above explains why a further submit + # would block this thread). Return without submitting AND without + # re-enqueueing -- ``_on_batch_done`` is the sole path that wakes the + # channel once an in-flight batch completes. Returning False leaves + # ``made_progress`` unset, so a park-only cycle does not busy-spin: + # Phase 3's idle tick engages instead. + if _is_channel_parked_locked(channel_id): + return False batch_indices = _epoch_batch_indices(state) if batch_indices is None: return False @@ -599,9 +698,11 @@ def _submit_one_batch(channel_id: int) -> bool: channel_input = input_by_channel_id[channel_id] current_epoch = state.epoch # Re-enqueue for the next round-robin pass if more batches remain. - if state.submitted_batches < state.total_batches and not state.cancelled: - runnable_channel_ids.append(channel_id) - runnable_channel_id_set.add(channel_id) + # Use the membership-guarded enqueue rather than an unconditional + # append: a concurrent ``_on_batch_done`` (running on a sampler + # thread) may have re-enqueued this channel between the pump's pop + # and here, and a second append would create a duplicate turn. + _enqueue_channel_if_runnable_locked(channel_id) sampler_input = channel_input[batch_indices] @@ -637,6 +738,10 @@ def _submit_one_batch(channel_id: int) -> bool: def _pump_runnable_channel_ids() -> bool: """Submit one batch per runnable channel in round-robin order. + Parked channels (see ``_submit_one_batch``) submit nothing, so + ``made_progress`` stays False and Phase 3's idle tick engages instead of + busy-spinning. + Returns True if at least one batch was submitted. """ made_progress = False @@ -803,6 +908,20 @@ def _handle_command(command: SharedMpCommand, payload: CommandPayload) -> bool: processed_command = True keep_running = _handle_command(command, payload) + # STOP drained above (``keep_running`` only ever goes False on STOP), + # so exit before Phase 2 rather than pumping once more. A final pump + # here could submit a fresh batch on a channel that a completion + # callback just re-enqueued (the callback runs concurrently and both + # re-enqueues the channel and wakes Phase 3), and that batch's + # ``channel.send`` can block forever on a full output channel whose + # consumer is already gone at teardown -- wedging ``finally``'s + # ``sampler.wait_all()``. Teardown needs no final pump: the + # ``finally`` block drains and shuts down every sampler, and any + # pending unregister finalizes through its own destroy/unregister + # protocol. + if not keep_running: + break + # Phase 2: Submit batches round-robin from runnable channels. made_progress = _pump_runnable_channel_ids() made_progress = _finish_ready_unregistrations() or made_progress @@ -810,13 +929,25 @@ def _handle_command(command: SharedMpCommand, payload: CommandPayload) -> bool: if not keep_running: break - # Phase 3: If idle (no commands, no batches), block until next command. + # Phase 3: Idle (no commands drained, no batches submitted), so wait + # on the completion wake-event with a fallback tick, then loop back + # to Phase 1. A batch completion (``_on_batch_done``) re-enqueues + # the freed channel and ``wake_event.set()``s, so this returns at + # once and the freed slot refills without idling ~25ms on average. + # ``wake_event`` is level-triggered: a ``set()`` racing ahead of this + # ``wait()`` returns True immediately (no lost wakeup); a ``set()`` + # seen before we reach the wait costs at most one extra non-sleeping + # spin. Commands are not handled here -- the ``continue`` returns to + # Phase 1, which drains ``task_queue`` via ``get_nowait()`` every + # iteration -- so worst-case command latency is one + # ``SCHEDULER_TICK_SECS`` tick, negligible against the per-epoch + # (~1e5 batches apart) command cadence. if not (processed_command or made_progress): - try: - command, payload = task_queue.get(timeout=SCHEDULER_TICK_SECS) - except queue.Empty: - continue - keep_running = _handle_command(command, payload) + if wake_event.wait(timeout=SCHEDULER_TICK_SECS): + # A completion woke us: clear so the next idle wait blocks + # again. + wake_event.clear() + continue except KeyboardInterrupt: pass finally: diff --git a/tests/unit/distributed/graph_store/shared_dist_sampling_producer_test.py b/tests/unit/distributed/graph_store/shared_dist_sampling_producer_test.py index 56e06532d..4b61b7eb5 100644 --- a/tests/unit/distributed/graph_store/shared_dist_sampling_producer_test.py +++ b/tests/unit/distributed/graph_store/shared_dist_sampling_producer_test.py @@ -1,6 +1,9 @@ +import contextlib import queue import threading -from collections.abc import Callable +import time +from collections import deque +from collections.abc import Callable, Iterator from typing import cast from unittest.mock import MagicMock, patch @@ -347,6 +350,7 @@ def test_worker_unregister_drains_buffered_output_and_waits_for_completion( worker_options.master_addr = "127.0.0.1" worker_options.master_port = 12345 worker_options.rpc_timeout = 30 + worker_options.worker_concurrency = 2 output_channel = _FakeOutputChannel() fake_sampler = _DeferredFakeSampler(output_channel) mock_create_dist_sampler.return_value = fake_sampler @@ -424,3 +428,845 @@ def run_callback() -> None: task_queue.put((SharedMpCommand.STOP, None)) worker_thread.join(timeout=5.0) self.assertFalse(worker_thread.is_alive()) + + +class _BoundedBlockingChannel: + """Output channel with GLT ``ShmChannel``-like bounded, blocking semantics. + + ``send`` blocks while the buffer is full: a paused consumer never drains, so + its channel saturates and the sampler coroutines wedge in ``send`` -- exactly + the condition that must park a channel instead of blocking the shared + scheduler thread. + + ``recv`` mirrors ``ShmChannel``: a positive ``timeout_ms`` raises + ``QueueTimeoutError`` on an empty channel, while a non-positive/None timeout + on an empty channel is a would-be production deadlock and asserts. + """ + + def __init__(self, capacity: int) -> None: + self._capacity = capacity + self._buffer: deque[object] = deque() + self._cond = threading.Condition() + self.total_sent = 0 + self.total_received = 0 + + def send(self, msg: object) -> None: + with self._cond: + while len(self._buffer) >= self._capacity: + self._cond.wait() + self._buffer.append(msg) + self.total_sent += 1 + self._cond.notify_all() + + def recv(self, timeout_ms: int | None = None, **_: object) -> object: + with self._cond: + if not self._buffer: + if timeout_ms is None or timeout_ms <= 0: + raise AssertionError( + "_BoundedBlockingChannel.recv called with blocking timeout " + "on empty channel -- production code is about to deadlock." + ) + deadline = time.monotonic() + timeout_ms / 1000.0 + while not self._buffer: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise QueueTimeoutError("Timeout: Queue is empty.") + self._cond.wait(remaining) + msg = self._buffer.popleft() + self.total_received += 1 + self._cond.notify_all() + return msg + + +class _GltOrderFakeSampler: + """Sampler that reproduces GLT ``ConcurrentEventLoop`` submit/complete order. + + ``sample_from_nodes`` acquires a per-channel ``BoundedSemaphore`` on the + CALLING (scheduler) thread -- the exact point where GLT's ``add_task`` blocks + once ``worker_concurrency`` coroutines are already pending -- and only then + hands the "coroutine" to a background thread that (1) sends the result on the + bounded channel (blocking while the consumer is paused), (2) fires the + completion callback, and (3) releases the semaphore, in that order. + + Firing the callback BEFORE the release mirrors ``event_loop.py`` ``on_done`` + (``result``/``callback`` then ``self._sem.release()``), so + ``completed_batches`` reflects a freed slot at the same instant the real code + would. The semaphore acquire on the scheduler thread means that if the + scheduler ever over-submits past the in-flight cap (i.e. the fix regresses), + the whole scheduler thread wedges here -- which the tests detect as a stall. + """ + + def __init__( + self, output_channel: _BoundedBlockingChannel, concurrency: int + ) -> None: + self._channel = output_channel + self._sem = threading.BoundedSemaphore(concurrency) + self._threads: list[threading.Thread] = [] + self._threads_lock = threading.Lock() + self.submit_count = 0 + + def start_loop(self) -> None: ... + + def wait_all(self) -> None: + with self._threads_lock: + threads = list(self._threads) + for thread in threads: + thread.join() + + def shutdown_loop(self) -> None: + self.wait_all() + + def sample_from_nodes( + self, _sampler_input: object, callback: Callable[[object], None] + ) -> None: + with self._threads_lock: + self.submit_count += 1 + # Acquired on the scheduler thread: blocks here iff the scheduler + # over-submits past the in-flight cap (the regression this fix prevents). + self._sem.acquire() + + def _coroutine() -> None: + try: + self._channel.send({"seed": torch.tensor([1], dtype=torch.long)}) + callback(None) + finally: + self._sem.release() + + thread = threading.Thread(target=_coroutine, daemon=True) + with self._threads_lock: + self._threads.append(thread) + thread.start() + + +class _CountingTaskQueue: + """Task queue that counts blocking vs non-blocking gets. + + Phase 3 idles on ``wake_event.wait()``, not on the task queue, so a correctly + parked scheduler issues ZERO blocking ``task_queue.get`` calls and only a + handful of ``get_nowait`` calls -- one per idle wake, each iteration draining + commands once before blocking on the event for a full ``SCHEDULER_TICK_SECS``. + The counts let a test catch a busy-spin regression: a scheduler that wrongly + reports progress on a park-only pump cycle skips the wait and burns thousands + of ``get_nowait`` calls at 100% CPU instead of engaging the idle wait. + """ + + def __init__(self) -> None: + self._queue: queue.Queue[tuple[SharedMpCommand, object]] = queue.Queue() + self._lock = threading.Lock() + self.blocking_get_calls = 0 + self.nowait_get_calls = 0 + + def put(self, item: tuple[SharedMpCommand, object]) -> None: + self._queue.put(item) + + def get_nowait(self) -> tuple[SharedMpCommand, object]: + with self._lock: + self.nowait_get_calls += 1 + return self._queue.get_nowait() + + def get(self, timeout: float | None = None) -> tuple[SharedMpCommand, object]: + with self._lock: + self.blocking_get_calls += 1 + return self._queue.get(timeout=timeout) + + +class StallFixWorkerLoopTest(TestCase): + """GLT-order-faithful regressions for the wait-free scheduler stall fix. + + Each drives the real ``_shared_sampling_worker_loop`` with a paused consumer + (its channel saturates and its coroutines wedge in ``send``, holding the + per-channel sampler semaphore) alongside an active consumer, and asserts the + active channel keeps progressing and commands keep draining -- i.e. the + single scheduler thread is never parked on the saturated channel. + + Without the in-flight cap the scheduler blocks in ``sample_from_nodes`` on the + saturated channel's exhausted semaphore, so these ``event_queue.get`` calls + time out (the pre-fix cross-rank deadlock). + """ + + @staticmethod + def _make_worker_options(worker_concurrency: int) -> MagicMock: + worker_options = MagicMock() + worker_options.worker_world_size = 1 + worker_options.worker_ranks = [0] + worker_options.use_all2all = False + worker_options.num_rpc_threads = 1 + worker_options.worker_devices = [torch.device("cpu")] + worker_options.master_addr = "127.0.0.1" + worker_options.master_port = 12345 + worker_options.rpc_timeout = 30 + worker_options.worker_concurrency = worker_concurrency + return worker_options + + @staticmethod + @contextlib.contextmanager + def _draining(*channels: _BoundedBlockingChannel) -> Iterator[None]: + """Continuously drain the given channels for the duration of the block. + + Releases any sends wedged on a full channel so the worker can always + reach Phase 1 and drain ``STOP``. Without this, a reverted in-flight cap + would leave the non-daemon worker blocked in the sampler semaphore, so a + red test would hang the whole suite instead of failing cleanly. + """ + stop = threading.Event() + + def _drain() -> None: + while not stop.is_set(): + for channel in channels: + try: + channel.recv(timeout_ms=20) + except QueueTimeoutError: + continue + + thread = threading.Thread(target=_drain, daemon=True) + thread.start() + try: + yield + finally: + stop.set() + thread.join(timeout=5.0) + + def _run_paused_and_active( + self, + *, + mock_create_dist_sampler: MagicMock, + ) -> None: + worker_concurrency = 2 + # channel_a's consumer drains continuously so its coroutines never wedge; + # channel_b has NO consumer, so it saturates at capacity and its + # coroutines block in send -- the paused-consumer condition. Under plain + # round-robin the scheduler must rotate past channel_b's saturated channel + # instead of parking the shared thread on it. + channel_a = _BoundedBlockingChannel(capacity=4) + channel_b = _BoundedBlockingChannel(capacity=worker_concurrency) + + mock_create_dist_sampler.side_effect = lambda **kwargs: _GltOrderFakeSampler( + kwargs["channel"], worker_concurrency + ) + + worker_options = self._make_worker_options(worker_concurrency) + task_queue: queue.Queue[tuple[SharedMpCommand, object]] = queue.Queue() + event_queue: queue.Queue[tuple[object, ...]] = queue.Queue() + barrier = MagicMock(wait=MagicMock()) + data = MagicMock(num_partitions=1) + sampling_config = _make_sampling_config() + + channel_a_id, channel_b_id = 1, 2 + active_batches, paused_batches = 6, 20 # batch_size 2 -> 12 vs 40 seeds + + stop_consumer = threading.Event() + + def _consume_active() -> None: + while not stop_consumer.is_set(): + try: + channel_a.recv(timeout_ms=20) + except QueueTimeoutError: + continue + + consumer_thread = threading.Thread(target=_consume_active, daemon=True) + consumer_thread.start() + + # create_dist_sampler picks up each channel by the ``channel`` kwarg, so + # channel_a's sampler sends to channel_a and channel_b's to channel_b. + for channel_id, channel, node_len in ( + (channel_a_id, channel_a, active_batches * 2), + (channel_b_id, channel_b, paused_batches * 2), + ): + task_queue.put( + ( + SharedMpCommand.REGISTER_INPUT, + RegisterInputCmd( + channel_id=channel_id, + worker_key=f"loader_compute_rank_{channel_id}", + sampler_input=NodeSamplerInput(node=torch.arange(node_len)), + sampling_config=sampling_config, + channel=channel, + ), + ) + ) + # Start the paused channel first so it begins saturating. + task_queue.put( + ( + SharedMpCommand.START_EPOCH, + StartEpochCmd( + channel_id=channel_b_id, + epoch=0, + seeds_index=torch.arange(paused_batches * 2), + ), + ) + ) + task_queue.put( + ( + SharedMpCommand.START_EPOCH, + StartEpochCmd( + channel_id=channel_a_id, + epoch=0, + seeds_index=torch.arange(active_batches * 2), + ), + ) + ) + + worker_thread = threading.Thread( + target=_shared_sampling_worker_loop, + args=( + 0, + data, + worker_options, + task_queue, + event_queue, + barrier, + KHopNeighborSamplerOptions(num_neighbors=[2]), + None, + ), + ) + worker_thread.start() + try: + # Active channel completes epoch 0 despite channel B's paused + # consumer: the scheduler rotated past B's saturated channel instead + # of parking on it. (Pre-fix: deadlock -> this get times out.) + done_epoch_0 = event_queue.get(timeout=10.0) + self.assertEqual(done_epoch_0, (EPOCH_DONE_EVENT, channel_a_id, 0, 0)) + + # Commands keep draining while B stays wedged: start a SECOND active + # epoch AFTER B has saturated and confirm it also completes (proves + # Phase-1 command draining + Phase-2 pumping never parked). + task_queue.put( + ( + SharedMpCommand.START_EPOCH, + StartEpochCmd( + channel_id=channel_a_id, + epoch=1, + seeds_index=torch.arange(active_batches * 2), + ), + ) + ) + done_epoch_1 = event_queue.get(timeout=10.0) + self.assertEqual(done_epoch_1, (EPOCH_DONE_EVENT, channel_a_id, 1, 0)) + + # The paused channel never streamed to a consumer while the active + # channel drained two full epochs -- asymmetric progress, no + # head-of-line blocking. + self.assertEqual(channel_b.total_received, 0) + self.assertGreaterEqual(channel_a.total_received, active_batches) + + # Unregister the wedged (parked) channel: draining its buffered output + # unblocks the stuck sends so cleanup finalizes without hanging. + task_queue.put((SharedMpCommand.UNREGISTER_INPUT, channel_b_id)) + finally: + # Drain the paused channel while stopping so the worker can always + # reach Phase 1 to process STOP -- see ``_draining``. + with self._draining(channel_b): + task_queue.put((SharedMpCommand.STOP, None)) + worker_thread.join(timeout=10.0) + stop_consumer.set() + consumer_thread.join(timeout=5.0) + + self.assertFalse(worker_thread.is_alive()) + + @patch("gigl.distributed.graph_store.shared_dist_sampling_producer.shutdown_rpc") + @patch("gigl.distributed.graph_store.shared_dist_sampling_producer.init_rpc") + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer.init_worker_group" + ) + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer._set_worker_signal_handlers" + ) + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer.torch.set_num_threads" + ) + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer.create_dist_sampler" + ) + def test_paused_consumer_does_not_stall_active_channel_unweighted( + self, + mock_create_dist_sampler: MagicMock, + _mock_set_num_threads: MagicMock, + _mock_signal_handlers: MagicMock, + _mock_init_worker_group: MagicMock, + _mock_init_rpc: MagicMock, + _mock_shutdown_rpc: MagicMock, + ) -> None: + self._run_paused_and_active( + mock_create_dist_sampler=mock_create_dist_sampler, + ) + + @patch("gigl.distributed.graph_store.shared_dist_sampling_producer.shutdown_rpc") + @patch("gigl.distributed.graph_store.shared_dist_sampling_producer.init_rpc") + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer.init_worker_group" + ) + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer._set_worker_signal_handlers" + ) + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer.torch.set_num_threads" + ) + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer.create_dist_sampler" + ) + def test_unregister_while_parked_tears_down_cleanly( + self, + mock_create_dist_sampler: MagicMock, + _mock_set_num_threads: MagicMock, + _mock_signal_handlers: MagicMock, + _mock_init_worker_group: MagicMock, + _mock_init_rpc: MagicMock, + _mock_shutdown_rpc: MagicMock, + ) -> None: + worker_concurrency = 2 + channel_b = _BoundedBlockingChannel(capacity=worker_concurrency) + created_samplers: list[_GltOrderFakeSampler] = [] + + def _make_sampler(**kwargs: object) -> _GltOrderFakeSampler: + sampler = _GltOrderFakeSampler( + cast(_BoundedBlockingChannel, kwargs["channel"]), worker_concurrency + ) + created_samplers.append(sampler) + return sampler + + mock_create_dist_sampler.side_effect = _make_sampler + + worker_options = self._make_worker_options(worker_concurrency) + task_queue = _CountingTaskQueue() + event_queue: queue.Queue[tuple[object, ...]] = queue.Queue() + barrier = MagicMock(wait=MagicMock()) + data = MagicMock(num_partitions=1) + sampling_config = _make_sampling_config() + channel_b_id, paused_batches = 5, 20 + + task_queue.put( + ( + SharedMpCommand.REGISTER_INPUT, + RegisterInputCmd( + channel_id=channel_b_id, + worker_key="loader_compute_rank_5", + sampler_input=NodeSamplerInput( + node=torch.arange(paused_batches * 2) + ), + sampling_config=sampling_config, + channel=channel_b, + ), + ) + ) + task_queue.put( + ( + SharedMpCommand.START_EPOCH, + StartEpochCmd( + channel_id=channel_b_id, + epoch=0, + seeds_index=torch.arange(paused_batches * 2), + ), + ) + ) + + worker_thread = threading.Thread( + target=_shared_sampling_worker_loop, + args=( + 0, + data, + worker_options, + task_queue, + event_queue, + barrier, + KHopNeighborSamplerOptions(num_neighbors=[2]), + None, + ), + ) + worker_thread.start() + try: + # The channel parks after submitting exactly capacity (completed + + # buffered) plus worker_concurrency (wedged in-flight) batches, and no + # more -- a busy-spinning or non-parking scheduler would keep climbing + # toward all 20 batches (or wedge the whole scheduler on the exhausted + # semaphore before reaching this count). + expected_parked_submits = worker_concurrency + channel_b._capacity + deadline = time.monotonic() + 10.0 + sampler: _GltOrderFakeSampler | None = None + while time.monotonic() < deadline: + if created_samplers: + sampler = created_samplers[0] + if sampler.submit_count >= expected_parked_submits: + break + time.sleep(0.01) + self.assertIsNotNone(sampler) + assert sampler is not None + self.assertGreaterEqual(sampler.submit_count, expected_parked_submits) + + # Give the scheduler ample time to (incorrectly) submit more or + # busy-spin; parked means it idles in Phase 3 instead. + blocking_gets_before = task_queue.blocking_get_calls + nowait_gets_before = task_queue.nowait_get_calls + time.sleep(0.3) + blocking_gets = task_queue.blocking_get_calls - blocking_gets_before + nowait_gets = task_queue.nowait_get_calls - nowait_gets_before + + # It stays put -- no submits past the cap. + self.assertEqual(sampler.submit_count, expected_parked_submits) + self.assertEqual(channel_b.total_received, 0) + # F3 event-wake: Phase 3 now waits on the completion ``wake_event``, NOT on + # ``task_queue``, so the scheduler issues ZERO blocking ``task_queue.get`` + # calls while parked. Still NOT busy-spinning: each idle iteration does one + # Phase-1 ``get_nowait()`` then blocks in ``wake_event.wait()`` for a full + # tick, so over ~0.3s / 0.05s tick only a handful of ``get_nowait`` calls + # accrue -- a regression that wrongly set ``made_progress`` on a park-only + # cycle would skip the wait and burn thousands of ``get_nowait`` at 100% CPU. + self.assertEqual(blocking_gets, 0) + self.assertGreaterEqual(nowait_gets, 1) + self.assertLess(nowait_gets, 100) + + # Unregister while parked: must drain buffered output, unblock the + # wedged sends, and finalize cleanup without hanging. + task_queue.put((SharedMpCommand.UNREGISTER_INPUT, channel_b_id)) + finally: + # Drain the wedged channel while stopping so the worker can always + # reach Phase 1 to process STOP -- see ``_draining``. + with self._draining(channel_b): + task_queue.put((SharedMpCommand.STOP, None)) + worker_thread.join(timeout=10.0) + + self.assertFalse(worker_thread.is_alive()) + # Every wedged coroutine was released during teardown. + for thread in sampler._threads: + self.assertFalse(thread.is_alive()) + + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer.SCHEDULER_TICK_SECS", + 30.0, + ) + @patch("gigl.distributed.graph_store.shared_dist_sampling_producer.shutdown_rpc") + @patch("gigl.distributed.graph_store.shared_dist_sampling_producer.init_rpc") + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer.init_worker_group" + ) + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer._set_worker_signal_handlers" + ) + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer.torch.set_num_threads" + ) + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer.create_dist_sampler" + ) + def test_completion_event_wakes_parked_channel_without_tick( + self, + mock_create_dist_sampler: MagicMock, + _mock_set_num_threads: MagicMock, + _mock_signal_handlers: MagicMock, + _mock_init_worker_group: MagicMock, + _mock_init_rpc: MagicMock, + _mock_shutdown_rpc: MagicMock, + ) -> None: + """A batch completion wakes the parked scheduler immediately via the F3 + ``wake_event`` -- it does NOT wait out the Phase-3 fallback tick. + + ``SCHEDULER_TICK_SECS`` is patched to 30s so the fallback tick is + effectively disabled: the ONLY way the parked channel's freed slot is + refilled within the test timeout is the event-wake. This makes the wake + path load-bearing (non-vacuous). Without the fix -- i.e. no + ``wake_event.set()`` in ``_on_batch_done`` and Phase 3 sleeping on the + tick -- the second submit would not issue for 30s and the + ``submit_signals.get(timeout=5.0)`` below would time out (this is exactly + the ~25ms-mean freed-slot idle the fix removes, amplified to 30s here). + + ``worker_concurrency=1`` so a single in-flight batch parks the channel; + the sampler withholds completion callbacks so parking is fully + controlled by the test. ``_CountingTaskQueue`` confirms Phase 3 waits on + the event rather than issuing any blocking ``task_queue.get``. + """ + worker_concurrency = 1 + channel_id = 3 + + class _WithheldSampler: + """Sampler that records each submit and withholds its callback. + + Withholding the callback keeps the single in-flight batch pending, so + at ``worker_concurrency=1`` the channel parks after one submit and + only the test can drive completion (and thus the wake). + """ + + def __init__(self) -> None: + self.callbacks: list[Callable[[object], None]] = [] + self.submit_signals: queue.Queue[int] = queue.Queue() + self.submit_count = 0 + self._lock = threading.Lock() + + def start_loop(self) -> None: ... + + def wait_all(self) -> None: ... + + def shutdown_loop(self) -> None: ... + + def sample_from_nodes( + self, _sampler_input: object, callback: Callable[[object], None] + ) -> None: + with self._lock: + self.callbacks.append(callback) + self.submit_count += 1 + index = self.submit_count + self.submit_signals.put(index) + + sampler = _WithheldSampler() + mock_create_dist_sampler.side_effect = lambda **_: sampler + + worker_options = self._make_worker_options(worker_concurrency) + task_queue = _CountingTaskQueue() + event_queue: queue.Queue[tuple[object, ...]] = queue.Queue() + barrier = MagicMock(wait=MagicMock()) + data = MagicMock(num_partitions=1) + sampling_config = _make_sampling_config() + + # 3 batches (6 seeds / batch_size 2): enough to submit, park, then resume. + task_queue.put( + ( + SharedMpCommand.REGISTER_INPUT, + RegisterInputCmd( + channel_id=channel_id, + worker_key="loader_compute_rank_3", + sampler_input=NodeSamplerInput(node=torch.arange(6)), + sampling_config=sampling_config, + channel=MagicMock(), + ), + ) + ) + task_queue.put( + ( + SharedMpCommand.START_EPOCH, + StartEpochCmd( + channel_id=channel_id, + epoch=0, + seeds_index=torch.arange(6), + ), + ) + ) + + worker_thread = threading.Thread( + target=_shared_sampling_worker_loop, + args=( + 0, + data, + worker_options, + cast(mp.Queue, task_queue), + cast(mp.Queue, event_queue), + barrier, + KHopNeighborSamplerOptions(num_neighbors=[2]), + None, + ), + ) + worker_thread.start() + try: + # First batch submits, then the channel parks (in-flight == cap with + # the callback withheld) and the scheduler enters Phase 3, blocking + # on wake_event with the (patched 30s) fallback tick. + first_submit = sampler.submit_signals.get(timeout=10.0) + self.assertEqual(first_submit, 1) + + # Parked: no further submit issues on its own. With the 30s tick + # this also confirms the scheduler is genuinely blocked in the wait, + # not spinning through ticks. + time.sleep(0.2) + self.assertTrue(sampler.submit_signals.empty()) + self.assertEqual(len(sampler.callbacks), 1) + # Phase 3 has NOT touched task_queue -- it waits on the event. + self.assertEqual(task_queue.blocking_get_calls, 0) + + # Fire the withheld completion from a separate thread (mirroring the + # sampler worker thread): _on_batch_done re-enqueues the channel and + # wake_event.set()s, waking Phase 3. + first_callback = sampler.callbacks[0] + wake_thread = threading.Thread(target=lambda: first_callback(None)) + wake_thread.start() + wake_thread.join(timeout=5.0) + self.assertFalse(wake_thread.is_alive()) + + # LOAD-BEARING: the freed slot is refilled PROMPTLY via the event -- + # within 5s, far under the 30s fallback tick. Without + # wake_event.set() the scheduler would sleep the full 30s and this + # get would time out. + second_submit = sampler.submit_signals.get(timeout=5.0) + self.assertEqual(second_submit, 2) + + # The wake came purely from the event: still zero blocking + # task_queue.get calls. + self.assertEqual(task_queue.blocking_get_calls, 0) + finally: + # The channel re-parked after the second submit (callback 2 withheld), + # so STOP alone would sit unseen until the 30s tick. Fire the second + # completion to set wake_event -> Phase 3 wakes -> Phase 1 drains STOP. + task_queue.put((SharedMpCommand.STOP, None)) + if len(sampler.callbacks) >= 2: + sampler.callbacks[1](None) + worker_thread.join(timeout=10.0) + + self.assertFalse(worker_thread.is_alive()) + # Phase 3 never blocked on the task queue across the whole run. + self.assertEqual(task_queue.blocking_get_calls, 0) + # No submit after STOP: the epoch has 3 batches, but the STOP iteration + # must break right after Phase 1 drains STOP -- BEFORE Phase 2 -- so the + # 3rd batch is never submitted. With event-wake command handling in + # Phase 1, firing the 2nd completion both re-enqueues the channel and + # wakes Phase 3; without the pre-Phase-2 break the scheduler would then + # pump that 3rd batch after teardown began, whose ``channel.send`` would + # block on a full output channel at real teardown. Exactly the two + # submits observed above (initial + event-woken refill) occurred. + self.assertEqual(sampler.submit_count, 2) + + @patch("gigl.distributed.graph_store.shared_dist_sampling_producer.shutdown_rpc") + @patch("gigl.distributed.graph_store.shared_dist_sampling_producer.init_rpc") + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer.init_worker_group" + ) + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer._set_worker_signal_handlers" + ) + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer.torch.set_num_threads" + ) + @patch( + "gigl.distributed.graph_store.shared_dist_sampling_producer.create_dist_sampler" + ) + def test_parked_channel_resumes_and_completes_when_consumer_drains( + self, + mock_create_dist_sampler: MagicMock, + _mock_set_num_threads: MagicMock, + _mock_signal_handlers: MagicMock, + _mock_init_worker_group: MagicMock, + _mock_init_rpc: MagicMock, + _mock_shutdown_rpc: MagicMock, + ) -> None: + """A parked channel wakes and finishes its epoch once its consumer drains. + + The other stall-fix tests only park a channel; none resumes it. This + pins the WAKE path directly: ``_on_batch_done`` re-enqueuing the parked + channel is the SOLE way a paused-then-resumed channel makes progress + again (the pump never revisits a channel it parked). If that re-enqueue + regressed, the epoch would hang after the consumer resumed and the + ``event_queue.get`` below would time out. + """ + worker_concurrency = 2 + channel = _BoundedBlockingChannel(capacity=worker_concurrency) + created_samplers: list[_GltOrderFakeSampler] = [] + + def _make_sampler(**kwargs: object) -> _GltOrderFakeSampler: + sampler = _GltOrderFakeSampler( + cast(_BoundedBlockingChannel, kwargs["channel"]), worker_concurrency + ) + created_samplers.append(sampler) + return sampler + + mock_create_dist_sampler.side_effect = _make_sampler + + worker_options = self._make_worker_options(worker_concurrency) + task_queue: queue.Queue[tuple[SharedMpCommand, object]] = queue.Queue() + event_queue: queue.Queue[tuple[object, ...]] = queue.Queue() + barrier = MagicMock(wait=MagicMock()) + data = MagicMock(num_partitions=1) + sampling_config = _make_sampling_config() + channel_id, total_batches = 7, 12 # batch_size 2 -> 24 seeds + + task_queue.put( + ( + SharedMpCommand.REGISTER_INPUT, + RegisterInputCmd( + channel_id=channel_id, + worker_key="loader_compute_rank_7", + sampler_input=NodeSamplerInput( + node=torch.arange(total_batches * 2) + ), + sampling_config=sampling_config, + channel=channel, + ), + ) + ) + task_queue.put( + ( + SharedMpCommand.START_EPOCH, + StartEpochCmd( + channel_id=channel_id, + epoch=0, + seeds_index=torch.arange(total_batches * 2), + ), + ) + ) + + # The consumer stays paused until ``resume`` is set, so the channel + # saturates and parks first; only then do we resume draining. + resume = threading.Event() + stop_consumer = threading.Event() + + def _consume() -> None: + resume.wait() + while not stop_consumer.is_set(): + try: + channel.recv(timeout_ms=20) + except QueueTimeoutError: + continue + + consumer_thread = threading.Thread(target=_consume, daemon=True) + consumer_thread.start() + + worker_thread = threading.Thread( + target=_shared_sampling_worker_loop, + args=( + 0, + data, + worker_options, + task_queue, + event_queue, + barrier, + KHopNeighborSamplerOptions(num_neighbors=[2]), + None, + ), + ) + worker_thread.start() + try: + # Wait for the channel to park at the in-flight cap: with the + # consumer paused, submit_count climbs to worker_concurrency + # (buffered) + capacity (wedged in send) and stops. + expected_parked_submits = worker_concurrency + channel._capacity + deadline = time.monotonic() + 10.0 + sampler: _GltOrderFakeSampler | None = None + while time.monotonic() < deadline: + if created_samplers: + sampler = created_samplers[0] + if sampler.submit_count >= expected_parked_submits: + break + time.sleep(0.01) + self.assertIsNotNone(sampler) + assert sampler is not None + + # Confirm it is genuinely parked (stuck at the cap, well short of the + # full epoch, nothing drained) before resuming -- not merely slow. + time.sleep(0.2) + self.assertEqual(sampler.submit_count, expected_parked_submits) + self.assertEqual(channel.total_received, 0) + self.assertLess(sampler.submit_count, total_batches) + + # Resume the consumer. Draining a batch completes a wedged send, + # which fires the callback -> _on_batch_done increments completed and + # WAKES the parked channel. Without that wake the epoch never + # finishes and this get times out. + resume.set() + done = event_queue.get(timeout=10.0) + self.assertEqual(done, (EPOCH_DONE_EVENT, channel_id, 0, 0)) + + # The woken channel submitted every remaining batch (deterministic: + # EPOCH_DONE implies all batches completed, so all were submitted). + self.assertEqual(sampler.submit_count, total_batches) + # EPOCH_DONE fires when the last batch is SENT (buffered); the + # consumer drains that tail a beat later, so poll rather than race it. + drain_deadline = time.monotonic() + 5.0 + while ( + channel.total_received < total_batches + and time.monotonic() < drain_deadline + ): + time.sleep(0.01) + self.assertEqual(channel.total_received, total_batches) + finally: + with self._draining(channel): + task_queue.put((SharedMpCommand.STOP, None)) + worker_thread.join(timeout=10.0) + stop_consumer.set() + resume.set() # unblock the consumer if we failed before resuming + consumer_thread.join(timeout=5.0) + + self.assertFalse(worker_thread.is_alive())