use server parameters

Author: pathfinder-fp

python/sgl_jax/srt/entrypoints/engine.py

@@ -49,7 +49,6 @@
     get_zmq_socket,
     kill_process_tree,
     launch_dummy_health_check_server,
-    pathways_available,
     prepare_model_and_tokenizer,
     set_ulimit,
 )
@@ -248,7 +247,7 @@ def rerank(
     def shutdown(self):
         """Shutdown the engine"""
         kill_process_tree(os.getpid(), include_parent=False)
-        if pathways_available():
+        if self.server_args.enable_single_process:
             self.send_to_rpc.close()
 
     def __enter__(self):
@@ -410,7 +409,7 @@ def get_default_sampling_params(self) -> SamplingParams:
         return SamplingParams()
 
 
-def _set_envs_and_config():
+def _set_envs_and_config(server_args):
     # Set ulimit
     set_ulimit()
 
@@ -434,7 +433,7 @@ def sigquit_handler(signum, frame):
         kill_process_tree(os.getpid())
 
     signal.signal(signal.SIGQUIT, sigquit_handler)
-    if not pathways_available():
+    if not server_args.enable_single_process:
         # Set mp start method
         mp.set_start_method("spawn", force=True)
     else:
@@ -450,7 +449,7 @@ def _launch_subprocesses(
     # Configure global environment
     configure_logger(server_args)
     server_args.check_server_args()
-    _set_envs_and_config()
+    _set_envs_and_config(server_args)
 
     # Allocate ports for inter-process communications
     if port_args is None:
@@ -632,7 +631,7 @@ def _launch_threads(
 def _launch_subprocesses_or_threads(
     server_args, port_args: Optional[PortArgs] = None
 ) -> Tuple[TokenizerManager, TemplateManager, Dict]:
-    if pathways_available():
+    if server_args.enable_single_process:
         return _launch_threads(server_args, port_args)
     else:
         return _launch_subprocesses(server_args, port_args)

python/sgl_jax/srt/managers/tp_worker_overlap_thread.py

@@ -18,7 +18,6 @@
 from sgl_jax.srt.managers.utils import resolve_future_token_ids, set_future_token_ids
 from sgl_jax.srt.sampling.sampling_batch_info import SamplingMetadata
 from sgl_jax.srt.server_args import ServerArgs
-from sgl_jax.srt.utils import pathways_available
 from sgl_jax.utils import get_exception_traceback
 
 logger = logging.getLogger(__name__)
@@ -52,7 +51,7 @@ def __init__(
         # JAX handles device execution automatically, no need for explicit streams
         self.forward_thread = threading.Thread(
             target=self.forward_thread_func,
-            daemon=True if pathways_available() else False,
+            daemon=True if server_args.enable_single_process else False,
         )
         self.forward_thread.start()
         self.parent_process = psutil.Process().parent()

python/sgl_jax/srt/server_args.py

@@ -132,6 +132,7 @@ class ServerArgs:
 
     # For deterministic sampling
     enable_deterministic_sampling: bool = False
+    enable_single_process: bool = False
 
     def __post_init__(self):
         # Set missing default values
@@ -771,6 +772,12 @@ def add_cli_args(parser: argparse.ArgumentParser):
             help="Enable deterministic sampling",
         )
 
+        parser.add_argument(
+            "--enable-single-process",
+            action = "store_true",
+            help = "Enable run the engine with single process.",
+        )
+
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace):
         args.tp_size = args.tensor_parallel_size

python/sgl_jax/srt/utils/__init__.py

@@ -13,4 +13,3 @@
     set_ulimit,
     set_uvicorn_logging_configs,
 )
-from .tunix_utils import pathways_available