From 0e08c2bfd3b6cb628bf6daf03063f95c1467be01 Mon Sep 17 00:00:00 2001 From: Alex Cheema Date: Thu, 12 Feb 2026 08:55:40 -0800 Subject: [PATCH] Add pipe-based JACCL SideChannel relay via exo control plane Replace fragile TCP SideChannel with anonymous pipes relayed through exo's event-sourced control plane. RunnerSupervisor creates pipe pairs for MlxJaccl instances, relays all_gather rounds via JacclSideChannelData/ JacclSideChannelGathered events through the master, eliminating errno=57 crashes from Thunderbolt RDMA driver instability. Also includes dashboard RDMA warning improvements and instance retry fixes. Co-Authored-By: Claude Opus 4.6 --- dashboard/src/routes/+page.svelte | 105 +++++++++--- src/exo/master/main.py | 48 ++++++ src/exo/master/placement.py | 4 +- .../master/process_managers/meta_instance.py | 5 +- src/exo/master/reconcile.py | 6 +- src/exo/master/tests/test_reconcile.py | 26 +-- src/exo/shared/apply.py | 13 +- src/exo/shared/types/events.py | 22 +++ src/exo/worker/main.py | 10 ++ src/exo/worker/runner/bootstrap.py | 7 + src/exo/worker/runner/runner_supervisor.py | 159 +++++++++++++++++- 11 files changed, 353 insertions(+), 52 deletions(-) diff --git a/dashboard/src/routes/+page.svelte b/dashboard/src/routes/+page.svelte index 7e27192f..9ed29b41 100644 --- a/dashboard/src/routes/+page.svelte +++ b/dashboard/src/routes/+page.svelte @@ -1274,7 +1274,8 @@ typeof runnerData === "object" ) { const rd = runnerData as { errorMessage?: string }; - if (rd.errorMessage) failedErrors.push(`${getNodeName(nodeId)}: ${rd.errorMessage}`); + if (rd.errorMessage) + failedErrors.push(`${getNodeName(nodeId)}: ${rd.errorMessage}`); } } if (status) { @@ -1305,7 +1306,12 @@ errorMessage: failedErrors.length > 0 ? failedErrors.join("; ") : null, }; if (has("Shutdown")) - return { statusText: "SHUTDOWN", statusClass: "inactive", perNodeStatus, errorMessage: null }; + return { + statusText: "SHUTDOWN", + statusClass: "inactive", + perNodeStatus, + errorMessage: null, + }; // For loading/warming states, show node progress when multi-node if (has("Loading")) { @@ -1314,7 +1320,12 @@ totalNodes > 1 ? `LOADING (${readyCount}/${totalNodes} nodes ready)` : "LOADING"; - return { statusText, statusClass: "starting", perNodeStatus, errorMessage: null }; + return { + statusText, + statusClass: "starting", + perNodeStatus, + errorMessage: null, + }; } if (has("WarmingUp")) { const readyCount = count("Ready") + count("Running"); @@ -1322,23 +1333,63 @@ totalNodes > 1 ? `WARMING UP (${readyCount}/${totalNodes} nodes ready)` : "WARMING UP"; - return { statusText, statusClass: "starting", perNodeStatus, errorMessage: null }; + return { + statusText, + statusClass: "starting", + perNodeStatus, + errorMessage: null, + }; } if (has("Running")) - return { statusText: "RUNNING", statusClass: "running", perNodeStatus, errorMessage: null }; + return { + statusText: "RUNNING", + statusClass: "running", + perNodeStatus, + errorMessage: null, + }; if (has("Ready")) - return { statusText: "READY", statusClass: "loaded", perNodeStatus, errorMessage: null }; + return { + statusText: "READY", + statusClass: "loaded", + perNodeStatus, + errorMessage: null, + }; if (has("Loaded")) - return { statusText: "LOADED", statusClass: "loaded", perNodeStatus, errorMessage: null }; + return { + statusText: "LOADED", + statusClass: "loaded", + perNodeStatus, + errorMessage: null, + }; if (has("WaitingForModel")) - return { statusText: "WAITING", statusClass: "starting", perNodeStatus, errorMessage: null }; + return { + statusText: "WAITING", + statusClass: "starting", + perNodeStatus, + errorMessage: null, + }; if (has("InitializingBackend")) - return { statusText: "INITIALIZING", statusClass: "starting", perNodeStatus, errorMessage: null }; + return { + statusText: "INITIALIZING", + statusClass: "starting", + perNodeStatus, + errorMessage: null, + }; if (has("WaitingForInitialization")) - return { statusText: "INITIALIZING", statusClass: "starting", perNodeStatus, errorMessage: null }; + return { + statusText: "INITIALIZING", + statusClass: "starting", + perNodeStatus, + errorMessage: null, + }; - return { statusText: "RUNNING", statusClass: "active", perNodeStatus, errorMessage: null }; + return { + statusText: "RUNNING", + statusClass: "active", + perNodeStatus, + errorMessage: null, + }; } function getBytes(value: unknown): number { @@ -2066,8 +2117,8 @@ class="absolute top-full left-0 mt-2 w-80 p-3 rounded border border-red-500/30 bg-exo-dark-gray/95 backdrop-blur-sm opacity-0 invisible group-hover:opacity-100 group-hover:visible transition-all duration-200 z-50 shadow-lg" >

- A macOS RDMA driver error was detected. This is a known issue - with the experimental RDMA driver in macOS. + A macOS RDMA driver error was detected. This is a known issue with + the experimental RDMA driver in macOS.

Error: @@ -2655,12 +2706,15 @@ {@const downloadInfo = instance ? getInstanceDownloadStatus(item.instanceId ?? id, instance) : getMetaInstancePlacingStatus(id)} - {@const metaData = item.isMetaInstance ? metaInstancesData[id] : null} - {@const retryError = metaData?.lastFailureError && !downloadInfo.isFailed - ? metaData.consecutiveFailures > 0 - ? `(${((metaData.consecutiveFailures - 1) % 3) + 1}/3) ${metaData.lastFailureError}` - : metaData.lastFailureError + {@const metaData = item.isMetaInstance + ? metaInstancesData[id] : null} + {@const retryError = + metaData?.lastFailureError && !downloadInfo.isFailed + ? metaData.consecutiveFailures > 0 + ? `(${((metaData.consecutiveFailures - 1) % 3) + 1}/3) ${metaData.lastFailureError}` + : metaData.lastFailureError + : null} {@const statusText = downloadInfo.statusText} {@const isDownloading = downloadInfo.isDownloading} {@const isFailed = @@ -3523,12 +3577,15 @@ instance, ) : getMetaInstancePlacingStatus(id)} - {@const metaData = item.isMetaInstance ? metaInstancesData[id] : null} - {@const retryError = metaData?.lastFailureError && !downloadInfo.isFailed - ? metaData.consecutiveFailures > 0 - ? `(${((metaData.consecutiveFailures - 1) % 3) + 1}/3) ${metaData.lastFailureError}` - : metaData.lastFailureError + {@const metaData = item.isMetaInstance + ? metaInstancesData[id] : null} + {@const retryError = + metaData?.lastFailureError && !downloadInfo.isFailed + ? metaData.consecutiveFailures > 0 + ? `(${((metaData.consecutiveFailures - 1) % 3) + 1}/3) ${metaData.lastFailureError}` + : metaData.lastFailureError + : null} {@const statusText = downloadInfo.statusText} {@const isDownloading = downloadInfo.isDownloading} {@const isFailed = @@ -3539,7 +3596,7 @@ statusText.startsWith("WARMING UP") || statusText === "WAITING" || statusText === "PLACING" || - statusText.startsWith("RETRYING")} + statusText.startsWith("RETRYING")} {@const isReady = statusText === "READY" || statusText === "LOADED"} {@const isRunning = statusText === "RUNNING"} diff --git a/src/exo/master/main.py b/src/exo/master/main.py index bdfeac14..7197f57d 100644 --- a/src/exo/master/main.py +++ b/src/exo/master/main.py @@ -43,6 +43,8 @@ from exo.shared.types.events import ( ForwarderEvent, IndexedEvent, InputChunkReceived, + JacclSideChannelData, + JacclSideChannelGathered, MetaInstanceCreated, MetaInstanceDeleted, NodeGatheredInfo, @@ -67,6 +69,7 @@ from exo.shared.types.tasks import ( TextGeneration as TextGenerationTask, ) from exo.shared.types.worker.instances import InstanceId +from exo.shared.types.worker.runners import RunnerId from exo.utils.channels import Receiver, Sender from exo.utils.event_buffer import MultiSourceBuffer @@ -95,6 +98,7 @@ class Master: self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master") self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {} self._expected_ranks: dict[TaskId, set[int]] = {} + self._jaccl_pending: dict[InstanceId, dict[int, dict[RunnerId, bytes]]] = {} self._process_managers: Sequence[ProcessManager] = [ InstanceHealthReconciler(), NodeTimeoutReconciler(), @@ -412,6 +416,11 @@ class Master: await self._handle_traces_collected(event) continue + if isinstance(event, JacclSideChannelData): + await self._apply_and_broadcast(event) + await self._handle_jaccl_side_channel(event) + continue + if isinstance(event, NodeGatheredInfo): event.when = str(datetime.now(tz=timezone.utc)) @@ -454,3 +463,42 @@ class Master: del self._pending_traces[task_id] if task_id in self._expected_ranks: del self._expected_ranks[task_id] + + async def _handle_jaccl_side_channel(self, event: JacclSideChannelData) -> None: + """Accumulate SideChannel contributions; when all runners for an instance + have submitted for the same sequence, emit JacclSideChannelGathered.""" + iid = event.instance_id + seq = event.sequence + + if iid not in self._jaccl_pending: + self._jaccl_pending[iid] = {} + if seq not in self._jaccl_pending[iid]: + self._jaccl_pending[iid][seq] = {} + self._jaccl_pending[iid][seq][event.runner_id] = event.data + + instance = self.state.instances.get(iid) + if instance is None: + logger.warning(f"JacclSideChannelData for unknown instance {iid}") + return + + expected_runners = set(instance.shard_assignments.runner_to_shard.keys()) + submitted = set(self._jaccl_pending[iid][seq].keys()) + + logger.info( + f"JACCL side channel: instance={iid} seq={seq} " + f"submitted={len(submitted)}/{len(expected_runners)}" + ) + + if submitted >= expected_runners: + gathered = dict(self._jaccl_pending[iid][seq]) + del self._jaccl_pending[iid][seq] + if not self._jaccl_pending[iid]: + del self._jaccl_pending[iid] + + await self._apply_and_broadcast( + JacclSideChannelGathered( + instance_id=iid, + sequence=seq, + gathered_data=gathered, + ) + ) diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py index 5157c5d7..a2b5201b 100644 --- a/src/exo/master/placement.py +++ b/src/exo/master/placement.py @@ -63,9 +63,7 @@ def place_instance( required_nodes: set[NodeId] | None = None, ) -> dict[InstanceId, Instance]: cycles = topology.get_cycles() - candidate_cycles = list( - filter(lambda it: len(it) >= command.min_nodes, cycles) - ) + candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles)) # Filter to cycles containing all required nodes (subset matching) if required_nodes: diff --git a/src/exo/master/process_managers/meta_instance.py b/src/exo/master/process_managers/meta_instance.py index 2a4c1c3d..4b92cc05 100644 --- a/src/exo/master/process_managers/meta_instance.py +++ b/src/exo/master/process_managers/meta_instance.py @@ -43,7 +43,10 @@ class MetaInstanceReconciler: all_events.extend(result.events) # Emit placement failure if error differs from what's already in state - if result.error is not None and meta_instance.placement_error != result.error: + if ( + result.error is not None + and meta_instance.placement_error != result.error + ): all_events.append( MetaInstancePlacementFailed( meta_instance_id=meta_instance.meta_instance_id, diff --git a/src/exo/master/reconcile.py b/src/exo/master/reconcile.py index e2362593..846ffca1 100644 --- a/src/exo/master/reconcile.py +++ b/src/exo/master/reconcile.py @@ -133,7 +133,11 @@ def instance_runners_failed( has_any_failed = True if status.error_message: node_id = runner_to_node.get(runner_id) - name = node_identities[node_id].friendly_name if node_id and node_id in node_identities else node_id or "unknown" + name = ( + node_identities[node_id].friendly_name + if node_id and node_id in node_identities + else node_id or "unknown" + ) error_messages.append(f"{name}: {status.error_message}") elif isinstance(status, RunnerShutdown): pass # Terminal but not a failure indicator on its own diff --git a/src/exo/master/tests/test_reconcile.py b/src/exo/master/tests/test_reconcile.py index 7a3a813f..e2d6e776 100644 --- a/src/exo/master/tests/test_reconcile.py +++ b/src/exo/master/tests/test_reconcile.py @@ -179,9 +179,7 @@ def test_not_satisfies_fewer_than_min_nodes(): def test_satisfies_with_node_ids_specified(): - meta = _meta_instance( - node_ids=[NodeId("node-a"), NodeId("node-b")], min_nodes=2 - ) + meta = _meta_instance(node_ids=[NodeId("node-a"), NodeId("node-b")], min_nodes=2) _, inst = _instance(node_ids=["node-a", "node-b", "node-c"]) assert instance_satisfies_meta_instance(meta, inst) is True @@ -494,8 +492,7 @@ def test_runners_not_failed_all_shutdown(): """All Shutdown (graceful) = not a failure.""" _, inst = _instance(node_ids=["node-a"]) runners = { - rid: RunnerShutdown() - for rid in inst.shard_assignments.node_to_runner.values() + rid: RunnerShutdown() for rid in inst.shard_assignments.node_to_runner.values() } is_failed, _ = instance_runners_failed(inst, runners, {}) assert is_failed is False @@ -524,8 +521,7 @@ def test_runners_not_failed_healthy(): """Runners in Ready state = not failed.""" _, inst = _instance(node_ids=["node-a"]) runners = { - rid: RunnerReady() - for rid in inst.shard_assignments.node_to_runner.values() + rid: RunnerReady() for rid in inst.shard_assignments.node_to_runner.values() } is_failed, _ = instance_runners_failed(inst, runners, {}) assert is_failed is False @@ -537,9 +533,7 @@ def test_runners_not_failed_healthy(): def test_apply_instance_deleted_tracks_failure(): """InstanceDeleted with failure_error increments meta instance failure count.""" meta = _meta_instance() - iid, inst = _instance( - node_ids=["node-a"], meta_instance_id=meta.meta_instance_id - ) + iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) state = State( meta_instances={meta.meta_instance_id: meta}, instances={iid: inst}, @@ -556,9 +550,7 @@ def test_apply_instance_deleted_increments_failure(): meta = _meta_instance().model_copy( update={"consecutive_failures": 2, "last_failure_error": "previous error"} ) - iid, inst = _instance( - node_ids=["node-a"], meta_instance_id=meta.meta_instance_id - ) + iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) state = State( meta_instances={meta.meta_instance_id: meta}, instances={iid: inst}, @@ -573,9 +565,7 @@ def test_apply_instance_deleted_increments_failure(): def test_apply_instance_deleted_no_failure_no_tracking(): """InstanceDeleted without failure_error does not track.""" meta = _meta_instance() - iid, inst = _instance( - node_ids=["node-a"], meta_instance_id=meta.meta_instance_id - ) + iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) state = State( meta_instances={meta.meta_instance_id: meta}, instances={iid: inst}, @@ -601,7 +591,9 @@ def test_apply_instance_deleted_orphan_no_tracking(): def test_apply_instance_retrying_removes_runners(): """InstanceRetrying removes the instance's runners from state but keeps the instance.""" meta = _meta_instance() - iid, inst = _instance(node_ids=["node-a", "node-b"], meta_instance_id=meta.meta_instance_id) + iid, inst = _instance( + node_ids=["node-a", "node-b"], meta_instance_id=meta.meta_instance_id + ) runner_ids = list(inst.shard_assignments.node_to_runner.values()) runners = { runner_ids[0]: RunnerFailed(error_message="OOM"), diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py index 62bb719e..5fb8a4d1 100644 --- a/src/exo/shared/apply.py +++ b/src/exo/shared/apply.py @@ -13,6 +13,8 @@ from exo.shared.types.events import ( InstanceCreated, InstanceDeleted, InstanceRetrying, + JacclSideChannelData, + JacclSideChannelGathered, MetaInstanceCreated, MetaInstanceDeleted, MetaInstancePlacementFailed, @@ -71,6 +73,8 @@ def event_apply(event: Event, state: State) -> State: | InputChunkReceived() | TracesCollected() | TracesMerged() + | JacclSideChannelData() + | JacclSideChannelGathered() ): # Pass-through events that don't modify state return state case InstanceCreated(): @@ -252,9 +256,7 @@ def apply_instance_retrying(event: InstanceRetrying, state: State) -> State: # Remove all runners belonging to this instance from state runner_ids_to_remove = set(instance.shard_assignments.node_to_runner.values()) new_runners: Mapping[RunnerId, RunnerStatus] = { - rid: rs - for rid, rs in state.runners.items() - if rid not in runner_ids_to_remove + rid: rs for rid, rs in state.runners.items() if rid not in runner_ids_to_remove } update: dict[str, object] = {"runners": new_runners} @@ -264,7 +266,10 @@ def apply_instance_retrying(event: InstanceRetrying, state: State) -> State: update["meta_instances"] = _update_meta_instance( state, event.meta_instance_id, - consecutive_failures=state.meta_instances[event.meta_instance_id].consecutive_failures + 1, + consecutive_failures=state.meta_instances[ + event.meta_instance_id + ].consecutive_failures + + 1, last_failure_error=event.failure_error, ) diff --git a/src/exo/shared/types/events.py b/src/exo/shared/types/events.py index 77e8f292..fabb367d 100644 --- a/src/exo/shared/types/events.py +++ b/src/exo/shared/types/events.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from datetime import datetime from typing import final @@ -157,6 +158,25 @@ class TracesMerged(BaseEvent): traces: list[TraceEventData] +@final +class JacclSideChannelData(BaseEvent): + """A runner's local contribution to a JACCL SideChannel all_gather round.""" + + instance_id: InstanceId + runner_id: RunnerId + sequence: int + data: bytes + + +@final +class JacclSideChannelGathered(BaseEvent): + """Gathered result of a JACCL SideChannel all_gather round.""" + + instance_id: InstanceId + sequence: int + gathered_data: Mapping[RunnerId, bytes] + + Event = ( TestEvent | TaskCreated @@ -181,6 +201,8 @@ Event = ( | TopologyEdgeDeleted | TracesCollected | TracesMerged + | JacclSideChannelData + | JacclSideChannelGathered ) diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index 4abd3c43..39fa22ff 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -24,6 +24,7 @@ from exo.shared.types.events import ( ForwarderEvent, IndexedEvent, InputChunkReceived, + JacclSideChannelGathered, NodeGatheredInfo, TaskCreated, TaskStatusUpdated, @@ -158,6 +159,15 @@ class Worker: for idx, event in indexed_events: self.state = apply(self.state, IndexedEvent(idx=idx, event=event)) + # Dispatch JACCL gathered events to the relevant RunnerSupervisor + if isinstance(event, JacclSideChannelGathered): + for runner in self.runners.values(): + if ( + runner.bound_instance.instance.instance_id + == event.instance_id + ): + runner.notify_gathered(event) + # Buffer input image chunks for image editing if isinstance(event, InputChunkReceived): cmd_id = event.command_id diff --git a/src/exo/worker/runner/bootstrap.py b/src/exo/worker/runner/bootstrap.py index bf08ab6c..62e3d61f 100644 --- a/src/exo/worker/runner/bootstrap.py +++ b/src/exo/worker/runner/bootstrap.py @@ -16,6 +16,7 @@ def entrypoint( event_sender: MpSender[Event], task_receiver: MpReceiver[Task], _logger: "loguru.Logger", + pipe_fds: tuple[int, int] | None = None, ) -> None: fast_synch_override = os.environ.get("EXO_FAST_SYNCH") if fast_synch_override == "on" or ( @@ -29,6 +30,12 @@ def entrypoint( else: os.environ["MLX_METAL_FAST_SYNCH"] = "0" + # Set up JACCL pipe-based SideChannel env vars if pipe fds were provided + if pipe_fds is not None: + mlx_reads, mlx_writes = pipe_fds + os.environ["MLX_JACCL_PIPE_IN"] = str(mlx_reads) + os.environ["MLX_JACCL_PIPE_OUT"] = str(mlx_writes) + global logger logger = _logger diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py index 3b859711..45499948 100644 --- a/src/exo/worker/runner/runner_supervisor.py +++ b/src/exo/worker/runner/runner_supervisor.py @@ -1,6 +1,9 @@ import contextlib +import os import signal +import struct from dataclasses import dataclass, field +from functools import partial from multiprocessing import Process from typing import Self @@ -14,12 +17,14 @@ from loguru import logger from exo.shared.types.events import ( Event, + JacclSideChannelData, + JacclSideChannelGathered, RunnerStatusUpdated, TaskAcknowledged, TaskStatusUpdated, ) from exo.shared.types.tasks import Task, TaskId, TaskStatus -from exo.shared.types.worker.instances import BoundInstance +from exo.shared.types.worker.instances import BoundInstance, MlxJacclInstance from exo.shared.types.worker.runners import ( RunnerConnecting, RunnerFailed, @@ -34,6 +39,26 @@ from exo.shared.types.worker.shards import ShardMetadata from exo.utils.channels import MpReceiver, MpSender, Sender, mp_channel from exo.worker.runner.bootstrap import entrypoint + +def _pipe_read_exact(fd: int, n: int) -> bytes | None: + """Read exactly n bytes from a file descriptor. Returns None on EOF.""" + data = b"" + while len(data) < n: + chunk = os.read(fd, n - len(data)) + if not chunk: + return None + data += chunk + return data + + +def _pipe_write_all(fd: int, data: bytes) -> None: + """Write all bytes to a file descriptor.""" + view = memoryview(data) + while view: + written = os.write(fd, view) + view = view[written:] + + PREFILL_TIMEOUT_SECONDS = 60 DECODE_TIMEOUT_SECONDS = 5 @@ -47,9 +72,15 @@ class RunnerSupervisor: _ev_recv: MpReceiver[Event] _task_sender: MpSender[Task] _event_sender: Sender[Event] + _pipe_read_fd: int | None = None # Python reads runner's pipe output + _pipe_write_fd: int | None = None # Python writes gathered data to runner + _child_pipe_fds: tuple[int, int] | None = None # fds to close after fork status: RunnerStatus = field(default_factory=RunnerIdle, init=False) pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False) completed: set[TaskId] = field(default_factory=set, init=False) + _gathered_waiters: dict[ + int, tuple[anyio.Event, JacclSideChannelGathered | None] + ] = field(default_factory=dict, init=False) @classmethod def create( @@ -63,6 +94,22 @@ class RunnerSupervisor: # A task is kind of a runner command task_sender, task_recv = mp_channel[Task]() + # For MlxJaccl instances, create pipe pairs for SideChannel relay. + # Pipe pair 1: C++ writes local data → Python reads it + # Pipe pair 2: Python writes gathered data → C++ reads it + pipe_read_fd: int | None = None + pipe_write_fd: int | None = None + child_pipe_fds: tuple[int, int] | None = None + pipe_fds_for_child: tuple[int, int] | None = None + + if isinstance(bound_instance.instance, MlxJacclInstance): + python_reads, mlx_writes = os.pipe() # C++ → Python + mlx_reads, python_writes = os.pipe() # Python → C++ + pipe_read_fd = python_reads + pipe_write_fd = python_writes + child_pipe_fds = (mlx_reads, mlx_writes) + pipe_fds_for_child = (mlx_reads, mlx_writes) + runner_process = Process( target=entrypoint, args=( @@ -70,6 +117,7 @@ class RunnerSupervisor: ev_send, task_recv, logger, + pipe_fds_for_child, ), daemon=True, ) @@ -84,19 +132,35 @@ class RunnerSupervisor: _ev_recv=ev_recv, _task_sender=task_sender, _event_sender=event_sender, + _pipe_read_fd=pipe_read_fd, + _pipe_write_fd=pipe_write_fd, + _child_pipe_fds=child_pipe_fds, ) return self async def run(self): self.runner_process.start() - await self._forward_events() + + # Close the child-side pipe fds in the parent process (child inherited them via fork) + if self._child_pipe_fds is not None: + for fd in self._child_pipe_fds: + os.close(fd) + self._child_pipe_fds = None + + if self._pipe_read_fd is not None: + async with anyio.create_task_group() as tg: + tg.start_soon(self._pipe_relay) + tg.start_soon(self._forward_events) + else: + await self._forward_events() def shutdown(self): logger.info("Runner supervisor shutting down") self._ev_recv.close() self._task_sender.close() self._event_sender.close() + self._close_pipe_fds() self.runner_process.join(1) if not self.runner_process.is_alive(): logger.info("Runner process succesfully terminated") @@ -172,6 +236,97 @@ class RunnerSupervisor: for tid in self.pending: self.pending[tid].set() + def _close_pipe_fds(self) -> None: + if self._pipe_read_fd is not None: + with contextlib.suppress(OSError): + os.close(self._pipe_read_fd) + self._pipe_read_fd = None + if self._pipe_write_fd is not None: + with contextlib.suppress(OSError): + os.close(self._pipe_write_fd) + self._pipe_write_fd = None + if self._child_pipe_fds is not None: + for fd in self._child_pipe_fds: + with contextlib.suppress(OSError): + os.close(fd) + self._child_pipe_fds = None + + async def _pipe_relay(self) -> None: + """Relay JACCL SideChannel all_gather rounds between runner pipes and exo events.""" + assert self._pipe_read_fd is not None + assert self._pipe_write_fd is not None + read_fd = self._pipe_read_fd + write_fd = self._pipe_write_fd + sequence = 0 + + try: + while True: + # 1. Read local data from runner: [uint32 size][size bytes] + header = await to_thread.run_sync(partial(_pipe_read_exact, read_fd, 4)) + if header is None: + logger.info("JACCL pipe relay: runner closed pipe (EOF)") + break + data_size: int = struct.unpack(" None: + """Called by the worker when a JacclSideChannelGathered event arrives.""" + seq = event.sequence + if seq not in self._gathered_waiters: + logger.warning(f"JACCL: received gathered event for unknown sequence {seq}") + return + waiter, _ = self._gathered_waiters[seq] + self._gathered_waiters[seq] = (waiter, event) + waiter.set() + def __del__(self) -> None: if self.runner_process.is_alive(): logger.warning("RunnerSupervisor was not stopped cleanly.")