1818import os
1919import signal
2020import socket
21+ import sys
2122import time
2223from multiprocessing import Process
2324from typing import Callable , Dict , List , Optional , Tuple
@@ -42,8 +43,31 @@ def _start_triton_server(
4243 model_path : Optional [str ] = None ,
4344) -> List [tuple ]:
4445 """Task to start Triton server process on a Spark executor."""
45- sig = inspect .signature (triton_server_fn )
46- params = sig .parameters
46+
47+ def _prepare_pytriton_env ():
48+ """Expose PyTriton to correct libpython3.11.so and Triton bundled libraries."""
49+ ld_library_paths = []
50+
51+ # Add nvidia_pytriton.libs to LD_LIBRARY_PATH
52+ for path in sys .path :
53+ if os .path .isdir (path ) and "site-packages" in path :
54+ libs_path = os .path .join (path , "nvidia_pytriton.libs" )
55+ if os .path .isdir (libs_path ):
56+ ld_library_paths .append (libs_path )
57+ break
58+
59+ # Add ${CONDA_PREFIX}/lib to LD_LIBRARY_PATH for conda environments
60+ if os .path .exists (os .path .join (sys .prefix , "conda-meta" )):
61+ conda_lib = os .path .join (sys .prefix , "lib" )
62+ if os .path .isdir (conda_lib ):
63+ ld_library_paths .append (conda_lib )
64+
65+ if "LD_LIBRARY_PATH" in os .environ :
66+ ld_library_paths .append (os .environ ["LD_LIBRARY_PATH" ])
67+
68+ os .environ ["LD_LIBRARY_PATH" ] = ":" .join (ld_library_paths )
69+
70+ return None
4771
4872 def _find_ports (start_port : int = 7000 ) -> List [int ]:
4973 """Find available ports for Triton's HTTP, gRPC, and metrics services."""
@@ -59,6 +83,8 @@ def _find_ports(start_port: int = 7000) -> List[int]:
5983 return ports
6084
6185 ports = _find_ports ()
86+ sig = inspect .signature (triton_server_fn )
87+ params = sig .parameters
6288
6389 if model_path is not None :
6490 assert (
@@ -69,6 +95,7 @@ def _find_ports(start_port: int = 7000) -> List[int]:
6995 assert len (params ) == 1 , "Server function must accept (ports) argument"
7096 args = (ports ,)
7197
98+ _prepare_pytriton_env ()
7299 hostname = socket .gethostname ()
73100 process = Process (target = triton_server_fn , args = args )
74101 process .start ()
@@ -83,6 +110,11 @@ def _find_ports(start_port: int = 7000) -> List[int]:
83110 except Exception :
84111 pass
85112
113+ client .close ()
114+ if process .is_alive ():
115+ # Terminate if timeout is exceeded to avoid dangling server processes
116+ process .terminate ()
117+
86118 raise TimeoutError (
87119 "Failure: server startup timeout exceeded. Check the executor logs for more info."
88120 )
@@ -98,14 +130,19 @@ def _stop_triton_server(
98130 pid , _ = server_pids_ports .get (hostname )
99131 assert pid is not None , f"No server PID found for host { hostname }
100132
101- for _ in range (wait_retries ):
133+ try :
134+ process = psutil .Process (pid )
135+ process .terminate ()
136+ process .wait (timeout = wait_timeout * wait_retries )
137+ return [True ]
138+ except psutil .NoSuchProcess :
139+ return [True ]
140+ except psutil .TimeoutExpired :
102141 try :
103- os .kill (pid , signal .SIGTERM )
104- except OSError :
142+ process .kill ()
105143 return [True ]
106- time .sleep (wait_timeout )
107-
108- return [False ] # Failed to terminate or timed out
144+ except :
145+ return [False ]
109146
110147
111148class TritonServerManager :
0 commit comments