Handle child run in run_in_task

pytorch-labs · Nov 21, 2024 · 502cea3 · 502cea3
1 parent b84cfee
commit 502cea3
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/benchmarks/nightly/run.py b/benchmarks/nightly/run.py
@@ -32,13 +32,7 @@ def setup_tritonbench_cwd():
 
 def run():
     setup_tritonbench_cwd()
-    from tritonbench.utils.parser import get_parser
-    from tritonbench.utils.runner import run_in_task, tritonbench_run
-    if "--op" in sys.argv:
-        parser = get_parser()
-        args, extra_args = parser.parse_known_args(sys.argv[1:])
-        tritonbench_run(args, extra_args)
-        return
+    from tritonbench.utils.runner import run_in_task
     for op in OPERATORS:
         run_in_task(op)
 

diff --git a/tritonbench/utils/runner.py b/tritonbench/utils/runner.py
@@ -69,11 +69,19 @@ def tritonbench_run(args: argparse.Namespace, extra_args: List[str]) -> Benchmar
         return metrics
 
 def run_in_task(op: str) -> None:
+    if "--child" in sys.argv:
+        sys.argv = remove_cmd_parameter(sys.argv, "--child")
+        from tritonbench.utils.parser import get_parser
+        parser = get_parser()
+        args, extra_args = parser.parse_known_args(sys.argv[1:])
+        tritonbench_run(args, extra_args)
+        return
     op_task_cmd = [] if IS_FBCODE else [sys.executable]
     copy_sys_argv = copy.deepcopy(sys.argv)
     copy_sys_argv = remove_cmd_parameter(copy_sys_argv, "--op")
     copy_sys_argv = remove_cmd_parameter(copy_sys_argv, "--isolate")
     add_cmd_parameter(copy_sys_argv, "--op", op)
+    add_cmd_parameter(copy_sys_argv, "--child")
     op_task_cmd.extend(copy_sys_argv)
     try:
         print("[tritonbench] running command: " + " ".join(op_task_cmd))