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 std::sync::Arc;
10
11use hyperactor::Instance;
12use hyperactor::context;
13use hyperactor_mesh::comm::multicast::CastInfo;
14use ndslice::Extent;
15use ndslice::Point;
16use pyo3::exceptions::PyRuntimeError;
17use pyo3::prelude::*;
18
19use crate::actor::ExecutionTracker;
20use crate::actor::PythonActor;
21use crate::actor::root_client_actor;
22use crate::mailbox::PyMailbox;
23use crate::proc::PyActorAddr;
24use crate::runtime;
25use crate::shape::PyPoint;
26
27#[pyclass(name = "Instance", module = "monarch._src.actor.actor_mesh")]
28pub struct PyInstance {
29    inner: Instance<PythonActor>,
30    #[pyo3(get, set)]
31    proc_mesh: Option<Py<PyAny>>,
32    #[pyo3(get, set, name = "_controller_controller")]
33    controller_controller: Option<Py<PyAny>>,
34    #[pyo3(get, set)]
35    pub(crate) rank: PyPoint,
36    #[pyo3(get, set, name = "_children")]
37    children: Option<Py<PyAny>>,
38
39    #[pyo3(get, set, name = "name")]
40    name: String,
41    #[pyo3(get, set, name = "class_name")]
42    class_name: Option<String>,
43    #[pyo3(get, set, name = "creator")]
44    creator: Option<Py<PyAny>>,
45
46    #[pyo3(get, set, name = "_mock_tensor_engine_factory")]
47    mock_tensor_engine_factory: Option<Py<PyAny>>,
48
49    /// Per-actor execution tracker -- a clone of the `PythonActor`'s `Arc`,
50    /// injected when the actor's own `PyInstance` is created. `None` for
51    /// root-client / `From`-built instances, which never bracket handlers.
52    /// Not exposed to Python.
53    execution_tracker: Option<Arc<ExecutionTracker>>,
54}
55
56impl Clone for PyInstance {
57    fn clone(&self) -> Self {
58        PyInstance {
59            inner: self.inner.clone_for_py(),
60            proc_mesh: self.proc_mesh.clone(),
61            controller_controller: self.controller_controller.clone(),
62            rank: self.rank.clone(),
63            children: self.children.clone(),
64            name: self.name.clone(),
65            class_name: self.class_name.clone(),
66            creator: self.creator.clone(),
67            mock_tensor_engine_factory: self.mock_tensor_engine_factory.clone(),
68            execution_tracker: self.execution_tracker.clone(),
69        }
70    }
71}
72
73impl std::ops::Deref for PyInstance {
74    type Target = Instance<PythonActor>;
75
76    fn deref(&self) -> &Self::Target {
77        &self.inner
78    }
79}
80
81#[pymethods]
82impl PyInstance {
83    #[getter]
84    pub(crate) fn _mailbox(&self) -> PyMailbox {
85        PyMailbox {
86            inner: self.inner.mailbox_for_py().clone(),
87        }
88    }
89
90    #[getter]
91    pub fn actor_id(&self) -> PyActorAddr {
92        let actor_id: hyperactor::ActorAddr = self.inner.self_addr().clone();
93        actor_id.into()
94    }
95
96    #[pyo3(signature = (reason = None))]
97    fn abort(&self, reason: Option<&str>) -> PyResult<()> {
98        let reason = reason.unwrap_or("(no reason provided)");
99        Ok(self.inner.abort(reason).map_err(anyhow::Error::from)?)
100    }
101
102    #[pyo3(signature = (reason = None))]
103    fn kill(&self, reason: Option<&str>) -> PyResult<()> {
104        let reason = reason.unwrap_or("(no reason provided)");
105        Ok(self.inner.kill(reason).map_err(anyhow::Error::from)?)
106    }
107
108    #[pyo3(signature = (reason = None))]
109    fn stop(&self, reason: Option<&str>) -> PyResult<()> {
110        tracing::info!(actor_id = %self.inner.self_addr(), "stopping PyInstance");
111        let reason = reason.unwrap_or("(no reason provided)");
112        self.inner
113            .stop(reason)
114            .map_err(|e| PyRuntimeError::new_err(e.to_string()))
115    }
116
117    /// Stop the actor and return a future that resolves when it reaches
118    /// a terminal status (stopped or failed). This ensures all pending
119    /// messages are drained and connections are flushed before returning.
120    #[pyo3(signature = (reason = None))]
121    fn stop_and_wait(&self, reason: Option<&str>) -> PyResult<crate::pytokio::PyPythonTask> {
122        let reason = reason.unwrap_or("shutdown").to_string();
123        let actor_id = self.inner.self_addr().clone();
124        let proc = self.inner.proc().clone();
125        crate::pytokio::PyPythonTask::new(async move {
126            let status_rx = proc.stop_actor(actor_id.id(), reason);
127            if let Some(mut rx) = status_rx {
128                let _ = rx.wait_for(|s| s.is_terminal()).await;
129            }
130            if let Err(e) = proc.flush().await {
131                tracing::warn!(%actor_id, "stop_and_wait: flush failed: {}", e);
132            }
133            Ok(())
134        })
135    }
136
137    /// Mark this actor as system/infrastructure.
138    ///
139    /// **PY-SYS-2:** Python actors use the `_is_system_actor = True`
140    /// class attribute so that this is called during actor init,
141    /// before ProcAgent publishes its first introspection snapshot.
142    fn set_system(&self) {
143        self.inner.set_system();
144    }
145
146    /// Reserve `count` ordering seqs against the receiver actor's
147    /// `PythonMessage` handler port. Subsequent fire-and-forget
148    /// endpoint sends to `receiver` pick up at the post-reservation
149    /// seq, creating a deterministic gap visible through the
150    /// receiver's `OrderedSender::snapshot` and
151    /// `/v1/{receiver}.inbound_ordering`.
152    ///
153    /// Test/demo only. Underscore-prefixed; production code must not
154    /// use this. The port computation matches the one used by normal
155    /// Python `.broadcast()` / `.call_one()` sends, so the gap is
156    /// observable without any extra plumbing.
157    #[pyo3(name = "_debug_skip_next_ordering_seq")]
158    fn debug_skip_next_ordering_seq(&self, receiver: &PyActorAddr, count: u64) {
159        use crate::actor::PythonMessage;
160        let port_addr = receiver
161            .inner
162            .port_addr(hyperactor::Port::handler::<PythonMessage>());
163        self.inner.debug_skip_next_ordering_seq(&port_addr, count);
164    }
165
166    /// Producer write-side for the mesh `execution` field: `_Actor.handle`
167    /// calls these around the real user-method invocation. `_execution_start`
168    /// returns the in-flight token; `_execution_finish` ends it. When this
169    /// instance has no tracker (root-client / `From`-built), start returns
170    /// the `0` no-op sentinel and finish ignores it (PE-4).
171    fn _execution_start(&self, name: String) -> u64 {
172        match &self.execution_tracker {
173            Some(tracker) => tracker.start(name),
174            None => 0,
175        }
176    }
177
178    fn _execution_finish(&self, token: u64) {
179        if let Some(tracker) = &self.execution_tracker {
180            tracker.finish(token);
181        }
182    }
183}
184
185impl PyInstance {
186    pub fn into_instance(self) -> Instance<PythonActor> {
187        self.inner
188    }
189
190    /// Inject the per-actor execution tracker. Called when the actor's own
191    /// `PyInstance` is first created (see `PythonActor::ensure_py_instance`).
192    pub(crate) fn set_execution_tracker(&mut self, tracker: Arc<ExecutionTracker>) {
193        self.execution_tracker = Some(tracker);
194    }
195}
196
197impl<I: context::Actor<A = PythonActor>> From<I> for PyInstance {
198    fn from(ins: I) -> Self {
199        PyInstance {
200            inner: ins.instance().clone_for_py(),
201            proc_mesh: None,
202            controller_controller: None,
203            rank: PyPoint::new(0, Extent::unity().into()),
204            children: None,
205            name: "root".to_string(),
206            class_name: None,
207            creator: None,
208            mock_tensor_engine_factory: None,
209            execution_tracker: None,
210        }
211    }
212}
213
214#[pyclass(name = "Context", module = "monarch._src.actor.actor_mesh")]
215pub struct PyContext {
216    instance: Py<PyInstance>,
217    rank: Point,
218    /// Cloneable handle to a span carrying the actor's recording key.
219    /// When entered, events emitted under this span are captured by
220    /// the per-actor flight recorder. `None` for bootstrap/client
221    /// contexts that are not actor handler execution paths.
222    recording_span: Option<tracing::Span>,
223}
224
225#[pymethods]
226impl PyContext {
227    #[getter]
228    fn actor_instance(&self) -> &Py<PyInstance> {
229        &self.instance
230    }
231
232    #[getter]
233    fn message_rank(&self) -> PyPoint {
234        self.rank.clone().into()
235    }
236
237    #[staticmethod]
238    fn _root_client_context(py: Python<'_>) -> PyResult<PyContext> {
239        let _guard = runtime::get_tokio_runtime().enter();
240        let instance: PyInstance = root_client_actor(py).into();
241        Ok(PyContext {
242            instance: instance.into_pyobject(py)?.into(),
243            rank: Extent::unity().point_of_rank(0).unwrap(),
244            recording_span: None,
245        })
246    }
247
248    /// Create a context from an existing instance.
249    /// This is used when the root client was bootstrapped via bootstrap_host()
250    /// instead of the default bootstrap_client().
251    #[staticmethod]
252    fn _from_instance(py: Python<'_>, instance: PyInstance) -> PyResult<PyContext> {
253        Ok(PyContext {
254            instance: instance.into_pyobject(py)?.into(),
255            rank: Extent::unity().point_of_rank(0).unwrap(),
256            recording_span: None,
257        })
258    }
259}
260
261impl PyContext {
262    pub(crate) fn new<T: hyperactor::actor::Actor>(
263        cx: &hyperactor::Context<T>,
264        instance: Py<PyInstance>,
265    ) -> PyContext {
266        PyContext {
267            instance,
268            rank: cx.cast_point(),
269            recording_span: Some(cx.recording_span()),
270        }
271    }
272
273    /// The actor's recording span, if this context is an actor handler
274    /// execution path. Used by `forward_to_tracing` to enter the
275    /// recording scope on the asyncio thread so that log events are
276    /// captured in the flight recorder.
277    pub(crate) fn recording_span(&self) -> Option<&tracing::Span> {
278        self.recording_span.as_ref()
279    }
280
281    /// Test-only: build a PyContext with a chosen recording span.
282    /// Uses the root client actor for the instance field (the test
283    /// only exercises the recording_span extraction path).
284    #[doc(hidden)]
285    pub fn for_test(py: Python<'_>, recording_span: Option<tracing::Span>) -> PyResult<PyContext> {
286        let mut ctx = Self::_root_client_context(py)?;
287        ctx.recording_span = recording_span;
288        Ok(ctx)
289    }
290}
291
292pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
293    hyperactor_mod.add_class::<PyInstance>()?;
294    hyperactor_mod.add_class::<PyContext>()?;
295    Ok(())
296}
monarch_hyperactor/context.rs

monarch_hyperactor/
context.rs
