From 1ed00408541b85632db072ad638a2c3db1010940 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 23 Jul 2026 17:15:29 +0000 Subject: [PATCH 1/7] graph-store: high-fan-out scheduler stall fix (in-flight submission cap) Add a per-channel in-flight submission cap to the shared graph-store sampling worker's scheduler so the single scheduler thread stays wait-free at high fan-out. - worker_concurrency loop-local mirrors the sampler's bounded semaphore. - _is_channel_parked_locked: a channel with worker_concurrency batches submitted-but-not-completed is "parked". - _submit_one_batch PARKs (returns False, no re-enqueue) instead of issuing a submit that could block the scheduler on a saturated output channel; re-enqueue now goes through the membership-guarded helper to avoid a duplicate round-robin turn. - _on_batch_done WAKEs a parked channel by re-enqueueing it once an in-flight slot frees -- the sole wake path. - Phase 3 idle wait, both worker/backend signatures, and the plain round-robin pump body are unchanged from main. Co-Authored-By: Claude Opus 4.8 --- .../shared_dist_sampling_producer.py | 109 ++++++++++++++++-- 1 file changed, 101 insertions(+), 8 deletions(-) diff --git a/gigl/distributed/graph_store/shared_dist_sampling_producer.py b/gigl/distributed/graph_store/shared_dist_sampling_producer.py index 608504d8b..78fd68ad0 100644 --- a/gigl/distributed/graph_store/shared_dist_sampling_producer.py +++ b/gigl/distributed/graph_store/shared_dist_sampling_producer.py @@ -395,6 +395,24 @@ def _shared_sampling_worker_loop( state_lock = threading.RLock() 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 the + # core of the stall fix. + # + # 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 +429,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 +591,14 @@ 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) if channel_id in removing_channel_ids: drained_messages = _drain_channel_locked(channel_id) if drained_messages > 0: @@ -580,16 +624,54 @@ 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 (fix step 4) 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: if the channel already has ``worker_concurrency`` batches in + # flight, submitting another could block this thread in the sampler's + # bounded semaphore on a saturated output channel. Return without + # submitting AND without re-enqueueing -- ``_on_batch_done`` is the + # sole path that re-enqueues (wakes) the channel once an in-flight + # batch completes. Returning False leaves ``made_progress`` unset in + # the pump, so a cycle that only parks 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 +681,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 +721,15 @@ def _submit_one_batch(channel_id: int) -> bool: def _pump_runnable_channel_ids() -> bool: """Submit one batch per runnable channel in round-robin order. + Channels that have hit their in-flight cap are parked by + ``_submit_one_batch`` (return False, no re-enqueue); ``_on_batch_done`` + re-enqueues them once a slot frees. A cycle whose only visited channels + park submits nothing, so ``made_progress`` stays False and Phase 3's idle + tick engages instead of busy-spinning. + + Channels re-enqueued during this pass are deferred to the next pump cycle + to preserve round-robin fairness. + Returns True if at least one batch was submitted. """ made_progress = False From 383b45cd7e9aaec0b8dff3d8a41e70acb5ad49e6 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 23 Jul 2026 17:21:19 +0000 Subject: [PATCH 2/7] graph-store: high-fan-out stall-fix regression tests (deadlock guard) Adds the GLT-order-faithful regression tests that guard the per-channel in-flight cap against the cross-rank deadlock it fixes. The shared harness (_BoundedBlockingChannel, _GltOrderFakeSampler, _CountingTaskQueue) drives the real _shared_sampling_worker_loop with a paused consumer whose channel saturates and whose coroutines wedge in send, holding the sampler semaphore. StallFixWorkerLoopTest then asserts the co-located active channel keeps completing epochs and commands keep draining, and that a parked channel idles on Phase 3's timed wait instead of busy-spinning. Also give the existing worker-loop base test a real int worker_concurrency so the parked-channel check does not compare against a MagicMock. Co-Authored-By: Claude Opus 4.8 --- .../shared_dist_sampling_producer_test.py | 501 +++++++++++++++++- 1 file changed, 500 insertions(+), 1 deletion(-) 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..a4f17503b 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,498 @@ 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. + + Lets a test distinguish a scheduler idling in Phase 3 (``get`` with a + timeout, ~one per ``SCHEDULER_TICK_SECS``) from one busy-spinning through the + phases (``get_nowait`` many times per tick, zero blocking ``get`` calls) -- + the signature of a regression that wrongly reports progress on a park-only + pump cycle and pins the core at 100% CPU instead of engaging the idle tick. + """ + + 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) + # Phase 3's timed wait is engaged (>=1 blocking get over a ~0.3s / + # 0.05s tick window) and the loop is NOT busy-spinning: a regression + # that wrongly set ``made_progress`` on a park-only cycle would skip + # the blocking get entirely and burn thousands of ``get_nowait`` + # calls at 100% CPU. + self.assertGreaterEqual(blocking_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()) From e6233c8aaff1c384cad5900373b8945bff820948 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 23 Jul 2026 17:37:24 +0000 Subject: [PATCH 3/7] graph-store: event-driven scheduler wake for Phase-3 idle refill The per-server scheduler's Phase-3 idle wait blocked on task_queue.get(timeout); a batch completion re-enqueued a parked channel but could not interrupt that sleep, so a freed sampler slot idled up to one SCHEDULER_TICK_SECS (~50ms, mean ~25ms). Add a loop-local threading.Event: _on_batch_done set()s it (under state_lock, non-blocking, never task_queue.put) right after the WAKE re-enqueue; Phase 3 waits on the event with the existing tick as fallback, then loops back so Phase 1 drains commands via get_nowait(). The Event is level-triggered and set() happens under the same state_lock hold as the WAKE re-enqueue, so no wakeup is lost. Command handling moves entirely to Phase 1; worst-case command latency is one fallback tick, negligible at per-epoch cadence. PARK/WAKE/in-flight-cap and the sampler-call-outside-state_lock invariant are unchanged; the stall fix is not reopened. Co-Authored-By: Claude Opus 4.8 --- .../shared_dist_sampling_producer.py | 48 +++++++++++++++---- 1 file changed, 39 insertions(+), 9 deletions(-) diff --git a/gigl/distributed/graph_store/shared_dist_sampling_producer.py b/gigl/distributed/graph_store/shared_dist_sampling_producer.py index 78fd68ad0..01afe8376 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,6 +397,12 @@ 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 @@ -599,6 +609,14 @@ def _on_batch_done(channel_id: int, epoch: int) -> None: # 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: @@ -903,13 +921,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: From bea066cdf86e5d16261ee63c4e18c9c424a2e29e Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 23 Jul 2026 17:42:23 +0000 Subject: [PATCH 4/7] graph-store: event-wake regression test + flip parked-idle assertion Phase 3 now waits on the completion wake_event instead of task_queue, so a parked scheduler issues ZERO blocking task_queue.get calls. Flip test_unregister_while_parked_tears_down_cleanly's Phase-3 assertion to the event-wake form (blocking_gets == 0; still a handful of get_nowait, not a busy-spin) and reframe the neighboring comment. Add test_completion_event_wakes_parked_channel_without_tick: with SCHEDULER_TICK_SECS patched to 30s and worker_concurrency=1, a withheld-callback sampler parks the single channel; firing the completion must refill the freed slot within 5s via the event (far under the 30s fallback tick), with task_queue.blocking_get_calls == 0 throughout. Non-vacuous only with the event-wake: without wake_event.set() the second submit waits the full 30s and times out. Co-Authored-By: Claude Opus 4.8 --- .../shared_dist_sampling_producer_test.py | 185 +++++++++++++++++- 1 file changed, 179 insertions(+), 6 deletions(-) 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 a4f17503b..4aa814751 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 @@ -901,12 +901,15 @@ def _make_sampler(**kwargs: object) -> _GltOrderFakeSampler: # It stays put -- no submits past the cap. self.assertEqual(sampler.submit_count, expected_parked_submits) self.assertEqual(channel_b.total_received, 0) - # Phase 3's timed wait is engaged (>=1 blocking get over a ~0.3s / - # 0.05s tick window) and the loop is NOT busy-spinning: a regression - # that wrongly set ``made_progress`` on a park-only cycle would skip - # the blocking get entirely and burn thousands of ``get_nowait`` - # calls at 100% CPU. - self.assertGreaterEqual(blocking_gets, 1) + # 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 @@ -923,3 +926,173 @@ def _make_sampler(**kwargs: object) -> _GltOrderFakeSampler: # 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._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) + index = len(self.callbacks) + 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) From d9e14b9199b131238a0066213044ad651601ce3b Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Thu, 23 Jul 2026 17:55:35 +0000 Subject: [PATCH 5/7] =?UTF-8?q?graph-store:=20event-wake=20shutdown=20fix?= =?UTF-8?q?=20=E2=80=94=20break=20before=20post-STOP=20pump?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix a shutdown hang reachable once event-wake moved command handling into Phase 1. A parked channel's in-flight batch could complete (re-enqueue + wake_event.set()), a STOP arrive, Phase 3 wake, Phase 1 drain STOP, and Phase 2 still submit a fresh batch on the unparked channel whose channel.send blocks on a full output channel at teardown — wedging finally's wait_all(). Add an early `if not keep_running: break` immediately after the Phase-1 drain and before Phase 2 so the STOP iteration never pumps again. Teardown is handled by the finally block and the separate unregister protocol, so no final pump is needed. Only the STOP path is affected (keep_running only goes False on STOP). Strengthen test_completion_event_wakes_parked_channel_without_tick: the fake sampler now tracks submit_count, and the test asserts exactly 2 submits after shutdown (no 3rd batch after STOP). This fails without the break (3 != 2). Update the stale _CountingTaskQueue docstring to describe PR-2's wake_event Phase-3 wait instead of the old blocking task_queue.get. Co-Authored-By: Claude Opus 4.8 --- .../shared_dist_sampling_producer.py | 14 +++++++++++ .../shared_dist_sampling_producer_test.py | 25 ++++++++++++++----- 2 files changed, 33 insertions(+), 6 deletions(-) diff --git a/gigl/distributed/graph_store/shared_dist_sampling_producer.py b/gigl/distributed/graph_store/shared_dist_sampling_producer.py index 01afe8376..163d74d31 100644 --- a/gigl/distributed/graph_store/shared_dist_sampling_producer.py +++ b/gigl/distributed/graph_store/shared_dist_sampling_producer.py @@ -914,6 +914,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 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 4aa814751..20150e4aa 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 @@ -541,11 +541,13 @@ def _coroutine() -> None: class _CountingTaskQueue: """Task queue that counts blocking vs non-blocking gets. - Lets a test distinguish a scheduler idling in Phase 3 (``get`` with a - timeout, ~one per ``SCHEDULER_TICK_SECS``) from one busy-spinning through the - phases (``get_nowait`` many times per tick, zero blocking ``get`` calls) -- - the signature of a regression that wrongly reports progress on a park-only - pump cycle and pins the core at 100% CPU instead of engaging the idle tick. + 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: @@ -985,6 +987,7 @@ class _WithheldSampler: 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: ... @@ -998,7 +1001,8 @@ def sample_from_nodes( ) -> None: with self._lock: self.callbacks.append(callback) - index = len(self.callbacks) + self.submit_count += 1 + index = self.submit_count self.submit_signals.put(index) sampler = _WithheldSampler() @@ -1096,3 +1100,12 @@ def sample_from_nodes( 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) From a762961609168cde201adf4a436890174f401461 Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Fri, 24 Jul 2026 17:35:51 +0000 Subject: [PATCH 6/7] graph-store stall fix: trim redundant comments + add WAKE-resume test Apply comment/docstring review feedback (humanify): reword the top in-flight-cap block to a standing property (drop the change-story "core of the stall fix"), remove the "(fix step 4)" plan reference from the finalization TODO, and cut the semaphore-block rationale that was restated in the PARK inline comment and the _pump_runnable_channel_ids docstring down to pointers at the canonical top block. No logic change. Add test_parked_channel_resumes_and_completes_when_consumer_drains (ace-it): the existing tests only park a channel; none resumes one. This pins the WAKE path directly -- park a channel by pausing its consumer, then resume and assert the epoch completes (EPOCH_DONE + every batch drained). Verified non-vacuous: with the _on_batch_done re-enqueue disabled the test fails at the epoch-done get (10s timeout), since the pump never revisits a channel it parked. Co-Authored-By: Claude Opus 4.8 --- .../shared_dist_sampling_producer.py | 32 ++-- .../shared_dist_sampling_producer_test.py | 152 ++++++++++++++++++ 2 files changed, 165 insertions(+), 19 deletions(-) diff --git a/gigl/distributed/graph_store/shared_dist_sampling_producer.py b/gigl/distributed/graph_store/shared_dist_sampling_producer.py index 78fd68ad0..3d793321e 100644 --- a/gigl/distributed/graph_store/shared_dist_sampling_producer.py +++ b/gigl/distributed/graph_store/shared_dist_sampling_producer.py @@ -404,8 +404,8 @@ def _shared_sampling_worker_loop( # ``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 the - # core of the stall fix. + # 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 @@ -631,7 +631,7 @@ def _submit_one_batch(channel_id: int) -> bool: 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 (fix step 4) is deferred. GLT's + 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 @@ -662,14 +662,13 @@ def _submit_one_batch(channel_id: int) -> bool: state = active_epoch_by_channel_id.get(channel_id) if state is None: return False - # PARK: if the channel already has ``worker_concurrency`` batches in - # flight, submitting another could block this thread in the sampler's - # bounded semaphore on a saturated output channel. Return without - # submitting AND without re-enqueueing -- ``_on_batch_done`` is the - # sole path that re-enqueues (wakes) the channel once an in-flight - # batch completes. Returning False leaves ``made_progress`` unset in - # the pump, so a cycle that only parks does not busy-spin: Phase 3's - # idle tick engages instead. + # 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) @@ -721,14 +720,9 @@ def _submit_one_batch(channel_id: int) -> bool: def _pump_runnable_channel_ids() -> bool: """Submit one batch per runnable channel in round-robin order. - Channels that have hit their in-flight cap are parked by - ``_submit_one_batch`` (return False, no re-enqueue); ``_on_batch_done`` - re-enqueues them once a slot frees. A cycle whose only visited channels - park submits nothing, so ``made_progress`` stays False and Phase 3's idle - tick engages instead of busy-spinning. - - Channels re-enqueued during this pass are deferred to the next pump cycle - to preserve round-robin fairness. + 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. """ 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 a4f17503b..b9c5c7d5b 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 @@ -923,3 +923,155 @@ def _make_sampler(**kwargs: object) -> _GltOrderFakeSampler: # 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.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 and drained every remaining batch. + self.assertEqual(sampler.submit_count, total_batches) + 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()) From 2195cb2195256a4f996f1d49ca39deb9e6fda45f Mon Sep 17 00:00:00 2001 From: kmontemayor Date: Fri, 24 Jul 2026 17:42:21 +0000 Subject: [PATCH 7/7] graph-store stall fix: de-flake WAKE-resume drain assertion EPOCH_DONE fires when the last batch is sent (buffered), not when the consumer has received it, so asserting channel.total_received == total_batches immediately after the epoch-done event races the consumer draining the tail. Poll for the count instead. submit_count == total_batches stays an immediate assertion -- it is deterministic once EPOCH_DONE fires (all batches completed => all submitted). Co-Authored-By: Claude Opus 4.8 --- .../graph_store/shared_dist_sampling_producer_test.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) 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 b9c5c7d5b..ed743bcbc 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 @@ -1063,8 +1063,17 @@ def _consume() -> None: done = event_queue.get(timeout=10.0) self.assertEqual(done, (EPOCH_DONE_EVENT, channel_id, 0, 0)) - # The woken channel submitted and drained every remaining batch. + # 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):