monarch_hyperactor/
context.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 hyperactor::Instance;
10use hyperactor::context;
11use hyperactor_mesh::comm::multicast::CastInfo;
12use ndslice::Extent;
13use ndslice::Point;
14use pyo3::exceptions::PyRuntimeError;
15use pyo3::prelude::*;
16
17use crate::actor::PythonActor;
18use crate::actor::root_client_actor;
19use crate::mailbox::PyMailbox;
20use crate::proc::PyActorId;
21use crate::runtime;
22use crate::shape::PyPoint;
23
24#[pyclass(name = "Instance", module = "monarch._src.actor.actor_mesh")]
25pub struct PyInstance {
26    inner: Instance<PythonActor>,
27    #[pyo3(get, set)]
28    proc_mesh: Option<Py<PyAny>>,
29    #[pyo3(get, set, name = "_controller_controller")]
30    controller_controller: Option<Py<PyAny>>,
31    #[pyo3(get, set)]
32    pub(crate) rank: PyPoint,
33    #[pyo3(get, set, name = "_children")]
34    children: Option<Py<PyAny>>,
35
36    #[pyo3(get, set, name = "name")]
37    name: String,
38    #[pyo3(get, set, name = "class_name")]
39    class_name: Option<String>,
40    #[pyo3(get, set, name = "creator")]
41    creator: Option<Py<PyAny>>,
42
43    #[pyo3(get, set, name = "_mock_tensor_engine_factory")]
44    mock_tensor_engine_factory: Option<Py<PyAny>>,
45}
46
47impl Clone for PyInstance {
48    fn clone(&self) -> Self {
49        PyInstance {
50            inner: self.inner.clone_for_py(),
51            proc_mesh: self.proc_mesh.clone(),
52            controller_controller: self.controller_controller.clone(),
53            rank: self.rank.clone(),
54            children: self.children.clone(),
55            name: self.name.clone(),
56            class_name: self.class_name.clone(),
57            creator: self.creator.clone(),
58            mock_tensor_engine_factory: self.mock_tensor_engine_factory.clone(),
59        }
60    }
61}
62
63impl std::ops::Deref for PyInstance {
64    type Target = Instance<PythonActor>;
65
66    fn deref(&self) -> &Self::Target {
67        &self.inner
68    }
69}
70
71#[pymethods]
72impl PyInstance {
73    #[getter]
74    pub(crate) fn _mailbox(&self) -> PyMailbox {
75        PyMailbox {
76            inner: self.inner.mailbox_for_py().clone(),
77        }
78    }
79
80    #[getter]
81    pub fn actor_id(&self) -> PyActorId {
82        self.inner.self_id().clone().into()
83    }
84
85    #[pyo3(signature = (reason = None))]
86    fn abort(&self, reason: Option<&str>) -> PyResult<()> {
87        let reason = reason.unwrap_or("(no reason provided)");
88        Ok(self.inner.abort(reason).map_err(anyhow::Error::from)?)
89    }
90
91    #[pyo3(signature = (reason = None))]
92    fn stop(&self, reason: Option<&str>) -> PyResult<()> {
93        tracing::info!(actor_id = %self.inner.self_id(), "stopping PyInstance");
94        let reason = reason.unwrap_or("(no reason provided)");
95        self.inner
96            .stop(reason)
97            .map_err(|e| PyRuntimeError::new_err(e.to_string()))
98    }
99
100    /// Deprecated alias for `stop`.
101    #[pyo3(signature = (reason = None))]
102    fn _stop_instance(&self, reason: Option<&str>) -> PyResult<()> {
103        self.stop(reason)
104    }
105
106    /// Mark this actor as system/infrastructure.
107    ///
108    /// **PY-SYS-2:** Python actors use the `_is_system_actor = True`
109    /// class attribute so that this is called during actor init,
110    /// before ProcAgent publishes its first introspection snapshot.
111    fn set_system(&self) {
112        self.inner.set_system();
113    }
114}
115
116impl PyInstance {
117    pub fn into_instance(self) -> Instance<PythonActor> {
118        self.inner
119    }
120}
121
122impl<I: context::Actor<A = PythonActor>> From<I> for PyInstance {
123    fn from(ins: I) -> Self {
124        PyInstance {
125            inner: ins.instance().clone_for_py(),
126            proc_mesh: None,
127            controller_controller: None,
128            rank: PyPoint::new(0, Extent::unity().into()),
129            children: None,
130            name: "root".to_string(),
131            class_name: None,
132            creator: None,
133            mock_tensor_engine_factory: None,
134        }
135    }
136}
137
138#[pyclass(name = "Context", module = "monarch._src.actor.actor_mesh")]
139pub struct PyContext {
140    instance: Py<PyInstance>,
141    rank: Point,
142}
143
144#[pymethods]
145impl PyContext {
146    #[getter]
147    fn actor_instance(&self) -> &Py<PyInstance> {
148        &self.instance
149    }
150
151    #[getter]
152    fn message_rank(&self) -> PyPoint {
153        self.rank.clone().into()
154    }
155
156    #[staticmethod]
157    fn _root_client_context(py: Python<'_>) -> PyResult<PyContext> {
158        let _guard = runtime::get_tokio_runtime().enter();
159        let instance: PyInstance = root_client_actor(py).into();
160        Ok(PyContext {
161            instance: instance.into_pyobject(py)?.into(),
162            rank: Extent::unity().point_of_rank(0).unwrap(),
163        })
164    }
165
166    /// Create a context from an existing instance.
167    /// This is used when the root client was bootstrapped via bootstrap_host()
168    /// instead of the default bootstrap_client().
169    #[staticmethod]
170    fn _from_instance(py: Python<'_>, instance: PyInstance) -> PyResult<PyContext> {
171        Ok(PyContext {
172            instance: instance.into_pyobject(py)?.into(),
173            rank: Extent::unity().point_of_rank(0).unwrap(),
174        })
175    }
176}
177
178impl PyContext {
179    pub(crate) fn new<T: hyperactor::actor::Actor>(
180        cx: &hyperactor::Context<T>,
181        instance: Py<PyInstance>,
182    ) -> PyContext {
183        PyContext {
184            instance,
185            rank: cx.cast_point(),
186        }
187    }
188}
189
190pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
191    hyperactor_mod.add_class::<PyInstance>()?;
192    hyperactor_mod.add_class::<PyContext>()?;
193    Ok(())
194}
monarch_hyperactor/context.rs

monarch_hyperactor/
context.rs
