tensor engine: Simplify CallFunctionError (#552)

Summary:
Pull Request resolved: #552

 The only case we every 'catch' is DependentError. Everything else is just a textural description. So this collapses those cases into just the two cases.

This retains the enum variants as function calls so that we can also reify them again if we need to catch them but having them obscures the places where DependentError needs to be handled differently.

Reviewed By: suo

Differential Revision: D78363663

fbshipit-source-id: f1d2ba31dfe65e1b1202ea4d2b67d224cfa76b5d

Author: zdevito

monarch_messages/src/worker.rs

@@ -506,74 +506,112 @@ pub enum StreamCreationMode {
     CreateNewStream,
 }
 
-/// The kinds of errors that a CallFunction message can return with.
-// TODO(agallagher): We should move most variants out into `ValueError`.
+/// When a worker runs any function, it may not succeed either because the function itself
+/// failed (Error) or because an input to the function already had an error value
+/// DependentError.
 #[derive(Error, Debug, Named)]
 pub enum CallFunctionError {
-    #[error("ref not found: {0}")]
-    RefNotFound(Ref),
+    #[error("{0}")]
+    Error(#[from] anyhow::Error),
+    #[error("Computation depended on an input that failed with errror: {0}")]
+    DependentError(Arc<CallFunctionError>),
+}
 
-    #[error("dependent error {0}")]
-    DependentError(#[from] Arc<CallFunctionError>),
+impl CallFunctionError {
+    /// Checks if the error is a `DependentError` and returns the underlying
+    /// error if so. Otherwise, returns `None`.
+    pub fn unwrap_dependent_error(&self) -> Option<Arc<CallFunctionError>> {
+        match self {
+            CallFunctionError::DependentError(e) => Some(e.clone()),
+            _ => None,
+        }
+    }
 
-    #[error("invalid remote function: {0}")]
-    InvalidRemoteFunction(String),
+    // Static functions for backward compatibility with existing enum cases
 
-    #[error("unsupported arg type for {0} function: {1}")]
-    UnsupportedArgType(String, String),
+    #[allow(non_snake_case)]
+    pub fn RefNotFound(r: Ref) -> Self {
+        Self::Error(anyhow::anyhow!("ref not found: {}", r))
+    }
 
-    #[error("remote function failed: {0}")]
-    RemoteFunctionFailed(#[from] SerializablePyErr),
+    #[allow(non_snake_case)]
+    pub fn InvalidRemoteFunction(msg: String) -> Self {
+        Self::Error(anyhow::anyhow!("invalid remote function: {}", msg))
+    }
 
-    #[error("borrow failed: {0}")]
-    BorrowError(#[from] BorrowError),
+    #[allow(non_snake_case)]
+    pub fn UnsupportedArgType(function_type: String, arg_type: String) -> Self {
+        Self::Error(anyhow::anyhow!(
+            "unsupported arg type for {} function: {}",
+            function_type,
+            arg_type
+        ))
+    }
 
-    #[error("torch operator failed: {0}")]
-    OperatorFailed(#[from] CallOpError),
+    #[allow(non_snake_case)]
+    pub fn RemoteFunctionFailed(err: SerializablePyErr) -> Self {
+        Self::Error(anyhow::anyhow!("remote function failed: {}", err))
+    }
 
-    #[error("unexpected number of returns from op, expected {expected}, got {actual}")]
-    UnexpectedNumberOfReturns { expected: usize, actual: usize },
+    #[allow(non_snake_case)]
+    pub fn BorrowError(err: BorrowError) -> Self {
+        Self::Error(anyhow::anyhow!("borrow failed: {}", err))
+    }
 
-    #[error(
-        "expected only a single arg (and no kwargs) when no function is given: {args}, {kwargs}"
-    )]
-    TooManyArgsForValue { args: String, kwargs: String },
+    #[allow(non_snake_case)]
+    pub fn OperatorFailed(err: CallOpError) -> Self {
+        Self::Error(anyhow::anyhow!("torch operator failed: {}", err))
+    }
 
-    #[error("error: {0}")]
-    Anyhow(#[from] anyhow::Error),
+    #[allow(non_snake_case)]
+    pub fn UnexpectedNumberOfReturns(expected: usize, actual: usize) -> Self {
+        Self::Error(anyhow::anyhow!(
+            "unexpected number of returns from op, expected {}, got {}",
+            expected,
+            actual
+        ))
+    }
 
-    #[error("recording failed at message {index}: ({message}). Error: {error}")]
-    RecordingFailed {
-        index: usize,
-        message: String,
-        error: Arc<CallFunctionError>,
-    },
+    #[allow(non_snake_case)]
+    pub fn TooManyArgsForValue(args: String, kwargs: String) -> Self {
+        Self::Error(anyhow::anyhow!(
+            "expected only a single arg (and no kwargs) when no function is given: {}, {}",
+            args,
+            kwargs
+        ))
+    }
+
+    #[allow(non_snake_case)]
+    pub fn Anyhow(err: anyhow::Error) -> Self {
+        Self::Error(err)
+    }
+
+    #[allow(non_snake_case)]
+    pub fn RecordingFailed(index: usize, message: String, error: Arc<CallFunctionError>) -> Self {
+        Self::Error(anyhow::anyhow!(
+            "recording failed at message {}: ({}). Error: {}",
+            index,
+            message,
+            error
+        ))
+    }
 }
 
-impl CallFunctionError {
-    /// Checks if the error is a `DependentError` and returns the underlying
-    /// error if so. Otherwise, returns `None`.
-    pub fn unwrap_dependent_error(&self) -> Option<Arc<CallFunctionError>> {
-        match self {
-            CallFunctionError::DependentError(e) => Some(e.clone()),
-            _ => None,
-        }
+impl From<SerializablePyErr> for CallFunctionError {
+    fn from(v: SerializablePyErr) -> CallFunctionError {
+        CallFunctionError::Error(v.into())
     }
 }
 
-/// Errors encountered during worker operations which get propagated in the
-/// refs-to-values maps used to store values.
-#[derive(Debug, Serialize, Deserialize, Error, Named, EnumAsInner)]
-pub enum ValueError {
-    // TODO(agallagher): Migrate to variants for each error (after we cleanup
-    // `CallFunctionError` to make it serializable).
-    #[error("call function error: {0}")]
-    CallFunctionError(String),
+impl From<BorrowError> for CallFunctionError {
+    fn from(v: BorrowError) -> CallFunctionError {
+        CallFunctionError::Error(v.into())
+    }
 }
 
-impl From<Arc<CallFunctionError>> for ValueError {
-    fn from(err: Arc<CallFunctionError>) -> Self {
-        ValueError::CallFunctionError(format!("{:?}", err))
+impl From<CallOpError> for CallFunctionError {
+    fn from(v: CallOpError) -> CallFunctionError {
+        CallFunctionError::Error(v.into())
     }
 }
 
@@ -874,7 +912,7 @@ pub enum WorkerMessage {
         /// The stream to retrieve from.
         stream: StreamRef,
         #[reply]
-        response_port: hyperactor::OncePortRef<Option<Result<WireValue, ValueError>>>,
+        response_port: hyperactor::OncePortRef<Option<Result<WireValue, String>>>,
     },
 }
 

monarch_simulator/src/worker.rs

@@ -639,7 +639,7 @@ impl WorkerMessageHandler for WorkerActor {
         _cx: &hyperactor::Context<Self>,
         _ref_id: Ref,
         _stream: StreamRef,
-    ) -> Result<Option<Result<WireValue, ValueError>>> {
+    ) -> Result<Option<Result<WireValue, String>>> {
         bail!("unimplemented: get_ref_unit_tests_only")
     }
 

monarch_tensor_worker/src/borrow.rs

@@ -422,10 +422,9 @@ mod tests {
         let error = result
             .context("no such ref")?
             .err()
-            .context("expected error")?
-            .into_call_function_error()?;
+            .context("expected error")?;
         assert!(
-            error.starts_with("DependentError"),
+            error.contains("Computation depended on an input that failed"),
             "If a borrowed value contains an error, downstream calls should propagate that error (unexpected error string: {})",
             error,
         );

monarch_tensor_worker/src/lib.rs

@@ -80,7 +80,6 @@ use monarch_messages::worker::Ref;
 use monarch_messages::worker::ResolvableFunction;
 use monarch_messages::worker::StreamCreationMode;
 use monarch_messages::worker::StreamRef;
-use monarch_messages::worker::ValueError;
 use monarch_messages::worker::WorkerMessage;
 use monarch_messages::worker::WorkerMessageHandler;
 use monarch_messages::worker::WorkerParams;
@@ -679,10 +678,9 @@ impl WorkerMessageHandler for WorkerActor {
                         (k, RValue::PyObject(object.into_py_object().unwrap()).into())
                     })
                     .collect();
-                let device_mesh = self
-                    .device_meshes
-                    .get(&device_mesh)
-                    .ok_or_else(|| CallFunctionError::RefNotFound(device_mesh))?;
+                let device_mesh = self.device_meshes.get(&device_mesh).ok_or_else(|| {
+                    CallFunctionError::Error(anyhow::anyhow!("ref not found: {}", device_mesh))
+                })?;
                 let pipe = PipeActor::spawn(
                     cx,
                     PipeParams {
@@ -989,10 +987,15 @@ impl WorkerMessageHandler for WorkerActor {
 
         // Get a port for the pipe
         let pipe = match self.pipes.get(&pipe) {
-            None => Err(Arc::new(CallFunctionError::RefNotFound(pipe))),
+            None => Err(Arc::new(CallFunctionError::Error(anyhow::anyhow!(
+                "ref not found: {}",
+                pipe
+            )))),
             Some(pipe) => match pipe.as_ref() {
                 Ok(pipe) => Ok(pipe.port()),
-                Err(e) => Err(Arc::new(CallFunctionError::DependentError(e.clone()))),
+                Err(e) => Err(Arc::new(CallFunctionError::DependentError(
+                    e.unwrap_dependent_error().unwrap_or(e.clone()),
+                ))),
             },
         };
 
@@ -1020,7 +1023,7 @@ impl WorkerMessageHandler for WorkerActor {
         cx: &hyperactor::Context<Self>,
         ref_id: Ref,
         stream: StreamRef,
-    ) -> Result<Option<Result<WireValue, ValueError>>> {
+    ) -> Result<Option<Result<WireValue, String>>> {
         let stream = self.try_get_stream(stream)?;
         Ok(stream
             .get_ref_unit_tests_only(cx, ref_id.clone())
@@ -1437,8 +1440,7 @@ mod tests {
         let mutated_ref = result
             .context("no such ref")?
             .err()
-            .context("expected error")?
-            .into_call_function_error()?;
+            .context("expected error")?;
         assert!(mutated_ref.contains("InvalidRemoteFunction"));
 
         let responses = controller_rx.drain();
@@ -2542,7 +2544,7 @@ mod tests {
             ]))
             .unwrap();
 
-        let value: Result<_, ValueError> = handle
+        let value: Result<_, String> = handle
             .get_ref_unit_tests_only(&client, 3.into(), 0.into())
             .await
             .unwrap()