monarch_hyperactor/
pywaker.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9//! Pywaker provides a GIL-free way to wake up Python event loops from Rust.
10//! This is accomplished by writing a single byte to a "self pipe" for every
11//! wake up. The corresponding read-side FD is added to the Python-side event
12//! loop.
13//!
14//! The approach does not add overhead in uncontended scenarios:
15//!
16//! ```ignore
17//!   | Approach             | Mean   | Notes                           |
18//!   |----------------------|--------|---------------------------------|
19//!   | FD-based             | ~48 µs | No GIL acquisition on Rust side |
20//!   | call_soon_threadsafe | ~56 µs | Acquires GIL on Rust side       |
21//!   | Pure Python          | ~44 µs | Baseline for comparison         |
22//! ```
23//!
24//! ([Full benchmarks](https://github.com/mariusae/tmp/tree/main/2025-12-16-wakerbench)).
25//!
26//! ## Implementation notes
27//!
28//! [`event`] creates a pipe(2), and wakeup events are implemented by writing a single
29//! byte into the pipe. The Python (read side) event registers the read fd into its
30//! asyncio event loop. Thus the pipe acts as a notification queue. When the fd becomes
31//! readable, the python event drains the pipe and then sets an `asyncio.Event` to notify
32//! any waiters. [`pipe`] supplies the same pipe to a reader that is not an event loop,
33//! such as a sync actor's inbox, which waits on the read fd with `poll`.
34//!
35//! Each read side has a single reader thread, and is thus free of race conditions: an
36//! event's reads all occur within its asyncio event loop, which processes events
37//! sequentially on one thread, and a sync actor's inbox is read only by its driver.
38
39use std::os::fd::FromRawFd;
40use std::os::fd::IntoRawFd;
41use std::os::fd::OwnedFd;
42use std::os::fd::RawFd;
43
44use monarch_types::MapPyErr;
45use nix::errno::Errno;
46use nix::fcntl::FcntlArg;
47use nix::fcntl::OFlag;
48use nix::fcntl::fcntl;
49use nix::unistd;
50use nix::unistd::write;
51use pyo3::Bound;
52use pyo3::Py;
53use pyo3::PyAny;
54use pyo3::PyResult;
55use pyo3::prelude::*;
56use pyo3::pyclass;
57use pyo3::types::PyModule;
58use pyo3::types::PyModuleMethods;
59
60/// Write side of a non-blocking self-pipe used to notify its receiver.
61#[derive(Debug)]
62pub struct Waker {
63    /// Pipe descriptor written by [`Waker::wake`]. The corresponding read end
64    /// is owned by a [`PyEvent`] or a blocking Python receiver.
65    write_fd: OwnedFd,
66}
67
68impl Waker {
69    /// Notify the receiver by writing one byte to the pipe.
70    ///
71    /// A full pipe already represents a pending notification and therefore
72    /// counts as success. A closed read end returns `false`.
73    pub fn wake(&self) -> Result<bool, nix::Error> {
74        static DATA: [u8; 1] = *b"w";
75
76        match retry_eintr(|| write(&self.write_fd, &DATA)) {
77            Ok(_) => Ok(true),
78            // Pipe is full. This is ok.
79            Err(Errno::EAGAIN) => Ok(true),
80            // The Python side closed the pipe. It is no longer listening.
81            Err(Errno::EPIPE) => Ok(false),
82            Err(e) => Err(e),
83        }
84    }
85}
86
87/// Wakers are not intended to be used from Python; TestWaker
88/// is provided to faciliate pure-Python unit testing.
89// NOTE: We can't use a Python calss name that starts with "Test" since
90// during Python testing, Pytest will inspect anything that starts with
91// "Test" and check if its callable which in pyo3 >= 0.26 will raise
92// a TypeError.
93#[pyclass(name = "PyTestWaker", module = "monarch._src.actor.waker")]
94struct TestPyWaker(Waker);
95
96#[pymethods]
97impl TestPyWaker {
98    fn wake(&self) -> PyResult<bool> {
99        self.0.wake().map_pyerr()
100    }
101
102    #[staticmethod]
103    fn create() -> PyResult<(TestPyWaker, PyEvent)> {
104        let (waker, event) = event().map_pyerr()?;
105        Ok((TestPyWaker(waker), event))
106    }
107}
108
109/// An event that is awoken by a [`Waker`].
110#[pyclass(name = "Event", module = "monarch._src.actor.waker")]
111pub struct PyEvent {
112    #[pyo3(get, name = "_read_fd")]
113    read_fd: RawFd,
114
115    #[pyo3(get, set, name = "_event_loop")]
116    event_loop: Option<Py<PyAny>>,
117
118    #[pyo3(get, set, name = "_event")]
119    event: Option<Py<PyAny>>,
120}
121
122impl Drop for PyEvent {
123    fn drop(&mut self) {
124        // SAFETY: fd was obtained via into_raw_fd() in the event() function
125        let _ = unsafe { OwnedFd::from_raw_fd(self.read_fd) };
126    }
127}
128
129/// Create a new event, returning the (Rust only) [`Waker`], and
130/// a Python [`PyEvent`], intended for passing to Python code.
131pub fn event() -> Result<(Waker, PyEvent), nix::Error> {
132    let (waker, read_fd) = pipe()?;
133
134    Ok((
135        waker,
136        PyEvent {
137            read_fd: read_fd.into_raw_fd(),
138            event_loop: None,
139            event: None,
140        },
141    ))
142}
143
144/// Create a non-blocking pipe, returning its write end as a [`Waker`] and
145/// its read end for the caller to wait on.
146pub(crate) fn pipe() -> Result<(Waker, OwnedFd), nix::Error> {
147    let (read_fd, write_fd) = unistd::pipe()?;
148
149    set_nonblocking(&read_fd)?;
150    set_nonblocking(&write_fd)?;
151
152    Ok((Waker { write_fd }, read_fd))
153}
154
155/// Call `f` again for as long as it fails with `EINTR`.
156pub(crate) fn retry_eintr<T>(mut f: impl FnMut() -> nix::Result<T>) -> nix::Result<T> {
157    loop {
158        match f() {
159            Err(Errno::EINTR) => continue,
160            result => return result,
161        }
162    }
163}
164
165fn set_nonblocking(fd: &OwnedFd) -> Result<(), nix::Error> {
166    let flags = OFlag::from_bits_truncate(fcntl(fd, FcntlArg::F_GETFL)?) | OFlag::O_NONBLOCK;
167    fcntl(fd, FcntlArg::F_SETFL(flags))?;
168    Ok(())
169}
170
171pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
172    hyperactor_mod.add_class::<PyEvent>()?;
173    hyperactor_mod.add_class::<TestPyWaker>()?;
174    Ok(())
175}
monarch_hyperactor/pywaker.rs

monarch_hyperactor/
pywaker.rs
