diff --git a/sgl-router/py_src/sglang_router/launch_router.py b/sgl-router/py_src/sglang_router/launch_router.py index 50d96da6e..8da5191c5 100644 --- a/sgl-router/py_src/sglang_router/launch_router.py +++ b/sgl-router/py_src/sglang_router/launch_router.py @@ -38,15 +38,6 @@ def launch_router(args: argparse.Namespace) -> Optional[Router]: router_args = args if router_args.mini_lb: - if router_args.enable_trace: - from sglang.srt.tracing.trace import ( - process_tracing_init, - trace_set_thread_info, - ) - - process_tracing_init(router_args.otlp_traces_endpoint, "sglang") - trace_set_thread_info("Mini lb") - mini_lb = MiniLoadBalancer(router_args) mini_lb.start() else: diff --git a/sgl-router/py_src/sglang_router/mini_lb.py b/sgl-router/py_src/sglang_router/mini_lb.py index 4fcd1cc17..7c3089b5e 100644 --- a/sgl-router/py_src/sglang_router/mini_lb.py +++ b/sgl-router/py_src/sglang_router/mini_lb.py @@ -18,13 +18,20 @@ from fastapi import FastAPI, HTTPException from fastapi.responses import ORJSONResponse, Response, StreamingResponse from sglang_router.router_args import RouterArgs -from sglang.srt.tracing.trace import ( - trace_get_remote_propagate_context, - trace_req_finish, - trace_req_start, - trace_slice_end, - trace_slice_start, -) +try: + from sglang.srt.tracing.trace import ( + process_tracing_init, + trace_get_remote_propagate_context, + trace_req_finish, + trace_req_start, + trace_set_thread_info, + trace_slice_end, + trace_slice_start, + ) + + trace_package_imported = True +except ImportError: + trace_package_imported = False logger = logging.getLogger(__name__) @@ -54,7 +61,13 @@ class MiniLoadBalancer: self.prefill_urls = [url[0] for url in router_args.prefill_urls] self.prefill_bootstrap_ports = [url[1] for url in router_args.prefill_urls] self.decode_urls = router_args.decode_urls + self.otlp_traces_endpoint = router_args.otlp_traces_endpoint self.enable_trace = router_args.enable_trace + if self.enable_trace and not trace_package_imported: + logger.warning( + "Tracing is not supported in this environment. Please install sglang." + ) + self.enable_trace = False def _validate_router_args(self, router_args: RouterArgs): logger.warning( @@ -77,6 +90,9 @@ class MiniLoadBalancer: def start(self): global lb lb = self + if self.enable_trace: + process_tracing_init(self.otlp_traces_endpoint, "sglang") + trace_set_thread_info("Mini lb") uvicorn.run(app, host=self.host, port=self.port) def select_pair(self): @@ -441,8 +457,9 @@ async def handle_completion_request(request_data: dict): def _generate_bootstrap_room(): bootstrap_room = random.randint(0, 2**63 - 1) - trace_req_start(bootstrap_room, bootstrap_room, role="router") - trace_slice_start("mini_lb_launch", bootstrap_room) + if lb.enable_trace: + trace_req_start(bootstrap_room, bootstrap_room, role="router") + trace_slice_start("mini_lb_launch", bootstrap_room) return bootstrap_room