From b295ca7ae20cf625caf72b4b928a5a2780b4379b Mon Sep 17 00:00:00 2001 From: Colin Gaffney Date: Wed, 1 Jul 2026 08:07:23 -0700 Subject: [PATCH] Internal change PiperOrigin-RevId: 941101071 --- .../checkpoint/_src/testing/benchmarks/xpk/launch_xpk.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/checkpoint/orbax/checkpoint/_src/testing/benchmarks/xpk/launch_xpk.py b/checkpoint/orbax/checkpoint/_src/testing/benchmarks/xpk/launch_xpk.py index 8a583c884f..76552f348d 100644 --- a/checkpoint/orbax/checkpoint/_src/testing/benchmarks/xpk/launch_xpk.py +++ b/checkpoint/orbax/checkpoint/_src/testing/benchmarks/xpk/launch_xpk.py @@ -721,8 +721,13 @@ def construct_workload_command( env_cmd = ' && '.join(env_vars) + ' && ' if env_vars else '' + if benchmark_binary_path.endswith('.py'): + cmd_executable = f'python3 {benchmark_binary_path}' + else: + cmd_executable = benchmark_binary_path + python_args = [ - f'python3 {benchmark_binary_path}', + cmd_executable, f'--config_file={config_file}', f'--output_directory={os.path.join(output_directory, run_id)}', '--alsologtostderr',