monarch_messages/
worker.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
9// NOTE: Until https://github.com/PyO3/pyo3/pull/4674, `pyo3::pymethods` trigger
10// and unsafe-op-in-unsafe-fn warnings.
11#![allow(unsafe_op_in_unsafe_fn)]
12// EnumAsInner generates code that triggers a false positive
13// unused_assignments lint on struct variant fields. #[allow] on the
14// enum itself doesn't propagate into derive-macro-generated code, so
15// the suppression must be at module scope.
16#![allow(unused_assignments)]
17
18use std::collections::HashMap;
19use std::fmt;
20use std::fmt::Display;
21use std::sync::Arc;
22
23use anyhow::Context;
24use derive_more::Display;
25use derive_more::From;
26use derive_more::TryInto;
27use enum_as_inner::EnumAsInner;
28use hyperactor as reference;
29use hyperactor::HandleClient;
30use hyperactor::Handler;
31use hyperactor::RefClient;
32use monarch_gil::GilSite;
33use monarch_gil::monarch_with_gil_blocking;
34use monarch_types::ReduceOp;
35use monarch_types::SerializablePyErr;
36use monarch_types::UniqueId;
37use monarch_types::py_global;
38use ndslice::Slice;
39use pyo3::exceptions::PyValueError;
40use pyo3::prelude::*;
41use pyo3::types::PyBytes;
42use pyo3::types::PyDict;
43use pyo3::types::PyTuple;
44use serde::Deserialize;
45use serde::Serialize;
46use thiserror::Error;
47use torch_sys2::BorrowError;
48use torch_sys2::Device;
49use torch_sys2::Layout;
50use torch_sys2::ScalarType;
51use typeuri::Named;
52
53use crate::controller::ControllerActor;
54use crate::controller::Seq;
55use crate::wire_value::WireValue;
56
57#[derive(
58    Serialize,
59    Deserialize,
60    Debug,
61    Clone,
62    Hash,
63    PartialEq,
64    Eq,
65    Copy,
66    PartialOrd,
67    Ord,
68    From
69)]
70#[pyo3::pyclass(
71    frozen,
72    module = "monarch._rust_bindings.monarch_extension.tensor_worker"
73)]
74pub struct StreamRef {
75    #[pyo3(get)]
76    pub id: u64,
77}
78
79#[pyo3::pymethods]
80impl StreamRef {
81    #[new]
82    #[pyo3(signature = (*, id))]
83    fn new(id: u64) -> Self {
84        Self { id }
85    }
86
87    fn __repr__(&self) -> String {
88        format!("StreamRef({})", self.id)
89    }
90
91    // TODO: Upgrade pyo3 to use eq, ord on pyclass
92    fn __richcmp__(&self, other: PyRef<Self>, op: pyo3::class::basic::CompareOp) -> PyResult<bool> {
93        Ok(match op {
94            pyo3::class::basic::CompareOp::Eq => self.id == other.id,
95            pyo3::class::basic::CompareOp::Ne => self.id != other.id,
96            pyo3::class::basic::CompareOp::Lt => self.id < other.id,
97            pyo3::class::basic::CompareOp::Le => self.id <= other.id,
98            pyo3::class::basic::CompareOp::Gt => self.id > other.id,
99            pyo3::class::basic::CompareOp::Ge => self.id >= other.id,
100        })
101    }
102
103    fn __hash__(&self) -> PyResult<u64> {
104        Ok(self.id)
105    }
106}
107
108// TODO: The Python implementation uses `Ref` to describe any worker value that
109// can be referenced by the controller, including: tensors, streams, pipes,
110// device meshes. We might be able to more explicitly type these, as they are
111// not generally interchangeable.
112#[derive(
113    Serialize,
114    Deserialize,
115    Debug,
116    Clone,
117    Hash,
118    PartialEq,
119    Eq,
120    Copy,
121    PartialOrd,
122    Ord,
123    From
124)]
125#[pyo3::pyclass(
126    frozen,
127    module = "monarch._rust_bindings.monarch_extension.tensor_worker"
128)]
129pub struct Ref {
130    #[pyo3(get)]
131    pub id: u64,
132}
133
134#[pyo3::pymethods]
135impl Ref {
136    #[new]
137    fn new(id: u64) -> Self {
138        Self { id }
139    }
140
141    #[getter]
142    fn r#ref(&self) -> u64 {
143        self.id
144    }
145
146    fn __repr__(&self) -> String {
147        format!("Ref({})", self.id)
148    }
149
150    // TODO: Upgrade pyo3 to use eq, ord on pyclass
151    fn __richcmp__(&self, other: PyRef<Self>, op: pyo3::class::basic::CompareOp) -> PyResult<bool> {
152        Ok(match op {
153            pyo3::class::basic::CompareOp::Eq => self.id == other.id,
154            pyo3::class::basic::CompareOp::Ne => self.id != other.id,
155            pyo3::class::basic::CompareOp::Lt => self.id < other.id,
156            pyo3::class::basic::CompareOp::Le => self.id <= other.id,
157            pyo3::class::basic::CompareOp::Gt => self.id > other.id,
158            pyo3::class::basic::CompareOp::Ge => self.id >= other.id,
159        })
160    }
161
162    fn __hash__(&self) -> PyResult<u64> {
163        Ok(self.id)
164    }
165
166    fn __getnewargs_ex__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
167        let kwargs = PyDict::new(py);
168        kwargs.set_item("id", self.id).unwrap();
169
170        PyTuple::new(
171            py,
172            vec![
173                PyTuple::empty(py).unbind().into_any(),
174                kwargs.unbind().into_any(),
175            ],
176        )
177    }
178}
179
180impl Ref {
181    // This is a function on ref instead of impl FromPyObject due to a bug in pyo3
182    // https://github.com/PyO3/pyo3/issues/4337
183    pub fn from_py_object(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
184        let attr_name = pyo3::intern!(obj.py(), "__monarch_ref__");
185        if let Ok(ref_) = obj.extract::<Ref>() {
186            return Ok(ref_);
187        }
188        if let Ok(func) = obj.getattr(attr_name)
189            && let Ok(Ok(val)) = func.call0().map(|val| val.extract::<u64>())
190        {
191            return Ok(val.into());
192        }
193        Err(PyValueError::new_err("Could not convert object to Ref"))
194    }
195}
196
197impl Display for Ref {
198    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
199        write!(f, "r{}", self.id)
200    }
201}
202
203/// Identifies a CallFunction target. Can either be a torch op or a Python
204/// global reference.
205// TODO: do some validation on the namespace/opname/overload
206#[derive(PartialEq, Serialize, Deserialize, Debug, Clone)]
207#[pyo3::pyclass(
208    frozen,
209    module = "monarch._rust_bindings.monarch_extension.tensor_worker"
210)]
211pub struct FunctionPath {
212    #[pyo3(get)]
213    pub path: String,
214}
215
216impl fmt::Display for FunctionPath {
217    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
218        write!(f, "<function \"{}\">", self.path)
219    }
220}
221
222impl<T: Into<String>> From<T> for FunctionPath {
223    fn from(val: T) -> Self {
224        Self { path: val.into() }
225    }
226}
227
228#[pyo3::pymethods]
229impl FunctionPath {
230    #[new]
231    #[pyo3(signature = (*, path))]
232    pub fn new(path: String) -> Self {
233        Self { path }
234    }
235
236    fn __repr__(&self) -> String {
237        self.path.clone()
238    }
239
240    pub fn resolve<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
241        let (start, rest) = self.path.split_once(".").with_context(|| {
242            format!(
243                "invalid function path {}: paths must be fully qualified",
244                self.path
245            )
246        })?;
247        if start == "torch" {
248            let mut cur = py.import("torch")?.into_any();
249            for p in rest.split(".") {
250                cur = cur.getattr(p)?;
251            }
252            Ok(cur)
253        } else {
254            let (module_fqn, function_name) = self.path.rsplit_once(".").with_context(|| {
255                format!(
256                    "invalid function path {}: paths must be fully qualified",
257                    self.path
258                )
259            })?;
260            let module = PyModule::import(py, module_fqn)?;
261            let mut function = module.getattr(function_name)?;
262            if function.hasattr("_remote_impl")? {
263                function = function.getattr("_remote_impl")?;
264            }
265            Ok(function.cast_into()?)
266        }
267    }
268}
269
270/// Identifies a CallFunction target. Can either be a torch op or a Python
271/// global reference.
272// TODO: do some validation on the namespace/opname/overload
273#[derive(PartialEq, Serialize, Deserialize, Debug, Clone, From)]
274#[pyo3::pyclass(
275    frozen,
276    module = "monarch._rust_bindings.monarch_extension.tensor_worker"
277)]
278pub struct Cloudpickle {
279    #[serde(with = "serde_bytes")]
280    bytes: Vec<u8>,
281}
282
283impl fmt::Display for Cloudpickle {
284    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
285        write!(f, "<cloud-pickle>")
286    }
287}
288
289py_global!(cloudpickle_dumps, "cloudpickle", "dumps");
290
291#[pyo3::pymethods]
292impl Cloudpickle {
293    #[new]
294    #[pyo3(signature = (*, bytes))]
295    pub fn new(bytes: Vec<u8>) -> Self {
296        Self { bytes }
297    }
298
299    fn __repr__(&self) -> String {
300        format!("Cloudpickle(bytes={:?})", self.bytes)
301    }
302
303    pub fn resolve<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
304        let module = PyModule::import(py, "cloudpickle")?;
305        let loads = module.getattr("loads")?;
306        loads.call1((PyBytes::new(py, &self.bytes),))
307    }
308}
309
310impl Cloudpickle {
311    pub fn dumps(obj: Bound<'_, PyAny>) -> PyResult<Self> {
312        let py = obj.py();
313        let dumps = cloudpickle_dumps(py);
314        let bytes_obj = dumps.call1((obj,))?;
315        let bytes = bytes_obj.cast::<PyBytes>()?.as_bytes().to_vec();
316        Ok(Self { bytes })
317    }
318}
319
320#[derive(
321    PartialEq,
322    Serialize,
323    Deserialize,
324    Debug,
325    Clone,
326    TryInto,
327    From,
328    FromPyObject,
329    Display
330)]
331pub enum ResolvableFunction {
332    #[pyo3(transparent)]
333    Cloudpickle(Cloudpickle),
334    #[pyo3(transparent)]
335    FunctionPath(FunctionPath),
336}
337
338#[derive(PartialEq, Serialize, Deserialize, Debug, Clone, From)]
339pub struct ArgsKwargs {
340    payload: Cloudpickle,
341}
342
343impl ArgsKwargs {
344    pub fn from_python<'py>(args: Bound<'py, PyAny>, kwargs: Bound<'py, PyAny>) -> PyResult<Self> {
345        // Create tuple (args, kwargs), then cloudpickle it
346        let py = args.py();
347        let tuple = PyTuple::new(py, vec![args, kwargs])?;
348        let payload = Cloudpickle::dumps(tuple.into_any())?;
349        Ok(Self { payload })
350    }
351
352    pub fn from_wire_values(
353        args: Vec<WireValue>,
354        kwargs: HashMap<String, WireValue>,
355    ) -> PyResult<Self> {
356        monarch_with_gil_blocking(GilSite::Convert, |py| {
357            // Convert WireValue args to Python objects
358            let py_args: Vec<Bound<'_, PyAny>> = args
359                .into_iter()
360                .map(|v| v.into_pyobject(py))
361                .collect::<PyResult<_>>()?;
362            let args_tuple = PyTuple::new(py, py_args)?;
363
364            // Convert WireValue kwargs to Python dict
365            let kwargs_dict = PyDict::new(py);
366            for (k, v) in kwargs {
367                kwargs_dict.set_item(k, v.into_pyobject(py)?)?;
368            }
369
370            Self::from_python(args_tuple.into_any(), kwargs_dict.into_any())
371        })
372    }
373
374    pub fn to_python<'py>(
375        &self,
376        py: Python<'py>,
377    ) -> PyResult<(Bound<'py, PyTuple>, Bound<'py, PyDict>)> {
378        let tuple = self.payload.resolve(py)?;
379        let tuple = tuple.cast::<PyTuple>()?;
380
381        // Extract args (first element)
382        let args = tuple.get_item(0)?;
383        let args_tuple = args.cast::<PyTuple>()?;
384
385        // Extract kwargs (second element)
386        let kwargs = tuple.get_item(1)?;
387        let kwargs_dict = kwargs.cast::<PyDict>()?;
388
389        Ok((args_tuple.clone(), kwargs_dict.clone()))
390    }
391}
392
393impl<'py> IntoPyObject<'py> for ResolvableFunction {
394    type Target = PyAny;
395    type Output = Bound<'py, Self::Target>;
396    type Error = PyErr;
397
398    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
399        Ok(match self {
400            Self::Cloudpickle(func) => func.into_pyobject(py)?.into_any(),
401            Self::FunctionPath(func) => func.into_pyobject(py)?.into_any(),
402        })
403    }
404}
405
406impl ResolvableFunction {
407    pub fn resolve<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
408        match self {
409            Self::Cloudpickle(func) => Ok(func.resolve(py)?.into_any()),
410            Self::FunctionPath(func) => func.resolve(py),
411        }
412    }
413
414    /// For testing: this is a special remote function path that induces a panic
415    /// when called.
416    pub fn panic_if_requested(&self) {
417        match self {
418            Self::FunctionPath(func) if func.path == "__test_panic" => {
419                panic!("__test_panic called");
420            }
421            _ => (),
422        }
423    }
424}
425
426impl<T: Into<String>> From<T> for ResolvableFunction {
427    fn from(val: T) -> Self {
428        FunctionPath::from(val).into()
429    }
430}
431
432#[derive(Serialize, Deserialize, Debug, Clone)]
433pub struct CallFunctionParams {
434    /// Sequence ID of the invocation.
435    pub seq: Seq,
436    /// The references of the results to set.
437    pub results: Vec<Option<Ref>>,
438    /// The references of the mutates to set.
439    pub mutates: Vec<Ref>,
440    /// The function to call.
441    pub function: ResolvableFunction,
442    /// The arguments and keyword arguments to the function.
443    pub args_kwargs: ArgsKwargs,
444    /// The stream to call the function on.
445    pub stream: StreamRef,
446    /// The process groups to execute the function on.
447    pub remote_process_groups: Vec<Ref>,
448}
449
450#[derive(Serialize, Deserialize, Debug, Clone)]
451pub struct ActorCallParams {
452    pub seq: Seq,
453    // The BrokerId but we do not depend on hyperactor in messages.
454    pub broker_id: (String, usize),
455    /// Referenceable objects to pass to the actor as LocalState,
456    /// these will be put into the PythonMessage
457    /// during its unpickling.
458    pub local_state: Vec<Ref>,
459    /// Tensors that will be mutated by the call.
460    pub mutates: Vec<Ref>,
461    pub stream: StreamRef,
462}
463
464#[derive(Serialize, Deserialize, Debug, Clone)]
465pub struct ActorMethodParams {
466    pub results: Vec<Option<Ref>>,
467    pub call: ActorCallParams,
468}
469
470/// Type of reduction for [`WorkerMessage::Reduce`].
471#[derive(Debug, Clone, Serialize, Deserialize)]
472pub enum Reduction {
473    /// A gather, concat'ing the values along the reduction dimension.
474    Stack,
475    /// A NCCL reduction type.
476    ReduceOp(ReduceOp),
477}
478
479#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
480#[pyo3::pyclass(
481    frozen,
482    name = "TensorFactory",
483    module = "monarch._rust_bindings.monarch_extension.tensor_worker"
484)]
485pub struct Factory {
486    pub size: Vec<i64>,
487    #[serde(with = "torch_sys2::ScalarTypeDef")]
488    pub dtype: ScalarType,
489    #[serde(with = "torch_sys2::LayoutDef")]
490    pub layout: Layout,
491    pub device: Device,
492}
493
494#[pyo3::pymethods]
495impl Factory {
496    #[new]
497    #[pyo3(signature = (*, size, dtype, layout, device))]
498    pub fn new(
499        py: Python<'_>,
500        size: Vec<i64>,
501        dtype: Py<PyAny>,
502        layout: Py<PyAny>,
503        device: Py<PyAny>,
504    ) -> PyResult<Self> {
505        // TODO: Add some validation around dtype / layout. We should have pyre types on
506        // the python side to help in the short term.
507        Ok(Self {
508            size,
509            dtype: dtype.extract::<ScalarType>(py)?,
510            layout: layout.extract::<Layout>(py)?,
511            device: device.extract::<Device>(py)?,
512        })
513    }
514
515    #[staticmethod]
516    pub fn from_py(obj: Bound<'_, PyAny>) -> PyResult<Self> {
517        Self::new(
518            obj.py(),
519            obj.getattr("size")?.extract()?,
520            obj.getattr("dtype")?.unbind(),
521            obj.getattr("layout")?.unbind(),
522            obj.getattr("device")?.unbind(),
523        )
524    }
525
526    #[getter]
527    fn size<'a>(&self, py: Python<'a>) -> PyResult<Bound<'a, PyTuple>> {
528        PyTuple::new(py, self.size.iter())
529    }
530
531    #[getter]
532    fn dtype<'a>(&self, py: Python<'a>) -> PyResult<Bound<'a, PyAny>> {
533        self.dtype.into_pyobject(py)
534    }
535
536    #[getter]
537    fn layout<'a>(&self, py: Python<'a>) -> PyResult<Bound<'a, PyAny>> {
538        self.layout.into_pyobject(py)
539    }
540
541    #[getter]
542    fn device(&self) -> String {
543        self.device.to_string()
544    }
545}
546
547/// Controls what CUDA stream an actor will use.
548#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
549#[pyo3::pyclass(
550    module = "monarch._rust_bindings.monarch_extension.tensor_worker",
551    eq,
552    eq_int
553)]
554pub enum StreamCreationMode {
555    /// Use the default stream for the current device.
556    UseDefaultStream,
557    /// Create a new stream for this actor.
558    CreateNewStream,
559}
560
561/// An error associated with a seq number that failed to execute.
562/// Any defined value that has an error value will have an assocated
563/// SeqError that is the root cause of why that value has an error.
564/// A value may have this error because it was directly defined by the
565/// action associated with the sequence number, or if it was defined by
566/// another action that dependend on the failing one.
567#[derive(Debug, Named)]
568pub struct SeqError {
569    pub seq: Seq,
570    pub error: anyhow::Error,
571}
572
573impl Display for SeqError {
574    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
575        self.error.fmt(f)
576    }
577}
578
579/// When a worker runs any function, it may not succeed either because the function itself
580/// failed (Error) or because an input to the function already had an error value
581/// DependentError.
582#[derive(Error, Debug, Named)]
583pub enum CallFunctionError {
584    #[error("{0}")]
585    Error(#[from] anyhow::Error),
586    #[error("Computation depended on an input that failed with error: {0}")]
587    DependentError(Arc<SeqError>),
588}
589
590impl CallFunctionError {
591    // Static functions for backward compatibility with existing enum cases
592    #[allow(non_snake_case)]
593    pub fn RefNotFound(r: Ref) -> Self {
594        Self::Error(anyhow::anyhow!("ref not found: {}", r))
595    }
596
597    #[allow(non_snake_case)]
598    pub fn InvalidRemoteFunction(msg: String) -> Self {
599        Self::Error(anyhow::anyhow!("invalid remote function: {}", msg))
600    }
601
602    #[allow(non_snake_case)]
603    pub fn UnsupportedArgType(function_type: String, arg_type: String) -> Self {
604        Self::Error(anyhow::anyhow!(
605            "unsupported arg type for {} function: {}",
606            function_type,
607            arg_type
608        ))
609    }
610
611    #[allow(non_snake_case)]
612    pub fn RemoteFunctionFailed(err: SerializablePyErr) -> Self {
613        Self::Error(anyhow::anyhow!("remote function failed: {}", err))
614    }
615
616    #[allow(non_snake_case)]
617    pub fn BorrowError(err: BorrowError) -> Self {
618        Self::Error(anyhow::anyhow!("borrow failed: {}", err))
619    }
620
621    #[allow(non_snake_case)]
622    pub fn UnexpectedNumberOfReturns(expected: usize, actual: usize) -> Self {
623        Self::Error(anyhow::anyhow!(
624            "unexpected number of returns from op, expected {}, got {}",
625            expected,
626            actual
627        ))
628    }
629
630    #[allow(non_snake_case)]
631    pub fn TooManyArgsForValue(args: String, kwargs: String) -> Self {
632        Self::Error(anyhow::anyhow!(
633            "expected only a single arg (and no kwargs) when no function is given: {}, {}",
634            args,
635            kwargs
636        ))
637    }
638
639    #[allow(non_snake_case)]
640    pub fn Anyhow(err: anyhow::Error) -> Self {
641        Self::Error(err)
642    }
643}
644
645impl From<SerializablePyErr> for CallFunctionError {
646    fn from(v: SerializablePyErr) -> CallFunctionError {
647        CallFunctionError::Error(v.into())
648    }
649}
650
651impl From<BorrowError> for CallFunctionError {
652    fn from(v: BorrowError) -> CallFunctionError {
653        CallFunctionError::Error(v.into())
654    }
655}
656
657/// Worker messages. These define the observable behavior of the worker, so the
658/// documentations here
659#[derive(
660    Handler,
661    HandleClient,
662    RefClient,
663    Clone,
664    Serialize,
665    Deserialize,
666    Debug,
667    Named,
668    EnumAsInner
669)]
670pub enum WorkerMessage {
671    /// Initialize backend network state.
672    BackendNetworkInit(UniqueId),
673
674    /// Initialize backend network state for point-to-point communication.
675    BackendNetworkPointToPointInit {
676        from_stream: StreamRef,
677        to_stream: StreamRef,
678    },
679
680    /// Call a function, either a torch op or a Python `remote_function`.
681    CallFunction(CallFunctionParams),
682
683    /// Groups commands together; these commands will be executed in order by
684    /// the worker.
685    CommandGroup(Vec<WorkerMessage>),
686
687    /// Create a [`Stream`] on the worker wih the provided id. Commands will be
688    /// generally be scheduled onto streams to run; different streams can
689    /// execute concurrently with one another.
690    CreateStream {
691        /// Id of the stream to create.
692        id: StreamRef,
693        /// Whether to use the default device stream or create a new one.
694        stream_creation: StreamCreationMode,
695    },
696
697    /// Create a [`DeviceMesh`] on the worker, which can be used to schedule
698    /// efficient inter-worker communication.
699    CreateDeviceMesh {
700        result: Ref,
701        names: Vec<String>,
702        ranks: Slice,
703    },
704
705    /// Create a PyTorch distributed process group on the worker, which can be
706    /// used to schedule collectives in UDFs using monarch communicators.
707    CreateRemoteProcessGroup {
708        result: Ref,
709        device_mesh: Ref,
710        dims: Vec<String>,
711    },
712
713    /// Create a borrow of a tensor from one stream to another.
714    ///
715    /// Borrows allows streams to access tensors on another stream. The runtime
716    /// will insert appropriate synchronization to ensure that cross-stream
717    /// usage is safe.
718    BorrowCreate {
719        /// Ref of the resulting borrowed tensor
720        result: Ref,
721        /// Id for the borrow
722        borrow: u64,
723        /// Tensor to borrow
724        tensor: Ref,
725        /// Stream to borrow from
726        from_stream: StreamRef,
727        /// Stream to borrow to
728        to_stream: StreamRef,
729    },
730
731    /// First use of the borrow on the receiving stream. This is a marker for
732    /// synchronization.
733    BorrowFirstUse {
734        borrow: u64,
735    },
736
737    /// Last use of the borrow on the receiving stream. This is a marker for
738    /// synchronization.
739    BorrowLastUse {
740        borrow: u64,
741    },
742
743    /// Drop the borrow and free the resources associated with it.
744    BorrowDrop {
745        borrow: u64,
746    },
747
748    /// Delete these refs from the worker state.
749    DeleteRefs(Vec<Ref>),
750
751    /// A [`ControllerMessage::Status`] will be send to the controller
752    /// when all streams have processed all the message sent before this one.
753    RequestStatus {
754        seq: Seq,
755        controller: bool,
756    },
757
758    /// Perform a reduction operation, using an efficient communication backend.
759    /// Only NCCL is supported for now.
760    Reduce {
761        /// Where to store the result of the reduction.
762        result: Ref,
763        /// The tensor to reduce.
764        tensor: Ref,
765        /// Tensor metadata for `tensor` that can be used to construct a
766        /// fresh tensor of appropriate size/shape. We use this if
767        /// `tensor` isn't accessible for some reason (like a previous
768        /// error on the worker).
769        factory: Factory,
770        /// The device mesh on which to perform the reduction.
771        mesh: Ref,
772        /// The stream to call the reduction on.
773        stream: StreamRef,
774        /// The dimensions of the device mesh to reduce over. The members of
775        /// these dimension will form the members of the reduction collective.
776        dims: Vec<String>,
777        /// What kind of reduction to perform.
778        reduction: Reduction,
779        /// If `true`, the reduced result will be evenly split across the tensors
780        /// of `dim`.
781        scatter: bool,
782        /// If `true`, the reduction will be performed in-place on `tensor`.
783        in_place: bool,
784        /// Pre-existing tensor that should be used as the output for the reduction.
785        out: Option<Ref>,
786    },
787
788    /// Create a new communicator on each rank in `ranks`, capable of
789    /// communicating with its peers along the specified dimensions.
790    SplitComm {
791        /// The device mesh dimensions along which the constructed communicator
792        /// should be able to exchange data.
793        dims: Vec<String>,
794        /// The device mesh associated with the new communicator. One communicator
795        /// will be created for every member of the mesh.
796        device_mesh: Ref,
797        /// The stream associated with the communicator. Communicator operations
798        /// will be ordered with respect to other operations scheduled on this
799        /// stream.
800        stream: StreamRef,
801    },
802
803    /// Create a new communicator on each rank in `ranks`, capable of
804    /// communicating with its peers along the specified dimensions.
805    SplitCommForProcessGroup {
806        /// The device mesh associated with the new communicator. One communicator
807        /// will be created for every member of the mesh.
808        remote_process_group: Ref,
809        /// The stream associated with the communicator. Communicator operations
810        /// will be ordered with respect to other operations scheduled on this
811        /// stream.
812        stream: StreamRef,
813    },
814
815    SendTensor {
816        result: Ref,
817        from_ranks: Slice,
818        to_ranks: Slice,
819        tensor: Ref,
820        factory: Factory,
821        from_stream: StreamRef,
822        to_stream: StreamRef,
823    },
824
825    SendValue {
826        seq: Seq,
827        /// Pipe to send value to.  If `None`, value is sent to controller.
828        destination: Option<Ref>,
829        mutates: Vec<Ref>,
830        /// Function to resolve the value to retrieve.  If `None`, then `args_kwargs`
831        /// must contain the value as the only element in args with no kwargs.
832        function: Option<ResolvableFunction>,
833        args_kwargs: ArgsKwargs,
834        /// The stream to retrieve from.
835        stream: StreamRef,
836    },
837
838    SendResultOfActorCall(ActorCallParams),
839    CallActorMethod(ActorMethodParams),
840    PipeRecv {
841        seq: Seq,
842        /// Result refs.
843        results: Vec<Option<Ref>>,
844        /// Pipe to receive value from.
845        pipe: Ref,
846        /// The stream to retrieve from.
847        stream: StreamRef,
848    },
849
850    /// Finish processing all messages previously sent to this worker and stop
851    /// the actor loop. Any streams will also be drained.
852    Exit {
853        /// Optional error reason if the exit is the result of an error, including
854        /// - optional actor id to indicate the source of the error
855        /// - error message or stacktrace
856        ///
857        /// The worker process will be stopped if the error is provided.
858        error: Option<(Option<reference::ActorAddr>, String)>,
859    },
860
861    /// Defines (part of) a new recording on the worker. This is a list of commands
862    /// representing the execution of a function that was defined using
863    /// monarch.compile. If there are too many commands to send in a single
864    /// DefineRecording message, the commands may be chunked into `ntotal_messages`,
865    /// with the `index` field indicating how to order the DefineRecording messages
866    /// for a single recording.
867    DefineRecording {
868        /// The ref associated with this recording that will be used to
869        /// call it in the future.
870        result: Ref,
871        /// The number of output tensors.
872        nresults: usize,
873        /// The number of input tensors.
874        nformals: usize,
875        /// The list of commands to run.
876        commands: Vec<WorkerMessage>,
877        /// How many total DefineRecording messages make up this recording.
878        ntotal_messages: usize,
879        /// This DefineRecording message's index in the set of messages
880        /// that make up this recording.
881        index: usize,
882    },
883
884    /// Defines an input tensor for a recording.
885    RecordingFormal {
886        /// The ref that will be used to pass the input tensor to the
887        /// recording.
888        result: Ref,
889        /// The index of the input tensor in the list of input tensors.
890        argument_index: usize,
891        /// The stream that this input tensor will be used on.
892        stream: StreamRef,
893    },
894
895    /// Defines an output tensor for a recording.
896    RecordingResult {
897        /// The ref that will be used to store the output tensor.
898        result: Ref,
899        /// The index of the output tensor in the list of output tensors.
900        output_index: usize,
901        /// The stream that this output tensor will come from.
902        stream: StreamRef,
903    },
904
905    /// Calls a recording that was previously defined using
906    /// DefineRecording.
907    CallRecording {
908        /// The sequence number of the invocation.
909        seq: Seq,
910        /// The ref of the recording to call.
911        recording: Ref,
912        /// The list of refs where the result tensors from the recording
913        /// will be stored.
914        results: Vec<Ref>,
915        /// The list of refs of input tensors to the recording.
916        actuals: Vec<Ref>,
917    },
918
919    SetRefUnitTestsOnly {
920        /// The reference to set.
921        reference: Ref,
922        /// The value to set it with.
923        value: WireValue,
924        /// The stream to set it on.
925        stream: StreamRef,
926    },
927
928    GetRefUnitTestsOnly {
929        /// The value to retrieve, expected to be a bool.
930        value: Ref,
931        /// The stream to retrieve from.
932        stream: StreamRef,
933        #[reply]
934        response_port: reference::OncePortRef<Option<Result<WireValue, String>>>,
935    },
936}
937
938/// The parameters to spawn a worker actor.
939#[derive(Debug, Clone, Serialize, Deserialize, Named)]
940pub struct WorkerParams {
941    // Global world size for this job
942    pub world_size: usize,
943
944    // Rank of the worker within the global world
945    pub rank: usize,
946
947    // Local cuda device that this worker represents. If None, we won't do CUDA
948    // synchronization.
949    pub device_index: Option<i8>,
950
951    // Actor Ref for the controller that the worker is associated with.
952    pub controller_actor: reference::ActorRef<ControllerActor>,
953}
954wirevalue::register_type!(WorkerParams);
955
956hyperactor::behavior!(WorkerActor, WorkerMessage,);
monarch_messages/worker.rs

monarch_messages/
worker.rs
