Skip to content

Commit 84facba

Browse files
committed
init
Signed-off-by: Superjomn <[email protected]>
1 parent 89d03d7 commit 84facba

File tree

2 files changed

+19
-0
lines changed

2 files changed

+19
-0
lines changed

tensorrt_llm/llmapi/llm_args.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from strenum import StrEnum
2222
from transformers import PreTrainedTokenizerBase
2323

24+
from tensorrt_llm.llmapi.utils import orchestrator_type_env
2425
from tensorrt_llm.lora_helper import (LoraConfig,
2526
get_default_trtllm_modules_to_hf_modules)
2627

@@ -2094,6 +2095,20 @@ def validate_peft_cache_config(self):
20942095
"while LoRA prefetch is not supported")
20952096
return self
20962097

2098+
@field_validator('orchestrator_type', mode='before')
2099+
def validate_orchestrator_config(v):
2100+
# The environment variable will override the orchestrator_type field.
2101+
# TODO: remove the environment variable after RPC path is stable, then
2102+
# there will be only two stable options: None(RPC) and 'ray'.
2103+
if (ev := orchestrator_type_env()) is not None:
2104+
if ev not in ['rpc', 'ray']:
2105+
raise ValueError(
2106+
f"Invalid orchestrator type: {ev}. Please set orchestrator_type to 'rpc' or 'ray'."
2107+
)
2108+
v = ev
2109+
2110+
return v
2111+
20972112
def _update_plugin_config(self, key: str, value: Any):
20982113
setattr(self.build_config.plugin_config, key, value)
20992114

tensorrt_llm/llmapi/utils.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -356,6 +356,10 @@ def enable_llmapi_debug() -> bool:
356356
return _enable_llmapi_debug_
357357

358358

359+
def orchestrator_type_env() -> Optional[str]:
360+
return os.environ.get("TLLM_ORCHESTRATOR_TYPE", None)
361+
362+
359363
@cache
360364
def enable_worker_single_process_for_tp1() -> bool:
361365
''' Tell whether to make worker use single process for TP1.

0 commit comments

Comments
 (0)