Add profiling capture support to the encoder server (#15730)

Signed-off-by: liuanqi <liuanqi6@xiaomi.com>
Co-authored-by: liusy58 <liusy58@linux.alibaba.com>

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,6 +464,14 @@ async def run_encoder(
     encoder = MMEncoder(server_args, schedule_path, dist_init_method, rank)
     while True:
         request = await encoder.schedule_socket.recv_pyobj()
+        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"],
@@ -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)