add a way to write python methods for rust-bound classes, round 2 (#913)

Summary:
Pull Request resolved: #913

This will be used to generate future python bindings for Context and Instance.

The implementation documents how it works.

This is the least bad way I could figure out how to do this since pyo3 does not natively let you subclass an existing python class.
ghstack-source-id: 304046443
exported-using-ghexport

Reviewed By: mariusae

Differential Revision: D80487153

fbshipit-source-id: 780376778ad7dcc36a10ae6b575d32ad4672914b

Author: zdevito

monarch_extension/src/lib.rs

@@ -28,6 +28,8 @@ mod tensor_worker;
 mod blocking;
 mod panic;
 mod trace;
+
+use monarch_types::py_global;
 use pyo3::prelude::*;
 
 #[pyfunction]

monarch_types/src/lib.rs

@@ -16,3 +16,22 @@ pub use pyobject::PickledPyObject;
 pub use python::SerializablePyErr;
 pub use python::TryIntoPyObjectUnsafe;
 pub use pytree::PyTree;
+
+/// Macro to generate a Python object lookup function with caching
+///
+/// # Arguments
+/// * `$fn_name` - Name of the Rust function to generate
+/// * `$python_path` - Path to the Python object as a string (e.g., "module.submodule.function")
+#[macro_export]
+macro_rules! py_global {
+    ($fn_name:ident, $python_module:literal, $python_class:literal) => {
+        fn $fn_name<'py>(py: ::pyo3::Python<'py>) -> ::pyo3::Bound<'py, ::pyo3::PyAny> {
+            static CACHE: ::pyo3::sync::GILOnceCell<::pyo3::PyObject> =
+                ::pyo3::sync::GILOnceCell::new();
+            CACHE
+                .import(py, $python_module, $python_class)
+                .unwrap()
+                .clone()
+        }
+    };
+}

python/monarch/_src/actor/python_extension_methods.py

@@ -0,0 +1,86 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import importlib
+
+from typing import cast, Type, TypeVar
+
+
+T = TypeVar("T")
+
+
+class PatchRustClass:
+    def __init__(self, rust_class: Type):
+        self.rust_class = rust_class
+
+    def __call__(self, python_class: Type[T]) -> Type[T]:
+        assert self.rust_class.__module__ == python_class.__module__
+        assert self.rust_class.__name__ == python_class.__name__
+        for name, implementation in python_class.__dict__.items():
+            if hasattr(self.rust_class, name):
+                # do not patch in the stub methods that
+                # are already defined by the rust implementation
+                continue
+            if not callable(implementation):
+                continue
+            setattr(self.rust_class, name, implementation)
+        return cast(Type[T], self.rust_class)
+
+
+def rust_struct(name: str) -> PatchRustClass:
+    """
+    When we bind a rust struct into Python, it is sometimes faster to implement
+    parts of the desired Python API in Python. It is also easier to understand
+    what the class does in terms of these methods.
+
+    We also want to avoid having to wrap rust objects in another layer of python objects
+    because:
+    * wrappers double the python overhead
+    * it is easy to confuse which level of wrappers and API takes, especially
+      along the python<->rust boundary.
+
+
+    To avoid wrappers we first define the class in pyo3. Lets say we add a class
+    monarch_hyperactor::actor_mesh::TestClass which we will want to extend with python methods in
+    the monarch/actor/_src/actor_mesh.py. In rust we will define the class as
+
+         #[pyclass(name = "TestClass", module = "monarch._src.actor_mesh")]
+         struct TestClass {}
+         #[pymethods]
+         impl TestClass {
+            fn hello(&self) {
+                println!("hello");
+            }
+         }
+
+    Then rather than writing typing stubs in a pyi file we write the stub code directly in
+    monarch/actor/_src/actor_mesh.py along with any helper methods:
+
+        @rust_struct("monarch_hyperactor::actor_mesh::TestClass")
+        class TestClass:
+            def hello(self) -> None:
+                ...
+            def hello_world(self) -> None:
+                self.hello()
+                print("world")
+
+    This class annotation then merges the python extension methods with the rust
+    class implementation. Any rust code that returns the TestClass will have the `hello_world`
+    extension method attached. Python typechecking always things TestClass is the python code,
+    so typing works.
+
+    It is ok to have the pyclass module not match where it is defined because (1) we patch it into the right place
+    to make sure pickling works, and (2) the rust_struct annotation points directly to where to find the rust code,
+    and will be discovered by goto line in the IDE.
+    """
+
+    *modules, name = name.split("::")
+    module_name = ".".join(modules)
+    module = importlib.import_module(f"monarch._rust_bindings.{module_name}")
+
+    rust_class = getattr(module, name)
+
+    return PatchRustClass(rust_class)

python/tests/test_python_actors.py

@@ -246,6 +246,16 @@ def test_value_mesh() -> None:
 
 @pytest.mark.timeout(60)
 def test_rust_binding_modules_correct() -> None:
+    """
+    This tests that rust bindings will survive pickling correctly.
+
+    To correctly define a rust binding, either
+
+    (1) Set its module to "monarch._rust_bindings.rust_crate.rust_module",
+        and make sure it is registered in monarch_extension/lib.rs
+    (2) Set its module to some existing python file, and use @rust_struct to install
+        the rust struct in that file and patch in any python extension methods.
+    """
     import monarch._rust_bindings as bindings
 
     def check(module, path):
@@ -255,8 +265,9 @@ def check(module, path):
             if isinstance(value, ModuleType):
                 check(value, f"{path}.{name}")
             elif hasattr(value, "__module__"):
-                assert value.__name__ == name
-                assert value.__module__ == path
+                value_module = importlib.import_module(value.__module__)
+                resolved_value = getattr(value_module, value.__name__)
+                assert value is resolved_value
 
     check(bindings, "monarch._rust_bindings")