diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index f21dcc5927cc..0f283973bee1 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -12,6 +12,7 @@ from tqdm import tqdm from transformers import PreTrainedTokenizerBase +from tensorrt_llm._utils import mpi_disabled from tensorrt_llm.inputs.data import TextPrompt from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.inputs.registry import DefaultInputProcessor @@ -132,8 +133,12 @@ def __init__(self, if backend == "pytorch": logger.info("Using LLM with PyTorch backend") llm_args_cls = TorchLlmArgs - if self._orchestrator_type == "ray": + if self._orchestrator_type == "ray" or mpi_disabled(): + self._orchestrator_type = "ray" os.environ["TLLM_DISABLE_MPI"] = "1" + # Propagate to args construction + kwargs["orchestrator_type"] = "ray" + elif backend == '_autodeploy': logger.info("Using LLM with AutoDeploy backend") from .._torch.auto_deploy.llm_args import \ diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index a04c24e38b69..f8fc108d0992 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -2031,6 +2031,13 @@ def pytest_addoption(parser): parser.addoption("--perf", action="store_true", help="'--perf' will run perf tests") + parser.addoption( + "--run-ray", + action="store_true", + default=False, + help= + "Enable Ray orchestrator path for integration tests (disables MPI).", + ) parser.addoption( "--perf-log-formats", help= @@ -2130,6 +2137,8 @@ def pytest_collection_modifyitems(session, config, items): def pytest_configure(config): # avoid thread leak of tqdm's TMonitor tqdm.tqdm.monitor_interval = 0 + if config.getoption("--run-ray"): + os.environ["TLLM_DISABLE_MPI"] = "1" def deselect_by_regex(regexp, items, test_prefix, config):