Fix multihost bugs and unify CLI entry point

Author: dipannita08

src/maxtext/eval/reporting/json_reporter.py

@@ -20,6 +20,7 @@
 import json
 import logging
 import os
+import tempfile
 
 logger = logging.getLogger(__name__)
 
@@ -50,13 +51,9 @@ def write_results(
       - results: The full results dict written to disk.
       - local_path: Absolute path of the written file.
   """
-  os.makedirs(results_path, exist_ok=True)
-
   timestamp = datetime.datetime.now(datetime.timezone.utc).strftime("%Y%m%dT%H%M%SZ")
-  # Create filename.
   safe_model = model_name.replace("/", "_").replace(":", "_")
   filename = f"{benchmark}_{safe_model}_{timestamp}.json"
-  local_path = os.path.join(results_path, filename)
 
   results = {
       "benchmark": benchmark,
@@ -67,7 +64,20 @@ def write_results(
       "config": config,
   }
 
-  with open(local_path, "w") as f:
-    json.dump(results, f, indent=2)
-  logger.info("Results written to %s", local_path)
+  if results_path.startswith("gs://"):
+    from maxtext.utils.gcs_utils import upload_blob  # pylint: disable=import-outside-toplevel
+    tmp_dir = tempfile.mkdtemp(prefix="eval_results_")
+    local_path = os.path.join(tmp_dir, filename)
+    with open(local_path, "w") as f:
+      json.dump(results, f, indent=2)
+    gcs_dest = f"{results_path.rstrip('/')}/{filename}"
+    upload_blob(gcs_dest, local_path)
+    logger.info("Results written to %s", gcs_dest)
+  else:
+    os.makedirs(results_path, exist_ok=True)
+    local_path = os.path.join(results_path, filename)
+    with open(local_path, "w") as f:
+      json.dump(results, f, indent=2)
+    logger.info("Results written to %s", local_path)
+
   return {"results": results, "local_path": os.path.abspath(local_path)}

src/maxtext/eval/runner/async_client.py

@@ -89,17 +89,17 @@ async def _generate_one(session: aiohttp.ClientSession, prompt: str) -> Generati
         "max_tokens": max_tokens,
         "temperature": temperature,
     }
-    t0 = time.monotonic()
     async with semaphore:
+      t0 = time.monotonic()
       try:
         async with session.post(api_url, json=payload) as resp:
           if resp.status != 200:
             body = await resp.text()
             return GenerationResult(error=f"HTTP {resp.status}: {body[:200]}")
           data = await resp.json()
-      except aiohttp.ClientError as exc:
+      except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
         return GenerationResult(error=str(exc))
-    latency = time.monotonic() - t0
+      latency = time.monotonic() - t0
 
     choice = data["choices"][0]
     usage = data.get("usage", {})

src/maxtext/eval/runner/common.py

@@ -0,0 +1,135 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared helpers for MaxText eval runners."""
+
+from __future__ import annotations
+
+import argparse
+import os
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+  from maxtext.eval.runner.server_manager import VllmServerManager
+
+ENABLE_EXPERT_PARALLEL_HELP = (
+    "Enable expert parallelism in vLLM. Required for MoE models such as "
+    "qwen3-30b-a3b, qwen3-235b-a22b, deepseek-v3, etc. Without this flag "
+    "tpu-inference omits the 'expert' mesh axis and MaxText's MoE sharding "
+    "raises KeyError."
+)
+
+
+def resolve_token(cfg: dict, hf_token: str | None) -> str | None:
+  """Return HF token from explicit arg or HF_TOKEN env var."""
+  return hf_token or os.environ.get("HF_TOKEN") or None
+
+
+def build_server_manager(cfg: dict, token: str | None) -> "VllmServerManager":
+  """Build a VllmServerManager from a merged config dict.
+
+  Handles token forwarding, MaxText adapter vs HF mode selection, and the
+  enable_expert_parallel to additional_config sharding injection.
+
+  Args:
+    cfg: Merged configuration dict. Required key: max_model_len. Common
+      optional keys: tensor_parallel_size, server_host, server_port,
+      max_num_batched_tokens, max_num_seqs, hf_mode, enable_expert_parallel.
+    token: HuggingFace token (or None).
+
+  Returns:
+    A VllmServerManager instance ready for use as a context manager (unstarted).
+  """
+  from maxtext.eval.runner.server_manager import VllmServerManager  # pylint: disable=import-outside-toplevel
+
+  hf_path = cfg["hf_path"]
+  model_name = cfg["model_name"]
+  checkpoint_path = cfg.get("checkpoint_path")
+  hf_mode = cfg.get("hf_mode", False)
+  use_maxtext_adapter = bool(checkpoint_path) and not hf_mode
+
+  tensor_parallel_size = int(cfg.get("tensor_parallel_size", 4))
+  max_model_len = int(cfg["max_model_len"])
+  server_host = cfg.get("server_host", "localhost")
+  server_port = int(cfg.get("server_port", 8000))
+
+  max_num_batched_tokens = cfg.get("max_num_batched_tokens")
+  if max_num_batched_tokens is not None:
+    max_num_batched_tokens = int(max_num_batched_tokens)
+  max_num_seqs = cfg.get("max_num_seqs")
+  if max_num_seqs is not None:
+    max_num_seqs = int(max_num_seqs)
+
+  server_env = {"HF_TOKEN": token} if token else None
+  additional_vllm_kwargs: dict = {}
+  if cfg.get("enable_expert_parallel"):
+    additional_vllm_kwargs["enable_expert_parallel"] = True
+
+  return VllmServerManager(
+      model_path=hf_path,
+      checkpoint_path=checkpoint_path if use_maxtext_adapter else None,
+      maxtext_model_name=model_name if use_maxtext_adapter else None,
+      host=server_host,
+      port=server_port,
+      tensor_parallel_size=tensor_parallel_size,
+      max_model_len=max_model_len,
+      max_num_batched_tokens=max_num_batched_tokens,
+      max_num_seqs=max_num_seqs,
+      env=server_env,
+      additional_vllm_kwargs=additional_vllm_kwargs or None,
+  )
+
+
+def maybe_upload_to_gcs(output: dict, gcs_results_path: str | None) -> None:
+  """Upload the results JSON to GCS if gcs_results_path is provided."""
+  if gcs_results_path:
+    from maxtext.eval.reporting.gcs_reporter import upload_results  # pylint: disable=import-outside-toplevel
+    upload_results(output["local_path"], gcs_results_path)
+
+
+def add_server_args(parser: argparse.ArgumentParser) -> None:
+  """Add the server/model CLI args shared by all eval runner parsers."""
+  parser.add_argument("--checkpoint_path", help="MaxText orbax checkpoint path (/0/items).")
+  parser.add_argument("--model_name", required=True, help="MaxText model name (e.g. llama3.1-8b).")
+  parser.add_argument("--hf_path", required=True, help="HF model ID or local tokenizer dir.")
+  parser.add_argument(
+      "--base_output_directory",
+      required=True,
+      help="Base output directory (local path or gs://<bucket>/).",
+  )
+  parser.add_argument("--run_name", required=True, help="Run name/identifier.")
+  parser.add_argument("--max_model_len", type=int, required=True, help="vLLM max context length.")
+  parser.add_argument(
+      "--tensor_parallel_size", type=int, default=4, help="vLLM tensor parallelism."
+  )
+  parser.add_argument("--server_host", default="localhost", help="vLLM server bind host.")
+  parser.add_argument("--server_port", type=int, default=8000, help="vLLM server port.")
+  parser.add_argument(
+      "--max_num_batched_tokens", type=int, help="vLLM tokens per scheduler step."
+  )
+  parser.add_argument("--max_num_seqs", type=int, help="vLLM max concurrent sequences.")
+  parser.add_argument("--hf_mode", action="store_true", help="HF safetensors mode.")
+  parser.add_argument(
+      "--enable_expert_parallel", action="store_true", help=ENABLE_EXPERT_PARALLEL_HELP
+  )
+  parser.add_argument("--hf_token", help="HuggingFace token for gated models.")
+  parser.add_argument(
+      "--gcs_results_path", help="Optional secondary GCS path to upload the results JSON."
+  )
+  parser.add_argument(
+      "--log_level",
+      default="INFO",
+      choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+      help="Logging level.",
+  )

src/maxtext/eval/runner/eval_runner.py

@@ -12,40 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""CLI entry point for model evaluation.
-
-MaxTextForCausalLM mode (preferred):
-Load weights directly from the MaxText checkpoint, no HuggingFace weight
-conversion required. Flag --hf_path supplies the tokenizer (HF model ID
-or local tokenizer dir).
-
-  python -m maxtext.eval.runner.eval_runner \
-      --config src/maxtext/eval/configs/mlperf.yml \
-      --base_config src/maxtext/configs/base.yml  \
-      --base_output_directory gs://<gcs_bucket>/ \
-      --run_name my_run \
-      --checkpoint_path gs://<gcs_bucket>/checkpoint/0/items \
-      --model_name llama3.1-8b \
-      --hf_path meta-llama/Llama-3.1-8B-Instruct
-
-HuggingFace safetensors mode:
-Use --hf_mode and point --hf_path to an existing HF model directory.
-
-  python -m maxtext.eval.runner.eval_runner \
-      --config src/maxtext/eval/configs/mlperf.yml \
-      --hf_path TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
-      --model_name tinyllama \
-      --hf_mode \
-      --base_output_directory /tmp/eval/ \
-      --run_name smoke_test \
-      --tensor_parallel_size 1
+"""Custom dataset eval runner (MLPerf OpenOrca, ROUGE scoring).
+
+Unified entry point:
+
+  python -m maxtext.eval.runner.run --runner eval ...
 """
 
 from __future__ import annotations
 
 import argparse
 import logging
-import os
 import time
 
 import yaml
@@ -116,7 +93,7 @@ def run_eval(cfg: dict, hf_token: str | None = None) -> dict:
   from maxtext.eval.datasets.registry import get_dataset
   from maxtext.eval.reporting.json_reporter import write_results
   from maxtext.eval.runner.async_client import generate_batch
-  from maxtext.eval.runner.server_manager import VllmServerManager
+  from maxtext.eval.runner.common import build_server_manager, maybe_upload_to_gcs, resolve_token
   from maxtext.eval.runner.warmup import warmup_server
   from maxtext.eval.scoring.registry import get_scorer
 
@@ -128,27 +105,10 @@ def run_eval(cfg: dict, hf_token: str | None = None) -> dict:
   max_tokens = int(cfg.get("max_tokens", 1024))
   temperature = float(cfg.get("temperature", 0.0))
   concurrency = int(cfg.get("concurrency", 64))
-  tensor_parallel_size = int(cfg.get("tensor_parallel_size", 4))
   if "max_model_len" not in cfg:
-    raise ValueError(
-        "Error: max_model_len is required."
-    )
-  max_model_len = int(cfg["max_model_len"])
-  server_host = cfg.get("server_host", "localhost")
-  server_port = int(cfg.get("server_port", 8000))
-  max_num_batched_tokens = cfg.get("max_num_batched_tokens")
-  if max_num_batched_tokens is not None:
-    max_num_batched_tokens = int(max_num_batched_tokens)
-  max_num_seqs = cfg.get("max_num_seqs")
-  if max_num_seqs is not None:
-    max_num_seqs = int(max_num_seqs)
+    raise ValueError("Error: max_model_len is required.")
   gcs_results_path = cfg.get("gcs_results_path")
-  token = hf_token or os.environ.get("HF_TOKEN") or None
-  checkpoint_path = cfg.get("checkpoint_path")
-  hf_mode = cfg.get("hf_mode", False)
-
-  # Determine loading mode.
-  use_maxtext_adapter = bool(checkpoint_path) and not hf_mode
+  token = resolve_token(cfg, hf_token)
 
   # Load tokenizer for prompt formatting.
   logger.info("Loading tokenizer from %s.", hf_path)
@@ -164,42 +124,40 @@ def run_eval(cfg: dict, hf_token: str | None = None) -> dict:
   references = [r.reference for r in requests]
 
   # Start vLLM server.
-  server_env = {"HF_TOKEN": token} if token else None
-  additional_vllm_kwargs = {}
-  if cfg.get("enable_expert_parallel"):
-    additional_vllm_kwargs["enable_expert_parallel"] = True
-
-  with VllmServerManager(
-      model_path=hf_path,
-      checkpoint_path=checkpoint_path if use_maxtext_adapter else None,
-      maxtext_model_name=model_name if use_maxtext_adapter else None,
-      host=server_host,
-      port=server_port,
-      tensor_parallel_size=tensor_parallel_size,
-      max_model_len=max_model_len,
-      max_num_batched_tokens=max_num_batched_tokens,
-      max_num_seqs=max_num_seqs,
-      env=server_env,
-      additional_vllm_kwargs=additional_vllm_kwargs or None,
-  ) as server:
-    base_url = server.base_url
-
-    # Warmup server.
-    warmup_server(base_url=base_url, model=model_name, sample_requests=requests)
-
-    # Generate responses.
-    logger.info("Generating responses for %d prompts.", len(prompts))
-    t0 = time.time()
-    results = generate_batch(
-        prompts=prompts,
-        base_url=base_url,
-        model=model_name,
-        max_tokens=max_tokens,
-        temperature=temperature,
-        concurrency=concurrency,
-    )
-    elapsed = time.time() - t0
-    logger.info("Generation completed in %.1fs (%.1f samples/s).", elapsed, len(prompts) / elapsed)
+  with build_server_manager(cfg, token) as server:
+    import jax as _jax  # pylint: disable=import-outside-toplevel
+    from jax.experimental import multihost_utils as _multihost_utils  # pylint: disable=import-outside-toplevel
+    is_rank0 = _jax.process_index() == 0
+
+    if is_rank0:
+      base_url = server.base_url
+
+      # Warmup server.
+      warmup_server(base_url=base_url, model=model_name, sample_requests=requests)
+
+      # Generate responses.
+      logger.info("Generating responses for %d prompts.", len(prompts))
+      t0 = time.time()
+      results = generate_batch(
+          prompts=prompts,
+          base_url=base_url,
+          model=model_name,
+          max_tokens=max_tokens,
+          temperature=temperature,
+          concurrency=concurrency,
+      )
+      elapsed = time.time() - t0
+      logger.info("Generation completed in %.1fs (%.1f samples/s).", elapsed, len(prompts) / elapsed)
+
+    # All ranks block here until rank-0 finishes generation. Non-rank-0 hosts
+    # keep their in-process LLM alive so rank-0's llm.generate() calls can
+    # complete their tensor-parallel collectives across all hosts.
+    _multihost_utils.sync_global_devices("eval_runner_complete")
+
+  # All ranks exit the context manager together above (LLM stopped on all).
+  # Only rank-0 has results/elapsed defined, non-rank-0 return early.
+  if not is_rank0:
+    return {}
 
   # Score.
   responses = [r.text for r in results]
@@ -229,11 +187,7 @@ def run_eval(cfg: dict, hf_token: str | None = None) -> dict:
       results_path=results_path,
   )
 
-  # Optional GCS Upload.
-  if gcs_results_path:
-    from maxtext.eval.reporting.gcs_reporter import upload_results  # pylint: disable=import-outside-toplevel
-    upload_results(output["local_path"], gcs_results_path)
-
+  maybe_upload_to_gcs(output, gcs_results_path)
   return output