Compare commits

...

1 Commits

Author SHA1 Message Date
Evan 16a3ab5155 sloppin it up 2026-03-03 10:49:36 +00:00
9 changed files with 634 additions and 118 deletions
+47
View File
@@ -32,8 +32,55 @@ def _maybe_run_inline_code(argv: Sequence[str]) -> bool:
return True
def _is_frozen() -> bool:
"""Check if running in a frozen environment (e.g., PyInstaller)."""
return getattr(sys, "frozen", False) or hasattr(sys, "_MEIPASS")
def _maybe_run_runner(argv: Sequence[str]) -> bool:
"""
Handle the --exo-runner flag for spawning runner subprocesses in frozen builds.
"""
if "--exo-runner" not in argv:
return False
try:
flag_index = argv.index("--exo-runner")
args_index = flag_index + 1
if args_index >= len(argv):
return False
# Set up the runner environment
import os
os.environ["EXO_FROZEN"] = "1"
# Import and run the bootstrap entrypoint
from exo.worker.runner.bootstrap import entrypoint
# Inject the JSON args into sys.argv as expected by entrypoint
sys.argv = [argv[0], argv[args_index]]
entrypoint()
return True
except Exception as e:
print(f"Failed to start runner: {e}", file=sys.stderr)
return False
if __name__ == "__main__":
if _maybe_run_inline_code(sys.argv):
sys.exit(0)
# Check if we should run as a runner subprocess (frozen builds)
if _maybe_run_runner(sys.argv):
sys.exit(0)
freeze_support()
# Store frozen status for subprocess spawning
if _is_frozen():
import os
os.environ["EXO_FROZEN"] = "1"
main()
@@ -0,0 +1,22 @@
"""CLI arguments for runner processes passed as JSON."""
from exo.shared.types.worker.instances import BoundInstance
from exo.utils.pydantic_ext import FrozenModel
class RunnerCliArgs(FrozenModel):
"""
Arguments passed to the runner process via JSON on the command line.
File descriptors for communication are passed separately via pass_fds.
Each logical channel uses a single pipe:
- Events: child -> parent (child writes, parent reads)
- Tasks: parent -> child (parent writes, child reads)
- Cancels: parent -> child (parent writes, child reads)
"""
bound_instance: BoundInstance
event_fd: int # Child writes events here
task_fd: int # Child reads tasks from here
cancel_fd: int # Child reads cancels from here
log_level: str = "INFO"
+320
View File
@@ -0,0 +1,320 @@
import contextlib
import os
import struct
from dataclasses import dataclass
from typing import Any, Self
from anyio import (
BrokenResourceError,
ClosedResourceError,
EndOfStream,
WouldBlock,
to_thread,
)
from anyio.streams.memory import (
MemoryObjectReceiveStream as AnyioReceiver,
)
from anyio.streams.memory import (
MemoryObjectSendStream as AnyioSender,
)
from pydantic import TypeAdapter
def _encode_message(data: bytes) -> bytes:
"""Encode a message with its length prefix."""
return struct.pack("!I", len(data)) + data
def _decode_message(data: bytes) -> tuple[bytes, bytes]:
"""Decode a message from a buffer, returning (message, remaining_buffer)."""
if len(data) < 4:
return b"", data
length = struct.unpack("!I", data[:4])[0]
if len(data) < 4 + length:
return b"", data
return data[4 : 4 + length], data[4 + length :]
class FdSender[T]:
"""A sender that writes JSON-encoded Pydantic models to a file descriptor."""
def __init__(
self,
fd: int,
type_adapter: TypeAdapter[T],
max_buffer_size: int = 1024 * 1024,
) -> None:
self._fd = fd
self._type_adapter = type_adapter
self._max_buffer_size = max_buffer_size
self._closed = False
self._write_buffer = bytearray()
def send_nowait(self, item: T) -> None:
"""Send an item without blocking."""
if self._closed:
raise ClosedResourceError
data = self._type_adapter.dump_json(item)
message = _encode_message(data)
if len(self._write_buffer) + len(message) > self._max_buffer_size:
raise WouldBlock
self._write_buffer.extend(message)
self._flush()
def send(self, item: T) -> None:
"""Send an item, blocking if necessary."""
if self._closed:
raise ClosedResourceError
data = self._type_adapter.dump_json(item)
message = _encode_message(data)
self._write_buffer.extend(message)
self._flush_blocking()
async def send_async(self, item: T) -> None:
"""Send an item asynchronously."""
await to_thread.run_sync(self.send, item)
def _flush(self) -> None:
"""Flush the write buffer non-blocking."""
if not self._write_buffer:
return
try:
written = os.write(self._fd, self._write_buffer)
self._write_buffer = self._write_buffer[written:]
except (OSError, BrokenPipeError) as e:
raise BrokenResourceError from e
def _flush_blocking(self) -> None:
"""Flush the write buffer, blocking until complete."""
while self._write_buffer:
self._flush()
def close(self) -> None:
"""Close the sender."""
if not self._closed:
self._closed = True
try:
self._flush_blocking()
os.close(self._fd)
except (OSError, BrokenPipeError):
pass
def __enter__(self) -> Self:
return self
def __exit__(self, *args: Any) -> None:
self.close()
class FdReceiver[T]:
"""A receiver that reads JSON-encoded Pydantic models from a file descriptor."""
def __init__(
self,
fd: int,
type_adapter: TypeAdapter[T],
max_buffer_size: int = 1024 * 1024,
) -> None:
self._fd = fd
self._type_adapter = type_adapter
self._max_buffer_size = max_buffer_size
self._closed = False
self._read_buffer = bytearray()
def receive_nowait(self) -> T:
"""Receive an item without blocking."""
if self._closed:
raise ClosedResourceError
# Try to read any available data
self._read_available()
# Try to decode a message
message, remaining = _decode_message(bytes(self._read_buffer))
if message:
self._read_buffer = bytearray(remaining)
return self._type_adapter.validate_json(message)
raise WouldBlock
def receive(self) -> T:
"""Receive an item, blocking if necessary."""
if self._closed:
raise ClosedResourceError
while True:
# Try to decode a message
message, remaining = _decode_message(bytes(self._read_buffer))
if message:
self._read_buffer = bytearray(remaining)
return self._type_adapter.validate_json(message)
# Read more data
try:
chunk = os.read(self._fd, 8192)
if not chunk:
# EOF
self.close()
raise EndOfStream
self._read_buffer.extend(chunk)
except (OSError, BrokenPipeError) as e:
self.close()
raise BrokenResourceError from e
async def receive_async(self) -> T:
"""Receive an item asynchronously."""
return await to_thread.run_sync(self.receive)
def _read_available(self) -> None:
"""Read any available data without blocking."""
try:
import fcntl
# Set non-blocking mode temporarily
flags = fcntl.fcntl(self._fd, fcntl.F_GETFL)
fcntl.fcntl(self._fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
try:
chunk = os.read(self._fd, 8192)
if chunk:
self._read_buffer.extend(chunk)
except BlockingIOError:
pass
finally:
fcntl.fcntl(self._fd, fcntl.F_SETFL, flags)
except ImportError:
# Windows doesn't have fcntl, just try reading
pass
def close(self) -> None:
"""Close the receiver."""
if not self._closed:
self._closed = True
with contextlib.suppress(OSError, BrokenPipeError):
os.close(self._fd)
def __iter__(self) -> Self:
return self
def __next__(self) -> T:
try:
return self.receive()
except EndOfStream:
raise StopIteration from None
def __aiter__(self) -> Self:
return self
async def __anext__(self) -> T:
try:
return await self.receive_async()
except EndOfStream:
raise StopAsyncIteration from None
def __enter__(self) -> Self:
return self
def __exit__(self, *args: Any) -> None:
self.close()
def collect(self) -> list[T]:
"""Collect all currently available items from this receiver."""
out: list[T] = []
while True:
try:
item = self.receive_nowait()
out.append(item)
except WouldBlock:
break
return out
@dataclass(eq=False)
class FdChannel:
"""A pair of file descriptors for bidirectional communication."""
sender_fd: int
receiver_fd: int
def close(self) -> None:
"""Close both file descriptors."""
for fd in (self.sender_fd, self.receiver_fd):
with contextlib.suppress(OSError, BrokenPipeError):
os.close(fd)
def fd_channel_pair() -> tuple[FdChannel, FdChannel]:
"""
Create a pair of connected file descriptor channels.
Returns (parent_channel, child_channel).
"""
# Create two pipes: one for parent->child, one for child->parent
parent_to_child_read, parent_to_child_write = os.pipe()
child_to_parent_read, child_to_parent_write = os.pipe()
# Parent writes to parent_to_child_write, reads from child_to_parent_read
parent_channel = FdChannel(
sender_fd=parent_to_child_write,
receiver_fd=child_to_parent_read,
)
# Child reads from parent_to_child_read, writes to child_to_parent_write
child_channel = FdChannel(
sender_fd=child_to_parent_write,
receiver_fd=parent_to_child_read,
)
return parent_channel, child_channel
class Sender[T](AnyioSender[T]):
def clone(self) -> "Sender[T]":
if self._closed:
raise ClosedResourceError
return Sender(_state=self._state)
def clone_receiver(self) -> "Receiver[T]":
"""Constructs a Receiver using a Senders shared state - similar to calling Receiver.clone() without needing the receiver"""
if self._closed:
raise ClosedResourceError
return Receiver(_state=self._state)
class Receiver[T](AnyioReceiver[T]):
def clone(self) -> "Receiver[T]":
if self._closed:
raise ClosedResourceError
return Receiver(_state=self._state)
def clone_sender(self) -> Sender[T]:
"""Constructs a Sender using a Receivers shared state - similar to calling Sender.clone() without needing the receiver"""
if self._closed:
raise ClosedResourceError
return Sender(_state=self._state)
def collect(self) -> list[T]:
"""Collect all currently available items from this receiver"""
out: list[T] = []
while True:
try:
item = self.receive_nowait()
out.append(item)
except WouldBlock:
break
return out
async def receive_at_least(self, n: int) -> list[T]:
out: list[T] = []
out.append(await self.receive())
out.extend(self.collect())
while len(out) < n:
out.append(await self.receive())
out.extend(self.collect())
return out
def __enter__(self) -> Self:
return self
+5 -5
View File
@@ -99,7 +99,7 @@ class Worker:
self.command_sender.close()
self.download_command_sender.close()
for runner in self.runners.values():
runner.shutdown()
await runner.shutdown()
async def _forward_info(self, recv: Receiver[GatheredInfo]):
with recv as info_stream:
@@ -160,7 +160,7 @@ class Worker:
# lets not kill the worker if a runner is unresponsive
match task:
case CreateRunner():
self._create_supervisor(task)
await self._create_supervisor(task)
await self.event_sender.send(
TaskStatusUpdated(
task_id=task.task_id, task_status=TaskStatus.Complete
@@ -220,7 +220,7 @@ class Worker:
)
)
finally:
runner.shutdown()
await runner.shutdown()
case CancelTask(
cancelled_task_id=cancelled_task_id, runner_id=runner_id
):
@@ -280,9 +280,9 @@ class Worker:
instance.shard_assignments.node_to_runner[self.node_id]
].start_task(task)
def _create_supervisor(self, task: CreateRunner) -> RunnerSupervisor:
async def _create_supervisor(self, task: CreateRunner) -> RunnerSupervisor:
"""Creates and stores a new AssignedRunner with initial downloading status."""
runner = RunnerSupervisor.create(
runner = await RunnerSupervisor.create(
bound_instance=task.bound_instance,
event_sender=self.event_sender.clone(),
)
+55 -24
View File
@@ -1,26 +1,44 @@
import contextlib
import os
import sys
import loguru
from pydantic import TypeAdapter
from exo.shared.types.events import Event, RunnerStatusUpdated
from exo.shared.types.tasks import Task, TaskId
from exo.shared.types.worker.instances import BoundInstance
from exo.shared.types.worker.runner_args import RunnerCliArgs
from exo.shared.types.worker.runners import RunnerFailed
from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
from exo.utils.fd_channels import FdReceiver, FdSender
logger: "loguru.Logger" = loguru.logger
# Type adapters for union types
event_adapter: TypeAdapter[Event] = TypeAdapter(Event)
task_adapter: TypeAdapter[Task] = TypeAdapter(Task)
task_id_adapter: TypeAdapter[TaskId] = TypeAdapter(TaskId)
def entrypoint(
bound_instance: BoundInstance,
event_sender: MpSender[Event],
task_receiver: MpReceiver[Task],
cancel_receiver: MpReceiver[TaskId],
_logger: "loguru.Logger",
) -> None:
global logger
logger = _logger
def entrypoint() -> None:
"""
Entrypoint for the runner process.
Expects a single command-line argument containing JSON-serialized RunnerCliArgs.
File descriptors for communication are inherited via pass_fds.
"""
if len(sys.argv) < 2:
print("Usage: runner <json_args>", file=sys.stderr)
sys.exit(1)
# Parse CLI arguments
args_json = sys.argv[1]
args = RunnerCliArgs.model_validate_json(args_json)
# Configure logging
logger.remove()
logger.add(sys.stderr, level=args.log_level)
# Set up environment
fast_synch_override = os.environ.get("EXO_FAST_SYNCH")
if fast_synch_override != "off":
os.environ["MLX_METAL_FAST_SYNCH"] = "1"
@@ -29,8 +47,16 @@ def entrypoint(
logger.info(f"Fast synch flag: {os.environ['MLX_METAL_FAST_SYNCH']}")
# Import main after setting global logger - this lets us just import logger from this module
# Create FD-based channels with TypeAdapters
# Child uses: event_fd for writing events, task_fd for reading tasks, cancel_fd for reading cancels
event_sender = FdSender(args.event_fd, event_adapter)
task_receiver = FdReceiver(args.task_fd, task_adapter)
cancel_receiver = FdReceiver(args.cancel_fd, task_id_adapter)
bound_instance = args.bound_instance
try:
# Import main after setting global logger - this lets us just import logger from this module
if bound_instance.is_image_model:
from exo.worker.runner.image_models.runner import main
else:
@@ -38,23 +64,28 @@ def entrypoint(
main(bound_instance, event_sender, task_receiver, cancel_receiver)
except ClosedResourceError:
logger.warning("Runner communication closed unexpectedly")
except Exception as e:
logger.opt(exception=e).warning(
f"Runner {bound_instance.bound_runner_id} crashed with critical exception {e}"
)
event_sender.send(
RunnerStatusUpdated(
runner_id=bound_instance.bound_runner_id,
runner_status=RunnerFailed(error_message=str(e)),
with contextlib.suppress(Exception):
event_sender.send(
RunnerStatusUpdated(
runner_id=bound_instance.bound_runner_id,
runner_status=RunnerFailed(error_message=str(e)),
)
)
)
finally:
try:
# Clean up channels
with contextlib.suppress(Exception):
event_sender.close()
with contextlib.suppress(Exception):
task_receiver.close()
finally:
event_sender.join()
task_receiver.join()
logger.info("bye from the runner")
with contextlib.suppress(Exception):
cancel_receiver.close()
logger.info("bye from the runner")
if __name__ == "__main__":
entrypoint()
+7 -7
View File
@@ -55,7 +55,7 @@ from exo.shared.types.worker.shards import (
PipelineShardMetadata,
ShardMetadata,
)
from exo.utils.channels import MpReceiver, MpSender
from exo.utils.fd_channels import FdReceiver, FdSender
from exo.worker.engines.image import (
DistributedImageModel,
generate_image,
@@ -88,7 +88,7 @@ def _process_image_response(
response: ImageGenerationResponse | PartialImageResponse,
command_id: CommandId,
shard_metadata: ShardMetadata,
event_sender: MpSender[Event],
event_sender: FdSender[Event],
image_index: int,
) -> None:
"""Process a single image response and send chunks."""
@@ -111,7 +111,7 @@ def _process_image_response(
def _send_traces_if_enabled(
event_sender: MpSender[Event],
event_sender: FdSender[Event],
task_id: TaskId,
rank: int,
) -> None:
@@ -144,7 +144,7 @@ def _send_image_chunk(
encoded_data: str,
command_id: CommandId,
model_id: ModelId,
event_sender: MpSender[Event],
event_sender: FdSender[Event],
image_index: int,
is_partial: bool,
partial_index: int | None = None,
@@ -184,9 +184,9 @@ def _send_image_chunk(
def main(
bound_instance: BoundInstance,
event_sender: MpSender[Event],
task_receiver: MpReceiver[Task],
cancel_receiver: MpReceiver[TaskId],
event_sender: FdSender[Event],
task_receiver: FdReceiver[Task],
cancel_receiver: FdReceiver[TaskId],
):
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (min(max(soft, 2048), hard), hard))
@@ -63,7 +63,7 @@ from exo.shared.types.worker.runners import (
RunnerStatus,
RunnerWarmingUp,
)
from exo.utils.channels import MpReceiver, MpSender
from exo.utils.fd_channels import FdReceiver, FdSender
from exo.worker.engines.mlx.cache import KVPrefixCache
from exo.worker.engines.mlx.generator.generate import (
PrefillCancelled,
@@ -85,9 +85,9 @@ from .tool_parsers import ToolParser, make_mlx_parser
def main(
bound_instance: BoundInstance,
event_sender: MpSender[Event],
task_receiver: MpReceiver[Task],
cancel_receiver: MpReceiver[TaskId],
event_sender: FdSender[Event],
task_receiver: FdReceiver[Task],
cancel_receiver: FdReceiver[TaskId],
):
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (min(max(soft, 2048), hard), hard))
+121 -55
View File
@@ -1,6 +1,8 @@
import contextlib
import multiprocessing as mp
import json
import os
import signal
import sys
from dataclasses import dataclass, field
from typing import Self
@@ -8,9 +10,11 @@ import anyio
from anyio import (
BrokenResourceError,
ClosedResourceError,
to_thread,
)
from anyio.abc import Process
from anyio.streams.text import TextReceiveStream
from loguru import logger
from pydantic import TypeAdapter
from exo.shared.types.events import (
Event,
@@ -20,6 +24,7 @@ from exo.shared.types.events import (
)
from exo.shared.types.tasks import Task, TaskId, TaskStatus
from exo.shared.types.worker.instances import BoundInstance
from exo.shared.types.worker.runner_args import RunnerCliArgs
from exo.shared.types.worker.runners import (
RunnerConnecting,
RunnerFailed,
@@ -31,24 +36,29 @@ from exo.shared.types.worker.runners import (
RunnerWarmingUp,
)
from exo.shared.types.worker.shards import ShardMetadata
from exo.utils.channels import MpReceiver, MpSender, Sender, mp_channel
from exo.utils.channels import Sender
from exo.utils.fd_channels import FdReceiver, FdSender
from exo.utils.task_group import TaskGroup
from exo.worker.runner.bootstrap import entrypoint
PREFILL_TIMEOUT_SECONDS = 60
DECODE_TIMEOUT_SECONDS = 5
# Type adapters for union types
event_adapter: TypeAdapter[Event] = TypeAdapter(Event)
task_adapter: TypeAdapter[Task] = TypeAdapter(Task)
task_id_adapter: TypeAdapter[TaskId] = TypeAdapter(TaskId)
@dataclass(eq=False)
class RunnerSupervisor:
shard_metadata: ShardMetadata
bound_instance: BoundInstance
runner_process: mp.Process
process: Process
initialize_timeout: float
_ev_recv: MpReceiver[Event]
_task_sender: MpSender[Task]
_ev_recv: FdReceiver[Event]
_task_sender: FdSender[Task]
_event_sender: Sender[Event]
_cancel_sender: MpSender[TaskId]
_cancel_sender: FdSender[TaskId]
_tg: TaskGroup = field(default_factory=TaskGroup, init=False)
status: RunnerStatus = field(default_factory=RunnerIdle, init=False)
pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False)
@@ -59,35 +69,75 @@ class RunnerSupervisor:
)
@classmethod
def create(
async def create(
cls,
*,
bound_instance: BoundInstance,
event_sender: Sender[Event],
initialize_timeout: float = 400,
) -> Self:
ev_send, ev_recv = mp_channel[Event]()
task_sender, task_recv = mp_channel[Task]()
cancel_sender, cancel_recv = mp_channel[TaskId]()
# Create pipes for each communication direction:
# - Events: child -> parent (child writes, parent reads)
# - Tasks: parent -> child (parent writes, child reads)
# - Cancels: parent -> child (parent writes, child reads)
event_read_fd, event_write_fd = os.pipe() # Child writes, parent reads
task_read_fd, task_write_fd = os.pipe() # Parent writes, child reads
cancel_read_fd, cancel_write_fd = os.pipe() # Parent writes, child reads
runner_process = mp.Process(
target=entrypoint,
args=(
bound_instance,
ev_send,
task_recv,
cancel_recv,
logger,
),
daemon=True,
# Build CLI args with only the FDs the child needs
args = RunnerCliArgs(
bound_instance=bound_instance,
event_fd=event_write_fd, # Child writes events here
task_fd=task_read_fd, # Child reads tasks from here
cancel_fd=cancel_read_fd, # Child reads cancels from here
log_level="INFO",
)
# File descriptors to pass to the child process
pass_fds = [event_write_fd, task_read_fd, cancel_read_fd]
shard_metadata = bound_instance.bound_shard
self = cls(
# Create FD-based channels for the parent side using TypeAdapters
ev_recv = FdReceiver[Event](event_read_fd, event_adapter)
task_sender = FdSender[Task](task_write_fd, task_adapter)
cancel_sender = FdSender[TaskId](cancel_write_fd, task_id_adapter)
# Close the child-side FDs in the parent process
with contextlib.suppress(OSError):
os.close(event_write_fd)
with contextlib.suppress(OSError):
os.close(task_read_fd)
with contextlib.suppress(OSError):
os.close(cancel_read_fd)
# Start the process using anyio.open_process
# Check if running in a frozen environment (e.g., PyInstaller)
is_frozen = os.environ.get("EXO_FROZEN") == "1" or getattr(sys, "frozen", False)
if is_frozen:
# In frozen builds, use the same executable with a special flag
# The __main__.py handles this by checking sys.argv
cmd = [
sys.executable,
"--exo-runner",
json.dumps(args.model_dump(mode="json")),
]
else:
# In normal Python, use -m to run the module
cmd = [
sys.executable,
"-m",
"exo.worker.runner.bootstrap",
json.dumps(args.model_dump(mode="json")),
]
process = await anyio.open_process(cmd, pass_fds=pass_fds)
return cls(
bound_instance=bound_instance,
shard_metadata=shard_metadata,
runner_process=runner_process,
process=process,
initialize_timeout=initialize_timeout,
_ev_recv=ev_recv,
_task_sender=task_sender,
@@ -95,15 +145,20 @@ class RunnerSupervisor:
_event_sender=event_sender,
)
return self
async def run(self):
self.runner_process.start()
async with self._tg as tg:
tg.start_soon(self._watch_runner)
tg.start_soon(self._forward_events)
tg.start_soon(self._read_stderr)
def shutdown(self):
async def _read_stderr(self):
"""Read and log stderr from the runner process."""
if self.process.stderr:
async with TextReceiveStream(self.process.stderr) as stderr_stream:
async for line in stderr_stream:
logger.debug(f"Runner stderr: {line.rstrip()}")
async def shutdown(self):
logger.info("Runner supervisor shutting down")
self._tg.cancel_tasks()
if not self._cancel_watch_runner.cancel_called:
@@ -118,20 +173,21 @@ class RunnerSupervisor:
self._cancel_sender.send(TaskId("CANCEL_CURRENT_TASK"))
with contextlib.suppress(ClosedResourceError):
self._cancel_sender.close()
self.runner_process.join(5)
if not self.runner_process.is_alive():
logger.info("Runner process succesfully terminated")
return
# This is overkill but it's not technically bad, just unnecessary.
logger.warning("Runner process didn't shutdown succesfully, terminating")
self.runner_process.terminate()
self.runner_process.join(1)
if not self.runner_process.is_alive():
return
logger.critical("Runner process didn't respond to SIGTERM, killing")
self.runner_process.kill()
# Try graceful termination
if self.process.returncode is None:
self.process.terminate()
try:
with anyio.fail_after(5):
await self.process.wait()
except TimeoutError:
logger.warning("Runner process didn't shutdown gracefully, killing")
self.process.kill()
try:
with anyio.fail_after(1):
await self.process.wait()
except TimeoutError:
logger.critical("Runner process didn't respond to SIGKILL")
async def start_task(self, task: Task):
if task.task_id in self.pending:
@@ -198,38 +254,48 @@ class RunnerSupervisor:
self.pending[tid].set()
def __del__(self) -> None:
if self.runner_process.is_alive():
if self.process.returncode is None:
logger.critical("RunnerSupervisor was not stopped cleanly.")
with contextlib.suppress(ValueError):
self.runner_process.kill()
with contextlib.suppress(Exception):
self.process.kill()
async def _watch_runner(self) -> None:
with self._cancel_watch_runner:
while True:
await anyio.sleep(5)
if not self.runner_process.is_alive():
if self.process.returncode is not None:
await self._check_runner(RuntimeError("Runner found to be dead"))
async def _check_runner(self, e: Exception) -> None:
if not self._cancel_watch_runner.cancel_called:
self._cancel_watch_runner.cancel()
logger.info("Checking runner's status")
if self.runner_process.is_alive():
logger.info("Runner was found to be alive, attempting to join process")
await to_thread.run_sync(self.runner_process.join, 5)
rc = self.runner_process.exitcode
logger.info(f"Runner exited with exit code {rc}")
if rc == 0:
returncode = self.process.returncode
if returncode is None:
# Process is still running, try to wait for it
self.process.terminate()
try:
with anyio.fail_after(5):
await self.process.wait()
returncode = self.process.returncode
except TimeoutError:
pass
logger.info(f"Runner exited with return code {returncode}")
if returncode == 0:
return
if isinstance(rc, int) and rc < 0:
sig = -rc
if isinstance(returncode, int) and returncode < 0:
sig = -returncode
try:
cause = f"signal={sig} ({signal.strsignal(sig)})"
except Exception:
cause = f"signal={sig}"
else:
cause = f"exitcode={rc}"
cause = f"returncode={returncode}"
logger.opt(exception=e).error(f"Runner terminated with {cause}")
@@ -248,4 +314,4 @@ class RunnerSupervisor:
logger.warning(
"Event sender already closed, unable to report runner failure"
)
self.shutdown()
await self.shutdown()
@@ -39,7 +39,6 @@ from exo.shared.types.worker.runners import (
RunnerShuttingDown,
RunnerWarmingUp,
)
from exo.utils.channels import mp_channel
from ...constants import (
CHAT_COMPLETION_TASK_ID,
@@ -139,8 +138,46 @@ class EventCollector:
def close(self) -> None:
pass
def join(self) -> None:
pass
class MockFdReceiver[T]:
"""Mock FdReceiver that iterates over a list of items."""
def __init__(self, items: list[T]) -> None:
self._items = items
self._index = 0
self._closed = False
def receive(self) -> T:
if self._closed:
raise StopIteration
if self._index >= len(self._items):
raise StopIteration
item = self._items[self._index]
self._index += 1
return item
def __iter__(self):
return self
def __next__(self) -> T:
return self.receive()
def close(self) -> None:
self._closed = True
def collect(self) -> list[T]:
"""Collect all remaining items."""
items = []
while self._index < len(self._items):
items.append(self._items[self._index])
self._index += 1
return items
def __enter__(self):
return self
def __exit__(self, *args: object) -> None:
self.close()
class MockTokenizer:
@@ -167,28 +204,21 @@ def _run(tasks: Iterable[Task]):
node_id=NODE_A,
)
task_sender, task_receiver = mp_channel[Task]()
_cancel_sender, cancel_receiver = mp_channel[TaskId]()
task_list = list(tasks)
task_receiver = MockFdReceiver[Task](task_list)
cancel_receiver = MockFdReceiver[TaskId]([])
event_sender = EventCollector()
with task_sender:
for t in tasks:
task_sender.send(t)
# worst monkeypatch known to man
# this is some c++ nonsense
task_receiver.close = nothin
task_receiver.join = nothin
with unittest.mock.patch(
"exo.worker.runner.llm_inference.runner.mx.distributed.all_gather",
make_nothin(mx.array([1])),
):
mlx_runner.main(
bound_instance,
event_sender, # pyright: ignore[reportArgumentType]
task_receiver,
cancel_receiver,
)
with unittest.mock.patch(
"exo.worker.runner.llm_inference.runner.mx.distributed.all_gather",
make_nothin(mx.array([1])),
):
mlx_runner.main(
bound_instance,
event_sender, # pyright: ignore[reportArgumentType]
task_receiver, # pyright: ignore[reportArgumentType]
cancel_receiver, # pyright: ignore[reportArgumentType]
)
return event_sender.events