[FEATURE] Add Profile Trace Merger for Distributed Traces (#11413)

python/sglang/profiler.py

@@ -25,6 +25,7 @@ def _run_profile(
     output_dir: Optional[str] = None,
     profile_name: Optional[str] = None,
     profile_by_stage: bool = False,
+    merge_profiles: bool = False,
 ) -> str:
     if output_dir is None:
         output_dir = PROFILER_DIR
@@ -60,6 +61,7 @@ def _run_profile(
         "num_steps": str(num_steps),
         "activities": activities,
         "profile_by_stage": profile_by_stage,
+        "merge_profiles": merge_profiles,
     }
 
     response = requests.post(url=url + "/start_profile", json=json_data)
@@ -76,10 +78,17 @@ def run_profile(
     output_dir: Optional[str] = None,
     profile_name: Optional[str] = None,
     profile_by_stage: bool = False,
+    merge_profiles: bool = False,
 ):
     # step based profile will self terminate on num_steps constraints
     link = _run_profile(
-        url, num_steps, activities, output_dir, profile_name, profile_by_stage
+        url,
+        num_steps,
+        activities,
+        output_dir,
+        profile_name,
+        profile_by_stage,
+        merge_profiles,
     )
     return link
 
@@ -145,6 +154,13 @@ if __name__ == "__main__":
         default=False,
         help="Whether to use rpd profiler (https://github.com/ROCm/rocmProfileData)",
     )
+    parser.add_argument(
+        "--merge-profiles",
+        action=argparse.BooleanOptionalAction,
+        type=bool,
+        default=False,
+        help="Whether to merge profiles from all ranks into a single trace file",
+    )
 
     args = parser.parse_args()
     activities = []
@@ -163,4 +179,5 @@ if __name__ == "__main__":
         args.output_dir,
         args.profile_name,
         args.profile_by_stage,
+        args.merge_profiles,
     )