chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: liuanqi <liuanqi6@xiaomi.com>
Co-authored-by: liusy58 <liusy58@linux.alibaba.com>
This commit is contained in:
Jumiar
2025-12-30 19:56:25 +08:00
committed by GitHub
parent 49adb37e37
commit 664f611e83
1 changed files with 134 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

140
python/sglang/srt/disaggregation/encode_server.py
View File

@@ -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)