Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions python/sglang/srt/managers/scheduler.py
Original file line number Diff line number Diff line change
Expand Up @@ -429,7 +429,7 @@ def __init__(
self.torch_profiler = None
self.torch_profiler_output_dir: Optional[str] = None
self.profiler_activities: Optional[List[str]] = None
self.profiler_id: Optional[str] = None
self.profile_id: Optional[str] = None
self.profiler_target_forward_ct: Optional[int] = None
self.profiler_target_prefill_ct: Optional[int] = None
self.profiler_target_decode_ct: Optional[int] = None
Expand Down Expand Up @@ -2144,6 +2144,7 @@ def profile(self, recv_req: ProfileReq):
recv_req.with_stack,
recv_req.record_shapes,
recv_req.profile_by_stage,
recv_req.profile_id,
)
else:
self.init_profile(
Expand All @@ -2153,6 +2154,7 @@ def profile(self, recv_req: ProfileReq):
recv_req.with_stack,
recv_req.record_shapes,
recv_req.profile_by_stage,
recv_req.profile_id,
)
return self.start_profile(True)
else:
Expand All @@ -2166,6 +2168,7 @@ def init_profile(
with_stack: Optional[bool],
record_shapes: Optional[bool],
profile_by_stage: bool,
profile_id: str,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The profile_id parameter is typed as str here, but it's derived from recv_req.profile_id which is defined as Optional[str] in the ProfileReq dataclass (in io_struct.py).

If recv_req.profile_id is None (which is its default value if not provided by the client), passing None to this parameter, which expects a str, could lead to unexpected behavior or type-related issues, especially if strict type checking were enforced or if subsequent code assumes it's always a string without further checks.

Consider changing the type hint to Optional[str] to accurately reflect that None might be passed. The handling of a None value should then be done explicitly in the method body (see related comment for line 2190).

) -> ProfileReqOutput:
if self.profile_in_progress:
return ProfileReqOutput(
Expand All @@ -2184,6 +2187,7 @@ def init_profile(
self.torch_profiler_with_stack = with_stack
self.torch_profiler_record_shapes = record_shapes
self.profiler_activities = activities
self.profile_id = profile_id
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

If profile_id (the parameter to init_profile) can be None (as suggested by its source ProfileReq.profile_id: Optional[str]), then assigning it directly to self.profile_id would result in self.profile_id being None.

This would subsequently cause a TypeError in the stop_profile method at line 2290, where self.profile_id is used in string concatenation to form the trace filename: self.profile_id + f"-TP-{self.tp_rank}" ....

To prevent this, it's crucial to ensure self.profile_id is always a string. A robust way to handle this is to assign a default generated ID if the input profile_id is None.

        if profile_id is None:
            # Generate a default profile ID if not provided.
            # Using a timestamp and a short random hex ensures reasonable uniqueness.
            self.profile_id = f"profile_{int(time.time())}_{os.urandom(4).hex()}"
        else:
            self.profile_id = profile_id


if num_steps:
self.profile_steps = num_steps
Expand Down Expand Up @@ -2283,7 +2287,7 @@ def stop_profile(
self.torch_profiler.export_chrome_trace(
os.path.join(
self.torch_profiler_output_dir,
str(time.time())
self.profile_id
+ f"-TP-{self.tp_rank}"
+ stage_suffix
+ ".trace.json.gz",
Expand Down
Loading