From 664f611e83deb319830f48ead72e413ec32aa851 Mon Sep 17 00:00:00 2001 From: Jumiar <48683774+Jumiar@users.noreply.github.com> Date: Tue, 30 Dec 2025 19:56:25 +0800 Subject: [PATCH] Add profiling capture support to the encoder server (#15730) Signed-off-by: liuanqi Co-authored-by: liusy58 --- .../srt/disaggregation/encode_server.py | 140 +++++++++++++++++- 1 file changed, 134 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index a9ad60f87..aa8e43be7 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -29,6 +29,7 @@ from sglang.srt.distributed.parallel_state import ( initialize_model_parallel, ) from sglang.srt.layers.dp_attention import initialize_dp_attention +from sglang.srt.managers.io_struct import ProfileReq, ProfileReqInput, ProfileReqType from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem from sglang.srt.mem_cache.multimodal_cache import MultiModalStaticCache from sglang.srt.model_loader import get_model @@ -106,6 +107,7 @@ class MMEncoder: self.server_args = server_args set_global_server_args_for_scheduler(server_args) self.rank = rank + self.profiler = EncoderProfiler(rank) self.image_processor = AutoImageProcessor.from_pretrained( server_args.model_path, @@ -223,6 +225,8 @@ class MMEncoder: logger.info( f"Vit time : {(end_time - start_time)*1000:.2f} ms {mm_embedding.shape = }" ) + if self.profiler is not None: + self.profiler.step() return _get_image_grid_dim(images_input), mm_embedding @@ -384,6 +388,71 @@ class MMEncoder: return response_json["embedding_port"] +class EncoderProfiler: + def __init__(self, rank: int): + self.rank = rank + self.profiler = None + self.steps_left = None + self.output_dir = None + self.prefix = None + self.profile_id = None + + def start(self, obj: ProfileReq): + if self.profiler is not None: + return False, "profiling already running" + + output_dir = obj.output_dir or os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp") + os.makedirs(output_dir, exist_ok=True) + self.output_dir = output_dir + self.prefix = obj.profile_prefix or "encoder" + self.profile_id = str(time.time()) + + activities = obj.activities or ["CPU", "GPU"] + torch_activities = [] + if "CPU" in activities: + torch_activities.append(torch.profiler.ProfilerActivity.CPU) + if "GPU" in activities: + torch_activities.append(torch.profiler.ProfilerActivity.CUDA) + + profile_memory = "MEM" in activities + if not torch_activities and not profile_memory: + return False, "no supported activities" + + self.profiler = torch.profiler.profile( + activities=torch_activities, + with_stack=True if obj.with_stack is None else obj.with_stack, + record_shapes=False if obj.record_shapes is None else obj.record_shapes, + profile_memory=profile_memory, + ) + self.profiler.start() + self.steps_left = obj.num_steps + logger.info( + f"Encoder profiling started. output_dir={self.output_dir} profile_id={self.profile_id}" + ) + return True, None + + def step(self): + if self.profiler is None: + return + self.profiler.step() + if self.steps_left is not None: + self.steps_left -= 1 + if self.steps_left <= 0: + self.stop() + + def stop(self): + if self.profiler is None: + return False, "profiling not running" + self.profiler.stop() + filename = f"{self.prefix}-rank{self.rank}-{self.profile_id}.trace.json" + trace_path = os.path.join(self.output_dir, filename) + self.profiler.export_chrome_trace(trace_path) + logger.info("Encoder profiling saved to: %s", trace_path) + self.profiler = None + self.steps_left = None + return True, None + + app = FastAPI() encoder: Optional[MMEncoder] = None send_sockets: List[zmq.Socket] = [] @@ -395,12 +464,20 @@ async def run_encoder( encoder = MMEncoder(server_args, schedule_path, dist_init_method, rank) while True: request = await encoder.schedule_socket.recv_pyobj() - await encoder.encode( - mm_items=request["mm_items"], - req_id=request["req_id"], - num_parts=request["num_parts"], - part_idx=request["part_idx"], - ) + if isinstance(request, ProfileReq): + if request.type == ProfileReqType.START_PROFILE: + if encoder.profiler is None: + encoder.profiler = EncoderProfiler(encoder.rank) + encoder.profiler.start(request) + else: + encoder.profiler.stop() + else: + await encoder.encode( + mm_items=request["mm_items"], + req_id=request["req_id"], + num_parts=request["num_parts"], + part_idx=request["part_idx"], + ) def launch_encoder(server_args, schedule_path, dist_init_method, rank): @@ -525,3 +602,54 @@ async def health_generate(): if encoder is None: return Response(status_code=503) return Response(status_code=200) + + +@app.api_route("/start_profile", methods=["GET", "POST"]) +async def start_profile_async(obj: Optional[ProfileReqInput] = None): + if encoder is None: + return Response(content="encoder not ready\n", status_code=503) + req = None + if obj is None: + req = ProfileReq(ProfileReqType.START_PROFILE) + else: + req = ProfileReq( + type=ProfileReqType.START_PROFILE, + output_dir=obj.output_dir, + start_step=obj.start_step, + num_steps=obj.num_steps, + activities=obj.activities, + with_stack=obj.with_stack, + record_shapes=obj.record_shapes, + profile_by_stage=obj.profile_by_stage, + profile_id=str(time.time()), + merge_profiles=obj.merge_profiles, + profile_prefix=obj.profile_prefix, + profile_stages=obj.profile_stages, + ) + for socket in send_sockets: + socket.send_pyobj(req) + if encoder.profiler is None: + encoder.profiler = EncoderProfiler(encoder.rank) + ok, msg = encoder.profiler.start(req) + if ok: + detail = ( + f"Start profiling. output_dir={encoder.profiler.output_dir} " + f"profile_id={encoder.profiler.profile_id}\n" + ) + return Response(content=detail, status_code=200) + return Response(content=(msg or "Start profiling failed.\n"), status_code=400) + + +@app.api_route("/stop_profile", methods=["GET", "POST"]) +async def stop_profile_async(): + if encoder is None: + return Response(content="encoder not ready\n", status_code=503) + if encoder.profiler is None: + return Response(content="profiling not initialized\n", status_code=400) + req = ProfileReq(ProfileReqType.STOP_PROFILE) + for socket in send_sockets: + socket.send_pyobj(req) + ok, msg = encoder.profiler.stop() + if ok: + return Response(content="Stop profiling.\n", status_code=200) + return Response(content=(msg or "Stop profiling failed.\n"), status_code=400)