tensor engine, implement rref on actor endpoints (#625)

Summary:
Pull Request resolved: #625

Allows actor endpoints to take tensors as arguments and return them as monarch.Tensor objects.

Actor endpoints that do this must have a propagator function defined.

"cached" does not currently work because the caching machinery cannot make the actor call.

Introduces `as_endpoint()` which can take something that was not annotated on a class and turn it into an endpoint. This is useful when needing to define a per-actor propagation function for rref calls because there otherwise is nowhere to put it.
ghstack-source-id: 298650755

Reviewed By: suo

Differential Revision: D78782587

fbshipit-source-id: 079fa58143fc32e29edb8147b5f226b810a787f7

Author: zdevito

monarch_extension/src/convert.rs

@@ -443,6 +443,18 @@ fn create_map(py: Python) -> HashMap<u64, FnType> {
             },
         ))
     });
+    m.insert(key("CallActorMethod"), |p| {
+        Ok(WorkerMessage::CallActorMethod(worker::ActorMethodParams {
+            call: worker::ActorCallParams {
+                seq: p.parseSeq("seq")?,
+                broker_id: p.parse("broker_id")?,
+                local_state: p.parseRefList("local_state")?,
+                mutates: p.parseRefList("mutates")?,
+                stream: p.parseStreamRef("stream")?,
+            },
+            results: p.parseFlatReferences("result")?,
+        }))
+    });
     m
 }
 

monarch_extension/src/tensor_worker.rs

@@ -1373,6 +1373,7 @@ pub(crate) fn worker_message_to_py(py: Python<'_>, message: &WorkerMessage) -> P
         WorkerMessage::SetRefUnitTestsOnly { .. } => unimplemented!(),
         WorkerMessage::GetRefUnitTestsOnly { .. } => unimplemented!(),
         WorkerMessage::SendResultOfActorCall { .. } => unimplemented!(),
+        WorkerMessage::CallActorMethod { .. } => unimplemented!(),
     }
 }
 

monarch_messages/src/worker.rs

@@ -415,6 +415,13 @@ pub struct ActorCallParams {
     pub mutates: Vec<Ref>,
     pub stream: StreamRef,
 }
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct ActorMethodParams {
+    pub results: Vec<Option<Ref>>,
+    pub call: ActorCallParams,
+}
+
 /// Type of reduction for [`WorkerMessage::Reduce`].
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum Reduction {
@@ -817,6 +824,7 @@ pub enum WorkerMessage {
     },
 
     SendResultOfActorCall(ActorCallParams),
+    CallActorMethod(ActorMethodParams),
     PipeRecv {
         seq: Seq,
         /// Result refs.

monarch_simulator/src/worker.rs

@@ -319,6 +319,14 @@ impl WorkerMessageHandler for WorkerActor {
         bail!("unimplemented: send_result_of_actor_call");
     }
 
+    async fn call_actor_method(
+        &mut self,
+        _cx: &hyperactor::Context<Self>,
+        _params: ActorMethodParams,
+    ) -> Result<()> {
+        bail!("unimplemented: call_actor_method");
+    }
+
     async fn command_group(
         &mut self,
         cx: &hyperactor::Context<Self>,

monarch_tensor_worker/src/lib.rs

@@ -72,6 +72,7 @@ use monarch_messages::controller::ControllerMessageClient;
 use monarch_messages::controller::Seq;
 use monarch_messages::wire_value::WireValue;
 use monarch_messages::worker::ActorCallParams;
+use monarch_messages::worker::ActorMethodParams;
 use monarch_messages::worker::CallFunctionError;
 use monarch_messages::worker::CallFunctionParams;
 use monarch_messages::worker::Factory;
@@ -860,6 +861,15 @@ impl WorkerMessageHandler for WorkerActor {
             .await?;
         Ok(())
     }
+    async fn call_actor_method(
+        &mut self,
+        cx: &hyperactor::Context<Self>,
+        params: ActorMethodParams,
+    ) -> Result<()> {
+        let stream = self.try_get_stream(params.call.stream)?;
+        stream.call_actor_method(cx, params).await?;
+        Ok(())
+    }
     async fn split_comm(
         &mut self,
         cx: &hyperactor::Context<Self>,

monarch_tensor_worker/src/stream.rs

@@ -45,6 +45,7 @@ use monarch_messages::controller::ControllerMessageClient;
 use monarch_messages::controller::Seq;
 use monarch_messages::controller::WorkerError;
 use monarch_messages::worker::ActorCallParams;
+use monarch_messages::worker::ActorMethodParams;
 use monarch_messages::worker::CallFunctionError;
 use monarch_messages::worker::CallFunctionParams;
 use monarch_messages::worker::SeqError;
@@ -265,6 +266,7 @@ pub enum StreamMessage {
     GetTensorRefUnitTestsOnly(Ref, #[reply] OncePortHandle<Option<TensorCellResult>>),
 
     SendResultOfActorCall(ActorId, ActorCallParams),
+    CallActorMethod(ActorMethodParams),
 }
 
 impl StreamMessage {
@@ -1062,6 +1064,44 @@ impl StreamActor {
         self.env.insert(dest, rvalue.clone());
         Ok(())
     }
+    async fn call_actor(
+        &mut self,
+        cx: &Context<'_, Self>,
+        params: ActorCallParams,
+    ) -> Result<PyObject, CallFunctionError> {
+        let local_state: Result<Vec<PyObject>> = Python::with_gil(|py| {
+            params
+                .local_state
+                .into_iter()
+                .map(|elem| {
+                    // SAFETY: python is gonna make unsafe copies of this stuff anyway
+                    unsafe {
+                        let x = self.ref_to_rvalue(&elem)?.try_to_object_unsafe(py)?.into();
+                        Ok(x)
+                    }
+                })
+                .collect()
+        });
+
+        let (send, recv) = cx.open_once_port();
+        let state = LocalState {
+            response_port: send,
+            state: local_state?,
+        };
+        let x: u64 = params.seq.into();
+        let message = LocalStateBrokerMessage::Set(x as usize, state);
+
+        let broker = BrokerId::new(params.broker_id).resolve(cx).unwrap();
+        broker
+            .send(message)
+            .map_err(|e| CallFunctionError::Error(e.into()))?;
+        let result = recv
+            .recv()
+            .await
+            .map_err(|e| CallFunctionError::Error(e.into()))?;
+
+        result.map_err(|pyerr| anyhow::Error::msg(pyerr.to_string()).into())
+    }
 }
 
 #[async_trait]
@@ -1669,62 +1709,38 @@ impl StreamMessageHandler for StreamActor {
         worker_actor_id: ActorId,
         params: ActorCallParams,
     ) -> anyhow::Result<()> {
-        // TODO: handle mutates
-        let local_state: Result<Vec<PyObject>> = Python::with_gil(|py| {
-            params
-                .local_state
-                .into_iter()
-                .map(|elem| {
-                    // SAFETY: python is gonna make unsafe copies of this stuff anyway
-                    unsafe {
-                        let x = self.ref_to_rvalue(&elem)?.try_to_object_unsafe(py)?.into();
-                        Ok(x)
-                    }
-                })
-                .collect()
-        });
-
-        let (send, recv) = cx.open_once_port();
-
-        let state = LocalState {
-            response_port: send,
-            state: local_state?,
-        };
-        let x: u64 = params.seq.into();
-        let message = LocalStateBrokerMessage::Set(x as usize, state);
+        let seq = params.seq;
+        let mutates = params.mutates.clone();
+        self.try_define(cx, seq, vec![], &mutates, async |self| {
+            let value = self.call_actor(cx, params).await?;
+            let result = Python::with_gil(|py| {
+                pickle_python_result(py, value.into_bound(py), worker_actor_id)
+            })?;
+            let result = Serialized::serialize(&result).unwrap();
+            self.controller_actor
+                .fetch_result(cx, seq, Ok(result))
+                .await?;
+            Ok(vec![])
+        })
+        .await
+    }
 
-        let broker = BrokerId::new(params.broker_id).resolve(cx).unwrap();
-        broker.send(message)?;
-        let result = recv.recv().await?;
-
-        match result {
-            Err(pyerr) => {
-                // If result has "exception" as its kind, then
-                // we need to unpickle and turn it into a WorkerError
-                // and call remote_function_failed otherwise the
-                // controller assumes the object is correct and doesn't handle
-                // dependency tracking correctly.
-                let err = Python::with_gil(|py| -> Result<WorkerError, SerializablePyErr> {
-                    Ok(WorkerError {
-                        worker_actor_id,
-                        backtrace: pyerr.to_string(),
-                    })
-                })?;
-                self.controller_actor
-                    .remote_function_failed(cx, params.seq, err)
-                    .await?;
-            }
-            Ok(value) => {
-                let result = Python::with_gil(|py| {
-                    pickle_python_result(py, value.into_bound(py), worker_actor_id)
-                })?;
-                let result = Serialized::serialize(&result).unwrap();
-                self.controller_actor
-                    .fetch_result(cx, params.seq, Ok(result))
-                    .await?;
-            }
-        }
-        Ok(())
+    async fn call_actor_method(
+        &mut self,
+        cx: &Context<Self>,
+        params: ActorMethodParams,
+    ) -> anyhow::Result<()> {
+        let seq = params.call.seq;
+        let mutates = params.call.mutates.clone();
+        self.try_define(cx, seq, params.results, &mutates, async |self| {
+            let result = self.call_actor(cx, params.call).await?;
+            let result = Python::with_gil(|py| {
+                PyTree::<RValue>::extract_bound(&result.into_bound(py))
+                    .map_err(SerializablePyErr::from_fn(py))
+            })?;
+            Ok(result.into_leaves())
+        })
+        .await
     }
 
     async fn set_value(

python/monarch/_src/actor/actor_mesh.py

@@ -65,6 +65,7 @@
     Endpoint,
     EndpointProperty,
     Extent,
+    NotAnEndpoint,
     Propagator,
     Selection,
 )
@@ -76,7 +77,7 @@
 from monarch._src.actor.shape import MeshTrait, NDSlice
 from monarch._src.actor.sync_state import fake_sync_state
 
-from monarch._src.actor.tensor_engine_shim import actor_send
+from monarch._src.actor.tensor_engine_shim import actor_rref, actor_send
 
 if TYPE_CHECKING:
     from monarch._src.actor.proc_mesh import ProcMesh
@@ -313,8 +314,7 @@ def _send(
         """
         self._signature.bind(None, *args, **kwargs)
         objects, bytes = flatten((args, kwargs), _is_ref_or_mailbox)
-        refs = [obj for obj in objects if hasattr(obj, "__monarch_ref__")]
-        if not refs:
+        if all(not hasattr(obj, "__monarch_ref__") for obj in objects):
             message = PythonMessage(
                 PythonMessageKind.CallMethod(
                     self._name, None if port is None else port._port_ref
@@ -323,7 +323,7 @@ def _send(
             )
             self._actor_mesh.cast(message, selection)
         else:
-            actor_send(self, bytes, refs, port, selection)
+            actor_send(self, bytes, objects, port, selection)
         shape = self._actor_mesh._shape
         return Extent(shape.labels, shape.ndslice.sizes)
 
@@ -335,6 +335,26 @@ def _port(self, once: bool = False) -> "PortTuple[R]":
             ), "unexpected receiver type"
         return PortTuple(p, PortReceiver(self._mailbox, self._supervise(r._receiver)))
 
+    def _rref(self, args, kwargs):
+        self._signature.bind(None, *args, **kwargs)
+        refs, bytes = flatten((args, kwargs), _is_ref_or_mailbox)
+
+        return actor_rref(self, bytes, refs)
+
+
+def as_endpoint(
+    not_an_endpoint: Callable[P, R], *, propagate: Propagator = None
+) -> Endpoint[P, R]:
+    if not isinstance(not_an_endpoint, NotAnEndpoint):
+        raise ValueError("expected an method of a spawned actor")
+    return ActorEndpoint(
+        not_an_endpoint._ref._actor_mesh_ref,
+        not_an_endpoint._name,
+        getattr(not_an_endpoint._ref, not_an_endpoint._name),
+        not_an_endpoint._ref._mailbox,
+        propagate,
+    )
+
 
 class Accumulator(Generic[P, R, A]):
     def __init__(
@@ -625,18 +645,23 @@ async def handle(
                         f" This is likely due to an earlier error: {self._saved_error}"
                     )
                 raise AssertionError(error_message)
-            the_method = getattr(self.instance, method)._method
+            the_method = getattr(self.instance, method)
+            if isinstance(the_method, EndpointProperty):
+                module = the_method._method.__module__
+                the_method = functools.partial(the_method._method, self.instance)
+            else:
+                module = the_method.__module__
 
             if inspect.iscoroutinefunction(the_method):
 
                 async def instrumented():
                     enter_span(
-                        the_method.__module__,
+                        module,
                         method,
                         str(ctx.mailbox.actor_id),
                     )
                     try:
-                        result = await the_method(self.instance, *args, **kwargs)
+                        result = await the_method(*args, **kwargs)
                         self._maybe_exit_debugger()
                     except Exception as e:
                         logging.critical(
@@ -649,9 +674,9 @@ async def instrumented():
 
                 result = await instrumented()
             else:
-                enter_span(the_method.__module__, method, str(ctx.mailbox.actor_id))
+                enter_span(module, method, str(ctx.mailbox.actor_id))
                 with fake_sync_state():
-                    result = the_method(self.instance, *args, **kwargs)
+                    result = the_method(*args, **kwargs)
                 self._maybe_exit_debugger()
                 exit_span()
 
@@ -758,35 +783,14 @@ def __init__(
                         attr_name,
                         attr_value._method,
                         self._mailbox,
+                        attr_value._propagator,
                     ),
                 )
 
-    def __getattr__(self, name: str) -> Any:
-        # This method is called when an attribute is not found
-        # For linting purposes, we need to tell the type checker that any attribute
-        # could be an endpoint that's dynamically added at runtime
-        # At runtime, we still want to raise AttributeError for truly missing attributes
-
-        # Check if this is a method on the underlying class
-        if hasattr(self._class, name):
-            attr = getattr(self._class, name)
-            if isinstance(attr, EndpointProperty):
-                # Dynamically create the endpoint
-                endpoint = ActorEndpoint(
-                    self._actor_mesh_ref,
-                    name,
-                    attr._method,
-                    self._mailbox,
-                    propagator=attr._propagator,
-                )
-                # Cache it for future use
-                setattr(self, name, endpoint)
-                return endpoint
-
-        # If we get here, it's truly not found
-        raise AttributeError(
-            f"'{self.__class__.__name__}' object has no attribute '{name}'"
-        )
+    def __getattr__(self, attr: str) -> NotAnEndpoint:
+        if attr in dir(self._class):
+            return NotAnEndpoint(self, attr)
+        raise AttributeError(attr)
 
     def _create(
         self,