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::ProcId;
17use hyperactor::RemoteMessage;
18use hyperactor::channel::ChannelAddr;
19use hyperactor::id::Label;
20use hyperactor::id::Uid;
21use hyperactor::mailbox::PortReceiver;
22use hyperactor::proc::Proc;
23use monarch_types::PickledPyObject;
24use pyo3::exceptions::PyRuntimeError;
25use pyo3::exceptions::PyValueError;
26use pyo3::prelude::*;
27use pyo3::types::PyList;
28use pyo3::types::PyType;
29
30use crate::actor::ActorKind;
31use crate::actor::PythonActor;
32use crate::actor::PythonActorHandle;
33use crate::runtime::signal_safe_block_on;
34
35#[pyclass(
36    frozen,
37    name = "Uid",
38    module = "monarch._rust_bindings.monarch_hyperactor.proc"
39)]
40#[derive(Clone)]
41pub struct PyUid {
42    inner: Uid,
43}
44
45#[pymethods]
46impl PyUid {
47    #[staticmethod]
48    fn anonymous() -> Self {
49        Self {
50            inner: Uid::anonymous(),
51        }
52    }
53
54    #[staticmethod]
55    #[pyo3(signature = (label=None))]
56    fn instance(label: Option<&str>) -> PyResult<Self> {
57        Ok(Self {
58            inner: match label {
59                Some(label) => Uid::instance(Label::new(label).map_err(|error| {
60                    PyValueError::new_err(format!("invalid uid label: {error}"))
61                })?),
62                None => Uid::anonymous(),
63            },
64        })
65    }
66
67    #[staticmethod]
68    #[pyo3(signature = (value, label=None))]
69    fn instance_from_value(value: u64, label: Option<&str>) -> PyResult<Self> {
70        Ok(Self {
71            inner: Uid::instance_from_value(
72                value,
73                label.map(Label::new).transpose().map_err(|error| {
74                    PyValueError::new_err(format!("invalid uid label: {error}"))
75                })?,
76            ),
77        })
78    }
79
80    #[staticmethod]
81    fn from_string(uid: &str) -> PyResult<Self> {
82        Ok(Self {
83            inner: uid
84                .parse()
85                .map_err(|error| PyValueError::new_err(format!("invalid uid: {error}")))?,
86        })
87    }
88
89    #[getter]
90    fn label(&self) -> Option<String> {
91        self.inner.label().map(ToString::to_string)
92    }
93
94    #[getter]
95    fn is_instance(&self) -> bool {
96        matches!(&self.inner, Uid::Instance(..))
97    }
98
99    #[getter]
100    fn is_singleton(&self) -> bool {
101        matches!(&self.inner, Uid::Singleton(..))
102    }
103
104    #[getter]
105    fn instance_value(&self) -> Option<u64> {
106        self.inner.instance_value()
107    }
108
109    fn __str__(&self) -> String {
110        self.inner.to_string()
111    }
112
113    fn __hash__(&self) -> u64 {
114        let mut hasher = DefaultHasher::new();
115        self.inner.hash(&mut hasher);
116        hasher.finish()
117    }
118
119    fn __eq__(&self, other: &Bound<'_, PyAny>) -> PyResult<bool> {
120        Ok(other
121            .extract::<PyRef<'_, PyUid>>()
122            .is_ok_and(|other| self.inner == other.inner))
123    }
124
125    fn __reduce__<'py>(slf: &Bound<'py, Self>) -> PyResult<(Bound<'py, PyAny>, (String,))> {
126        Ok((slf.getattr("from_string")?, (slf.borrow().__str__(),)))
127    }
128}
129
130#[pyclass(
131    frozen,
132    name = "ProcId",
133    module = "monarch._rust_bindings.monarch_hyperactor.proc"
134)]
135#[derive(Clone)]
136pub struct PyProcId {
137    pub(crate) inner: ProcId,
138}
139
140#[pymethods]
141impl PyProcId {
142    #[new]
143    fn new(uid: PyRef<'_, PyUid>) -> Self {
144        Self {
145            inner: ProcId::new(uid.inner.clone(), None),
146        }
147    }
148
149    #[staticmethod]
150    fn from_string(proc_id: &str) -> PyResult<Self> {
151        Ok(Self {
152            inner: proc_id
153                .parse()
154                .map_err(|error| PyValueError::new_err(format!("invalid proc id: {error}")))?,
155        })
156    }
157
158    #[getter]
159    fn uid(&self) -> PyUid {
160        PyUid {
161            inner: self.inner.uid().clone(),
162        }
163    }
164
165    #[getter]
166    fn label(&self) -> Option<String> {
167        self.inner.label().map(ToString::to_string)
168    }
169
170    fn __str__(&self) -> String {
171        self.inner.to_string()
172    }
173
174    fn __hash__(&self) -> u64 {
175        let mut hasher = DefaultHasher::new();
176        self.inner.hash(&mut hasher);
177        hasher.finish()
178    }
179
180    fn __eq__(&self, other: &Bound<'_, PyAny>) -> PyResult<bool> {
181        Ok(other
182            .extract::<PyRef<'_, PyProcId>>()
183            .is_ok_and(|other| self.inner == other.inner))
184    }
185
186    fn __reduce__<'py>(slf: &Bound<'py, Self>) -> PyResult<(Bound<'py, PyAny>, (String,))> {
187        Ok((slf.getattr("from_string")?, (slf.borrow().__str__(),)))
188    }
189}
190
191/// Wrapper around a proc that provides utilities to implement a python actor.
192#[derive(Clone, Debug)]
193#[pyclass(
194    name = "Proc",
195    module = "monarch._rust_bindings.monarch_hyperactor.proc"
196)]
197pub struct PyProc {
198    pub(super) inner: Proc,
199}
200
201#[pymethods]
202impl PyProc {
203    #[new]
204    #[pyo3(signature = ())]
205    fn new() -> PyResult<Self> {
206        Ok(Self {
207            inner: Proc::isolated(),
208        })
209    }
210
211    #[getter]
212    fn addr(&self) -> String {
213        self.inner.proc_addr().addr().to_string()
214    }
215
216    #[getter]
217    fn name(&self) -> String {
218        self.inner
219            .proc_addr()
220            .label()
221            .map(|l: &hyperactor::id::Label| l.as_str().to_string())
222            .unwrap_or_else(|| self.inner.proc_addr().id().to_string())
223    }
224
225    #[getter]
226    fn id(&self) -> String {
227        self.inner.proc_addr().to_string()
228    }
229
230    fn destroy<'py>(
231        &mut self,
232        timeout_in_secs: u64,
233        py: Python<'py>,
234    ) -> PyResult<Bound<'py, PyList>> {
235        let mut inner = self.inner.clone();
236        let (_stopped, aborted) = signal_safe_block_on(py, async move {
237            inner
238                .destroy_and_wait(Duration::from_secs(timeout_in_secs), "destroy")
239                .await
240                .map_err(|e| PyRuntimeError::new_err(e.to_string()))
241        })??;
242        let aborted_actors = aborted
243            .into_iter()
244            .map(|actor_id| format!("{}", actor_id))
245            .collect::<Vec<_>>();
246        // TODO: i don't think returning this list is of much use for
247        // anything?
248        PyList::new(py, aborted_actors)
249    }
250
251    /// Spawn a low-level Python dispatcher actor.
252    ///
253    /// `actor` implements the Rust-binding `Actor` protocol, whose entry point
254    /// is `async handle(...)`; it is not a user `monarch.actor.Actor` class with
255    /// decorated endpoints. This route carries no Init message or user class to
256    /// classify, so it always retains the async event-loop topology.
257    ///
258    /// User actors go through `ProcMesh.spawn`, which classifies their endpoint
259    /// methods before native spawn and passes the resulting [`ActorKind`].
260    #[pyo3(signature = (actor, name=None))]
261    fn spawn<'py>(
262        &self,
263        py: Python<'py>,
264        actor: &Bound<'py, PyType>,
265        name: Option<String>,
266    ) -> PyResult<Bound<'py, PyAny>> {
267        let proc = self.inner.clone();
268        let pickled_type = PickledPyObject::pickle(actor.as_any())?;
269        crate::runtime::future_into_py(py, async move {
270            // The low-level Actor protocol has one async `handle` entry point
271            // and no user endpoint class to classify.
272            let actor = PythonActor::new(pickled_type, None, None, ActorKind::Async)?;
273            Ok(PythonActorHandle {
274                inner: proc.spawn_with_label(name.as_deref().unwrap_or("anon"), actor),
275            })
276        })
277    }
278
279    /// Blocking form of [`Self::spawn`], with the same low-level async-dispatch
280    /// contract.
281    #[pyo3(signature = (actor, name=None))]
282    fn spawn_blocking<'py>(
283        &self,
284        py: Python<'py>,
285        actor: &Bound<'py, PyType>,
286        name: Option<String>,
287    ) -> PyResult<PythonActorHandle> {
288        let proc = self.inner.clone();
289        let pickled_type = PickledPyObject::pickle(actor.as_any())?;
290        Ok(PythonActorHandle {
291            inner: signal_safe_block_on(py, async move {
292                // The low-level Actor protocol has one async `handle` entry
293                // point and no user endpoint class to classify.
294                let actor = PythonActor::new(pickled_type, None, None, ActorKind::Async)?;
295                Ok(proc.spawn_with_label(name.as_deref().unwrap_or("anon"), actor))
296            })
297            .map_err(|e| PyRuntimeError::new_err(e.to_string()))?
298            .map_err(|e: anyhow::Error| PyRuntimeError::new_err(e.to_string()))?,
299        })
300    }
301}
302
303impl PyProc {
304    pub fn new_from_proc(proc: Proc) -> Self {
305        Self { inner: proc }
306    }
307}
308
309#[pyclass(
310    frozen,
311    name = "ActorAddr",
312    module = "monarch._rust_bindings.monarch_hyperactor.proc"
313)]
314#[derive(Clone)]
315pub struct PyActorAddr {
316    pub(super) inner: hyperactor::ActorAddr,
317}
318
319impl From<hyperactor::ActorAddr> for PyActorAddr {
320    fn from(actor_id: hyperactor::ActorAddr) -> Self {
321        Self { inner: actor_id }
322    }
323}
324
325impl From<PyActorAddr> for hyperactor::ActorAddr {
326    fn from(val: PyActorAddr) -> Self {
327        val.inner
328    }
329}
330
331#[pymethods]
332impl PyActorAddr {
333    #[new]
334    #[pyo3(signature = (*, addr, proc_name, actor_name))]
335    fn new(addr: &str, proc_name: &str, actor_name: &str) -> PyResult<Self> {
336        let addr: ChannelAddr = addr.parse().map_err(|e| {
337            PyValueError::new_err(format!("Failed to parse channel address '{}': {}", addr, e))
338        })?;
339        Ok(Self {
340            inner: hyperactor::ProcAddr::singleton(addr, proc_name).actor_addr(actor_name),
341        })
342    }
343
344    #[staticmethod]
345    fn from_string(actor_id: &str) -> PyResult<Self> {
346        Ok(Self {
347            inner: actor_id.parse().map_err(|e| {
348                PyValueError::new_err(format!(
349                    "Failed to extract actor id from {}: {}",
350                    actor_id, e
351                ))
352            })?,
353        })
354    }
355
356    #[getter]
357    fn addr(&self) -> String {
358        self.inner.proc_addr().addr().to_string()
359    }
360
361    #[getter]
362    fn proc_name(&self) -> String {
363        self.inner
364            .proc_addr()
365            .label()
366            .map(|l: &hyperactor::id::Label| l.as_str().to_string())
367            .unwrap_or_else(|| self.inner.proc_addr().id().to_string())
368    }
369
370    #[getter]
371    fn actor_name(&self) -> String {
372        self.inner
373            .label()
374            .map(|l: &hyperactor::id::Label| l.as_str().to_string())
375            .unwrap_or_else(|| self.inner.uid().to_string())
376    }
377
378    #[getter]
379    fn label(&self) -> Option<String> {
380        self.inner
381            .label()
382            .map(|l: &hyperactor::id::Label| l.as_str().to_string())
383    }
384
385    #[getter]
386    fn proc_label(&self) -> Option<String> {
387        self.inner
388            .proc_addr()
389            .label()
390            .map(|l: &hyperactor::id::Label| l.as_str().to_string())
391    }
392
393    #[getter]
394    fn uid(&self) -> String {
395        self.inner.uid().to_string()
396    }
397
398    #[getter]
399    fn pid(&self) -> String {
400        self.uid()
401    }
402
403    #[getter]
404    fn proc_id(&self) -> String {
405        self.inner.proc_addr().to_string()
406    }
407
408    #[getter]
409    fn is_root(&self) -> bool {
410        self.inner.is_root()
411    }
412
413    fn __str__(&self) -> String {
414        self.inner.to_string()
415    }
416
417    fn __hash__(&self) -> u64 {
418        let mut hasher = DefaultHasher::new();
419        self.inner.to_string().hash(&mut hasher);
420        hasher.finish()
421    }
422
423    fn __eq__(&self, other: &Bound<'_, PyAny>) -> PyResult<bool> {
424        if let Ok(other) = other.extract::<PyActorAddr>() {
425            Ok(self.inner == other.inner)
426        } else {
427            Ok(false)
428        }
429    }
430
431    fn __reduce__<'py>(slf: &Bound<'py, Self>) -> PyResult<(Bound<'py, PyAny>, (String,))> {
432        Ok((slf.getattr("from_string")?, (slf.borrow().__str__(),)))
433    }
434}
435
436impl From<&PyActorAddr> for hyperactor::ActorAddr {
437    fn from(actor_id: &PyActorAddr) -> Self {
438        actor_id.inner.clone()
439    }
440}
441
442impl std::fmt::Debug for PyActorAddr {
443    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
444        self.inner.fmt(f)
445    }
446}
447
448#[derive(Clone, Copy, PartialEq, Eq)]
449enum InstanceStatus {
450    Running,
451    Stopped,
452}
453
454/// Wrapper around a [`Any`] that allows returning it to python and
455/// passed to python based detached actors to send to other actors.
456#[pyclass(
457    frozen,
458    name = "Serialized",
459    module = "monarch._rust_bindings.monarch_hyperactor.proc"
460)]
461#[derive(Debug)]
462pub struct PySerialized {
463    inner: wirevalue::Any,
464    /// The handler port for this message type.
465    port: u64,
466}
467
468impl PySerialized {
469    pub fn new<M: RemoteMessage>(message: &M) -> PyResult<Self> {
470        Ok(Self {
471            inner: wirevalue::Any::serialize(message).map_err(|err| {
472                PyRuntimeError::new_err(format!(
473                    "failed to serialize message of type {} to Any: {}",
474                    std::any::type_name::<M>(),
475                    err
476                ))
477            })?,
478            port: M::port(),
479        })
480    }
481
482    pub fn deserialized<M: RemoteMessage>(&self) -> PyResult<M> {
483        self.inner.deserialized().map_err(|err| {
484            PyRuntimeError::new_err(format!("failed to deserialize message: {}", err))
485        })
486    }
487
488    /// The handler port for this message type.
489    pub fn port(&self) -> u64 {
490        self.port
491    }
492}
493
494/// Wrapper around an instance of an actor that provides utilities to implement
495/// a python actor. This helps by allowing users to specialize the actor to the
496/// message type they want to handle.
497pub struct InstanceWrapper<M: RemoteMessage> {
498    instance: Client,
499    message_receiver: PortReceiver<M>,
500    status: InstanceStatus,
501    actor_id: hyperactor::ActorAddr,
502}
503
504impl<M: RemoteMessage> InstanceWrapper<M> {
505    pub fn new(proc: &PyProc, actor_name: &str) -> Result<Self> {
506        let instance = proc.inner.client(actor_name);
507        // TEMPORARY: remove after using fixed handler ports.
508        let (_handler_port, message_receiver) = instance.bind_handler_port::<M>();
509
510        let actor_id = instance.self_addr().clone();
511
512        Ok(Self {
513            instance,
514            message_receiver,
515            status: InstanceStatus::Running,
516            actor_id,
517        })
518    }
519
520    /// Send a message to any actor. It is the responsibility of the caller to ensure the right
521    /// payload accepted by the target actor has been serialized and provided to this function.
522    pub fn send(&self, actor_id: &PyActorAddr, message: &PySerialized) -> PyResult<()> {
523        hyperactor::internal_macro_support::tracing::debug!(
524            name = "py_send_message",
525            actor_id =
526                hyperactor::internal_macro_support::tracing::field::display(self.actor_addr()),
527            receiver_actor_id = tracing::field::display(&actor_id.inner),
528            ?message,
529        );
530        actor_id
531            .inner
532            .port_addr(hyperactor::Port::handler_id(message.port(), None))
533            .send(&self.instance, message.inner.clone());
534        Ok(())
535    }
536
537    /// Make sure the actor is still alive (in the `Running` state).
538    fn ensure_alive(&self) -> Result<()> {
539        anyhow::ensure!(
540            self.status == InstanceStatus::Running,
541            "actor is not running"
542        );
543        Ok(())
544    }
545
546    /// Get the next message from the queue. It will wait until a message is received
547    /// or the timeout is reached in which case it will return None.
548    #[hyperactor::instrument(level = "trace", fields(actor_id = hyperactor::internal_macro_support::tracing::field::display(self.actor_addr())))]
549    pub async fn next_message(&mut self, timeout_msec: Option<u64>) -> Result<Option<M>> {
550        hyperactor::declare_static_timer!(
551            PY_NEXT_MESSAGE_TIMER,
552            "py_next_message",
553            hyperactor_telemetry::TimeUnit::Nanos
554        );
555        let metric_pairs = hyperactor::kv_pairs!("actor_id" => self.actor_addr().to_string(), "mode" => match timeout_msec{
556            None => "blocking",
557            Some(0) => "polling",
558            Some(_) => "blocking_with_timeout",
559        });
560        let _timer = PY_NEXT_MESSAGE_TIMER.start(metric_pairs);
561        self.ensure_alive()?;
562        match timeout_msec {
563            // Blocking wait for next message.
564            None => {
565                self.message_receiver.recv().await.map(Some)},
566            Some(0) => {
567                // Non-blocking.
568                // Try to get next message without waiting.
569                self.message_receiver.try_recv()
570            }
571            Some(timeout_msec) => {
572                // Blocking wait with a timeout.
573                match tokio::time::timeout(
574                    Duration::from_millis(timeout_msec),
575                    self.message_receiver.recv(),
576                )
577                .await
578                {
579                    Ok(output) => output.map(Some),
580                    Err(_) => Ok(None), // Timeout reached
581                }
582            }
583        }
584        .map_err(|err| err.into())
585        .inspect_err(|err| {
586            hyperactor::metrics::ACTOR_MESSAGE_RECEIVE_ERRORS.add(1, hyperactor::kv_pairs!("actor_id" => self.actor_addr().to_string()));
587            tracing::error!(err=?err, actor_id=%self.actor_addr(), "unable to receive next py message");
588        })
589        .inspect(|_|{
590            hyperactor::metrics::ACTOR_MESSAGES_RECEIVED.add(1, hyperactor::kv_pairs!("actor_id" => self.actor_addr().to_string()));
591        })
592    }
593
594    /// Put the actor in stopped mode and return any messages that were received.
595    #[hyperactor::instrument(fields(actor_id=hyperactor::internal_macro_support::tracing::field::display(self.actor_addr())))]
596    pub fn drain_and_stop(&mut self) -> Result<Vec<M>> {
597        self.ensure_alive()?;
598        let messages: Vec<M> = self.message_receiver.drain().into_iter().collect();
599        tracing::info!("stopping the client actor in Python client");
600        self.status = InstanceStatus::Stopped;
601        Ok(messages)
602    }
603
604    pub fn instance(&self) -> &Client {
605        &self.instance
606    }
607
608    pub fn actor_addr(&self) -> &hyperactor::ActorAddr {
609        &self.actor_id
610    }
611}
612
613pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
614    hyperactor_mod.add_class::<PyUid>()?;
615    hyperactor_mod.add_class::<PyProcId>()?;
616    hyperactor_mod.add_class::<PyProc>()?;
617    hyperactor_mod.add_class::<PyActorAddr>()?;
618    hyperactor_mod.add_class::<PySerialized>()?;
619    Ok(())
620}
monarch_hyperactor/proc.rs

monarch_hyperactor/
proc.rs
