[feature] Add start step profile argument in /start_profile (#7608)

This commit is contained in:
kyleliang-nv
2025-07-09 18:42:15 -07:00
committed by GitHub
parent 7590f5224b
commit dd445a41f5
6 changed files with 142 additions and 1 deletions

View File

@@ -905,6 +905,7 @@ class ProfileReqInput:
# If set, it profile as many as this number of steps.
# If it is set, profiling is automatically stopped after this step, and
# the caller doesn't need to run stop_profile.
start_step: Optional[int] = None
num_steps: Optional[int] = None
activities: Optional[List[str]] = None
profile_by_stage: bool = False
@@ -932,6 +933,7 @@ class ExpertDistributionReqOutput:
class ProfileReq:
type: ProfileReqType
output_dir: Optional[str] = None
start_step: Optional[int] = None
num_steps: Optional[int] = None
activities: Optional[List[str]] = None
profile_by_stage: bool = False