Merge pull request #1753 from SamuelMarks:inference_microbenchmark-EXIT_CODE

PiperOrigin-RevId: 761574119

Author: maxtext authors

MaxText/inference_microbenchmark.py

@@ -19,6 +19,7 @@
 import jax
 import json
 import os
+import sys
 
 from absl import app
 from collections.abc import MutableMapping
@@ -426,10 +427,14 @@ def run_benchmarks(config):
   return results
 
 
-def main(config, **kwargs):
+def run_benchmarks_with_unsafe_rbg(config, **kwargs):
   jax.config.update("jax_default_prng_impl", "unsafe_rbg")
   return run_benchmarks(pyconfig.initialize(config, **kwargs))
 
 
+def main(config, **kwargs):
+  json.dump(run_benchmarks_with_unsafe_rbg(config, **kwargs), sys.stdout)
+
+
 if __name__ == "__main__":
   app.run(main)

MaxText/inference_microbenchmark_sweep.py

@@ -123,7 +123,9 @@ def main():
         **inference_metadata,
     }
     try:
-      microbenchmark_results = inference_microbenchmark.main(config, inference_metadata=inference_metadata)
+      microbenchmark_results = inference_microbenchmark.run_benchmarks_with_unsafe_rbg(
+          config, inference_metadata=inference_metadata
+      )
       if microbenchmark_results:
         metrics = microbenchmark_results["flattened_results"]
         metrics = {k.lower(): v for k, v in metrics.items()}