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

Method __init__

tensorrt_llm/executor/base_worker.py:74–117  ·  view source on GitHub ↗
(
        self,
        engine: Union[Path, Engine],
        executor_config: Optional[tllm.ExecutorConfig] = None,
        batched_logits_processor: Optional[BatchedLogitsProcessor] = None,
        postproc_worker_config: Optional[PostprocWorkerConfig] = None,
        is_llm_executor: Optional[bool] = None,
        hf_model_dir: Optional[Path] = None,
        tokenizer: Optional[TokenizerBase] = None,
        llm_args: Optional[BaseLlmArgs] = None,
    )

Source from the content-addressed store, hash-verified

72 pass
73
74 def __init__(
75 self,
76 engine: Union[Path, Engine],
77 executor_config: Optional[tllm.ExecutorConfig] = None,
78 batched_logits_processor: Optional[BatchedLogitsProcessor] = None,
79 postproc_worker_config: Optional[PostprocWorkerConfig] = None,
80 is_llm_executor: Optional[bool] = None,
81 hf_model_dir: Optional[Path] = None,
82 tokenizer: Optional[TokenizerBase] = None,
83 llm_args: Optional[BaseLlmArgs] = None,
84 ) -> None:
85 postproc_config = postproc_worker_config or PostprocWorkerConfig()
86 super().__init__(
87 num_postprocess_workers=postproc_config.num_postprocess_workers,
88 postprocess_tokenizer_dir=postproc_config.postprocess_tokenizer_dir,
89 is_llm_executor=is_llm_executor,
90 )
91
92 # inputs
93 self._engine = engine
94 self._executor_config = executor_config
95 self._batched_logits_processor = batched_logits_processor
96 self._postproc_worker_config = postproc_worker_config
97 self._is_llm_executor = is_llm_executor
98 self._hf_model_dir = hf_model_dir
99 self._tokenizer = tokenizer
100 self.llm_args = llm_args
101
102 self.engine = None
103 self.result_queue: Optional[IpcQueue] = None
104 self.postproc_queues: Optional[List[IpcQueue]] = None
105 self.rank = mpi_rank()
106 self.global_rank = global_mpi_rank()
107 # mapping: client_id -> GenerationResult
108 self._results: Dict[int, GenerationResult] = {}
109 # mapping: client_id from Proxy -> request_id returned from runtime backend
110 self._client_id_to_request_id: Dict[int, int] = {}
111 self._await_response_helper = AwaitResponseHelper(weakref.proxy(self))
112 self._backend = None if llm_args is None else llm_args.backend
113 self._is_pytorch_backend = self._backend in ["pytorch", "_autodeploy"]
114 self._lora_config = llm_args.lora_config if self._is_pytorch_backend else None
115
116 if global_mpi_size() > 1:
117 logger.set_rank(self.global_rank)
118
119 def _configure_affinity(self, device_id):
120 '''Probe and configure the CPU affinity of the worker based on NUMA topology.

Callers

nothing calls this directly

Calls 6

global_mpi_rankFunction · 0.85
AwaitResponseHelperClass · 0.85
global_mpi_sizeFunction · 0.85
mpi_rankFunction · 0.50
set_rankMethod · 0.45

Tested by

no test coverage detected