(self,
n_workers: int = 0,
addr: str = f'tcp://127.0.0.1:*',
hmac_key: Optional[bytes] = None,
comm=None,
is_comm: bool = False)
| 409 | ''' |
| 410 | |
| 411 | def __init__(self, |
| 412 | n_workers: int = 0, |
| 413 | addr: str = f'tcp://127.0.0.1:*', |
| 414 | hmac_key: Optional[bytes] = None, |
| 415 | comm=None, |
| 416 | is_comm: bool = False): |
| 417 | # FIXME: this is a hack to avoid circular import, resolve later |
| 418 | from tensorrt_llm.executor.ipc import ZeroMqQueue |
| 419 | self.addr = addr |
| 420 | self.queue = ZeroMqQueue((addr, hmac_key), |
| 421 | is_server=True, |
| 422 | socket_type=zmq.PAIR, |
| 423 | use_hmac_encryption=bool(hmac_key)) |
| 424 | self.comm = comm |
| 425 | self.results = [] # the results may arrive in any order |
| 426 | |
| 427 | if self.comm is not None: |
| 428 | self.session = MpiCommSession(n_workers=self.comm.Get_size(), |
| 429 | comm=self.comm) |
| 430 | else: |
| 431 | self.session = MpiCommSession( |
| 432 | n_workers=n_workers) if is_comm else MpiPoolSession( |
| 433 | n_workers=n_workers) |
| 434 | |
| 435 | @staticmethod |
| 436 | def task_wrapper(task: Callable[..., T], *args, **kwargs) -> T: |
nothing calls this directly
no test coverage detected