finish all send requests before quitting pp event-loop to avoid mpi deadlock; synchronize sampler right after async calls to avoid hang

Signed-off-by: Lizhi Zhou <[email protected]>

Author: reasonsolo

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -627,6 +627,12 @@ def _process_iter_stats(self, finished_requests: list[LlmRequest],
                 batch_state.sample_state.scheduled_requests), req_stats)
 
     def _executor_loop_cleanup(self):
+        # Unblock receiving processes. When second-last rank quits before last rank,
+        # last rank will never return from recv_object.
+        for req in self.send_handles:
+            if req is not None:
+                req.wait()
+
         with self.response_cv:
             self.is_shutdown = True
             self.response_cv.notify_all()
@@ -750,8 +756,10 @@ def _executor_loop_pp(self):
 
                             sample_state = self._sample_async(
                                 scheduled_batch, batch_outputs)
+                            assert sample_state is not None, "Sampling failed"
                             sample_state.host.logits = logits_host
                             self._update_request_states(scheduled_batch)
+                            sample_state.sampler_event.synchronize()
 
                     if self.enable_iter_perf_stats:
                         iter_stats.inflight_batching_stats.num_ctx_tokens = self.model_engine.iter_states[

tests/integration/defs/disaggregated/test_disaggregated.py

@@ -156,10 +156,22 @@ def run_disaggregated_test(example_dir,
     run_env = env.copy()
     run_env["UCX_TLS"] = "^ib"
 
+    nsys_path = os.getenv("NSYS_PATH", None)
+    nsys_file = os.getenv("NSYS_FILE", None)
+    nsys_cmd = [
+        "nsys",
+        "profile",
+        "--trace",
+        "cuda,cublas,nvtx",
+        "--output",
+        nsys_file,
+        "--force-overwrite=true",
+    ] if nsys_path and nsys_file else []
+
     num_ranks, config_file = get_test_config(test_desc, example_dir,
                                              os.path.dirname(__file__))
 
-    workers_cmd = [
+    workers_cmd = nsys_cmd + [
         'mpirun', '--allow-run-as-root', '--oversubscribe', '-n',
         str(num_ranks), 'trtllm-serve', 'disaggregated_mpi_worker', '-c',
         config_file