[nvbug/5308432] fix: extend triton exit time for test_llava (NVIDIA#5971)

Signed-off-by: Chang Liu <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: chang-l

tests/integration/defs/triton_server/test_triton_llm.py

@@ -2491,7 +2491,7 @@ def test_llava(
     # NOTE
     # Due to mpi init error, manually set PMIX_MCA_gds=hash (ref: https://github.com/open-mpi/ompi/issues/6981)
     check_call(
-        f"PMIX_MCA_gds=hash python3 {launch_server_py} --world_size=1 --model_repo={new_model_repo}",
+        f"PMIX_MCA_gds=hash python3 {launch_server_py} --world_size=1 --model_repo={new_model_repo} --exit_timeout=300",
         shell=True)
     check_server_ready()
     # Run Test

triton_backend/scripts/launch_triton_server.py

@@ -103,7 +103,12 @@ def parse_arguments():
         help='Launch tritonserver with trtllm-llmapi-launch',
         default=False,
     )
-
+    parser.add_argument(
+        '--exit_timeout',
+        type=int,
+        help='Exit timeout in seconds',
+        default=None,
+    )
     return parser.parse_args()
 
 
@@ -152,9 +157,20 @@ def add_port_config(cmd, grpc_port, http_port, metrics_port):
     return cmd
 
 
-def get_cmd(world_size, tritonserver, grpc_port, http_port, metrics_port,
-            model_repo, log, log_file, tensorrt_llm_model_name, oversubscribe,
-            multimodal_gpu0_cuda_mem_pool_bytes, no_mpi, trtllm_llmapi_launch):
+def get_cmd(world_size,
+            tritonserver,
+            grpc_port,
+            http_port,
+            metrics_port,
+            model_repo,
+            log,
+            log_file,
+            tensorrt_llm_model_name,
+            oversubscribe,
+            multimodal_gpu0_cuda_mem_pool_bytes,
+            no_mpi,
+            trtllm_llmapi_launch,
+            exit_timeout=None):
     if no_mpi:
         assert world_size == 1, "world size must be 1 when using no-mpi"
 
@@ -172,6 +188,8 @@ def get_cmd(world_size, tritonserver, grpc_port, http_port, metrics_port,
         if trtllm_llmapi_launch:
             cmd += ['trtllm-llmapi-launch']
         cmd += [tritonserver, f'--model-repository={model_repo}']
+        if exit_timeout:
+            cmd += [f'--exit-timeout-secs={exit_timeout}']
 
         # Add port configuration
         cmd = add_port_config(cmd, grpc_port, http_port, metrics_port)
@@ -221,7 +239,7 @@ def get_cmd(world_size, tritonserver, grpc_port, http_port, metrics_port,
                   args.http_port, args.metrics_port, args.model_repo, args.log,
                   args.log_file, args.tensorrt_llm_model_name,
                   args.oversubscribe, args.multimodal_gpu0_cuda_mem_pool_bytes,
-                  args.no_mpi, args.trtllm_llmapi_launch)
+                  args.no_mpi, args.trtllm_llmapi_launch, args.exit_timeout)
     env = os.environ.copy()
     if args.multi_model:
         if not args.disable_spawn_processes: