SplitwiseConnector class for managing and scheduling Splitwise tasks.
| 29 | |
| 30 | |
| 31 | class SplitwiseConnector: |
| 32 | """ |
| 33 | SplitwiseConnector class for managing and scheduling Splitwise tasks. |
| 34 | """ |
| 35 | |
| 36 | def __init__(self, cfg, worker_queue, resource_manager): |
| 37 | """ |
| 38 | Initialize the SplitwiseConnector instance. |
| 39 | |
| 40 | Parameters: |
| 41 | cfg (dict): Configuration information. |
| 42 | worker_queue (object): Worker queue object. |
| 43 | resource_manager (object): Resource manager object. |
| 44 | """ |
| 45 | self.cfg = cfg |
| 46 | self.local_data_parallel_id = self.cfg.parallel_config.local_data_parallel_id |
| 47 | if self.cfg.parallel_config.data_parallel_size > 1: |
| 48 | self.logger = get_logger( |
| 49 | "splitwise_connector", f"splitwise_connector_dprank{self.local_data_parallel_id}.log" |
| 50 | ) |
| 51 | else: |
| 52 | self.logger = get_logger("splitwise_connector", "splitwise_connector.log") |
| 53 | self.engine_worker_queue = worker_queue |
| 54 | self.resource_manager = resource_manager |
| 55 | self.current_request_ids = dict() |
| 56 | self.enable_decode_cache_task = envs.FD_ENABLE_CACHE_TASK == "1" |
| 57 | |
| 58 | if self.cfg.scheduler_config.splitwise_role != "mixed": |
| 59 | self.zmq_ctx = zmq.Context() |
| 60 | self.push_sockets: Dict[str, zmq.Socket] = {} |
| 61 | self.pull_socket = None |
| 62 | self.io_executor = ThreadPoolExecutor(max_workers=4) |
| 63 | self._init_network() |
| 64 | |
| 65 | def _init_network(self): |
| 66 | """ |
| 67 | init network for splitwise |
| 68 | """ |
| 69 | |
| 70 | self.router_socket = self.zmq_ctx.socket(zmq.ROUTER) |
| 71 | self.router_socket.setsockopt(zmq.LINGER, 0) |
| 72 | self.router_socket.setsockopt(zmq.SNDHWM, 1000) |
| 73 | self.router_socket.setsockopt(zmq.ROUTER_MANDATORY, 1) |
| 74 | self.logger.info(f"_init_network: bind {self.cfg.cache_config.local_pd_comm_port}") |
| 75 | self.router_socket.bind(f"tcp://*:{self.cfg.cache_config.local_pd_comm_port}") |
| 76 | |
| 77 | self.poller = zmq.Poller() |
| 78 | self.poller.register(self.router_socket, zmq.POLLIN) |
| 79 | |
| 80 | self.push_sockets = {} |
| 81 | |
| 82 | self.prefill_cache_info = [] |
| 83 | |
| 84 | def start_receiver(self): |
| 85 | """ |
| 86 | start receiver thread |
| 87 | """ |
| 88 | while True: |
no outgoing calls