"tests/pytorch/vscode:/vscode.git/clone" did not exist on "06438d70335f3de7fd2be0fa7196414a4ff786d4"
Unverified Commit 1e90fe2e authored by Feng Su's avatar Feng Su Committed by GitHub
Browse files

[Bug fix] trace: fix import error in mini_lb if sgl-router image does not install sglang (#12338)

parent caa5d296
......@@ -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:
......
......@@ -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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment