init

Signed-off-by: Superjomn <[email protected]>

Author: Superjomn

tensorrt_llm/llmapi/llm_args.py

@@ -21,6 +21,7 @@
 from strenum import StrEnum
 from transformers import PreTrainedTokenizerBase
 
+from tensorrt_llm.llmapi.utils import orchestrator_type_env
 from tensorrt_llm.lora_helper import (LoraConfig,
                                       get_default_trtllm_modules_to_hf_modules)
 
@@ -2025,6 +2026,20 @@ def validate_peft_cache_config(self):
                 "while LoRA prefetch is not supported")
         return self
 
+    @field_validator('orchestrator_type', mode='before')
+    def validate_orchestrator_config(v):
+        # The environment variable will override the orchestrator_type field.
+        # TODO: remove the environment variable after RPC path is stable, then
+        # there will be only two stable options: None(RPC) and 'ray'.
+        if (ev := orchestrator_type_env()) is not None:
+            if ev not in ['rpc', 'ray']:
+                raise ValueError(
+                    f"Invalid orchestrator type: {ev}. Please set orchestrator_type to 'rpc' or 'ray'."
+                )
+            v = ev
+
+        return v
+
     def _update_plugin_config(self, key: str, value: Any):
         setattr(self.build_config.plugin_config, key, value)
 

tensorrt_llm/llmapi/utils.py

@@ -356,6 +356,10 @@ def enable_llmapi_debug() -> bool:
     return _enable_llmapi_debug_
 
 
+def orchestrator_type_env() -> Optional[str]:
+    return os.environ.get("TLLM_ORCHESTRATOR_TYPE", None)
+
+
 @cache
 def enable_worker_single_process_for_tp1() -> bool:
     ''' Tell whether to make worker use single process for TP1.