monarch_hyperactor/
proc.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
9use std::hash::DefaultHasher;
10use std::hash::Hash;
11use std::hash::Hasher;
12use std::time::Duration;
13
14use anyhow::Result;
15use hyperactor::Client;
16use hyperactor::RemoteMessage;
17use hyperactor::channel::ChannelAddr;
18use hyperactor::mailbox::PortReceiver;
19use hyperactor::proc::Proc;
20use monarch_types::PickledPyObject;
21use pyo3::exceptions::PyRuntimeError;
22use pyo3::exceptions::PyValueError;
23use pyo3::prelude::*;
24use pyo3::types::PyList;
25use pyo3::types::PyType;
26
27use crate::actor::PythonActor;
28use crate::actor::PythonActorHandle;
29use crate::runtime::signal_safe_block_on;
30
31/// Wrapper around a proc that provides utilities to implement a python actor.
32#[derive(Clone, Debug)]
33#[pyclass(
34    name = "Proc",
35    module = "monarch._rust_bindings.monarch_hyperactor.proc"
36)]
37pub struct PyProc {
38    pub(super) inner: Proc,
39}
40
41#[pymethods]
42impl PyProc {
43    #[new]
44    #[pyo3(signature = ())]
45    fn new() -> PyResult<Self> {
46        Ok(Self {
47            inner: Proc::isolated(),
48        })
49    }
50
51    #[getter]
52    fn addr(&self) -> String {
53        self.inner.proc_addr().addr().to_string()
54    }
55
56    #[getter]
57    fn name(&self) -> String {
58        self.inner
59            .proc_addr()
60            .label()
61            .map(|l: &hyperactor::id::Label| l.as_str().to_string())
62            .unwrap_or_else(|| self.inner.proc_addr().id().to_string())
63    }
64
65    #[getter]
66    fn id(&self) -> String {
67        self.inner.proc_addr().to_string()
68    }
69
70    fn destroy<'py>(
71        &mut self,
72        timeout_in_secs: u64,
73        py: Python<'py>,
74    ) -> PyResult<Bound<'py, PyList>> {
75        let mut inner = self.inner.clone();
76        let (_stopped, aborted) = signal_safe_block_on(py, async move {
77            inner
78                .destroy_and_wait(Duration::from_secs(timeout_in_secs), "destroy")
79                .await
80                .map_err(|e| PyRuntimeError::new_err(e.to_string()))
81        })??;
82        let aborted_actors = aborted
83            .into_iter()
84            .map(|actor_id| format!("{}", actor_id))
85            .collect::<Vec<_>>();
86        // TODO: i don't think returning this list is of much use for
87        // anything?
88        PyList::new(py, aborted_actors)
89    }
90
91    #[pyo3(signature = (actor, name=None))]
92    fn spawn<'py>(
93        &self,
94        py: Python<'py>,
95        actor: &Bound<'py, PyType>,
96        name: Option<String>,
97    ) -> PyResult<Bound<'py, PyAny>> {
98        let proc = self.inner.clone();
99        let pickled_type = PickledPyObject::pickle(actor.as_any())?;
100        crate::runtime::future_into_py(py, async move {
101            let actor = PythonActor::new(pickled_type, None, None, None)?;
102            Ok(PythonActorHandle {
103                inner: proc.spawn_with_label(name.as_deref().unwrap_or("anon"), actor),
104            })
105        })
106    }
107
108    #[pyo3(signature = (actor, name=None))]
109    fn spawn_blocking<'py>(
110        &self,
111        py: Python<'py>,
112        actor: &Bound<'py, PyType>,
113        name: Option<String>,
114    ) -> PyResult<PythonActorHandle> {
115        let proc = self.inner.clone();
116        let pickled_type = PickledPyObject::pickle(actor.as_any())?;
117        Ok(PythonActorHandle {
118            inner: signal_safe_block_on(py, async move {
119                let actor = PythonActor::new(pickled_type, None, None, None)?;
120                Ok(proc.spawn_with_label(name.as_deref().unwrap_or("anon"), actor))
121            })
122            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?
123            .map_err(|e: anyhow::Error| PyRuntimeError::new_err(e.to_string()))?,
124        })
125    }
126}
127
128impl PyProc {
129    pub fn new_from_proc(proc: Proc) -> Self {
130        Self { inner: proc }
131    }
132}
133
134#[pyclass(
135    frozen,
136    name = "ActorAddr",
137    module = "monarch._rust_bindings.monarch_hyperactor.proc"
138)]
139#[derive(Clone)]
140pub struct PyActorAddr {
141    pub(super) inner: hyperactor::ActorAddr,
142}
143
144impl From<hyperactor::ActorAddr> for PyActorAddr {
145    fn from(actor_id: hyperactor::ActorAddr) -> Self {
146        Self { inner: actor_id }
147    }
148}
149
150impl From<PyActorAddr> for hyperactor::ActorAddr {
151    fn from(val: PyActorAddr) -> Self {
152        val.inner
153    }
154}
155
156#[pymethods]
157impl PyActorAddr {
158    #[new]
159    #[pyo3(signature = (*, addr, proc_name, actor_name))]
160    fn new(addr: &str, proc_name: &str, actor_name: &str) -> PyResult<Self> {
161        let addr: ChannelAddr = addr.parse().map_err(|e| {
162            PyValueError::new_err(format!("Failed to parse channel address '{}': {}", addr, e))
163        })?;
164        Ok(Self {
165            inner: hyperactor::ProcAddr::singleton(addr, proc_name).actor_addr(actor_name),
166        })
167    }
168
169    #[staticmethod]
170    fn from_string(actor_id: &str) -> PyResult<Self> {
171        Ok(Self {
172            inner: actor_id.parse().map_err(|e| {
173                PyValueError::new_err(format!(
174                    "Failed to extract actor id from {}: {}",
175                    actor_id, e
176                ))
177            })?,
178        })
179    }
180
181    #[getter]
182    fn addr(&self) -> String {
183        self.inner.proc_addr().addr().to_string()
184    }
185
186    #[getter]
187    fn proc_name(&self) -> String {
188        self.inner
189            .proc_addr()
190            .label()
191            .map(|l: &hyperactor::id::Label| l.as_str().to_string())
192            .unwrap_or_else(|| self.inner.proc_addr().id().to_string())
193    }
194
195    #[getter]
196    fn actor_name(&self) -> String {
197        self.inner
198            .label()
199            .map(|l: &hyperactor::id::Label| l.as_str().to_string())
200            .unwrap_or_else(|| self.inner.uid().to_string())
201    }
202
203    #[getter]
204    fn label(&self) -> Option<String> {
205        self.inner
206            .label()
207            .map(|l: &hyperactor::id::Label| l.as_str().to_string())
208    }
209
210    #[getter]
211    fn proc_label(&self) -> Option<String> {
212        self.inner
213            .proc_addr()
214            .label()
215            .map(|l: &hyperactor::id::Label| l.as_str().to_string())
216    }
217
218    #[getter]
219    fn uid(&self) -> String {
220        self.inner.uid().to_string()
221    }
222
223    #[getter]
224    fn pid(&self) -> String {
225        self.uid()
226    }
227
228    #[getter]
229    fn proc_id(&self) -> String {
230        self.inner.proc_addr().to_string()
231    }
232
233    #[getter]
234    fn is_root(&self) -> bool {
235        self.inner.is_root()
236    }
237
238    fn __str__(&self) -> String {
239        self.inner.to_string()
240    }
241
242    fn __hash__(&self) -> u64 {
243        let mut hasher = DefaultHasher::new();
244        self.inner.to_string().hash(&mut hasher);
245        hasher.finish()
246    }
247
248    fn __eq__(&self, other: &Bound<'_, PyAny>) -> PyResult<bool> {
249        if let Ok(other) = other.extract::<PyActorAddr>() {
250            Ok(self.inner == other.inner)
251        } else {
252            Ok(false)
253        }
254    }
255
256    fn __reduce__<'py>(slf: &Bound<'py, Self>) -> PyResult<(Bound<'py, PyAny>, (String,))> {
257        Ok((slf.getattr("from_string")?, (slf.borrow().__str__(),)))
258    }
259}
260
261impl From<&PyActorAddr> for hyperactor::ActorAddr {
262    fn from(actor_id: &PyActorAddr) -> Self {
263        actor_id.inner.clone()
264    }
265}
266
267impl std::fmt::Debug for PyActorAddr {
268    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
269        self.inner.fmt(f)
270    }
271}
272
273#[derive(Clone, Copy, PartialEq, Eq)]
274enum InstanceStatus {
275    Running,
276    Stopped,
277}
278
279/// Wrapper around a [`Any`] that allows returning it to python and
280/// passed to python based detached actors to send to other actors.
281#[pyclass(
282    frozen,
283    name = "Serialized",
284    module = "monarch._rust_bindings.monarch_hyperactor.proc"
285)]
286#[derive(Debug)]
287pub struct PySerialized {
288    inner: wirevalue::Any,
289    /// The handler port for this message type.
290    port: u64,
291}
292
293impl PySerialized {
294    pub fn new<M: RemoteMessage>(message: &M) -> PyResult<Self> {
295        Ok(Self {
296            inner: wirevalue::Any::serialize(message).map_err(|err| {
297                PyRuntimeError::new_err(format!(
298                    "failed to serialize message of type {} to Any: {}",
299                    std::any::type_name::<M>(),
300                    err
301                ))
302            })?,
303            port: M::port(),
304        })
305    }
306
307    pub fn deserialized<M: RemoteMessage>(&self) -> PyResult<M> {
308        self.inner.deserialized().map_err(|err| {
309            PyRuntimeError::new_err(format!("failed to deserialize message: {}", err))
310        })
311    }
312
313    /// The handler port for this message type.
314    pub fn port(&self) -> u64 {
315        self.port
316    }
317}
318
319/// Wrapper around an instance of an actor that provides utilities to implement
320/// a python actor. This helps by allowing users to specialize the actor to the
321/// message type they want to handle.
322pub struct InstanceWrapper<M: RemoteMessage> {
323    instance: Client,
324    message_receiver: PortReceiver<M>,
325    status: InstanceStatus,
326    actor_id: hyperactor::ActorAddr,
327}
328
329impl<M: RemoteMessage> InstanceWrapper<M> {
330    pub fn new(proc: &PyProc, actor_name: &str) -> Result<Self> {
331        let instance = proc.inner.client(actor_name);
332        // TEMPORARY: remove after using fixed handler ports.
333        let (_handler_port, message_receiver) = instance.bind_handler_port::<M>();
334
335        let actor_id = instance.self_addr().clone();
336
337        Ok(Self {
338            instance,
339            message_receiver,
340            status: InstanceStatus::Running,
341            actor_id,
342        })
343    }
344
345    /// Send a message to any actor. It is the responsibility of the caller to ensure the right
346    /// payload accepted by the target actor has been serialized and provided to this function.
347    pub fn send(&self, actor_id: &PyActorAddr, message: &PySerialized) -> PyResult<()> {
348        hyperactor::internal_macro_support::tracing::debug!(
349            name = "py_send_message",
350            actor_id =
351                hyperactor::internal_macro_support::tracing::field::display(self.actor_addr()),
352            receiver_actor_id = tracing::field::display(&actor_id.inner),
353            ?message,
354        );
355        actor_id
356            .inner
357            .port_addr(hyperactor::Port::handler_id(message.port(), None))
358            .send(&self.instance, message.inner.clone());
359        Ok(())
360    }
361
362    /// Make sure the actor is still alive (in the `Running` state).
363    fn ensure_alive(&self) -> Result<()> {
364        anyhow::ensure!(
365            self.status == InstanceStatus::Running,
366            "actor is not running"
367        );
368        Ok(())
369    }
370
371    /// Get the next message from the queue. It will wait until a message is received
372    /// or the timeout is reached in which case it will return None.
373    #[hyperactor::instrument(level = "trace", fields(actor_id = hyperactor::internal_macro_support::tracing::field::display(self.actor_addr())))]
374    pub async fn next_message(&mut self, timeout_msec: Option<u64>) -> Result<Option<M>> {
375        hyperactor::declare_static_timer!(
376            PY_NEXT_MESSAGE_TIMER,
377            "py_next_message",
378            hyperactor_telemetry::TimeUnit::Nanos
379        );
380        let _ = PY_NEXT_MESSAGE_TIMER
381            .start(hyperactor::kv_pairs!("actor_id" => self.actor_addr().to_string(), "mode" => match timeout_msec{
382                None => "blocking",
383                Some(0) => "polling",
384                Some(_) => "blocking_with_timeout",
385            }));
386        self.ensure_alive()?;
387        match timeout_msec {
388            // Blocking wait for next message.
389            None => {
390                self.message_receiver.recv().await.map(Some)},
391            Some(0) => {
392                // Non-blocking.
393                // Try to get next message without waiting.
394                self.message_receiver.try_recv()
395            }
396            Some(timeout_msec) => {
397                // Blocking wait with a timeout.
398                match tokio::time::timeout(
399                    Duration::from_millis(timeout_msec),
400                    self.message_receiver.recv(),
401                )
402                .await
403                {
404                    Ok(output) => output.map(Some),
405                    Err(_) => Ok(None), // Timeout reached
406                }
407            }
408        }
409        .map_err(|err| err.into())
410        .inspect_err(|err| {
411            hyperactor::metrics::ACTOR_MESSAGE_RECEIVE_ERRORS.add(1, hyperactor::kv_pairs!("actor_id" => self.actor_addr().to_string()));
412            tracing::error!(err=?err, actor_id=%self.actor_addr(), "unable to receive next py message");
413        })
414        .inspect(|_|{
415            hyperactor::metrics::ACTOR_MESSAGES_RECEIVED.add(1, hyperactor::kv_pairs!("actor_id" => self.actor_addr().to_string()));
416        })
417    }
418
419    /// Put the actor in stopped mode and return any messages that were received.
420    #[hyperactor::instrument(fields(actor_id=hyperactor::internal_macro_support::tracing::field::display(self.actor_addr())))]
421    pub fn drain_and_stop(&mut self) -> Result<Vec<M>> {
422        self.ensure_alive()?;
423        let messages: Vec<M> = self.message_receiver.drain().into_iter().collect();
424        tracing::info!("stopping the client actor in Python client");
425        self.status = InstanceStatus::Stopped;
426        Ok(messages)
427    }
428
429    pub fn instance(&self) -> &Client {
430        &self.instance
431    }
432
433    pub fn actor_addr(&self) -> &hyperactor::ActorAddr {
434        &self.actor_id
435    }
436}
437
438pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
439    hyperactor_mod.add_class::<PyProc>()?;
440    hyperactor_mod.add_class::<PyActorAddr>()?;
441    hyperactor_mod.add_class::<PySerialized>()?;
442    Ok(())
443}
monarch_hyperactor/proc.rs

monarch_hyperactor/
proc.rs
