diff --git a/configs/examples/dummy.yaml b/configs/examples/dummy.yaml index a1bc684d9f..922bde5cb8 100644 --- a/configs/examples/dummy.yaml +++ b/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 diff --git a/mminf/communication/communicator.py b/mminf/communication/communicator.py new file mode 100644 index 0000000000..d47fbda4d1 --- /dev/null +++ b/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 + \ No newline at end of file diff --git a/mminf/communication/tensors.py b/mminf/communication/tensors.py new file mode 100644 index 0000000000..d15e7efb1a --- /dev/null +++ b/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 \ No newline at end of file diff --git a/mminf/conductor/dummy_conductor.py b/mminf/conductor/dummy_conductor.py index 1571c632fa..691fc7e4d8 100644 --- a/mminf/conductor/dummy_conductor.py +++ b/mminf/conductor/dummy_conductor.py @@ -1,33 +1,36 @@ +from copy import deepcopy from dataclasses import dataclass, field import time import numpy as np -import zmq +from mminf.communication.communicator import ZMQCommunicator +from mminf.graph.base import GraphPointer, TensorPointerInfo from mminf.ipc_formats import ( - ConductorMessage, ConductorMessageType, ConductorTensors, InputTensors, - NewRequest, NewRequestConductor, SubgraphsDone, WorkerMessage, + ConductorMessageType, InputSignals, + NewRequest, NewRequestConductor, PersistSignals, SubgraphsDone, WorkerMessage, WorkerMessageType ) -from mminf.model.base import CurrentForwardMetadata, ForwardPassInputs, Model, TensorData +from mminf.model.base import CurrentForwardMetadata, Model @dataclass class RequestData: current_forward_metadata: CurrentForwardMetadata - tensors: dict[str, TensorData] - new_outputs: dict[str, TensorData] # tensors passed back to conductor + fwd_inputs: list[GraphPointer] + persist_signals: dict[str, TensorPointerInfo] # signals passed back to conductor subgraph_to_worker: dict[str, str] + new_tokens: list[int] # for tracking progress all_subgraph_ids: set[str] + current_subgraph_ids: set[str] completed_subgraph_ids: set[str] = field(default_factory=set) # TODO: will need to add to this as we build things out - class DummyConductor: """ Initial in-progress conductor implementation. TODO: this is extremely @@ -39,8 +42,7 @@ def __init__( worker_ids: list[str], model: Model, model_config_file: str, - worker_socket_path_prefix: str="/tmp/mminf/workers/", - conductor_socket_path: str="/tmp/mminf/conductor.ipc" + socket_path_prefix: str="/tmp/mminf" ): self.requests: dict[str, RequestData] = {} self.worker_ids = worker_ids @@ -50,20 +52,13 @@ def __init__( subgraph.subgraph_id: subgraph \ for subgraph in model.get_subgraphs(model_config_file) } - # TODO: properly launch workers via Ray - - self.context = zmq.Context() - self.result_socket = self.context.socket(zmq.PULL) - self.result_socket.connect(f"ipc://{conductor_socket_path}") - self.result_socket.setsockopt(zmq.LINGER, 0) - - self.worker_sockets: dict[str, zmq.SyncSocket] = {} - for id in worker_ids: - self.worker_sockets[id] = self.context.socket(zmq.PUSH) - self.worker_sockets[id].connect( - f"ipc://{worker_socket_path_prefix}/{id}.ipc" - ) - self.worker_sockets[id].setsockopt(zmq.LINGER, 0) + # TODO: properly launch workers + + self.communicator = ZMQCommunicator( + my_id="conductor", + push_ids=worker_ids + ["api_server"], + ipc_socket_path_prefix=socket_path_prefix + ) def _assign_subgraphs_to_workers(self) -> dict[str, str]: """ @@ -79,30 +74,20 @@ def _assign_subgraphs_to_workers(self) -> dict[str, str]: def _split_inputs_to_workers( self, subgraph_to_worker: dict[str, str], - inputs: ForwardPassInputs - ) -> dict[str, ForwardPassInputs]: + inputs: list[GraphPointer] + ) -> dict[str, list[GraphPointer]]: """ Given the full ForwardPassInputs for kicking off a new forward pass, return a mapping of worker_id to the ForwardPassInputs that are routed to that worker. ForwardPassInputs consists of graph pointers and tensors. """ - inputs_per_worker: dict[str, ForwardPassInputs] = {} + inputs_per_worker: dict[str, list[GraphPointer]] = {} for subgraph_id, worker_id in subgraph_to_worker.items(): stages = set(self.subgraphs[subgraph_id].section.get_stage_names()) - pointers = {} - tensors = {} - - for signal, dests in inputs.pointers.items(): - dests = list(filter(lambda dest: dest.next_stage in stages, dests)) - if len(dests) == 0: - continue - pointers[signal] = dests - tensors[signal] = inputs.tensors[signal] - inputs_per_worker[worker_id] = ForwardPassInputs( - tensors=tensors, - pointers=pointers - ) - + + inputs_per_worker[worker_id] = [ + ptr for ptr in inputs if ptr.next_stage in stages + ] return inputs_per_worker @@ -121,15 +106,22 @@ def _ingest_request( body.initial_input_modalities, body.initial_output_modalities ), - tensors=body.initial_inputs, + fwd_inputs=[], + persist_signals=body.initial_signals, subgraph_to_worker=subgraph_to_worker, - all_subgraph_ids=set(subgraph_to_worker.keys()) + all_subgraph_ids=set(subgraph_to_worker.keys()), + current_subgraph_ids=set(), + new_tokens=[], ) self.requests[body.request_id] = request_data first_forward_inputs = self.model.get_forward_pass_inputs( - request_data.tensors, - request_data.current_forward_metadata + request_data.current_forward_metadata, + persist_signals=body.initial_signals + ) + self._set_current_subgraph_ids( + body.request_id, + request_data.current_forward_metadata.phase ) # send data to appropriate workers @@ -149,14 +141,23 @@ def _ingest_request( subgraph_ids=subgraph_ids, subgraph_to_worker=subgraph_to_worker, initial_phase=request_data.current_forward_metadata.phase, - initial_inputs=inputs_per_worker[worker].pointers, - initial_tensors=inputs_per_worker[worker].tensors + initial_inputs=inputs_per_worker[worker], + ) + self.communicator.send( + worker, WorkerMessage( + message_type=WorkerMessageType.NEW_REQUEST, + body=message + ) ) - self.worker_sockets[worker].send_pyobj(WorkerMessage( - message_type=WorkerMessageType.NEW_REQUEST, - body=message - )) + def _set_current_subgraph_ids( + self, request_id: str, phase: str + ): + self.requests[request_id].current_subgraph_ids = set([ + sg for sg in self.requests[request_id].all_subgraph_ids \ + if phase in self.subgraphs[sg].phases + ]) + def _process_subgraphs_done( self, body: SubgraphsDone ): @@ -173,22 +174,28 @@ def _process_subgraphs_done( body.subgraph_ids ) - done_with_forward = request_data.all_subgraph_ids.issubset( + done_with_forward = request_data.current_subgraph_ids.issubset( request_data.completed_subgraph_ids ) if done_with_forward: # start a new forward pass # TODO: look for EOS - self.model.update_for_next_forward( - metadata=request_data.current_forward_metadata, - input_tensors=request_data.tensors, - new_outputs=request_data.new_outputs + prev_forward_meta = deepcopy(request_data.current_forward_metadata) + request_data.current_forward_metadata = \ + self.model.update_for_next_forward( + metadata=request_data.current_forward_metadata, + new_tokens=request_data.new_tokens, + ) + self._set_current_subgraph_ids( + body.request_id, + request_data.current_forward_metadata.phase ) fwd_inputs = self.model.get_forward_pass_inputs( - input_tensors=request_data.tensors, - metadata=request_data.current_forward_metadata + metadata=request_data.current_forward_metadata, + persist_signals=request_data.persist_signals, + prev_forward_metadata=prev_forward_meta ) inputs_per_worker = self._split_inputs_to_workers( @@ -198,32 +205,39 @@ def _process_subgraphs_done( for worker, inputs in inputs_per_worker.items(): message = WorkerMessage( - message_type=WorkerMessageType.INPUT_TENSORS, - body=InputTensors( + message_type=WorkerMessageType.INPUT_SIGNALS, + body=InputSignals( request_id=body.request_id, phase=request_data.current_forward_metadata.phase, - inputs=inputs.pointers, - tensors=inputs.tensors + inputs=inputs, ) ) - self.worker_sockets[worker].send_pyobj(message) + self.communicator.send(worker, message) + + request_data.fwd_inputs = fwd_inputs + request_data.new_tokens = [] + request_data.completed_subgraph_ids = set() - def _process_new_tensors(self, body: ConductorTensors): + def _process_new_tensors(self, body: PersistSignals): """ If worker has sent tensors back to the conductor, process those. """ - self.requests[body.request_id].tensors.update(body.tensors) + self.requests[body.request_id].persist_signals.update( + body.signals + ) + self.requests[body.request_id].new_tokens.extend(body.new_tokens) + def run(self): while True: - message: ConductorMessage = self.result_socket.recv_pyobj() - if message.message_type == ConductorMessageType.NEW_REQUEST: - self._ingest_request(message.body) - elif message.message_type == ConductorMessageType.SUBGRAPHS_DONE: - self._process_subgraphs_done(message.body) - elif message.message_type == ConductorMessageType.TENSORS: - self._process_new_tensors(message.body) - else: - raise ValueError(f"Unknown message type: {message.message_type}") + for message in self.communicator.get_all_new_messages(): + if message.message_type == ConductorMessageType.NEW_REQUEST: + self._ingest_request(message.body) + elif message.message_type == ConductorMessageType.SUBGRAPHS_DONE: + self._process_subgraphs_done(message.body) + elif message.message_type == ConductorMessageType.PERSIST_SIGNALS: + self._process_new_tensors(message.body) + else: + raise ValueError(f"Unknown message type: {message.message_type}") time.sleep(0.1) # just for dummy conductor! \ No newline at end of file diff --git a/mminf/graph/base.py b/mminf/graph/base.py index 559560744b..2b751af784 100644 --- a/mminf/graph/base.py +++ b/mminf/graph/base.py @@ -11,49 +11,41 @@ def update_list_dicts(signals: dict[str, list], new_signals: dict[str, list]): signals[key] = new_signals[key] +@dataclass +class TensorPointerInfo: + dims: list[int] + dtype: str + nbytes: int + address: int + source_session_id: str # e.g., f"{HOSTNAME}:{client_engine.get_rpc_port()}" + source_entity: str # which {worker, api_server} the tensor is on + + @dataclass class GraphPointer: next_stage: str - back_to_conductor: bool = field(default=False) - + name: str + tensor_info: TensorPointerInfo | None = field(default=None) -# Three different ways of defining graph edges -SignalToDests = dict[str, list[str]] -SignalToDestsAndFlags = dict[str, list[GraphPointer]] -DestToInputs = dict[str, list[str]] + # Flags + back_to_conductor: bool = field(default=False) + is_new_token: bool = field(default=False) -def remove_flags( - new_inputs: SignalToDestsAndFlags -) -> SignalToDests: - return { - id: [ptr.next_stage for ptr in new_inputs[id]] \ - for id in new_inputs - } +# Two different ways of defining graph edges +DestToGraphPointers = dict[str, list[GraphPointer]] def get_stage_to_inputs_mapping( - new_inputs: SignalToDests -)-> DestToInputs: - result: DestToInputs = {} - for input_id, ptrs in new_inputs.items(): - for ptr in ptrs: - if ptr not in result: - result[ptr] = [] - result[ptr].append(input_id) + new_inputs: list[GraphPointer] +)-> DestToGraphPointers: + result: DestToGraphPointers = {} + for ptr in new_inputs: + if ptr.next_stage not in result: + result[ptr.next_stage] = [] + result[ptr.next_stage].append(ptr) return result -def get_signal_to_dest_mapping( - dest_to_inputs: DestToInputs -) -> SignalToDests: - result: SignalToDests = {} - for dest, inputs in dest_to_inputs.items(): - for input in inputs: - if input not in result: - result[input] = [] - result[input].append(dest) - return result - @dataclass class GraphSection(ABC): @abstractmethod @@ -61,21 +53,21 @@ def get_stage_names(self) -> list[str]: pass @abstractmethod - def get_inputs(self) -> SignalToDests: + def get_inputs(self) -> list[GraphPointer]: """ All external or "loop-back" inputs into a subgraph """ pass @abstractmethod - def get_outputs(self) -> SignalToDests: + def get_outputs(self) -> list[GraphPointer]: """ All external or "loop-back" outputs from a subgraph """ pass @abstractmethod - def ingest_inputs(self, stage_to_inputs: DestToInputs) -> set[str]: + def ingest_inputs(self, stage_to_inputs: DestToGraphPointers) -> set[str]: """ Adds inputs to the appropriate "ready_input_ids", and **mutates stage_to_inputs** to remove the inputs that were able to be added @@ -98,7 +90,7 @@ def split_off_ready(self) -> tuple[list["GraphStage"], "GraphSection"]: class GraphStage(GraphSection): name: str input_ids: set[str] - outputs: SignalToDestsAndFlags + outputs: list[GraphPointer] consumes_stream: bool = field(default=False) # populated as previous stages complete @@ -114,30 +106,28 @@ def is_ready(self): def get_stage_names(self) -> list[str]: return [self.name] - def get_inputs(self) -> SignalToDests: - return { - inp: [self.name] for inp in self.input_ids - } - - def get_outputs(self) -> SignalToDests: - return remove_flags(self.outputs) + def get_inputs(self) -> list[GraphPointer]: + return [ + GraphPointer(next_stage=self.name, name=id) for id in self.input_ids + ] + + def get_outputs(self) -> list[GraphPointer]: + return self.outputs - def ingest_inputs(self, stage_to_inputs: DestToInputs): + def ingest_inputs(self, stage_to_inputs: DestToGraphPointers): if self.name not in stage_to_inputs: return set() ingested_ids = set([ - id for id in stage_to_inputs[self.name] \ - if id in self.input_ids and id not in self.ready_input_ids + ptr.name for ptr in stage_to_inputs[self.name] \ + if ptr.name in self.input_ids and ptr.name not in self.ready_input_ids ]) # ingest ids that this stage takes in, and are not already ready self.ready_input_ids.update(ingested_ids) - # if ingested_ids: - # print(f"Stage {self.name} ingesting inputs {ingested_ids}") - # remove the ingested ids from the stage_to_inputs stage_to_inputs[self.name] = [ - id for id in stage_to_inputs[self.name] if id not in ingested_ids + ptr for ptr in stage_to_inputs[self.name] \ + if ptr.name not in ingested_ids ] if len(stage_to_inputs[self.name]) == 0: del stage_to_inputs[self.name] @@ -161,27 +151,33 @@ def get_stage_names(self) -> list[str]: def _get_inputs_outputs(self): # In the case that this section is part of a loop, "loop-back" # variables are included in the list of inputs and outputs - inputs: SignalToDests = {} - outputs: SignalToDests = {} + inputs: list[GraphPointer] = [] + outputs: list[GraphPointer] = [] + output_names = set() for s in self.sections: stage_names = s.get_stage_names() - new_inputs = s.get_inputs() - for i in new_inputs: - # filters out the internal signals from the input list - if i not in outputs: - inputs[i] = new_inputs[i] - else: - outputs[i] = [ptr for ptr in outputs[i] if ptr not in stage_names] - outputs.update(s.get_outputs()) + inputs_of_s = s.get_inputs() + + # filter out internal signals from the new inputs + inputs += ([inp for inp in inputs_of_s if inp.name not in output_names]) + # filter internal output signals from the output list + outputs = [ptr for ptr in outputs if not ( + ptr.name not in inputs_of_s and ptr.next_stage in stage_names + )] + + # add new outputs to the output list + outputs_of_s = s.get_outputs() + outputs += outputs_of_s + output_names.update([ptr.name for ptr in outputs_of_s]) return inputs, outputs - def get_inputs(self): + def get_inputs(self) -> list[GraphPointer]: return self._get_inputs_outputs()[0] - def get_outputs(self): + def get_outputs(self) -> list[GraphPointer]: return self._get_inputs_outputs()[1] - def ingest_inputs(self, stage_to_inputs: DestToInputs): + def ingest_inputs(self, stage_to_inputs: DestToGraphPointers): ingested = set() for s in self.sections: ingested.update(s.ingest_inputs(stage_to_inputs)) @@ -210,18 +206,16 @@ def get_stage_names(self) -> list[str]: ], start=[]) def get_inputs(self): - inputs = {} - for s in self.sections: - update_list_dicts(inputs, s.get_inputs()) - return inputs + return sum( + [s.get_inputs() for s in self.sections], start=[] + ) def get_outputs(self): - outputs = {} - for s in self.sections: - update_list_dicts(outputs, s.get_outputs()) - return outputs + return sum( + [s.get_outputs() for s in self.sections], start=[] + ) - def ingest_inputs(self, stage_to_inputs: DestToInputs): + def ingest_inputs(self, stage_to_inputs: DestToGraphPointers): ingested = set() for s in self.sections: ingested.update(s.ingest_inputs(stage_to_inputs)) @@ -246,23 +240,23 @@ class Loop(GraphSection): section: GraphSection # this is used to populate next_section and # in-progress section; it remains "clean" n_iters: int - outputs: SignalToDestsAndFlags + outputs: list[GraphPointer] curr_iter: int = field(default=0) - external_inputs: SignalToDests = field(default=None) - loop_back_signals: SignalToDests = field(default=None) + external_inputs: list[GraphPointer] = field(default=None) + loop_back_signals: list[GraphPointer] = field(default=None) curr_iter_section: GraphSection = field(default=None) nxt_iter_section: GraphSection = field(default=None) - def get_outputs(self) -> SignalToDests: - return remove_flags(self.outputs) + def get_outputs(self): + return self.outputs - def get_inputs(self) -> SignalToDests: + def get_inputs(self): return self.section.get_inputs() def get_stage_names(self): return self.section.get_stage_names() - def ingest_inputs(self, stage_to_inputs: DestToInputs): + def ingest_inputs(self, stage_to_inputs: DestToGraphPointers): # Populate the current iteration first, then populate the next iteration # if there are any leftover inputs (which would signal either inputs that # are not for this section, or loop-back inputs) @@ -275,9 +269,13 @@ def ingest_inputs(self, stage_to_inputs: DestToInputs): # we should only be populating the nxt_iter_section with loop-back inputs, # so exclude external inputs from populating nxt_iter_section. This logic # is required to make nested loops work. + ext_inp_name_dest = [ + (ptr.name, ptr.next_stage) for ptr in self.external_inputs + ] external_inputs = { - dest: [i for i in inputs if dest in self.external_inputs.get(i, [])] \ - for dest, inputs in stage_to_inputs.items() + dest: [ + ptr for ptr in inputs if (ptr.name, dest) in ext_inp_name_dest + ] for dest, inputs in stage_to_inputs.items() } for dest in stage_to_inputs: stage_to_inputs[dest] = [ @@ -290,26 +288,30 @@ def ingest_inputs(self, stage_to_inputs: DestToInputs): ) update_list_dicts(stage_to_inputs, external_inputs) return ingested - + def _get_external_inputs(self): inputs = self.section.get_inputs() internal_outputs = self.section.get_outputs() + output_names_dests = set([(ptr.name, ptr.next_stage) for ptr in internal_outputs]) # compute "external inputs", i.e., ones that don't come from looping # back, and make sure those are populated for the next loop iter - return { - i: inputs[i] for i in inputs if i not in internal_outputs - } + return [ + inp for inp in inputs if (inp.name, inp.next_stage) not in output_names_dests + ] - def _get_loop_back_signals(self) -> SignalToDests: + def _get_loop_back_signals(self) -> list[GraphPointer]: # these inputs and outputs only include external and loop-back signals; # they do not include signals that are purely internal to the section inputs = self.section.get_inputs() + input_names_dests = [ + (ptr.name, ptr.next_stage) for ptr in inputs + ] outputs = self.section.get_outputs() - return { - i: inputs[i] for i in inputs if i in outputs - } + return [ + ptr for ptr in outputs if (ptr.name, ptr.next_stage) in input_names_dests + ] def _replace_outputs_for_final_iter( self, section: GraphSection, @@ -319,18 +321,25 @@ def _replace_outputs_for_final_iter( from the graph, and (2) add in self.outputs where appropriate """ loop_back_signals = self.loop_back_signals - def replace_outputs(stage_outputs: SignalToDestsAndFlags): - for output in stage_outputs: - # remove loop backs - stage_outputs[output] = [ - pointer for pointer in stage_outputs[output] \ - if pointer.next_stage not in loop_back_signals.get(output, []) - ] - if output in self.outputs: - # add in final output - stage_outputs[output].extend(self.outputs[output]) + loop_back_name_dests = set([ + (ptr.name, ptr.next_stage) for ptr in loop_back_signals + ]) + full_outputs = self.outputs + + def replace_outputs(stage_outputs: list[GraphPointer]): + stage_output_names = set([ + ptr.name for ptr in stage_outputs + ]) + outputs_to_add = [ + ptr for ptr in full_outputs if ptr.name in stage_output_names + ] + + return outputs_to_add + [ + ptr for ptr in stage_outputs \ + if (ptr.name, ptr.next_stage) not in loop_back_name_dests + ] if isinstance(section, GraphStage) or isinstance(section, Loop): - replace_outputs(section.outputs) + section.outputs = replace_outputs(section.outputs) elif isinstance(section, Sequential) or isinstance(section, Parallel): for sec in section.sections: self._replace_outputs_for_final_iter(sec) diff --git a/mminf/graph/request_queues.py b/mminf/graph/request_queues.py index e2ba3366a3..809fe814c7 100644 --- a/mminf/graph/request_queues.py +++ b/mminf/graph/request_queues.py @@ -3,14 +3,14 @@ from dataclasses import dataclass, field from mminf.graph.base import ( - GraphSection, GraphStage, SignalToDests, SignalToDestsAndFlags, - get_signal_to_dest_mapping, get_stage_to_inputs_mapping, remove_flags + DestToGraphPointers, GraphPointer, GraphSection, GraphStage, + get_stage_to_inputs_mapping ) @dataclass class ProcessedInputs: routed_to_this_subgraph: set[str] - for_other_subgraphs: SignalToDests + for_other_subgraphs: list[GraphPointer] @dataclass @@ -37,7 +37,7 @@ def _update_ready_waiting(self): def process_new_inputs( self, - new_inputs: SignalToDests + new_inputs: list[GraphPointer] ) -> ProcessedInputs: """ Processes all outputs that feed into the waiting graph section, and @@ -47,13 +47,15 @@ def process_new_inputs( if self.waiting is None: return new_inputs - new_inputs = get_stage_to_inputs_mapping(new_inputs) + new_inputs: DestToGraphPointers = get_stage_to_inputs_mapping(new_inputs) ingested = self.waiting.ingest_inputs(new_inputs) - external_outputs = new_inputs + external_outputs = sum( + new_inputs.values(), start=[] + ) self._update_ready_waiting() return ProcessedInputs( - for_other_subgraphs=get_signal_to_dest_mapping(external_outputs), + for_other_subgraphs=external_outputs, routed_to_this_subgraph=ingested ) diff --git a/mminf/ipc_formats.py b/mminf/ipc_formats.py index 5bdb2ee86f..9dda608cf6 100644 --- a/mminf/ipc_formats.py +++ b/mminf/ipc_formats.py @@ -2,8 +2,10 @@ from dataclasses import asdict, dataclass from enum import Enum -from mminf.graph.base import SignalToDests, SignalToDestsAndFlags -from mminf.model.base import TensorData +from attrs import field + +from mminf.graph.base import GraphPointer, TensorPointerInfo + class Status(Enum): WAITING = "waiting" @@ -28,7 +30,8 @@ def from_dict(self, input: dict): class WorkerMessageType(Enum): NEW_REQUEST = "new_request" REMOVE_REQUEST = "remove_request" - INPUT_TENSORS = "input_tensors" + INPUT_SIGNALS = "input_signals" + TENSOR_RECEIVED = "tensor_received" @dataclass @@ -37,8 +40,8 @@ class NewRequest(MessageBody): subgraph_ids: list[str] subgraph_to_worker: dict[str, str] initial_phase: str - initial_inputs: SignalToDestsAndFlags - initial_tensors: dict[str, TensorData] + initial_inputs: list[GraphPointer] + @dataclass class RemoveRequest(MessageBody): @@ -46,11 +49,18 @@ class RemoveRequest(MessageBody): @dataclass -class InputTensors(MessageBody): +class InputSignals(MessageBody): request_id: str phase: str - inputs: SignalToDests - tensors: dict[str, TensorData] + inputs: list[GraphPointer] + + +@dataclass +class TensorReceived(MessageBody): + request_id: str + receiving_entity: str + successful_tensor_ids: list[str] + failed_tensor_ids: list[str] @dataclass @@ -65,28 +75,29 @@ class WorkerMessage: class ConductorMessageType(Enum): NEW_REQUEST = "new_request" - TENSORS = "tensors" + PERSIST_SIGNALS = "persist_signal" SUBGRAPHS_DONE = "subgraphs_done" @dataclass class NewRequestConductor(MessageBody): request_id: str - initial_inputs: dict[str, TensorData] + initial_signals: dict[str, TensorPointerInfo] initial_input_modalities: list[str] initial_output_modalities: list[str] @dataclass -class ConductorTensors(MessageBody): +class SubgraphsDone(MessageBody): request_id: str - pointers: SignalToDestsAndFlags - tensors: dict[str, TensorData] + subgraph_ids: list[str] + @dataclass -class SubgraphsDone(MessageBody): +class PersistSignals(MessageBody): request_id: str - subgraph_ids: list[str] + signals: dict[str, TensorPointerInfo] + new_tokens: list[int] = field(default=None) @dataclass diff --git a/mminf/model/base.py b/mminf/model/base.py index 5f3589ca31..a689d12fb1 100644 --- a/mminf/model/base.py +++ b/mminf/model/base.py @@ -8,7 +8,7 @@ import torch import yaml -from mminf.graph.base import GraphSection, GraphStage, Loop, Parallel, Sequential, SignalToDestsAndFlags +from mminf.graph.base import GraphPointer, GraphSection, GraphStage, Loop, Parallel, Sequential, TensorPointerInfo STREAM_OUT = "stream_out" @@ -24,15 +24,6 @@ class Subgraph: subgraph_id: str = field(default_factory=lambda: str(uuid4())) -@dataclass -class TensorData: - tensor: torch.Tensor | None # this should be replaced by a Ray reference - - # list of segment boundaries (e.g., [(0, 10), (50, 100)] means tokens - # 0 (inclusive) to 10 (exclusive) and 50 to 100. - token_ranges: list[tuple[int, int]] - - def _combine_sections_sequential_or_parallel( section: GraphSection, other: GraphSection, comb_type: Type[Sequential] | Type[Parallel] @@ -158,16 +149,6 @@ class CurrentForwardMetadata: kwargs: dict = field(default_factory=dict) -@dataclass -class ForwardPassInputs: - """ - Inputs for the current full-model forward pass, with where each input - is routed to. - """ - tensors: dict[str, TensorData] - pointers: SignalToDestsAndFlags - - class Model(ABC): def _get_subgraphs_for_phase( self, phase_name: str, graph: GraphSection, @@ -215,9 +196,10 @@ def get_initial_forward_metadata( @abstractmethod def get_forward_pass_inputs( - self, input_tensors: dict[str, TensorData], - metadata: CurrentForwardMetadata, - ) -> ForwardPassInputs: + self, metadata: CurrentForwardMetadata, + persist_signals: dict[str, TensorPointerInfo], + prev_forward_metadata: CurrentForwardMetadata=None, + ) -> list[GraphPointer]: """ Called by the conductor. @@ -230,26 +212,25 @@ def get_forward_pass_inputs( @abstractmethod def update_for_next_forward( self, metadata: CurrentForwardMetadata, - input_tensors: dict[str, TensorData], - new_outputs: dict[str, TensorData] - ) -> tuple[CurrentForwardMetadata, dict[str, TensorData]]: + new_tokens: list[int], + ) -> CurrentForwardMetadata: """ Called by the conductor at the end of a full model fwd pass. """ # e.g., check for BOI token, check if image was generated and should # be added to the input modalities and input tensors, adds new token # to the input text, etc... - # Returns the metadata and input tensors for the new forward pass + # Returns the metadata for the new forward pass pass @abstractmethod def step( self, stage_name: str, phase: str, - input_tensors: dict[str, TensorData], + input_tensors: dict[str, torch.Tensor], state, # TODO: figure out state **kwargs - ) -> dict[str, TensorData]: + ) -> dict[str, torch.Tensor]: """ Called by the worker to execute a stage """ diff --git a/mminf/model/dummy_model.py b/mminf/model/dummy_model.py index 5ecf5c4652..09a68f0f32 100644 --- a/mminf/model/dummy_model.py +++ b/mminf/model/dummy_model.py @@ -1,9 +1,10 @@ from copy import deepcopy import numpy as np +import torch -from mminf.graph.base import GraphPointer, GraphStage, Loop, Parallel, Sequential, SignalToDestsAndFlags -from mminf.model.base import STREAM_OUT, CurrentForwardMetadata, ForwardPassInputs, Model, TensorData +from mminf.graph.base import GraphPointer, GraphStage, Loop, Parallel, Sequential, TensorPointerInfo +from mminf.model.base import STREAM_OUT, CurrentForwardMetadata, Model class DummyModel(Model): @@ -11,22 +12,40 @@ class DummyModel(Model): Show-o2-inspired model that does nothing, for testing and example purposes. """ def _get_text_emb(self): - return GraphStage( - name="text_emb", - input_ids=["text"], - outputs={ - "text_emb": "LLM" - } - ) + return Sequential([ + GraphStage( + name="text_emb", + input_ids=["text"], + outputs=[ + GraphPointer(next_stage="concat_text", name="new_text_emb") + ] + ), + GraphStage( + name="concat_text", + input_ids=["new_text_emb", "existing_text_emb"], + outputs=[ + GraphPointer(next_stage="LLM", name="text_emb", back_to_conductor=True) + ] + ) + ]) def _get_img_emb(self): - return GraphStage( - name="image_emb", - input_ids=["images"], - outputs={ - "img_emb": "LLM" - } - ) + return Sequential([ + GraphStage( + name="image_emb", + input_ids=["images"], + outputs=[ + GraphPointer(next_stage="concat_img", name="new_image_emb") + ] + ), + GraphStage( + name="concat_img", + input_ids=["new_image_emb", "existing_image_emb"], + outputs=[ + GraphPointer(next_stage="LLM", name="img_emb", back_to_conductor=True) + ] + ) + ]) def get_phase_graphs(self): prefill = Sequential([ @@ -34,9 +53,13 @@ def get_phase_graphs(self): GraphStage( name="LLM", input_ids=["text_emb", "img_emb"], - outputs={ - "new_token": GraphPointer(STREAM_OUT, back_to_conductor=True) - } + outputs=[ + GraphPointer( + next_stage=STREAM_OUT, + name="new_token", + is_new_token=True + ) + ] ) ]) decode = deepcopy(prefill) @@ -47,29 +70,33 @@ def get_phase_graphs(self): GraphStage( name="LLM", input_ids=["text_emb", "img_emb", "latents"], - outputs={ - "hidden_states": GraphPointer("flow") - } + outputs=[ + GraphPointer(next_stage="flow", name="hidden_states") + ] ), GraphStage( "flow", input_ids=["hidden_states"], - outputs={ - "latents": GraphPointer("LLM") - } + outputs=[ + GraphPointer(next_stage="LLM", name="latents") + ] ) ]), n_iters=10, - outputs={ - "latents": GraphPointer("VAE_dec") - } + outputs=[ + GraphPointer(next_stage="VAE_dec", name="latents") + ] ), GraphStage( name="VAE_dec", input_ids=["latents"], - outputs={ - "image_output": GraphPointer(STREAM_OUT) - } + outputs=[ + GraphPointer( + next_stage=STREAM_OUT, + name="image_output", + back_to_conductor=True + ) + ] ) ]) @@ -90,34 +117,55 @@ def get_initial_forward_metadata( ) def get_forward_pass_inputs( - self, input_tensors: dict[str, TensorData], - metadata: CurrentForwardMetadata, - ) -> ForwardPassInputs: - pointers = { - "images": [GraphPointer("image_emb")], - "text": [GraphPointer("text_emb")] - } - if "images" not in input_tensors: # add null image for routing purposes - input_tensors["images"] = TensorData(tensor=None, token_ranges=[]) - if "text" not in input_tensors: # add null text for routing purposes - input_tensors["text"] = TensorData(tensor=None, token_ranges=[]) + self, metadata: CurrentForwardMetadata, + persist_signals: dict[str, TensorPointerInfo], + prev_forward_metadata: CurrentForwardMetadata=None, + ) -> list[GraphPointer]: + text_inp = GraphPointer( + next_stage="text_emb", + name="text", + ) + img_inp = GraphPointer( + next_stage="image_emb", + name="images", + ) + existing_text = GraphPointer( + next_stage="concat_text", + name="existing_text_emb", + ) + existing_img = GraphPointer( + next_stage="concat_img", + name="existing_image_emb", + ) + + pointers = [ + text_inp, img_inp, existing_text, existing_img + ] + + if metadata.is_prefill: # first forward + text_inp.tensor_info = persist_signals.get("text", None) + img_inp.tensor_info = persist_signals.get("images", None) + else: + existing_text.tensor_info = persist_signals.get("text_emb", None) + existing_img.tensor_info = persist_signals.get("img_emb", None) + if prev_forward_metadata.phase == "image_gen": + img_inp.tensor_info = persist_signals.get("image_output", None) + text_inp.tensor_info = persist_signals.get("new_token", None) - if metadata.phase == "image_gen": - # maybe this should be a random tensor in real life, or the latents - # will be initialized on the worker level... - input_tensors["latents"] = TensorData( tensor=None, token_ranges=[]) - pointers["latents"] = [GraphPointer("LLM")] + if metadata.phase == "image_gen": + pointers.append( + GraphPointer( + next_stage="LLM", + name="latents", + tensor_info=persist_signals.get("latents", None) + ) + ) + return pointers - return ForwardPassInputs( - tensors=input_tensors, - pointers=pointers - ) - def update_for_next_forward( self, metadata: CurrentForwardMetadata, - input_tensors: dict[str, TensorData], - new_outputs: dict[str, TensorData] - ): + new_tokens: list[int], + ) -> CurrentForwardMetadata: # dummy model doesn't actually do anything, so this function will just # randomly select a phase metadata.phase = str(np.random.choice(["decode", "image_gen"])) @@ -125,12 +173,12 @@ def update_for_next_forward( metadata.output_modalities = ["text"] else: metadata.output_modalities = ["image"] - return metadata, input_tensors + return metadata def step( self, stage_name: str, phase: str, - input_tensors: dict[str, TensorData], + input_tensors: dict[str, torch.Tensor], state, # TODO: figure out state **kwargs ): diff --git a/mminf/worker/dummy_worker.py b/mminf/worker/dummy_worker.py index da67375895..0f1d31f67c 100644 --- a/mminf/worker/dummy_worker.py +++ b/mminf/worker/dummy_worker.py @@ -1,13 +1,14 @@ from copy import deepcopy from dataclasses import dataclass, field import time -import zmq -from mminf.model.base import Subgraph, TensorData -from mminf.graph.base import GraphPointer, GraphStage, SignalToDests, SignalToDestsAndFlags, remove_flags +from mminf.communication.communicator import CommProtocol, ZMQCommunicator +from mminf.communication.tensors import MooncakeCommunicationManager +from mminf.model.base import Subgraph +from mminf.graph.base import GraphPointer, GraphStage from mminf.graph.request_queues import PerRequestStageQueues, ProcessedInputs from mminf.ipc_formats import ( - ConductorMessage, ConductorMessageType, ConductorTensors, InputTensors, + ConductorMessage, ConductorMessageType, InputSignals, NewRequest, RemoveRequest, SubgraphsDone, WorkerMessage, WorkerMessageType ) @@ -15,8 +16,7 @@ @dataclass(frozen=True) class StageAndPhase: """ - Syntactic sugar for a tuple of stage name and phase, e.g., (LLM, decode) - or (flow, image_gen) + Tuple of stage name and phase, e.g., (LLM, decode) or (flow, image_gen) """ stage: str phase: str @@ -25,8 +25,8 @@ class StageAndPhase: @dataclass class StageOutputRouting: routed_to_this_subgraph: set[str] # set of tensor ids - to_conductor: SignalToDestsAndFlags # outputs that are going back to the conductor - to_workers: dict[str, SignalToDests] # worker id to signals + to_conductor: list[GraphPointer] # outputs that are going back to the conductor + to_workers: dict[str, list[GraphPointer]] # worker id to signals completed_subgraphs: list[str] = field(default_factory=[]) # list of subgraph IDs @@ -42,7 +42,7 @@ class SubgraphQueues: subgraph: Subgraph per_request_queues: dict[str, PerRequestStageQueues] # request_id -> queue - def process_new_inputs(self, request_id: str, inputs: SignalToDests) -> ProcessedInputs: + def process_new_inputs(self, request_id: str, inputs: list[GraphPointer]) -> ProcessedInputs: """ Add new inputs for a request, and update waiting/ready stages accordingly. Returns any signals that should be sent to other subgraphs. @@ -122,9 +122,9 @@ class PerRequestInfo: stage_to_worker: dict[StageAndPhase, str] # for all stages subgraph_ids: list[str] # for this worker current_phase: str = field(default=None) + # phase_subgraph_ids = subgraphs for the current phase phase_subgraph_ids: list[str] = field(default_factory=list) # for this worker - tensors: dict[str, TensorData] = field(default_factory=dict) @dataclass @@ -156,7 +156,7 @@ def get_phase(self, request_id: str): def process_new_inputs( self, request_id: str, - inputs: SignalToDests + inputs: list[GraphPointer] ): """ Updates queues with new inputs for a request @@ -168,7 +168,7 @@ def process_new_inputs( def process_stage_outputs( self, request_id: str, - outputs: SignalToDestsAndFlags + outputs: list[GraphPointer] ) -> StageOutputRouting: """ After a stage has finished processing, use its outputs to update @@ -180,13 +180,9 @@ def process_stage_outputs( (different) worker. """ # (1) find back_to_conductor flags - to_conductor = { - signal: [dest for dest in dests if dest.back_to_conductor] \ - for signal, dests in outputs.items() - } + to_conductor = [ptr for ptr in outputs if ptr.back_to_conductor] # (2) process all internal-facing outputs - outputs = remove_flags(outputs) subgraph_ids = self.per_request_info[request_id].phase_subgraph_ids completed_subgraphs = [] @@ -206,24 +202,14 @@ def process_stage_outputs( # in different workers) # (3) get mapping of worker to external outputs - to_workers: dict[str, SignalToDests] = {} - for signal, dests in outputs.items(): - # to_workers_update is what we're going to add to the to_workers dit - # for the outputs from the current loop - to_workers_update: dict[str, list[GraphPointer]] = {} # worker: [graph_pointer] - for dest in dests: - worker_id = self.per_request_info[request_id].stage_to_worker[StageAndPhase( - stage=dest, phase=self.get_phase(request_id) - )] - if worker_id not in to_workers_update: - to_workers_update[worker_id] = [] - to_workers_update[worker_id].append(dest) - - # update the to_workers dict with results from to_workers_update - for worker_id, pointers in to_workers_update.items(): - if worker_id not in to_workers: - to_workers[worker_id] = {} - to_workers[worker_id][signal] = pointers + to_workers: dict[str, list[GraphPointer]] = {} + for ptr in outputs: + worker_id = self.per_request_info[request_id].stage_to_worker[StageAndPhase( + stage=ptr.next_stage, phase=self.get_phase(request_id) + )] + if worker_id not in to_workers: + to_workers[worker_id] = [] + to_workers[worker_id].append(ptr) return StageOutputRouting( routed_to_this_subgraph=routed_to_this_subgraph, @@ -274,8 +260,9 @@ def __init__( my_subgraphs: list[Subgraph], all_subgraph_ids_to_phases: dict[str, set[str]], # for all subgraphs all_subgraph_ids_to_stages: dict[str, list[str]], # for all subgraphs - worker_socket_path_prefix: str="/tmp/mminf/workers/", - conductor_socket_path: str="/tmp/mminf/conductor.ipc" + hostname: str="localhost", # TODO: figure this out + socket_path_prefix: str="/tmp/mminf", + tensor_comm_protocol=CommProtocol.RDMA, ): """ Initial in-progress worker implementation. This worker cannnot actually @@ -297,25 +284,17 @@ def __init__( all_subgraph_ids_to_stages=all_subgraph_ids_to_stages ) - # set up sockets between the other workers - self.context = zmq.Context() - self.request_socket = self.context.socket(zmq.PULL) - self.request_socket.bind(f"ipc://{worker_socket_path_prefix}/{worker_id}.ipc") - self.request_socket.setsockopt(zmq.LINGER, 0) - - self.result_socket = self.context.socket(zmq.PUSH) - self.result_socket.connect(f"ipc://{conductor_socket_path}") - self.result_socket.setsockopt(zmq.LINGER, 0) - - self.inter_worker_sockets: dict[str, zmq.SyncSocket] = {} - for id in worker_ids: - if id == worker_id: - continue - self.inter_worker_sockets[id] = self.context.socket(zmq.PUSH) - self.inter_worker_sockets[id].connect( - f"ipc://{worker_socket_path_prefix}/{id}.ipc" - ) - self.inter_worker_sockets[id].setsockopt(zmq.LINGER, 0) + self.communicator = ZMQCommunicator( + my_id=worker_id, + push_ids=worker_ids + ["conductor", "api_server"], + ipc_socket_path_prefix=socket_path_prefix + ) + self.tensor_manager = MooncakeCommunicationManager( + my_entity_id=worker_id, + hostname=hostname, + communicator=self.communicator, + protocol=tensor_comm_protocol, + ) def _add_new_request( self, body: NewRequest @@ -328,14 +307,15 @@ def _add_new_request( subgraph_ids=body.subgraph_ids, subgraph_to_worker_id=body.subgraph_to_worker ) - self.subgraphs_manager.per_request_info[ - body.request_id].tensors.update(body.initial_tensors) + + # TODO Atindra: start reading in tensors from body.initial_inputs + self.subgraphs_manager.update_phase( body.request_id, body.initial_phase ) self.subgraphs_manager.process_new_inputs( request_id=body.request_id, - inputs=remove_flags(body.initial_inputs) + inputs=body.initial_inputs ) def _remove_request(self, body: RemoveRequest): @@ -346,7 +326,7 @@ def _remove_request(self, body: RemoveRequest): self.subgraphs_manager.remove_request(body.request_id) def _process_new_inputs( - self, body: InputTensors + self, body: InputSignals ): """ When either the conductor or other workers send tensors to this worker, @@ -356,31 +336,24 @@ def _process_new_inputs( self.subgraphs_manager.update_phase( body.request_id, body.phase ) - self.subgraphs_manager.per_request_info[ - body.request_id].tensors.update(body.tensors) - self.subgraphs_manager.process_new_inputs( - request_id=body.request_id, - inputs=remove_flags(body.inputs) - ) + + # TODO Atindra: start reading in tensors from body.initial_inputs + # Also, somewhere else, we will have to call self.subgraphs_manager.process_new_inputs + # when the tensors are actually ready def _process_messages(self): """ Processes all pending messages (communication from conductor and other workers to this worker) """ - while True: - try: - message: WorkerMessage = self.request_socket.recv_pyobj( - flags=zmq.NOBLOCK - ) - if message.message_type == WorkerMessageType.NEW_REQUEST: - self._add_new_request(message.body) - elif message.message_type == WorkerMessageType.REMOVE_REQUEST: - self._remove_request(message.body) - elif message.message_type == WorkerMessageType.INPUT_TENSORS: - self._process_new_inputs(message.body) - except zmq.Again: - break + for message in self.communicator.get_all_new_messages(): + if message.message_type == WorkerMessageType.NEW_REQUEST: + self._add_new_request(message.body) + elif message.message_type == WorkerMessageType.REMOVE_REQUEST: + self._remove_request(message.body) + elif message.message_type == WorkerMessageType.INPUT_SIGNALS: + self._process_new_inputs(message.body) + # TODO: handle the tensor_received message def _send_outputs( self, @@ -391,40 +364,37 @@ def _send_outputs( Sends outputs to other workers and to the conductor """ for worker in outputs.to_workers: - request = WorkerMessage( - message_type=WorkerMessageType.INPUT_TENSORS, - body=InputTensors( + message = WorkerMessage( + message_type=WorkerMessageType.INPUT_SIGNALS, + body=InputSignals( request_id=request_id, phase=self.subgraphs_manager.get_phase(request_id), inputs=outputs.to_workers[worker] ) ) - self.inter_worker_sockets[worker].send_pyobj(request) + self.communicator.send(worker, message) # to conductor if outputs.to_conductor: - tensors = self.subgraphs_manager.per_request_info[ - request_id].tensors - request = ConductorMessage( - message_type=ConductorMessageType.TENSORS, - body=ConductorTensors( + message = ConductorMessage( + message_type=ConductorMessageType.PERSIST_SIGNAL, + body=InputSignals( request_id=request_id, inputs=outputs.to_conductor, - tensors={ - id: tensors[id] for id in outputs.to_conductor - } + phase=self.subgraphs_manager.get_phase(request_id) ) ) - self.result_socket.send_pyobj(request) + self.communicator.send("conductor", message) if outputs.completed_subgraphs: - request = ConductorMessage( + message = ConductorMessage( message_type=ConductorMessageType.SUBGRAPHS_DONE, body=SubgraphsDone( request_id=request_id, subgraph_ids=outputs.completed_subgraphs ) ) + self.communicator.send("conductor", message) def run(self): # TODO: this is just a dummy version diff --git a/test/test_request_queues.py b/test/test_request_queues.py index 98cc2a6ab6..a7a9c7b7a1 100644 --- a/test/test_request_queues.py +++ b/test/test_request_queues.py @@ -3,7 +3,7 @@ from mminf.graph.request_queues import PerRequestStageQueues sys.path.insert(0, ".") -from mminf.graph.base import GraphPointer, GraphStage, Loop, Parallel, Sequential, remove_flags +from mminf.graph.base import GraphPointer, GraphStage, Loop, Parallel, Sequential import numpy as np @@ -16,57 +16,58 @@ GraphStage( name="LLM", input_ids=["text_emb", "img_emb", "latents"], - outputs={ - "hidden_states": [GraphPointer("flow")], - "some_random_external_output": [GraphPointer("f")] - } + outputs=[ + GraphPointer(name="hidden_states", next_stage="flow"), + GraphPointer(name="some_random_external_output", next_stage="f") + ] ), Loop( section=Sequential([ GraphStage( name="flow", input_ids=["hidden_states", "mystery", "mystery2"], - outputs={ - "mystery2": [GraphPointer("flow")], - "partial_latents": [GraphPointer("flow2")] - } + outputs=[ + GraphPointer(name="partial_mystery2", next_stage="flow2"), + GraphPointer(name="partial_latents", next_stage="flow2") + ] ), GraphStage( name="flow2", input_ids=["partial_latents"], - outputs={ - "latents": [] - } + outputs=[ + GraphPointer(name="latents", next_stage="") + ] ), ]), n_iters=2, - outputs={ - "latents": [GraphPointer("LLM")], - } + outputs=[ + GraphPointer(name="latents", next_stage="LLM") + ] ) ]), Sequential([ GraphStage( name="f", input_ids=["mystery", "some_random_external_output"], - outputs={ - "xyz": [GraphPointer("g")] - } + outputs=[ + GraphPointer(name="xyz", next_stage="g") + ] ), GraphStage( name="g", input_ids=["xyz"], - outputs={ - "mystery": [GraphPointer("f"), GraphPointer("flow")] - } + outputs=[ + GraphPointer(name="mystery", next_stage="f"), + GraphPointer(name="mystery", next_stage="flow") + ] ) ]) ]), n_iters=3, - outputs={ - "latents": [GraphPointer("VAE_decoder")], - "some_random_external_output": [GraphPointer("STREAM_OUT")] - } + outputs=[ + GraphPointer(name="latents", next_stage="VAE_decoder"), + GraphPointer(name="some_random_external_output", next_stage="STREAM_OUT") + ] ) loop = Parallel([ @@ -104,18 +105,19 @@ GraphStage( name="text_emb", input_ids=["text"], - outputs={ - "text_emb": [GraphPointer("LLM")], - "mystery": [GraphPointer("flow"), GraphPointer("f")] - } + outputs=[ + GraphPointer(next_stage="LLM", name="text_emb"), + GraphPointer(next_stage="f", name="mystery"), + GraphPointer(next_stage="flow", name="mystery") + ] ), GraphStage( name="vit_encoder", input_ids=["image"], - outputs={ - "img_emb": [GraphPointer("LLM")], - "some_random_external_output": [GraphPointer("f")] - } + outputs=[ + GraphPointer(next_stage="LLM", name="img_emb"), + GraphPointer(next_stage="f", name="some_random_external_output"), + ] ) ]), Loop( @@ -124,74 +126,75 @@ GraphStage( name="LLM", input_ids=["text_emb", "img_emb", "latents"], - outputs={ - "hidden_states": [GraphPointer("flow")], - "some_random_external_output": [GraphPointer("f")] - } + outputs=[ + GraphPointer(next_stage="flow", name="hidden_states"), + GraphPointer(next_stage="f", name="some_random_external_output") + ] ), Loop( section=Sequential([ GraphStage( name="flow", input_ids=["hidden_states", "mystery", "mystery2"], - outputs={ - "partial_mystery2": [GraphPointer("flow2")], - "partial_latents": [GraphPointer("flow2")] - } + outputs=[ + GraphPointer(next_stage="flow2", name="partial_mystery2"), + GraphPointer(next_stage="flow2", name="partial_latents") + ] ), GraphStage( name="flow2", input_ids=["partial_latents", "partial_mystery2"], - outputs={ - "latents": [], - "mystery2": [GraphPointer("flow")] - } + outputs=[ + GraphPointer(next_stage="", name="latents"), + GraphPointer(next_stage="flow", name="mystery2") + ] ), ]), n_iters=2, - outputs={ - "latents": [GraphPointer("LLM")], - } + outputs=[ + GraphPointer(next_stage="LLM", name="latents") + ] ) ]), Sequential([ GraphStage( name="f", input_ids=["mystery", "some_random_external_output"], - outputs={ - "xyz": [GraphPointer("g")] - } + outputs=[ + GraphPointer(next_stage="g", name="xyz") + ] ), GraphStage( name="g", input_ids=["xyz"], - outputs={ - "mystery": [GraphPointer("f"), GraphPointer("flow")] - } + outputs=[ + GraphPointer(next_stage="f", name="mystery"), + GraphPointer(next_stage="flow", name="mystery") + ] ) ]) ]), n_iters=3, - outputs={ - "latents": [GraphPointer("VAE_decoder")], - "some_random_external_output": [GraphPointer("STREAM_OUT")] - } + outputs=[ + GraphPointer(next_stage="VAE_decoder", name="latents"), + GraphPointer(next_stage="STREAM_OUT", name="some_random_external_output") + ] ), GraphStage( name="VAE_decoder", input_ids=["latents"], - outputs={ - "generated_image": [GraphPointer("STREAM_OUT")] - } + outputs=[ + GraphPointer(next_stage="STREAM_OUT", name="generated_image") + ] ) ]) - provided_inputs = { - "text": [GraphPointer("text_emb")], - "image": [GraphPointer("vit_encoder")], - "latents": [GraphPointer("LLM")], - "mystery2": [GraphPointer("flow")] - } + provided_inputs = [ + GraphPointer(name="text", next_stage="text_emb"), + GraphPointer(name="image", next_stage="vit_encoder"), + GraphPointer(name="latents", next_stage="LLM"), + GraphPointer(name="mystery2", next_stage="flow") + ] queues = PerRequestStageQueues( ready=[], @@ -199,7 +202,7 @@ ) tic = time.perf_counter() - queues.process_new_inputs(remove_flags(provided_inputs)) + queues.process_new_inputs(provided_inputs) # loop until all stages are done and print out while len(queues.ready) > 0 or queues.waiting is not None: print("\n" + "="*60) @@ -208,16 +211,16 @@ print("Waiting stages:", queues.waiting.get_stage_names()) if len(queues.ready) == 0: - print(queues.waiting) + # print(queues.waiting) raise Exception("No ready stages but still waiting stages, something's wrong") print() # pop a random ready stage and process it stage = queues.ready.pop(np.random.randint(0, len(queues.ready))) print(f"Processing stage {stage.name} with inputs {stage.input_ids}") new_inputs = stage.outputs - external_outputs = queues.process_new_inputs(remove_flags(new_inputs)) - if external_outputs: - print(f"External outputs: {external_outputs}") + print(f"New inputs: {[f'{ptr.name} -> {ptr.next_stage}' for ptr in new_inputs]}") + external_outputs = queues.process_new_inputs(new_inputs) + print(f"Outputs: {external_outputs}") toc = time.perf_counter() print(toc - tic)