MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT-LLM / __init__

Method __init__

tensorrt_llm/executor/postproc_worker.py:75–108  ·  view source on GitHub ↗

Args: pull_pipe_addr (tuple[str, Optional[bytes]]): The address and HMAC key of the input IPC. push_pipe_addr (tuple[str, Optional[bytes]]): The address and HMAC key of the output IPC. tokenizer_dir (str): The directory to load tokenizer. reco

(
        self,
        pull_pipe_addr: tuple[str, Optional[bytes]],
        push_pipe_addr: tuple[str, Optional[bytes]],
        tokenizer_dir: str,
        record_creator: Callable[
            ["PostprocWorker.Input", TransformersTokenizer], Any],
    )

Source from the content-addressed store, hash-verified

73 disaggregated_params: Any = None
74
75 def __init__(
76 self,
77 pull_pipe_addr: tuple[str, Optional[bytes]],
78 push_pipe_addr: tuple[str, Optional[bytes]],
79 tokenizer_dir: str,
80 record_creator: Callable[
81 ["PostprocWorker.Input", TransformersTokenizer], Any],
82 ):
83 '''
84 Args:
85 pull_pipe_addr (tuple[str, Optional[bytes]]): The address and HMAC key of the input IPC.
86 push_pipe_addr (tuple[str, Optional[bytes]]): The address and HMAC key of the output IPC.
87 tokenizer_dir (str): The directory to load tokenizer.
88 record_creator (Callable[["ResponsePostprocessWorker.Input"], Any]): A creator for creating a record for a request.
89 result_handler (Optional[Callable[[GenerationResultBase], Any]]): A callback handles the final result.
90 '''
91
92 self._records: Dict[int, GenerationResult] = {}
93 self._record_creator = record_creator
94 self._pull_pipe = ZeroMqQueue(address=pull_pipe_addr,
95 is_async=True,
96 is_server=False,
97 name="postprocess_pull_pipe")
98 self._push_pipe = ZeroMqQueue(address=push_pipe_addr,
99 is_async=True,
100 is_server=False,
101 socket_type=zmq.PUSH,
102 name="postprocess_push_pipe")
103 self._to_stop = asyncio.Event()
104
105 self._q = deque()
106
107 # Load the tokenizer and share in all records
108 self._tokenizer = load_hf_tokenizer(tokenizer_dir)
109
110 @staticmethod
111 def default_record_creator(

Callers

nothing calls this directly

Calls 2

ZeroMqQueueClass · 0.85
load_hf_tokenizerFunction · 0.85

Tested by

no test coverage detected