diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 0b1ab7b99..d72babdb8 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -178,7 +178,6 @@ from sglang.srt.managers.scheduler_output_processor_mixin import ( SchedulerOutputProcessorMixin, ) from sglang.srt.managers.scheduler_pp_mixin import SchedulerPPMixin -from sglang.srt.managers.scheduler_profiler_mixin import SchedulerProfilerMixin from sglang.srt.managers.scheduler_recv_skipper import SchedulerRecvSkipper from sglang.srt.managers.scheduler_runtime_checker_mixin import ( SchedulerRuntimeCheckerMixin, @@ -323,7 +322,6 @@ def validate_dflash_request(req: Req) -> Optional[str]: class Scheduler( SchedulerOutputProcessorMixin, SchedulerUpdateWeightsMixin, - SchedulerProfilerMixin, SchedulerMetricsMixin, SchedulerDisaggregationDecodeMixin, SchedulerDisaggregationPrefillMixin, @@ -1325,7 +1323,7 @@ class Scheduler( (SlowDownReqInput, self.slow_down), ( ProfileReq, - lambda req: self._profile(self.profiler_manager, req), + lambda req: self.profiler_manager._profile(req), ), (FreezeGCReq, self.handle_freeze_gc), (GetInternalStateReq, self.get_internal_state), @@ -2725,7 +2723,7 @@ class Scheduler( batch.forward_iter = self.forward_ct # Whether to run the profiler - self._profile_batch_predicate(self.profiler_manager, batch) + self.profiler_manager._profile_batch_predicate(batch) if self.forward_sleep_time is not None: logger.info(f"Scheduler.run_batch sleep {self.forward_sleep_time}s") time.sleep(self.forward_sleep_time) diff --git a/python/sglang/srt/managers/scheduler_components/profiler_manager.py b/python/sglang/srt/managers/scheduler_components/profiler_manager.py index e3d21d3f7..8b446fe28 100644 --- a/python/sglang/srt/managers/scheduler_components/profiler_manager.py +++ b/python/sglang/srt/managers/scheduler_components/profiler_manager.py @@ -1,10 +1,44 @@ from __future__ import annotations +import logging +import os +import time from dataclasses import dataclass from pathlib import Path -from typing import Any, Callable, List, Optional +from typing import ( + TYPE_CHECKING, + Any, + Callable, + List, + Optional, +) + +import torch from sglang.srt.environ import envs +from sglang.srt.managers.io_struct import ProfileReq, ProfileReqOutput, ProfileReqType +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import is_npu +from sglang.srt.utils.profile_merger import ProfileMerger + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import ScheduleBatch + +_is_npu = is_npu() +if _is_npu: + import torch_npu + + patches = [ + ["profiler.profile", torch_npu.profiler.profile], + ["profiler.ProfilerActivity.CUDA", torch_npu.profiler.ProfilerActivity.NPU], + ["profiler.ProfilerActivity.CPU", torch_npu.profiler.ProfilerActivity.CPU], + ] + torch_npu._apply_patches(patches) + +logger = logging.getLogger(__name__) + + from sglang.srt.utils.profile_utils import ProfileManager @@ -41,3 +75,348 @@ class SchedulerProfilerManager: # For ROCM self.rpd_profiler = None + + def _init_profile( + self, + output_dir: Optional[str], + start_step: Optional[int], + num_steps: Optional[int], + activities: Optional[List[str]], + with_stack: Optional[bool], + record_shapes: Optional[bool], + profile_by_stage: bool, + profile_id: str, + merge_profiles: bool = False, + profile_prefix: str = "", + profile_stages: Optional[List[str]] = None, + ) -> ProfileReqOutput: + if envs.SGLANG_PROFILE_V2.get(): + return self._profile_manager.configure( + output_dir=output_dir, + start_step=start_step, + num_steps=num_steps, + activities=activities, + with_stack=with_stack, + record_shapes=record_shapes, + profile_by_stage=profile_by_stage, + profile_id=profile_id, + merge_profiles=merge_profiles, + profile_prefix=profile_prefix, + profile_stages=profile_stages, + ) + + if self.profile_in_progress: + return ProfileReqOutput( + success=False, + message="Profiling is already in progress. Call /stop_profile first.", + ) + + self.profile_by_stage = profile_by_stage + self.merge_profiles = merge_profiles + + if output_dir is None: + output_dir = os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp") + if activities is None: + activities = ["CPU", "GPU"] + + self.torch_profiler_output_dir = Path(output_dir).expanduser() + self.torch_profiler_with_stack = with_stack + self.torch_profiler_record_shapes = record_shapes + self.profiler_activities = activities + self.profile_id = profile_id + self.profile_prefix = profile_prefix + + if start_step: + self.profiler_start_forward_ct = max(start_step, self.get_forward_ct() + 1) + + if num_steps: + if self.profile_by_stage: + self.profiler_prefill_ct = 0 + self.profiler_decode_ct = 0 + self.profiler_target_prefill_ct = num_steps + self.profiler_target_decode_ct = num_steps + elif start_step: + self.profiler_target_forward_ct = ( + self.profiler_start_forward_ct + num_steps + ) + else: + self.profiler_target_forward_ct = self.get_forward_ct() + num_steps + # The caller will be notified when reaching profiler_target_forward_ct + else: + self.profiler_target_forward_ct = None + + return ProfileReqOutput(success=True, message="Succeeded") + + def _start_profile( + self, stage: Optional[ForwardMode] = None + ) -> ProfileReqOutput | None: + if envs.SGLANG_PROFILE_V2.get(): + return self._profile_manager.manual_start() + + stage_str = f" for {stage.name}" if stage else "" + logger.info( + f"Profiling starts{stage_str}. Traces will be saved to: {self.torch_profiler_output_dir} (with profile id: {self.profile_id})", + ) + + activities = self.profiler_activities + with_stack = self.torch_profiler_with_stack + record_shapes = self.torch_profiler_record_shapes + + activity_map = { + "CPU": torch.profiler.ProfilerActivity.CPU, + "GPU": torch.profiler.ProfilerActivity.CUDA, + } + if hasattr(torch.profiler.ProfilerActivity, "XPU"): + activity_map["XPU"] = torch.profiler.ProfilerActivity.XPU + torchprof_activities = [ + activity_map[a] for a in activities if a in activity_map + ] + + if "RPD" in activities: # for ROCM + from rpdTracerControl import rpdTracerControl + + rpdTracerControl.skipCreate() + + self.rpd_profile_path = os.path.join( + self.torch_profiler_output_dir, + "rpd-" + str(time.time()) + f"-TP-{self.ps.tp_rank}" + ".trace.json.gz", + ) + + if self.ps.tp_rank == 0: + import sqlite3 + + from rocpd.schema import RocpdSchema + + if os.path.exists("trace.rpd"): + os.unlink("trace.rpd") + schema = RocpdSchema() + connection = sqlite3.connect("trace.rpd") + schema.writeSchema(connection) + connection.commit() + del connection + torch.distributed.barrier(self.dp_tp_cpu_group) + + self.rpd_profiler = rpdTracerControl() + self.rpd_profiler.setPythonTrace(True) + self.rpd_profiler.start() + self.rpd_profiler.rangePush("", "rpd profile range", "") + self.profile_in_progress = True + elif torchprof_activities: + self.torch_profiler = torch.profiler.profile( + activities=torchprof_activities, + with_stack=with_stack if with_stack is not None else True, + record_shapes=record_shapes if record_shapes is not None else False, + on_trace_ready=( + None + if not _is_npu + else torch_npu.profiler.tensorboard_trace_handler( + str(self.torch_profiler_output_dir) + ) + ), + ) + self.torch_profiler.start() + self.profile_in_progress = True + + if "MEM" in activities: + torch.cuda.memory._record_memory_history(max_entries=100000) + self.profile_in_progress = True + + if "CUDA_PROFILER" in activities: + if self.ps.gpu_id == get_global_server_args().base_gpu_id: + torch.cuda.cudart().cudaProfilerStart() + self.profile_in_progress = True + + return ProfileReqOutput(success=True, message="Succeeded") + + def _merge_profile_traces(self) -> str: + if not self.merge_profiles: + return "" + + if self.ps.tp_rank != 0: + return "" + if self.ps.dp_size > 1 and self.ps.dp_rank != 0: + return "" + if self.ps.pp_size > 1 and self.ps.pp_rank != 0: + return "" + if self.ps.moe_ep_size > 1 and self.ps.moe_ep_rank != 0: + return "" + + try: + logger.info("Starting profile merge...") + merger = ProfileMerger(self.torch_profiler_output_dir, self.profile_id) + merged_path = merger.merge_chrome_traces() + + summary = merger.get_merge_summary() + merge_message = ( + f" Merged trace: {merged_path} " + f"(Events: {summary.get('total_events', '?')}, " + f"Files: {summary.get('total_files', '?')})" + ) + + logger.info(f"Profile merge completed: {merged_path}") + except Exception as e: + logger.error(f"Failed to merge profiles: {e}", exc_info=True) + return f" Merge failed: {e!s}" + else: + return merge_message + + def _stop_profile( + self, stage: Optional[ForwardMode] = None + ) -> ProfileReqOutput | None: + if envs.SGLANG_PROFILE_V2.get(): + return self._profile_manager.manual_stop() + + if not self.profile_in_progress: + return ProfileReqOutput( + success=False, + message="Profiling is not in progress. Call /start_profile first.", + ) + + self.torch_profiler_output_dir.mkdir(parents=True, exist_ok=True) + + if self.profile_prefix: + stage_prefix = self.profile_prefix + "-" + else: + stage_prefix = "" + + stage_suffix = f"-{stage.name}" if stage else "" + logger.info("Stop profiling" + stage_suffix + "...") + if self.torch_profiler is not None: + self.torch_profiler.stop() + if not _is_npu: + # Build filename with only non-zero ranks to maintain backward compatibility + filename_parts = [self.profile_id, f"TP-{self.ps.tp_rank}"] + + # Only add other ranks if parallelism is enabled (size > 1) + if self.ps.dp_size > 1: + filename_parts.append(f"DP-{self.ps.dp_rank}") + if self.ps.pp_size > 1: + filename_parts.append(f"PP-{self.ps.pp_rank}") + if self.ps.moe_ep_size > 1: + filename_parts.append(f"EP-{self.ps.moe_ep_rank}") + + filename = ( + stage_prefix + + "-".join(filename_parts) + + stage_suffix + + ".trace.json.gz" + ) + + self.torch_profiler.export_chrome_trace( + os.path.join(self.torch_profiler_output_dir, filename) + ) + torch.distributed.barrier(self.dp_tp_cpu_group) + + if self.rpd_profiler is not None: + self.rpd_profiler.rangePop() + self.rpd_profiler.stop() + self.rpd_profiler.flush() + + torch.distributed.barrier(self.dp_tp_cpu_group) + if self.ps.tp_rank == 0: + from sglang.srt.utils.rpd_utils import rpd_to_chrome_trace + + rpd_to_chrome_trace("trace.rpd", self.rpd_profile_path) + self.rpd_profiler = None + self.rpd_profile_path = None + + if self.profiler_activities is not None and "MEM" in self.profiler_activities: + memory_profile_path = os.path.join( + self.torch_profiler_output_dir, + str(time.time()) + + f"-TP-{self.ps.tp_rank}-memory" + + stage_suffix + + ".pickle", + ) + torch.cuda.memory._dump_snapshot(memory_profile_path) + torch.cuda.memory._record_memory_history(enabled=None) + + if "CUDA_PROFILER" in self.profiler_activities: + if self.ps.gpu_id == get_global_server_args().base_gpu_id: + torch.cuda.cudart().cudaProfilerStop() + + merge_message = self._merge_profile_traces() + + logger.info( + "Profiling done. Traces are saved to: %s%s", + self.torch_profiler_output_dir, + merge_message, + ) + self.torch_profiler = None + self.profile_in_progress = False + self.profiler_start_forward_ct = None + + return ProfileReqOutput(success=True, message=f"Succeeded.{merge_message}") + + def _profile_batch_predicate(self, batch: ScheduleBatch): + if envs.SGLANG_PROFILE_V2.get(): + self._profile_manager.step(forward_mode=batch.forward_mode) + return + + if self.profile_by_stage: + if batch.forward_mode.is_prefill(): + if self.profiler_prefill_ct == 0: + self._start_profile(batch.forward_mode) + self.profiler_prefill_ct += 1 + if self.profiler_prefill_ct > self.profiler_target_prefill_ct: + if self.profile_in_progress: + self._stop_profile(stage=ForwardMode.EXTEND) + elif batch.forward_mode.is_decode(): + if self.profiler_decode_ct == 0: + if self.profile_in_progress: + # force trace flush + self._stop_profile(stage=ForwardMode.EXTEND) + self._start_profile(batch.forward_mode) + self.profiler_decode_ct += 1 + if self.profiler_decode_ct > self.profiler_target_decode_ct: + if self.profile_in_progress: + self._stop_profile(stage=ForwardMode.DECODE) + elif batch.forward_mode.is_idle(): + pass + else: + raise RuntimeError(f"unsupported profile stage: {batch.forward_mode}") + else: + # Check profiler + if ( + self.profiler_target_forward_ct + and self.profiler_target_forward_ct <= self.get_forward_ct() + ): + self._stop_profile() + if ( + self.profiler_start_forward_ct + and self.profiler_start_forward_ct == self.get_forward_ct() + ): + self._start_profile() + + def _profile(self, recv_req: ProfileReq): + if recv_req.type == ProfileReqType.START_PROFILE: + if recv_req.profile_by_stage or recv_req.start_step: + return self._init_profile( + recv_req.output_dir, + recv_req.start_step, + recv_req.num_steps, + recv_req.activities, + recv_req.with_stack, + recv_req.record_shapes, + recv_req.profile_by_stage, + recv_req.profile_id, + recv_req.merge_profiles, + recv_req.profile_prefix, + recv_req.profile_stages, + ) + else: + self._init_profile( + recv_req.output_dir, + recv_req.start_step, + recv_req.num_steps, + recv_req.activities, + recv_req.with_stack, + recv_req.record_shapes, + recv_req.profile_by_stage, + recv_req.profile_id, + recv_req.merge_profiles, + recv_req.profile_prefix, + ) + return self._start_profile() + else: + return self._stop_profile() diff --git a/python/sglang/srt/managers/scheduler_profiler_mixin.py b/python/sglang/srt/managers/scheduler_profiler_mixin.py deleted file mode 100644 index c1c714ab2..000000000 --- a/python/sglang/srt/managers/scheduler_profiler_mixin.py +++ /dev/null @@ -1,398 +0,0 @@ -from __future__ import annotations - -import logging -import os -import time -from pathlib import Path -from typing import TYPE_CHECKING, List, Optional - -import torch - -from sglang.srt.environ import envs -from sglang.srt.managers.io_struct import ProfileReq, ProfileReqOutput, ProfileReqType -from sglang.srt.model_executor.forward_batch_info import ForwardMode -from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils import is_npu -from sglang.srt.utils.profile_merger import ProfileMerger - -if TYPE_CHECKING: - from sglang.srt.managers.schedule_batch import ScheduleBatch - from sglang.srt.managers.scheduler_components.profiler_manager import ( - SchedulerProfilerManager, - ) - -_is_npu = is_npu() -if _is_npu: - import torch_npu - - patches = [ - ["profiler.profile", torch_npu.profiler.profile], - ["profiler.ProfilerActivity.CUDA", torch_npu.profiler.ProfilerActivity.NPU], - ["profiler.ProfilerActivity.CPU", torch_npu.profiler.ProfilerActivity.CPU], - ] - torch_npu._apply_patches(patches) - -logger = logging.getLogger(__name__) - - -class SchedulerProfilerMixin: - @staticmethod - def _init_profile( - self: "SchedulerProfilerManager", - output_dir: Optional[str], - start_step: Optional[int], - num_steps: Optional[int], - activities: Optional[List[str]], - with_stack: Optional[bool], - record_shapes: Optional[bool], - profile_by_stage: bool, - profile_id: str, - merge_profiles: bool = False, - profile_prefix: str = "", - profile_stages: Optional[List[str]] = None, - ) -> ProfileReqOutput: - if envs.SGLANG_PROFILE_V2.get(): - return self._profile_manager.configure( - output_dir=output_dir, - start_step=start_step, - num_steps=num_steps, - activities=activities, - with_stack=with_stack, - record_shapes=record_shapes, - profile_by_stage=profile_by_stage, - profile_id=profile_id, - merge_profiles=merge_profiles, - profile_prefix=profile_prefix, - profile_stages=profile_stages, - ) - - if self.profile_in_progress: - return ProfileReqOutput( - success=False, - message="Profiling is already in progress. Call /stop_profile first.", - ) - - self.profile_by_stage = profile_by_stage - self.merge_profiles = merge_profiles - - if output_dir is None: - output_dir = os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp") - if activities is None: - activities = ["CPU", "GPU"] - - self.torch_profiler_output_dir = Path(output_dir).expanduser() - self.torch_profiler_with_stack = with_stack - self.torch_profiler_record_shapes = record_shapes - self.profiler_activities = activities - self.profile_id = profile_id - self.profile_prefix = profile_prefix - - if start_step: - self.profiler_start_forward_ct = max(start_step, self.get_forward_ct() + 1) - - if num_steps: - if self.profile_by_stage: - self.profiler_prefill_ct = 0 - self.profiler_decode_ct = 0 - self.profiler_target_prefill_ct = num_steps - self.profiler_target_decode_ct = num_steps - elif start_step: - self.profiler_target_forward_ct = ( - self.profiler_start_forward_ct + num_steps - ) - else: - self.profiler_target_forward_ct = self.get_forward_ct() + num_steps - # The caller will be notified when reaching profiler_target_forward_ct - else: - self.profiler_target_forward_ct = None - - return ProfileReqOutput(success=True, message="Succeeded") - - @staticmethod - def _start_profile( - self: "SchedulerProfilerManager", stage: Optional[ForwardMode] = None - ) -> ProfileReqOutput | None: - if envs.SGLANG_PROFILE_V2.get(): - return self._profile_manager.manual_start() - - stage_str = f" for {stage.name}" if stage else "" - logger.info( - f"Profiling starts{stage_str}. Traces will be saved to: {self.torch_profiler_output_dir} (with profile id: {self.profile_id})", - ) - - activities = self.profiler_activities - with_stack = self.torch_profiler_with_stack - record_shapes = self.torch_profiler_record_shapes - - activity_map = { - "CPU": torch.profiler.ProfilerActivity.CPU, - "GPU": torch.profiler.ProfilerActivity.CUDA, - } - if hasattr(torch.profiler.ProfilerActivity, "XPU"): - activity_map["XPU"] = torch.profiler.ProfilerActivity.XPU - torchprof_activities = [ - activity_map[a] for a in activities if a in activity_map - ] - - if "RPD" in activities: # for ROCM - from rpdTracerControl import rpdTracerControl - - rpdTracerControl.skipCreate() - - self.rpd_profile_path = os.path.join( - self.torch_profiler_output_dir, - "rpd-" + str(time.time()) + f"-TP-{self.ps.tp_rank}" + ".trace.json.gz", - ) - - if self.ps.tp_rank == 0: - import sqlite3 - - from rocpd.schema import RocpdSchema - - if os.path.exists("trace.rpd"): - os.unlink("trace.rpd") - schema = RocpdSchema() - connection = sqlite3.connect("trace.rpd") - schema.writeSchema(connection) - connection.commit() - del connection - torch.distributed.barrier(self.dp_tp_cpu_group) - - self.rpd_profiler = rpdTracerControl() - self.rpd_profiler.setPythonTrace(True) - self.rpd_profiler.start() - self.rpd_profiler.rangePush("", "rpd profile range", "") - self.profile_in_progress = True - elif torchprof_activities: - self.torch_profiler = torch.profiler.profile( - activities=torchprof_activities, - with_stack=with_stack if with_stack is not None else True, - record_shapes=record_shapes if record_shapes is not None else False, - on_trace_ready=( - None - if not _is_npu - else torch_npu.profiler.tensorboard_trace_handler( - str(self.torch_profiler_output_dir) - ) - ), - ) - self.torch_profiler.start() - self.profile_in_progress = True - - if "MEM" in activities: - torch.cuda.memory._record_memory_history(max_entries=100000) - self.profile_in_progress = True - - if "CUDA_PROFILER" in activities: - if self.ps.gpu_id == get_global_server_args().base_gpu_id: - torch.cuda.cudart().cudaProfilerStart() - self.profile_in_progress = True - - return ProfileReqOutput(success=True, message="Succeeded") - - @staticmethod - def _merge_profile_traces(self: "SchedulerProfilerManager") -> str: - if not self.merge_profiles: - return "" - - if self.ps.tp_rank != 0: - return "" - if self.ps.dp_size > 1 and self.ps.dp_rank != 0: - return "" - if self.ps.pp_size > 1 and self.ps.pp_rank != 0: - return "" - if self.ps.moe_ep_size > 1 and self.ps.moe_ep_rank != 0: - return "" - - try: - logger.info("Starting profile merge...") - merger = ProfileMerger(self.torch_profiler_output_dir, self.profile_id) - merged_path = merger.merge_chrome_traces() - - summary = merger.get_merge_summary() - merge_message = ( - f" Merged trace: {merged_path} " - f"(Events: {summary.get('total_events', '?')}, " - f"Files: {summary.get('total_files', '?')})" - ) - - logger.info(f"Profile merge completed: {merged_path}") - except Exception as e: - logger.error(f"Failed to merge profiles: {e}", exc_info=True) - return f" Merge failed: {e!s}" - else: - return merge_message - - @staticmethod - def _stop_profile( - self: "SchedulerProfilerManager", stage: Optional[ForwardMode] = None - ) -> ProfileReqOutput | None: - if envs.SGLANG_PROFILE_V2.get(): - return self._profile_manager.manual_stop() - - if not self.profile_in_progress: - return ProfileReqOutput( - success=False, - message="Profiling is not in progress. Call /start_profile first.", - ) - - self.torch_profiler_output_dir.mkdir(parents=True, exist_ok=True) - - if self.profile_prefix: - stage_prefix = self.profile_prefix + "-" - else: - stage_prefix = "" - - stage_suffix = f"-{stage.name}" if stage else "" - logger.info("Stop profiling" + stage_suffix + "...") - if self.torch_profiler is not None: - self.torch_profiler.stop() - if not _is_npu: - # Build filename with only non-zero ranks to maintain backward compatibility - filename_parts = [self.profile_id, f"TP-{self.ps.tp_rank}"] - - # Only add other ranks if parallelism is enabled (size > 1) - if self.ps.dp_size > 1: - filename_parts.append(f"DP-{self.ps.dp_rank}") - if self.ps.pp_size > 1: - filename_parts.append(f"PP-{self.ps.pp_rank}") - if self.ps.moe_ep_size > 1: - filename_parts.append(f"EP-{self.ps.moe_ep_rank}") - - filename = ( - stage_prefix - + "-".join(filename_parts) - + stage_suffix - + ".trace.json.gz" - ) - - self.torch_profiler.export_chrome_trace( - os.path.join(self.torch_profiler_output_dir, filename) - ) - torch.distributed.barrier(self.dp_tp_cpu_group) - - if self.rpd_profiler is not None: - self.rpd_profiler.rangePop() - self.rpd_profiler.stop() - self.rpd_profiler.flush() - - torch.distributed.barrier(self.dp_tp_cpu_group) - if self.ps.tp_rank == 0: - from sglang.srt.utils.rpd_utils import rpd_to_chrome_trace - - rpd_to_chrome_trace("trace.rpd", self.rpd_profile_path) - self.rpd_profiler = None - self.rpd_profile_path = None - - if self.profiler_activities is not None and "MEM" in self.profiler_activities: - memory_profile_path = os.path.join( - self.torch_profiler_output_dir, - str(time.time()) - + f"-TP-{self.ps.tp_rank}-memory" - + stage_suffix - + ".pickle", - ) - torch.cuda.memory._dump_snapshot(memory_profile_path) - torch.cuda.memory._record_memory_history(enabled=None) - - if "CUDA_PROFILER" in self.profiler_activities: - if self.ps.gpu_id == get_global_server_args().base_gpu_id: - torch.cuda.cudart().cudaProfilerStop() - - merge_message = SchedulerProfilerMixin._merge_profile_traces(self) - - logger.info( - "Profiling done. Traces are saved to: %s%s", - self.torch_profiler_output_dir, - merge_message, - ) - self.torch_profiler = None - self.profile_in_progress = False - self.profiler_start_forward_ct = None - - return ProfileReqOutput(success=True, message=f"Succeeded.{merge_message}") - - @staticmethod - def _profile_batch_predicate( - self: "SchedulerProfilerManager", batch: ScheduleBatch - ): - if envs.SGLANG_PROFILE_V2.get(): - self._profile_manager.step(forward_mode=batch.forward_mode) - return - - if self.profile_by_stage: - if batch.forward_mode.is_prefill(): - if self.profiler_prefill_ct == 0: - SchedulerProfilerMixin._start_profile(self, batch.forward_mode) - self.profiler_prefill_ct += 1 - if self.profiler_prefill_ct > self.profiler_target_prefill_ct: - if self.profile_in_progress: - SchedulerProfilerMixin._stop_profile( - self, stage=ForwardMode.EXTEND - ) - elif batch.forward_mode.is_decode(): - if self.profiler_decode_ct == 0: - if self.profile_in_progress: - # force trace flush - SchedulerProfilerMixin._stop_profile( - self, stage=ForwardMode.EXTEND - ) - SchedulerProfilerMixin._start_profile(self, batch.forward_mode) - self.profiler_decode_ct += 1 - if self.profiler_decode_ct > self.profiler_target_decode_ct: - if self.profile_in_progress: - SchedulerProfilerMixin._stop_profile( - self, stage=ForwardMode.DECODE - ) - elif batch.forward_mode.is_idle(): - pass - else: - raise RuntimeError(f"unsupported profile stage: {batch.forward_mode}") - else: - # Check profiler - if ( - self.profiler_target_forward_ct - and self.profiler_target_forward_ct <= self.get_forward_ct() - ): - SchedulerProfilerMixin._stop_profile(self) - if ( - self.profiler_start_forward_ct - and self.profiler_start_forward_ct == self.get_forward_ct() - ): - SchedulerProfilerMixin._start_profile(self) - - @staticmethod - def _profile(self: "SchedulerProfilerManager", recv_req: ProfileReq): - if recv_req.type == ProfileReqType.START_PROFILE: - if recv_req.profile_by_stage or recv_req.start_step: - return SchedulerProfilerMixin._init_profile( - self, - recv_req.output_dir, - recv_req.start_step, - recv_req.num_steps, - recv_req.activities, - recv_req.with_stack, - recv_req.record_shapes, - recv_req.profile_by_stage, - recv_req.profile_id, - recv_req.merge_profiles, - recv_req.profile_prefix, - recv_req.profile_stages, - ) - else: - SchedulerProfilerMixin._init_profile( - self, - recv_req.output_dir, - recv_req.start_step, - recv_req.num_steps, - recv_req.activities, - recv_req.with_stack, - recv_req.record_shapes, - recv_req.profile_by_stage, - recv_req.profile_id, - recv_req.merge_profiles, - recv_req.profile_prefix, - ) - return SchedulerProfilerMixin._start_profile(self) - else: - return SchedulerProfilerMixin._stop_profile(self) diff --git a/test/registered/unit/utils/test_profile_merger.py b/test/registered/unit/utils/test_profile_merger.py index 9f211dbc0..709a34114 100644 --- a/test/registered/unit/utils/test_profile_merger.py +++ b/test/registered/unit/utils/test_profile_merger.py @@ -229,9 +229,11 @@ class TestProfileMergerIntegration(unittest.TestCase): self.assertIn("merge_profiles", sig.parameters) # Test SchedulerProfilerMixin - from sglang.srt.managers.scheduler_profiler_mixin import SchedulerProfilerMixin + from sglang.srt.managers.scheduler_components.profiler_manager import ( + SchedulerProfilerManager, + ) - sig = inspect.signature(SchedulerProfilerMixin._init_profile) + sig = inspect.signature(SchedulerProfilerManager._init_profile) self.assertIn("merge_profiles", sig.parameters) # Test CLI profiler