[6.5/n] Small refactor of _call_on_shard_and_fetch

python/monarch/common/client.py

@@ -41,6 +41,8 @@
 from monarch.common.borrows import Borrow, StorageAliases
 from monarch.common.controller_api import LogMessage, MessageResult, TController
 from monarch.common.device_mesh import DeviceMesh
+
+from monarch.common.future import Future
 from monarch.common.invocation import DeviceException, RemoteException, Seq
 from monarch.common.recording import flatten_messages, Recording
 
@@ -52,9 +54,6 @@
 
 from . import _coalescing
 
-if TYPE_CHECKING:
-    from monarch.common.future import Future
-
 
 logger = logging.getLogger(__name__)
 
@@ -447,6 +446,39 @@ def no_coalescing(self, reason):
     def mesh_state(self) -> WorldState:
         return self.inner.worker_world_state()
 
+    def fetch(
+        self,
+        mesh: "DeviceMesh",
+        stream: "StreamRef",
+        shard,
+        preprocess_message,
+        args,
+        kwargs,
+        defs: Tuple["Tensor", ...],
+        uses: Tuple["Tensor", ...],
+    ) -> "Future":
+        fut = Future(self)
+        ident = self.new_node(defs, uses, fut)
+        process = mesh._process(shard)
+        self.send(
+            process,
+            messages.SendValue(
+                ident,
+                None,
+                defs,
+                preprocess_message,
+                args,
+                kwargs,
+                stream,
+            ),
+        )
+        # we have to ask for status updates
+        # from workers to be sure they have finished
+        # enough work to count this future as finished,
+        # and all potential errors have been reported
+        self._request_status()
+        return fut
+
 
 def tree_map_refs(first_ref: int, tree):
     def translate_id(ref: int) -> int:

python/monarch/common/remote.py

@@ -21,6 +21,7 @@
     overload,
     Protocol,
     Tuple,
+    TYPE_CHECKING,
     TypeVar,
 )
 
@@ -30,6 +31,9 @@
 
 from monarch.common import _coalescing, device_mesh, messages, stream
 
+if TYPE_CHECKING:
+    from monarch.common.client import Client
+
 from monarch.common.device_mesh import RemoteProcessGroup
 from monarch.common.fake import fake_call
 
@@ -173,30 +177,19 @@ def _call_on_shard_and_fetch(
         propagator, rfunction, args, kwargs, ambient_mesh, stream._active
     )
 
-    client = mesh.client
+    client: "Client" = mesh.client
     if _coalescing.is_active(client):
         raise NotImplementedError("NYI: fetching results during a coalescing block")
-    fut = Future(client)
-    ident = client.new_node(mutates, dtensors, fut)
-    process = mesh._process(shard)
-    client.send(
-        process,
-        messages.SendValue(
-            ident,
-            None,
-            mutates,
-            preprocess_message,
-            args,
-            kwargs,
-            stream._active._to_ref(client),
-        ),
+    return client.fetch(
+        mesh,
+        stream._active._to_ref(client),
+        shard,
+        preprocess_message,
+        args,
+        kwargs,
+        mutates,
+        dtensors,
     )
-    # we have to ask for status updates
-    # from workers to be sure they have finished
-    # enough work to count this future as finished,
-    # and all potential errors have been reported
-    client._request_status()
-    return fut
 
 
 @remote