fix the scheduler could start on rank > 0

Signed-off-by: richardhuo-nv <[email protected]>

Author: richardhuo-nv

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -404,14 +404,15 @@ def create_py_executor(
             scheduler_cls = getattr(
                 module, kv_connector_config.connector_scheduler_class)
 
+            rank = tensorrt_llm.mpi_rank()
             # Some connector API implementations may need to establish out-of-band communication between the scheduler and workers.
             # In this case, the worker may be dependent on the scheduler, or vice-versa.
             # To deal with cases like this, we instantiate them both concurrently.
             with ThreadPoolExecutor(max_workers=2) as executor:
                 connector_worker_task = executor.submit(worker_cls,
                                                         executor_config)
 
-                if scheduler_cls is not None:
+                if scheduler_cls is not None and rank == 0:
                     connector_scheduler_task = executor.submit(
                         scheduler_cls, executor_config)
                     connector_scheduler = connector_scheduler_task.result()