diff --git a/src/exo/worker/runner/bootstrap.py b/src/exo/worker/runner/bootstrap.py index 62e3d61f..fbe449f8 100644 --- a/src/exo/worker/runner/bootstrap.py +++ b/src/exo/worker/runner/bootstrap.py @@ -16,7 +16,7 @@ def entrypoint( event_sender: MpSender[Event], task_receiver: MpReceiver[Task], _logger: "loguru.Logger", - pipe_fds: tuple[int, int] | None = None, + pipe_fifo_paths: tuple[str, str] | None = None, ) -> None: fast_synch_override = os.environ.get("EXO_FAST_SYNCH") if fast_synch_override == "on" or ( @@ -30,11 +30,15 @@ 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) + # Open JACCL FIFOs by path and set env vars for C++ SideChannel. + # Named pipes (FIFOs) work across multiprocessing spawn (macOS default). + if pipe_fifo_paths is not None: + fifo_c2p, fifo_p2c = pipe_fifo_paths + # C++ reads gathered data from p2c (PIPE_IN), writes local data to c2p (PIPE_OUT) + pipe_in_fd = os.open(fifo_p2c, os.O_RDONLY) + pipe_out_fd = os.open(fifo_c2p, os.O_WRONLY) + os.environ["MLX_JACCL_PIPE_IN"] = str(pipe_in_fd) + os.environ["MLX_JACCL_PIPE_OUT"] = str(pipe_out_fd) global logger logger = _logger diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py index 45499948..d2c7573f 100644 --- a/src/exo/worker/runner/runner_supervisor.py +++ b/src/exo/worker/runner/runner_supervisor.py @@ -2,6 +2,7 @@ import contextlib import os import signal import struct +import tempfile from dataclasses import dataclass, field from functools import partial from multiprocessing import Process @@ -75,6 +76,9 @@ class RunnerSupervisor: _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 + _fifo_dir: str | None = None # Temp dir for FIFO files (for cleanup) + _fifo_c2p: str | None = None # FIFO path: C++ writes → Python reads + _fifo_p2c: str | None = None # FIFO path: Python writes → C++ reads 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) @@ -94,21 +98,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 + # For MlxJaccl instances, create named pipes (FIFOs) for SideChannel relay. + # Named pipes work across multiprocessing.Process spawn (macOS default). + # FIFO c2p: C++ writes local data → Python reads it + # FIFO p2c: Python writes gathered data → C++ reads it + fifo_dir: str | None = None + fifo_c2p: str | None = None + fifo_p2c: str | None = None + pipe_fifo_paths: tuple[str, str] | 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) + fifo_dir = tempfile.mkdtemp(prefix="exo_jaccl_") + fifo_c2p = os.path.join(fifo_dir, "c2p") # C++ → Python + fifo_p2c = os.path.join(fifo_dir, "p2c") # Python → C++ + os.mkfifo(fifo_c2p) + os.mkfifo(fifo_p2c) + pipe_fifo_paths = (fifo_c2p, fifo_p2c) runner_process = Process( target=entrypoint, @@ -117,7 +122,7 @@ class RunnerSupervisor: ev_send, task_recv, logger, - pipe_fds_for_child, + pipe_fifo_paths, ), daemon=True, ) @@ -132,9 +137,9 @@ 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, + _fifo_dir=fifo_dir, + _fifo_c2p=fifo_c2p, + _fifo_p2c=fifo_p2c, ) return self @@ -142,13 +147,30 @@ class RunnerSupervisor: async def run(self): self.runner_process.start() - # 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._fifo_c2p is not None and self._fifo_p2c is not None: + # Open FIFOs from parent side. These block until child opens the other end, + # so we run them in threads concurrently to avoid deadlock. + fifo_c2p = self._fifo_c2p + fifo_p2c = self._fifo_p2c + + async def open_read() -> None: + self._pipe_read_fd = await to_thread.run_sync( + partial(os.open, fifo_c2p, os.O_RDONLY) + ) + + async def open_write() -> None: + self._pipe_write_fd = await to_thread.run_sync( + partial(os.open, fifo_p2c, os.O_WRONLY) + ) + + async with anyio.create_task_group() as open_tg: + open_tg.start_soon(open_read) + open_tg.start_soon(open_write) + + logger.info( + f"JACCL pipe relay: FIFOs opened (read_fd={self._pipe_read_fd}, write_fd={self._pipe_write_fd})" + ) - 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) @@ -250,6 +272,19 @@ class RunnerSupervisor: with contextlib.suppress(OSError): os.close(fd) self._child_pipe_fds = None + # Clean up FIFO files + if self._fifo_c2p is not None: + with contextlib.suppress(OSError): + os.unlink(self._fifo_c2p) + self._fifo_c2p = None + if self._fifo_p2c is not None: + with contextlib.suppress(OSError): + os.unlink(self._fifo_p2c) + self._fifo_p2c = None + if self._fifo_dir is not None: + with contextlib.suppress(OSError): + os.rmdir(self._fifo_dir) + self._fifo_dir = None async def _pipe_relay(self) -> None: """Relay JACCL SideChannel all_gather rounds between runner pipes and exo events."""