Merge pull request #1708 from AI-Hypercomputer:sujinesh/support_headless_mode_recipe

PiperOrigin-RevId: 771265392

Author: maxtext authors

benchmarks/maxtext_xpk_runner.py

@@ -78,6 +78,7 @@ class PathwaysConfig:
   server_flags: str = ''
   proxy_flags: str = ''
   worker_flags: str = ''
+  headless: bool = False
 
 
 # TODO(@vbarr): Split out parameters related to XPK workload and a General workload
@@ -446,7 +447,7 @@ def _get_pathways_proxy_flags(wl_config: WorkloadConfig):
 
   # Get proxy and xla flag string from model config
   proxy_flags_string = pw_config.proxy_flags
-  xla_flags_string = wl_config.model.xla_flags
+  xla_flags_string = wl_config.model.xla_flags if not pw_config.headless else ''
 
   # Split both proxy_flags_string and xla_flags_string into lists of flags
   proxy_flags_list = proxy_flags_string.strip().split()
@@ -457,8 +458,8 @@ def _get_pathways_proxy_flags(wl_config: WorkloadConfig):
 
   # Remove the flags that are specified to be removed.
   if (
-      wl_config.model.pathways_xla_flag_options
-      and xla_flags.REMOVE in wl_config.model.pathways_xla_flag_options
+      not pw_config.headless and (wl_config.model.pathways_xla_flag_options
+      and xla_flags.REMOVE in wl_config.model.pathways_xla_flag_options)
   ):
     flags_to_remove = wl_config.model.pathways_xla_flag_options[
         xla_flags.REMOVE
@@ -471,8 +472,8 @@ def _get_pathways_proxy_flags(wl_config: WorkloadConfig):
 
   # Add the flags that are specified to be added.
   if (
-      wl_config.model.pathways_xla_flag_options
-      and xla_flags.ADD_PROXY in wl_config.model.pathways_xla_flag_options
+      not pw_config.headless and (wl_config.model.pathways_xla_flag_options
+      and xla_flags.ADD_PROXY in wl_config.model.pathways_xla_flag_options)
   ):
     flags_to_add = wl_config.model.pathways_xla_flag_options[
         xla_flags.ADD_PROXY
@@ -500,8 +501,8 @@ def _get_pathways_worker_flags(wl_config: WorkloadConfig):
 
   # Add the flags that are specified to be added.
   if (
-      wl_config.model.pathways_xla_flag_options
-      and xla_flags.ADD_WORKER in wl_config.model.pathways_xla_flag_options
+      not pw_config.headless and (wl_config.model.pathways_xla_flag_options
+      and xla_flags.ADD_WORKER in wl_config.model.pathways_xla_flag_options)
   ):
     flags_to_add = wl_config.model.pathways_xla_flag_options[
         xla_flags.ADD_WORKER
@@ -523,8 +524,8 @@ def _get_pathways_server_flags(wl_config: WorkloadConfig):
 
   # Add the flags that are specified to be added.
   if (
-      wl_config.model.pathways_xla_flag_options
-      and xla_flags.ADD_SERVER in wl_config.model.pathways_xla_flag_options
+      not pw_config.headless and (wl_config.model.pathways_xla_flag_options
+      and xla_flags.ADD_SERVER in wl_config.model.pathways_xla_flag_options)
   ):
     flags_to_add = wl_config.model.pathways_xla_flag_options[
         xla_flags.ADD_SERVER
@@ -569,6 +570,7 @@ def _get_pathways_specific_flags(wl_config: WorkloadConfig):
       f' --custom-pathways-server-args="{server_flags}" '
       f' --custom-pathways-proxy-server-args="{proxy_flags}" '
       f' --custom-pathways-worker-args="{worker_flags}" '
+      f' {"--headless" if pw_config.headless else ""}'
   )
   return pathways_specific_flags
 
@@ -582,6 +584,7 @@ def generate_xpk_workload_cmd(
   """Generates a command to run a maxtext model on XPK."""
 
   is_pathways_enabled = wl_config.pathways_config is not None
+  is_pathways_headless_enabled = wl_config.pathways_config and wl_config.pathways_config.headless
 
   time.localtime()
   length_of_random_str = 3
@@ -614,10 +617,12 @@ def generate_xpk_workload_cmd(
                                             wl_config.run_name,
                                             'metrics')
 
-  user_command = build_user_command(
-      name=name,
-      wl_config=wl_config
-  )
+  user_command = ''
+  if not is_pathways_headless_enabled:
+    user_command = build_user_command(
+        name=name,
+        wl_config=wl_config
+    )
 
   additional_flags = ''
   if not is_pathways_enabled and wl_config.libtpu_type == LibTpuType.CUSTOM:
@@ -641,7 +646,7 @@ def generate_xpk_workload_cmd(
     docker_image_flag = f'--docker-image="{wl_config.base_docker_image}"'
 
   upload_metrics_to_bq_cmd = ""
-  if wl_config.generate_metrics_and_upload_to_big_query:
+  if wl_config.generate_metrics_and_upload_to_big_query and not is_pathways_headless_enabled:
     # TODO (optionally) make it so that this upload step is done on local device instead of within the workload.
     args = _build_args_from_config(wl_config)
     args_str = ""

benchmarks/recipes/pw_headless_mode.py

@@ -0,0 +1,63 @@
+"""Copyright 2025 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.
+"""
+
+import os
+import sys
+
+import benchmarks.recipes.args_helper as helper
+import maxtext_xpk_runner as mxr
+from recipes.user_configs import cluster_config, xpk_path, pathways_config, base_output_directory, headless_workload_name
+
+
+def main() -> int:
+  # Handle command line arguments using args_helper
+  should_continue = helper.handle_cmd_args(
+      cluster_config, helper.DELETE, xpk_path=xpk_path
+  )
+
+  if not should_continue:
+    return 0
+
+  num_slices = 2
+
+  # Run workloads in the following slice configurations
+  wl_config = mxr.WorkloadConfig(
+      model=None,
+      num_slices=num_slices,
+      device_type=cluster_config.device_type,
+      base_output_directory=base_output_directory,
+      max_restarts=0,
+      libtpu_type=None,
+      libtpu_nightly_version="",
+      base_docker_image="",
+      pathways_config=pathways_config,
+      xpk_path=xpk_path,
+  )
+  command, name = mxr.generate_xpk_workload_cmd(
+      cluster_config=cluster_config,
+      wl_config=wl_config,
+      workload_name=headless_workload_name,
+  )
+
+  print(f"Name of the workload is: {name} \n")
+  print(f"XPK command to be used is: {command} \n")
+
+  return_code = mxr.run_command_with_updates(command, name)
+  if return_code != 0:
+    print(f"Unable to run xpk workload: {name}")
+
+
+if __name__ == "__main__":
+  main()

benchmarks/recipes/user_configs.py

@@ -0,0 +1,50 @@
+"""Copyright 2025 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.
+"""
+
+"""Define user specific configurations for recipes here."""
+
+import os
+
+import maxtext_xpk_runner as mxr
+from xpk_configs import XpkClusterConfig
+
+cluster_config = XpkClusterConfig(
+    cluster_name="test-v5e-32-cluster",
+    project="cloud-tpu-cluster",
+    zone="us-south1-a",
+    device_type="v5litepod-32",
+)
+xpk_path = "~/xpk"
+
+user = os.environ["USER"]
+region = "-".join(cluster_config.zone.split("-")[:-1])
+proxy_image = (
+    f"us-docker.pkg.dev/path/to/{user}/proxy_server"
+)
+server_image = (
+    f"us-docker.pkg.dev/path/to/{user}/server"
+)
+colocated_python_image = f"gcr.io/{cluster_config.project}/path/to/{user}/colocated_python_sidecar"
+runner = f"gcr.io/{cluster_config.project}/{user}_maxtext_latest:latest"
+base_output_directory = f"gs://{user}-{region}/{user}"
+headless = True
+pathways_config = mxr.PathwaysConfig(
+    server_image=server_image,
+    proxy_server_image=proxy_image,
+    runner_image=runner,
+    colocated_python_sidecar_image=colocated_python_image,
+    headless=headless,
+)
+headless_workload_name = f"{user[:3]}-headless"