Merge pull request #3299 from CIeNET-International:maxtxt/user/dora/benchmarks_workload_id

PiperOrigin-RevId: 882606227

Author: Google-ML-Automation

benchmarks/maxtext_xpk_runner.py

@@ -746,6 +746,7 @@ def xpk_benchmark_runner(
     command, name = generate_xpk_workload_cmd(
         cluster_config=cluster_config,
         wl_config=wl_config,
+        workload_name=wl_config.run_name,
         user=user,
         exp_name=exp_name,
     )

benchmarks/recipes/parser_utils.py

@@ -156,3 +156,4 @@ def add_arguments(parser: argparse.ArgumentParser):
   parser.add_argument("--delete", action="store_true", help="Delete the cluster workload")
   parser.add_argument("--max_restarts", type=int, default=0, help="Maximum number of restarts")
   parser.add_argument("--temp_key", type=str, default=None, help="Temporary placeholder code")
+  parser.add_argument("--workload_id", type=str, default=None, help="Optional custom name for the XPK workload.")

benchmarks/recipes/runner_utils.py

@@ -42,6 +42,7 @@ def _create_workload_config(
       "base_docker_image": (user_config.runner if Framework(framework) == Framework.MCJAX else None),
       "pathways_config": (user_config.pathways_config if Framework(framework) == Framework.PATHWAYS else None),
       "xpk_path": user_config.xpk_path,
+      "run_name": user_config.workload_id,
       "num_steps": num_steps,
       "priority": priority,
       "generate_metrics_and_upload_to_big_query": user_config.bq_enable,

benchmarks/recipes/user_configs.py

@@ -81,6 +81,7 @@ class UserConfig:
   delete: bool = False
   max_restarts: int = 0
   temp_key: str = None
+  workload_id: str = None
 
   def __post_init__(self):
     """Automatically generate derived attributes after the object is created."""