Tensor communication

configs/examples/dummy.yaml

@@ -1,11 +1,13 @@
 stage_groups:
   - stage_names:
       - text_emb
+      - concat_text
     ranks:
       - 0
 
   - stage_names:
       - image_emb
+      - concat_img
     ranks:
       - 0
 

mminf/communication/communicator.py

@@ -0,0 +1,83 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+
+import zmq
+
+
+class BaseCommunicator(ABC):
+    @abstractmethod
+    def send(self, entity_id: str, msg):
+        """
+        entity_id: worker_xyz, conductor, or api_server
+        """
+        pass
+
+    @abstractmethod
+    def get_all_new_messages(self) -> list:
+        pass
+
+    @abstractmethod
+    def get_session_id(self) -> str:
+        pass
+
+
+class CommProtocol(Enum):
+    IPC = "IPC"
+    RDMA = "RDMA"
+
+
+class ZMQCommunicator(BaseCommunicator):
+    def __init__(
+        self,
+        my_id: str,
+        push_ids: list[str],
+        protocol: CommProtocol=CommProtocol.IPC,
+        ipc_socket_path_prefix: str="/tmp/mminf/",
+        # TODO: for TCP
+    ):
+        self.context = zmq.Context()
+        self.protocol = protocol
+        self.pull_socket = self.context.socket(zmq.PULL)
+
+        # TODO: maybe only open sockets as we need them, and close sockets
+        # when we no longer need them
+        self.push_sockets: dict[str, zmq.SyncSocket] = {}
+        self.session_id = f"ipc://{ipc_socket_path_prefix}/{my_id}.ipc"
+
+        if protocol == CommProtocol.IPC:
+            self.pull_socket.bind(f"ipc://{ipc_socket_path_prefix}/{my_id}.ipc")
+            self.pull_socket.setsockopt(zmq.LINGER, 0)
+        else:
+            raise NotImplementedError(f"Protocol {protocol} not yet supported yet")
+
+        for id in push_ids:
+            if id == my_id:
+                continue
+            self.push_sockets[id] = self.context.socket(zmq.PUSH)
+            self.push_sockets[id].connect(
+                f"ipc://{ipc_socket_path_prefix}/{id}.ipc"
+            )
+            self.push_sockets[id].setsockopt(zmq.LINGER, 0)
+
+    def get_session_id(self) -> str:
+        return self.session_id
+
+    def send(self, entity_id: str, msg):
+        # TODO: maybe serialize to JSON instead if more efficient
+        self.push_sockets[entity_id].send_pyobj(msg)
+
+    def get_all_new_messages(self) -> list:
+        messages = []
+        while True:
+            try:
+                # zmq.NOBLOCK means zmq doesn't wait for a new message to be
+                # available, it returns a message if it exists or raises an error 
+                # if no messages are available (error is caught below)
+                messages.append(self.pull_socket.recv_pyobj(
+                    flags=zmq.NOBLOCK
+                ))
+            except zmq.Again:
+                # zmq.Again actually means no messages left to read
+                break 
+        return messages
+

mminf/communication/tensors.py

@@ -0,0 +1,164 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+import gc
+from mooncake.engine import TransferEngine
+import torch
+
+
+from mminf.communication.communicator import BaseCommunicator, CommProtocol
+from mminf.graph.base import GraphPointer, TensorPointerInfo
+
+
+class TensorCommunicationManager(ABC):
+    @abstractmethod
+    def register_and_prepare_to_send(
+        self, request_id: str, tensors: dict[str, torch.Tensor],
+        graph_pointers: list[GraphPointer]
+    ):
+        """
+        Updates graph_pointers with required tensor info (addresses, datatypes,
+        num bytes, etc.).
+        If relevant (e.g., mooncake rdma), registers buffers.
+        """
+        pass
+
+    @abstractmethod
+    def get_tensor(self, request_id: str, tensor_name: str) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def cleanup(self, request_id: str, tensor_name: str):
+        """
+        Removes buffer if exists. Unregisters buffers if relevant
+        """
+        pass
+
+    @abstractmethod
+    def start_read_tensors(self, graph_pointers: list[GraphPointer]):
+        """
+        Initializes empty buffer, initializes a read. May return immediately. 
+        """
+        pass
+
+    @abstractmethod
+    def get_ready_tensors(self) -> dict[str, GraphPointer]:
+        """
+        Returns request_id: list of the GraphPointers that are currently
+        ready for that request
+        """
+        pass
+
+
+@dataclass
+class NameAndRequestId:
+    tensor_name: str
+    request_id: str
+
+
+@dataclass
+class EventAndPointers:
+    event: torch.cuda.Event
+    pointers: list[GraphPointer]
+
+
+class MooncakeCommunicationManager(TensorCommunicationManager):
+    def __init__(
+        self,
+        my_entity_id: str,
+        hostname: str,
+        communicator: BaseCommunicator,
+        protocol: CommProtocol=CommProtocol.RDMA,
+        metadata_server: str="P2PHANDSHAKE" # [ETCD_SERVER_URL, P2PHANDSHAKE, ...]
+
+    ):
+        self.my_entity_id = my_entity_id
+        self.tensors: dict[NameAndRequestId, torch.Tensor]
+        self.communicator = communicator
+        self.protocol = protocol
+        self.my_session_id = communicator.get_session_id()
+
+        self.engine = TransferEngine()
+        self.engine.initialize(
+            hostname,
+            metadata_server,
+            protocol,
+            ""
+        )        
+        # request_id: EventAndPointers
+        self.pending: dict[str, EventAndPointers] = {}
+
+    def register_and_prepare_to_send(
+        self, request_id: str, tensors: dict[str, torch.Tensor],
+        graph_pointers: list[GraphPointer]
+    ):
+        # get tensor name to graph pointers
+        name_to_pointers: dict[str, list[GraphPointer]] = {}
+        for pointer in graph_pointers:
+            if pointer.name not in name_to_pointers:
+                name_to_pointers[pointer.name] = []
+            name_to_pointers[pointer.name].append(pointer)
+
+        for name in tensors:
+            self.tensors[NameAndRequestId(
+                name, request_id=request_id
+            )] = tensors[name]
+
+            pointer_info = TensorPointerInfo(
+                dims=tensors[name].shape,
+                dtype=tensors[name].dtype,
+                nbytes=tensors[name].element_size() * tensors[name].nelement(),
+                address=tensors[name].data_ptr(),
+                source_session_id=self.my_session_id,
+                source_entity=self.my_entity_id
+            )
+            for pointer in name_to_pointers[name]:
+                pointer.tensor_info = pointer_info
+
+            if self.protocol == CommProtocol.RDMA:
+                ret_value = self.engine.register_memory(
+                    pointer_info.address, pointer_info.nbytes
+                )
+                if ret_value != 0:
+                    # TODO: error handling
+                    raise RuntimeError("Mooncake memory registration failed.")
+
+    def cleanup(self, request_id: str, tensor_name: str):
+        if self.protocol == CommProtocol.RDMA:
+            ret_value = self.engine.unregister_memory(
+                self.tensors[NameAndRequestId(
+                    tensor_name, request_id
+                )].data_ptr()
+            )
+            if ret_value != 0:
+                # TODO: error handling
+                raise RuntimeError("Mooncake memory unregistration failed.")
+
+        del self.tensors[NameAndRequestId(
+            tensor_name, request_id
+        )]
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def get_tensor(self, request_id: str, tensor_name: str) -> torch.Tensor:
+        return self.tensors[NameAndRequestId(
+            tensor_name, request_id
+        )]
+
+    def get_ready_tensors(self) -> dict[str, GraphPointer]:
+        """
+        Returns request_id: list of the GraphPointers that are currently
+        ready for that request (by querying the events in self.pending)
+        """
+        pass # TODO Atindra
+
+    def start_read_tensors(self, graph_pointers: list[GraphPointer]):
+        """
+        Initializes empty buffer, initializes a read. May return immediately. 
+
+        Kicks of the engine transfer read on cuda. Updates self.pending.
+        """
+        # This also probably needs to handle None pointer.tensor_info, in which
+        # case the tensor becomes ready immediately and is available in
+        # self.tensors. Maybe this function should also read the corresponding
+        # tensor names
+        pass # TODO Atindra