Fix errant cuda stream syncs and tensor serializations in remote function impl (#14)

Summary:
Pull Request resolved: #14

This diff fixes two issues related to the remote function implementation inside `StreamActor`:
1. There were two calls to `tracing::debug!` that use the function inputs and outputs as arguments. Even if the log level was set such that the message wasn't printed to the console, this still incurred the overhead of serialization (and therefore, when tensors were involved, synchronizing the host with the cuda stream).
2. The code to extract an `RValue` from a `PyObject` worked by attempting to do a conversion in C++, and if that failed, handling the exception via rust and trying the next `RValue` variant. Each time the conversion failed, the code created a `PyValueError` that serialized the `PyObject` as part of its message. Importantly, the case for tensors was only handled after the cases for `ScalarType`, `Layout`, and `MemoryFormat`. So if the `RValue` was a tensor, we would incur 3 C++ exceptions, 3 cuda stream syncs, and 3 tensor serializations. And this would happen for every tensor in the output of a remote function.

The fixes in this diff are:
1. Remove the `tracing::debug!` calls.
2. Instead of trying an unchecked type conversion and handling the exception, actually check the type before trying to do the conversion.

Reviewed By: colin2328

Differential Revision: D75035191

fbshipit-source-id: 9acf397c0ea9a9af0ffef215f39f2ee5c50b7395

Author: Sam Lurye

monarch_worker/src/stream.rs

@@ -830,9 +830,6 @@ impl StreamActor {
             // Execute the borrow.
             let _borrow = multiborrow.borrow()?;
 
-            tracing::debug!(
-                "calling python function: {function:?} with args: {py_args:?} and kwargs: {py_kwargs:?}"
-            );
             // Call function.
             // Use custom subscriber to route Worker messages to stdout.
             let scoped_subscriber = Subscriber::builder().with_writer(std::io::stdout).finish();
@@ -858,7 +855,6 @@ impl StreamActor {
                         )
                         .map_err(SerializablePyErr::from_fn(py))
                 })?;
-            tracing::debug!("python function: {function:?} result: {result:?}");
 
             // Parse the python result as an `Object`, which should preserve the
             // original Python object structure, while providing access to the

torch-sys/src/bridge.cpp

@@ -408,6 +408,18 @@ IValue ivalue_from_arbitrary_py_object(PyObject* unowned) {
   return torch::jit::toIValue(obj, inferredType.type());
 }
 
+bool py_object_is_ivalue(PyObject* unowned) {
+  auto obj = py::reinterpret_steal<py::object>(unowned);
+  auto inferredType = torch::jit::tryToInferType(obj);
+  if (!inferredType.success()) {
+    return false;
+  }
+  // TODO(agallagher): Arbitrary Python objects -- which we can't and don't want
+  // to package into an `IValue` -- will be inferred as a `Class` type.  Throw
+  // here so that we'll fallback to parsing into `RValue::PyObject`.
+  return !inferredType.type()->cast<at::ClassType>();
+}
+
 c10::Device device_from_py_object(PyObject* unowned) {
   auto obj = py::reinterpret_steal<py::object>(unowned);
   if (!THPDevice_Check(obj.ptr())) {
@@ -435,6 +447,11 @@ PyObject* scalar_type_to_py_object(c10::ScalarType scalar_type) {
   return Py_NewRef(dtype);
 }
 
+bool py_object_is_scalar_type(PyObject* unowned) {
+  auto obj = py::reinterpret_steal<py::object>(unowned);
+  return THPDtype_Check(obj.ptr());
+}
+
 c10::Layout layout_from_py_object(PyObject* unowned) {
   auto obj = py::reinterpret_steal<py::object>(unowned);
   if (!THPLayout_Check(obj.ptr())) {
@@ -449,6 +466,11 @@ PyObject* layout_to_py_object(c10::Layout layout) {
   return Py_NewRef(thp_layout);
 }
 
+bool py_object_is_layout(PyObject* unowned) {
+  auto obj = py::reinterpret_steal<py::object>(unowned);
+  return THPLayout_Check(obj.ptr());
+}
+
 c10::MemoryFormat memory_format_from_py_object(PyObject* unowned) {
   auto obj = py::reinterpret_steal<py::object>(unowned);
   if (!THPMemoryFormat_Check(obj.ptr())) {
@@ -463,6 +485,11 @@ PyObject* memory_format_to_py_object(c10::MemoryFormat memory_format) {
   return Py_NewRef(thp_memory_format);
 }
 
+bool py_object_is_memory_format(PyObject* unowned) {
+  auto obj = py::reinterpret_steal<py::object>(unowned);
+  return THPMemoryFormat_Check(obj.ptr());
+}
+
 PyObject* tensor_to_py_object(Tensor tensor) {
   torch::jit::guardAgainstNamedTensor<Tensor>(tensor);
   return py::cast(std::move(tensor)).release().ptr();

torch-sys/src/bridge.h

@@ -92,6 +92,7 @@ rust::Vec<int32_t> sizes(const Tensor& tensor);
 
 FFIPyObject arbitrary_ivalue_to_py_object(IValue val);
 IValue ivalue_from_arbitrary_py_object(FFIPyObject obj);
+bool py_object_is_ivalue(FFIPyObject obj);
 
 inline IValue ivalue_from_py_object_with_type(
     FFIPyObject obj,
@@ -111,12 +112,15 @@ FFIPyObject device_to_py_object(c10::Device device);
 
 c10::ScalarType scalar_type_from_py_object(FFIPyObject obj);
 FFIPyObject scalar_type_to_py_object(c10::ScalarType scalar_type);
+bool py_object_is_scalar_type(FFIPyObject obj);
 
 c10::Layout layout_from_py_object(FFIPyObject obj);
 FFIPyObject layout_to_py_object(c10::Layout layout);
+bool py_object_is_layout(FFIPyObject obj);
 
 c10::MemoryFormat memory_format_from_py_object(FFIPyObject obj);
 FFIPyObject memory_format_to_py_object(c10::MemoryFormat memory_format);
+bool py_object_is_memory_format(FFIPyObject obj);
 
 FFIPyObject tensor_to_py_object(Tensor tensor);
 Tensor tensor_from_py_object(FFIPyObject obj);

torch-sys/src/bridge.rs

@@ -217,10 +217,12 @@ pub(crate) mod ffi {
         // Layout
         fn layout_from_py_object(obj: FFIPyObject) -> Result<Layout>;
         fn layout_to_py_object(layout: Layout) -> FFIPyObject;
+        fn py_object_is_layout(obj: FFIPyObject) -> bool;
 
         // MemoryFormat
         fn memory_format_from_py_object(obj: FFIPyObject) -> Result<MemoryFormat>;
         fn memory_format_to_py_object(memory_format: MemoryFormat) -> FFIPyObject;
+        fn py_object_is_memory_format(obj: FFIPyObject) -> bool;
 
         // Tensor
         fn tensor_from_py_object(obj: FFIPyObject) -> Result<Tensor>;
@@ -281,6 +283,7 @@ pub(crate) mod ffi {
         // Convert to Python object.
         fn scalar_type_from_py_object(obj: FFIPyObject) -> Result<ScalarType>;
         fn scalar_type_to_py_object(scalar_type: ScalarType) -> FFIPyObject;
+        fn py_object_is_scalar_type(obj: FFIPyObject) -> bool;
 
         /// # Safety
         /// - **Mutability**:
@@ -330,6 +333,7 @@ pub(crate) mod ffi {
         // Interop with Python object.
         fn arbitrary_ivalue_to_py_object(val: IValue) -> Result<FFIPyObject>;
         fn ivalue_from_arbitrary_py_object(obj: FFIPyObject) -> Result<IValue>;
+        fn py_object_is_ivalue(obj: FFIPyObject) -> bool;
         /// Converts the provided Python object to an `IValue` with the provided
         /// type. If the object is not convertible to the provided type, an
         /// exception will be thrown.

torch-sys/src/ivalue.rs

@@ -309,6 +309,11 @@ impl IValue {
                 ))
             })
     }
+
+    pub(crate) fn from_py_object_or_none(obj: &Bound<'_, PyAny>) -> Option<IValue> {
+        ffi::py_object_is_ivalue(obj.clone().into())
+            .then(|| ffi::ivalue_from_arbitrary_py_object(obj.into()).unwrap())
+    }
 }
 
 // impl `From` for all IValue kinds.

torch-sys/src/layout.rs

@@ -22,6 +22,13 @@ unsafe impl ExternType for Layout {
     type Kind = cxx::kind::Trivial;
 }
 
+impl Layout {
+    pub(crate) fn from_py_object_or_none(obj: &Bound<'_, PyAny>) -> Option<Self> {
+        ffi::py_object_is_layout(obj.clone().into())
+            .then(|| ffi::layout_from_py_object(obj.into()).unwrap())
+    }
+}
+
 impl FromPyObject<'_> for Layout {
     fn extract_bound(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
         ffi::layout_from_py_object(obj.into()).map_err(|e| {

torch-sys/src/memory_format.rs

@@ -22,6 +22,13 @@ unsafe impl ExternType for MemoryFormat {
     type Kind = cxx::kind::Trivial;
 }
 
+impl MemoryFormat {
+    pub(crate) fn from_py_object_or_none(obj: &Bound<'_, PyAny>) -> Option<Self> {
+        ffi::py_object_is_memory_format(obj.clone().into())
+            .then(|| ffi::memory_format_from_py_object(obj.into()).unwrap())
+    }
+}
+
 impl FromPyObject<'_> for MemoryFormat {
     fn extract_bound(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
         ffi::memory_format_from_py_object(obj.into()).map_err(|e| {

torch-sys/src/rvalue.rs

@@ -167,13 +167,16 @@ impl TryIntoPyObjectUnsafe<PyAny> for &RValue {
 
 impl FromPyObject<'_> for RValue {
     fn extract_bound(obj: &pyo3::Bound<'_, pyo3::PyAny>) -> pyo3::PyResult<Self> {
-        if let Ok(val) = obj.extract::<ScalarType>() {
+        // It's crucial for correctness to try converting to IValue after we've
+        // tried the other non-PyObject variants, because the IValue conversion
+        // will actually succeed when obj is a ScalarType, Layout, or MemoryFormat.
+        if let Some(val) = ScalarType::from_py_object_or_none(obj) {
             Ok(RValue::ScalarType(val))
-        } else if let Ok(val) = obj.extract::<Layout>() {
+        } else if let Some(val) = Layout::from_py_object_or_none(obj) {
             Ok(RValue::Layout(val))
-        } else if let Ok(val) = obj.extract::<MemoryFormat>() {
+        } else if let Some(val) = MemoryFormat::from_py_object_or_none(obj) {
             Ok(RValue::MemoryFormat(val))
-        } else if let Ok(val) = IValue::extract_bound(obj) {
+        } else if let Some(val) = IValue::from_py_object_or_none(obj) {
             Ok(val.into())
         } else {
             Ok(RValue::PyObject(PickledPyObject::pickle(obj)?))

torch-sys/src/scalar_type.rs

@@ -23,6 +23,13 @@ unsafe impl ExternType for ScalarType {
     type Kind = cxx::kind::Trivial;
 }
 
+impl ScalarType {
+    pub(crate) fn from_py_object_or_none(obj: &Bound<'_, PyAny>) -> Option<Self> {
+        ffi::py_object_is_scalar_type(obj.clone().into())
+            .then(|| ffi::scalar_type_from_py_object(obj.into()).unwrap())
+    }
+}
+
 impl FromPyObject<'_> for ScalarType {
     fn extract_bound(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
         ffi::scalar_type_from_py_object(obj.into()).map_err(|e| {