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

monarch_hyperactor/
context.rs
