TCPBuffer (#753)

Summary:

X-link: #753

Adds a TCPBuffer implementation to tensor_engine, as a fallback example for RDMABuffer.

The main change from the original TCP-based buffer is that it uses ZMQ for communications to 1) decouple from Monarch's message passing system and 2) to mimic RDMABuffer's design, and showcase how other backends can be easily added

Differential Revision: D79588454

Author: allenwang28

python/monarch/_src/actor/proc_mesh.py

@@ -78,6 +78,9 @@
         _RdmaManager,
     )
 
+    # type: ignore[import]
+    from monarch._src.tensor_engine.tcp import TCPManager  # @manual
+
     # type: ignore[16]
     HAS_TENSOR_ENGINE = torch.cuda.is_available()
 except ImportError:
@@ -132,6 +135,9 @@ def __init__(
         self._slice = False
         # type: ignore[21]
         self._rdma_manager: Optional["_RdmaManager"] = None
+        # type: ignore[21]
+        self._tcp_manager: Optional["TCPManager"] = None
+
         self._debug_manager: Optional[DebugManager] = None
         self._code_sync_client: Optional[CodeSyncMeshClient] = None
         self._logging_mesh_client: Optional[LoggingMeshClient] = None
@@ -182,12 +188,20 @@ async def _init_manager_actors_coro(
             else None
         )
 
+        _tcp_manager = (
+            # type: ignore[16]
+            await self._spawn_nonblocking_on(proc_mesh, "tcp_manager", TCPManager)
+            if HAS_TENSOR_ENGINE
+            else None
+        )
+
         _debug_manager = await self._spawn_nonblocking_on(
             proc_mesh, _DEBUG_MANAGER_ACTOR_NAME, DebugManager, await _debug_client()
         )
 
         self._debug_manager = _debug_manager
         self._rdma_manager = _rdma_manager
+        self._tcp_manager = _tcp_manager
 
         if setup is not None:
             # If the user has passed the setup lambda, we need to call

python/monarch/_src/tensor_engine/rdma.py

@@ -101,6 +101,14 @@ def read_into(
 
         Returns an ActorFuture that can be awaited or called with .get() for blocking operation.
         """
+        try:
+            MonarchContext.get()
+        except LookupError:
+            raise RuntimeError(
+                "RDMABuffer.read_into() can only be called from within a Monarch actor context. "
+                "Make sure you're calling this from within an actor method."
+            )
+
         _assert_tensor_is_1d_contiguous_uint8(dst)
         dst_gpu = None
         if dst.device.type != "cpu":
@@ -148,6 +156,14 @@ def write_from(
 
         Returns an ActorFuture that can be awaited or called with .get() for blocking operation.
         """
+        try:
+            MonarchContext.get()
+        except LookupError:
+            raise RuntimeError(
+                "RDMABuffer.write_from() can only be called from within a Monarch actor context. "
+                "Make sure you're calling this from within an actor method."
+            )
+
         _assert_tensor_is_1d_contiguous_uint8(src)
         src_gpu = None
         if src.device.type != "cpu":