[12/n] tensor engine, demote call_on_shard_and_fetch

Pull Request resolved: #470

remote is about to become an endpoint, along with the endpoint 'adverbs'.  By convention adverbs can't take additional arguments besides the endpoints args/kwargs, so like the `send` method `call_on_shard_and_fetch` needs to become a free function. Devmate says more

**Task Completed**

Successfully converted the [`**call_on_shard_and_fetch**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fpython%2Fmonarch%2Fcommon%2Fremote.py%22%2C179%5D "call_on_shard_and_fetch") method to a free function! Here's a summary of all the changes made:

Changes Made:
-------------

### 1. **Removed the method from the [`**Remote**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fpython%2Fmonarch%2Fcommon%2Fremote.py%22%2C65%5D "Remote") class**

*   Deleted the [`**call_on_shard_and_fetch**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fpython%2Fmonarch%2Fcommon%2Fremote.py%22%2C179%5D "call_on_shard_and_fetch") method from the [`**Remote**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fpython%2Fmonarch%2Fcommon%2Fremote.py%22%2C65%5D "Remote") class in [`**remote.py**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fpython%2Fmonarch%2Fcommon%2Fremote.py%22%2Cnull%5D "/data/users/zdevito/fbsource/fbcode/monarch/python/monarch/common/remote.py")

### 2. **Created a new free function**

*   Added `call_on_shard_and_fetch(remote_obj: Remote, *args, shard: Dict[str, int] | None = None, **kwargs) -> Future[Any]` as a standalone function
*   The function takes the [`**Remote**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fpython%2Fmonarch%2Fcommon%2Fremote.py%22%2C65%5D "Remote") instance as its first parameter, followed by the same arguments as the original method
*   Updated the function to call the internal [`**_call_on_shard_and_fetch**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fpython%2Fmonarch%2Fcommon%2Fremote.py%22%2C144%5D "_call_on_shard_and_fetch") with the appropriate parameters from the [`**Remote**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fpython%2Fmonarch%2Fcommon%2Fremote.py%22%2C65%5D "Remote") object

### 3. **Updated all call sites across the codebase**

*   **Test files**: Updated `test_remote_functions.py` and `test_rust_backend.py` to use the new free function syntax
*   **Example files**: Updated controller example and FSDP examples to use the new function
*   **Module files**: Updated `opaque_module.py` and `opaque_object.py` to use the new function
*   **Internal usage**: Updated the cached propagation logic in [`**remote.py**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fpython%2Fmonarch%2Fcommon%2Fremote.py%22%2Cnull%5D "remote.py") to use the new function

### 4. **Updated imports**

*   Added `from monarch.common.remote import call_on_shard_and_fetch` to all files that needed it
*   The function is also available through the main `monarch` module via the existing export in `__init__.py`

Syntax Change:
--------------

**Before (method):**

`remote_obj.call_on_shard_and_fetch(arg1, arg2, shard={"host": 0})`

**After (free function):**

`call_on_shard_and_fetch(remote_obj, arg1, arg2, shard={"host": 0})`

Files Modified:
---------------

*   [`**remote.py**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fpython%2Fmonarch%2Fcommon%2Fremote.py%22%2Cnull%5D "/data/users/zdevito/fbsource/fbcode/monarch/python/monarch/common/remote.py") (main changes)
*   [`**test_remote_functions.py**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fpython%2Ftests%2Ftest_remote_functions.py%22%2Cnull%5D "/data/users/zdevito/fbsource/fbcode/monarch/python/tests/test_remote_functions.py")
*   [`**test_rust_backend.py**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fpython%2Ftests%2Ftest_rust_backend.py%22%2Cnull%5D "/data/users/zdevito/fbsource/fbcode/monarch/python/tests/test_rust_backend.py")
*   [`**example.py**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fexamples%2Fmeta%2Fcontroller%2Fexample.py%22%2Cnull%5D "/data/users/zdevito/fbsource/fbcode/monarch/examples/meta/controller/example.py")
*   [`**opaque_module.py**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fpython%2Fmonarch%2Fopaque_module.py%22%2Cnull%5D "/data/users/zdevito/fbsource/fbcode/monarch/python/monarch/opaque_module.py")
*   [`**opaque_object.py**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fpython%2Fmonarch%2Fopaque_object.py%22%2Cnull%5D "/data/users/zdevito/fbsource/fbcode/monarch/python/monarch/opaque_object.py")
*   [`**fsdp_parametrization.py**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fexamples%2Fmeta%2Ffsdp%2Ffsdp_parametrization.py%22%2Cnull%5D "/data/users/zdevito/fbsource/fbcode/monarch/examples/meta/fsdp/fsdp_parametrization.py")
*   [`**fsdp_setattr.py**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fexamples%2Fmeta%2Ffsdp%2Ffsdp_setattr.py%22%2Cnull%5D "/data/users/zdevito/fbsource/fbcode/monarch/examples/meta/fsdp/fsdp_setattr.py")
*   [`**fetch.py**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fpython%2Fmonarch%2Ffetch.py%22%2Cnull%5D "/data/users/zdevito/fbsource/fbcode/monarch/python/monarch/fetch.py")

All changes have been validated and no errors were introduced. The function maintains the same behavior and API, just with a different calling convention where the [`**Remote**`](command:code-compose.open?%5B%22%2Fdata%2Fusers%2Fzdevito%2Ffbsource%2Ffbcode%2Fmonarch%2Fpython%2Fmonarch%2Fcommon%2Fremote.py%22%2C65%5D "Remote") object is passed as the first argument instead of being the method receiver.
ghstack-source-id: 295034203

Differential Revision: [D77967708](https://our.internmc.facebook.com/intern/diff/D77967708/)

Author: zdevito

python/monarch/common/remote.py

@@ -72,6 +72,14 @@ def __init__(self, impl: Any, propagator_arg: Propagator):
     def _resolvable(self):
         return resolvable_function(self._remote_impl)
 
+    @property
+    def _maybe_resolvable(self):
+        return (
+            None
+            if self._remote_impl is None
+            else resolvable_function(self._remote_impl)
+        )
+
     def _propagate(self, args, kwargs, fake_args, fake_kwargs):
         if self._propagator_arg is None or self._propagator_arg == "cached":
             if self._cache is None:
@@ -104,13 +112,6 @@ def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
             stream._active,
         )
 
-    def call_on_shard_and_fetch(
-        self, *args, shard: Dict[str, int] | None = None, **kwargs
-    ) -> Future[R]:
-        return _call_on_shard_and_fetch(
-            self._resolvable, self._fetch_propagate, *args, shard=shard, **kwargs
-        )
-
 
 # This can't just be Callable because otherwise we are not
 # allowed to use type arguments in the return value.
@@ -148,21 +149,26 @@ def remote(function: Any = None, *, propagate: Propagator = None) -> Any:
     return Remote(function, propagate)
 
 
-def _call_on_shard_and_fetch(
-    rfunction: ResolvableFunction | None,
-    propagator: Any,
+remote_identity = Remote(None, lambda x: x)
+
+
+def call_on_shard_and_fetch(
+    remote_obj: Remote[P, R],
     /,
     *args: object,
     shard: dict[str, int] | None = None,
     **kwargs: object,
-) -> Future:
+) -> Future[R]:
     """
     Call `function` at the coordinates `shard` of the current device mesh, and retrieve the result as a Future.
         function - the remote function to call
         *args/**kwargs - arguments to the function
         shard - a dictionary from mesh dimension name to coordinate of the shard
                 If None, this will fetch from coordinate 0 for all dimensions (useful after all_reduce/all_gather)
     """
+
+    rfunction = remote_obj._maybe_resolvable
+    propagator = remote_obj._fetch_propagate
     ambient_mesh = device_mesh._active
 
     if rfunction is None:
@@ -271,8 +277,8 @@ def _cached_propagation(_cache, rfunction, args, kwargs):
     if key not in _cache:
         _miss += 1
         args_no_pg, kwargs_no_pg = tree_map(_mock_pgs, (args, kwargs))
-        result_with_placeholders, output_pattern = _propagate.call_on_shard_and_fetch(
-            function=rfunction, args=args_no_pg, kwargs=kwargs_no_pg
+        result_with_placeholders, output_pattern = call_on_shard_and_fetch(
+            _propagate, function=rfunction, args=args_no_pg, kwargs=kwargs_no_pg
         ).result()
 
         _, unflatten_result = flatten(

python/monarch/fetch.py

@@ -15,7 +15,7 @@
 
 from monarch.common.future import Future
 
-from monarch.common.remote import _call_on_shard_and_fetch
+from monarch.common.remote import call_on_shard_and_fetch, remote_identity
 
 T = TypeVar("T")
 
@@ -37,9 +37,7 @@ def fetch_shard(
             shard = {}
         shard.update(kwargs)
 
-    return _call_on_shard_and_fetch(
-        None, lambda *args, **kwargs: None, obj, shard=shard
-    )
+    return call_on_shard_and_fetch(remote_identity, obj, shard=shard)
 
 
 def show(obj: T, shard: dict[str, int] | None = None, **kwargs: int) -> object:

python/monarch/opaque_module.py

@@ -9,7 +9,7 @@
 import torch
 from monarch.common.function_caching import TensorGroup, TensorGroupPattern
 from monarch.common.opaque_ref import OpaqueRef
-from monarch.common.remote import remote
+from monarch.common.remote import call_on_shard_and_fetch, remote
 from monarch.common.tensor_factory import TensorFactory
 from monarch.common.tree import flatten
 from monarch.opaque_object import _fresh_opaque_ref, OpaqueObject
@@ -144,11 +144,9 @@ def __init__(self, *args, **kwargs):
 
     def parameters(self):
         if self._parameters is None:
-            tensor_group_pattern = (
-                remote(_get_parameters_shape)
-                .call_on_shard_and_fetch(self._object)
-                .result()
-            )
+            tensor_group_pattern = call_on_shard_and_fetch(
+                remote(_get_parameters_shape), self._object
+            ).result()
             self._parameters = [
                 p.requires_grad_(True)
                 for p in remote(

python/monarch/opaque_object.py

@@ -14,7 +14,7 @@
 )
 
 from monarch.common.opaque_ref import OpaqueRef
-from monarch.common.remote import remote
+from monarch.common.remote import call_on_shard_and_fetch, remote
 
 
 def _invoke_method(obj: OpaqueRef, method_name: str, *args, **kwargs):
@@ -83,6 +83,6 @@ def call_method(self, method_name, propagation, *args, **kwargs):
         return endpoint(self, method_name, *args, **kwargs)
 
     def call_method_on_shard_and_fetch(self, method_name, *args, **kwargs):
-        return remote(_invoke_method).call_on_shard_and_fetch(
-            self, method_name, *args, **kwargs
+        return call_on_shard_and_fetch(
+            remote(_invoke_method), self, method_name, *args, **kwargs
         )

python/tests/test_remote_functions.py

@@ -35,7 +35,7 @@
 from monarch.cached_remote_function import remote_autograd_function
 from monarch.common import remote as remote_module
 from monarch.common.device_mesh import DeviceMesh
-from monarch.common.remote import Remote
+from monarch.common.remote import call_on_shard_and_fetch, Remote
 from monarch.mesh_controller import RemoteException as NewRemoteException
 
 from monarch.opaque_module import OpaqueModule
@@ -634,11 +634,10 @@ def test_fetch_preprocess(self, backend_type):
         with self.local_device_mesh(2, 2, backend_type):
             assert (
                 "an argument processed"
-                == remote("monarch.worker._testing_function.do_some_processing")
-                .call_on_shard_and_fetch(
+                == call_on_shard_and_fetch(
+                    remote("monarch.worker._testing_function.do_some_processing"),
                     "an argument",
-                )
-                .result()
+                ).result()
             )
 
     def test_cached_remote_function(self, backend_type):
@@ -733,7 +732,7 @@ def cuda_works(x):
 
         with self.local_device_mesh(2, 2, backend_type):
             a = torch.ones(())
-            assert check.call_on_shard_and_fetch(bar(a, a)).result()
+            assert call_on_shard_and_fetch(check, bar(a, a)).result()
             # ensure we do not attempt to pickle closures
             close()
 
@@ -776,7 +775,7 @@ def simple():
 
         with self.local_device_mesh(1, 1, backend_type):
             # This should be a valid return than an exception to raise
-            simple.call_on_shard_and_fetch().result()
+            call_on_shard_and_fetch(simple).result()
 
     def test_opaque_object(self, backend_type):
         with self.local_device_mesh(2, 2, backend_type):

python/tests/test_rust_backend.py

@@ -17,6 +17,7 @@
 import torch.utils._python_dispatch
 from monarch import fetch_shard, no_mesh, remote, Stream
 from monarch.common.device_mesh import DeviceMesh
+from monarch.common.remote import call_on_shard_and_fetch
 from monarch.rust_local_mesh import local_meshes, LoggingLocation, SocketType
 from torch.nn.attention import sdpa_kernel, SDPBackend
 from torch.nn.functional import scaled_dot_product_attention
@@ -111,9 +112,10 @@ def test_fetch_preprocess(self):
         with local_mesh():
             assert (
                 "an argument processed"
-                == remote("monarch.worker._testing_function.do_some_processing")
-                .call_on_shard_and_fetch("an argument")
-                .result()
+                == call_on_shard_and_fetch(
+                    remote("monarch.worker._testing_function.do_some_processing"),
+                    "an argument",
+                ).result()
             )
 
     def test_brutal_shutdown(self):
@@ -143,8 +145,8 @@ def has_nan(t):
                 return torch.isnan(t).any().item()
 
             t = torch.rand(3, 4)
-            res = has_nan.call_on_shard_and_fetch(
-                t, shard={"host": 0, "gpu": 0}
+            res = call_on_shard_and_fetch(
+                has_nan, t, shard={"host": 0, "gpu": 0}
             ).result()
 
         self.assertFalse(res)