From 502cea3f2d85a8bf72936ae3b586dec99bf215b0 Mon Sep 17 00:00:00 2001 From: Xu Zhao Date: Wed, 6 Nov 2024 11:07:40 -0500 Subject: [PATCH] Handle child run in run_in_task --- benchmarks/nightly/run.py | 8 +------- tritonbench/utils/runner.py | 8 ++++++++ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/benchmarks/nightly/run.py b/benchmarks/nightly/run.py index 8702c37f..30804875 100644 --- 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 index ced7a13e..596ada7d 100644 --- 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))