Unverified Commit ef21729c authored by fzyzcjy's avatar fzyzcjy Committed by GitHub
Browse files

Fix profiles do not have consistent names (#6811)

parent f5159315
...@@ -429,7 +429,7 @@ class Scheduler( ...@@ -429,7 +429,7 @@ class Scheduler(
self.torch_profiler = None self.torch_profiler = None
self.torch_profiler_output_dir: Optional[str] = None self.torch_profiler_output_dir: Optional[str] = None
self.profiler_activities: Optional[List[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_forward_ct: Optional[int] = None
self.profiler_target_prefill_ct: Optional[int] = None self.profiler_target_prefill_ct: Optional[int] = None
self.profiler_target_decode_ct: Optional[int] = None self.profiler_target_decode_ct: Optional[int] = None
...@@ -2145,6 +2145,7 @@ class Scheduler( ...@@ -2145,6 +2145,7 @@ class Scheduler(
recv_req.with_stack, recv_req.with_stack,
recv_req.record_shapes, recv_req.record_shapes,
recv_req.profile_by_stage, recv_req.profile_by_stage,
recv_req.profile_id,
) )
else: else:
self.init_profile( self.init_profile(
...@@ -2154,6 +2155,7 @@ class Scheduler( ...@@ -2154,6 +2155,7 @@ class Scheduler(
recv_req.with_stack, recv_req.with_stack,
recv_req.record_shapes, recv_req.record_shapes,
recv_req.profile_by_stage, recv_req.profile_by_stage,
recv_req.profile_id,
) )
return self.start_profile(True) return self.start_profile(True)
else: else:
...@@ -2167,6 +2169,7 @@ class Scheduler( ...@@ -2167,6 +2169,7 @@ class Scheduler(
with_stack: Optional[bool], with_stack: Optional[bool],
record_shapes: Optional[bool], record_shapes: Optional[bool],
profile_by_stage: bool, profile_by_stage: bool,
profile_id: str,
) -> ProfileReqOutput: ) -> ProfileReqOutput:
if self.profile_in_progress: if self.profile_in_progress:
return ProfileReqOutput( return ProfileReqOutput(
...@@ -2185,6 +2188,7 @@ class Scheduler( ...@@ -2185,6 +2188,7 @@ class Scheduler(
self.torch_profiler_with_stack = with_stack self.torch_profiler_with_stack = with_stack
self.torch_profiler_record_shapes = record_shapes self.torch_profiler_record_shapes = record_shapes
self.profiler_activities = activities self.profiler_activities = activities
self.profile_id = profile_id
if num_steps: if num_steps:
self.profile_steps = num_steps self.profile_steps = num_steps
...@@ -2284,7 +2288,7 @@ class Scheduler( ...@@ -2284,7 +2288,7 @@ class Scheduler(
self.torch_profiler.export_chrome_trace( self.torch_profiler.export_chrome_trace(
os.path.join( os.path.join(
self.torch_profiler_output_dir, self.torch_profiler_output_dir,
str(time.time()) self.profile_id
+ f"-TP-{self.tp_rank}" + f"-TP-{self.tp_rank}"
+ stage_suffix + stage_suffix
+ ".trace.json.gz", + ".trace.json.gz",
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment