From fd2bc906c1e843adbfa087fc77e8d2a442ff6533 Mon Sep 17 00:00:00 2001 From: Pate Motter Date: Mon, 20 Jul 2026 10:16:57 -0700 Subject: [PATCH] Fix pyconfig access regression in inference_microbenchmark_sweep --- .../inference_microbenchmark_sweep.py | 24 +++++++++++-------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/src/maxtext/inference/inference_microbenchmark_sweep.py b/src/maxtext/inference/inference_microbenchmark_sweep.py index afd8dd8317..c940e7d694 100644 --- a/src/maxtext/inference/inference_microbenchmark_sweep.py +++ b/src/maxtext/inference/inference_microbenchmark_sweep.py @@ -46,10 +46,10 @@ def main(): - flatten_microbenchmark_results: Whether or not to flatten results. Should be true """ - config = pyconfig.initialize(sys.argv) - base_run_name = config.run_name + base_config = pyconfig.initialize(sys.argv) + base_run_name = base_config.run_name - with open(config.inference_metadata_file, "rt", encoding="utf-8") as json_file: + with open(base_config.inference_metadata_file, "rt", encoding="utf-8") as json_file: inference_metadata = json.load(json_file) print(f"inference_metadata: {inference_metadata}") @@ -78,12 +78,16 @@ def main(): f"{two_axis_order_product_id}-{prefill_cache_axis_order.replace(',','')}-{ar_cache_axis_order.replace(',','')}" ) run_name = f"{base_run_name}/{run_tag}" - - tensorboard_dir = os.path.join(config.base_output_directory, run_name, "tensorboard", "") - pyconfig._config.keys["prefill_cache_axis_order"] = prefill_cache_axis_order # pylint: disable=protected-access - pyconfig._config.keys["ar_cache_axis_order"] = ar_cache_axis_order # pylint: disable=protected-access - pyconfig._config.keys["tensorboard_dir"] = tensorboard_dir # pylint: disable=protected-access - pyconfig._config.keys["run_name"] = run_name # pylint: disable=protected-access + tensorboard_dir = os.path.join(base_config.base_output_directory, run_name, "tensorboard", "") + + # Re-initialize configuration per sweep step using valid kwargs overrides + config = pyconfig.initialize( + sys.argv, + prefill_cache_axis_order=prefill_cache_axis_order, + ar_cache_axis_order=ar_cache_axis_order, + tensorboard_dir=tensorboard_dir, + run_name=run_name, + ) # Prepare metadata (dimensions) json for XLML dimensions_json = { @@ -114,7 +118,7 @@ def main(): "run_name": f"{run_name}", "run_tag": f"{run_tag}", "config_json_string": json.dumps( - pyconfig._config.keys, # pylint: disable=protected-access + config.get_keys(), default=lambda x: f"<>", ), }