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

Method __init__

tensorrt_llm/executor/ray_executor.py:35–108  ·  view source on GitHub ↗
(self,
                 worker_kwargs: Dict,
                 model_world_size: int,
                 postproc_worker_config: PostprocWorkerConfig,
                 is_llm_executor: bool,
                 tp_size=1)

Source from the content-addressed store, hash-verified

33class RayExecutor(RpcExecutorMixin, GenerationExecutor):
34
35 def __init__(self,
36 worker_kwargs: Dict,
37 model_world_size: int,
38 postproc_worker_config: PostprocWorkerConfig,
39 is_llm_executor: bool,
40 tp_size=1):
41 os.environ['RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES'] = '1'
42 os.environ["RAY_DEDUP_LOGS"] = "0" # for debug
43
44 super().__init__(model_world_size, postproc_worker_config,
45 is_llm_executor)
46
47 self.has_start_local_cluser = False
48 runtime_env = {
49 "env_vars": {
50 "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"
51 }
52 }
53
54 ray_init_args = {
55 "include_dashboard": False,
56 "namespace": "trtllm",
57 "ignore_reinit_error": True,
58 "runtime_env": runtime_env
59 }
60
61 try:
62 if os.environ.get("TLLM_RAY_FORCE_LOCAL_CLUSTER", "0") != "1":
63 try:
64 ray.init(address="auto", **ray_init_args)
65 logger.info(f"Attached to an existing Ray cluster.")
66 except ConnectionError:
67 logger.info(f"Ray cluster not found, starting a new one.")
68
69 if not ray.is_initialized():
70 ray.init(**ray_init_args)
71 self.has_start_local_cluser = True
72 else:
73 ray.init(address="local", **ray_init_args)
74 self.has_start_local_cluser = True
75
76 self.world_size = model_world_size
77 self.tp_size = tp_size
78 self.master_address = ray.util.get_node_ip_address()
79
80 self.worker_kwargs = dict(
81 **worker_kwargs,
82 postproc_worker_config=postproc_worker_config,
83 is_llm_executor=is_llm_executor)
84
85 self.init_rpc_executor()
86 # Inject the generated HMAC key into worker_kwargs for workers
87 self.worker_kwargs['hmac_key'] = self.hmac_key
88 self.worker_kwargs['rpc_addr'] = self.rpc_addr
89
90 placement_config = getattr(self.worker_kwargs['llm_args'],
91 'ray_placement_config', None)
92 defer_workers_init = placement_config.defer_workers_init if placement_config else False

Callers

nothing calls this directly

Calls 10

init_workers_syncMethod · 0.95
setup_engine_remoteMethod · 0.95
shutdownMethod · 0.95
has_event_loopFunction · 0.85
is_initializedMethod · 0.80
init_rpc_executorMethod · 0.80
setup_mainloopMethod · 0.80
getMethod · 0.45
infoMethod · 0.45
errorMethod · 0.45

Tested by

no test coverage detected