Merge

SujeethJinesh · SujeethJinesh · commit ddfe758b33fc · 2025-01-30T12:52:16.000-08:00
Remote Python Support for Benchmark Runner

Add Remote Python Recipe

Clean up

Add Delete Capabilities

Minor edits

Fix

Fix problems

Changes

Changes

Fixes

Undo Weird Merge

Save work

Break out args helper

Fix file imports

Move around stuff

cleanup

Fix

Clean up comment

whitespace

Fix issues

Fix lint
diff --git a/benchmarks/benchmark_runner.py b/benchmarks/benchmark_runner.py
@@ -126,6 +126,11 @@ def add_xpk_runner_arguments(custom_parser: argparse.ArgumentParser):
       default='us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest',
       help='version of pathways runner image to be benchmarked command.',
   )
+  custom_parser.add_argument(
+      '--remote_python_sidecar_image',
+      type=str,
+      help='version of remote python sidecar image to be benchmarked command.',
+  )
   custom_parser.add_argument(
       '--use_pathways',
       type=bool,
@@ -246,6 +251,7 @@ def main() -> None:
         server_image=options.pathways_server_image,
         proxy_image=options.pathways_proxy_image,
         runner_image=options.pathways_runner_image,
+        remote_python_sidecar_image=options.remote_python_sidecar_image,
       )
 
     workload_config = WorkloadConfig(
diff --git a/benchmarks/maxtext_trillium_model_configs.py b/benchmarks/maxtext_trillium_model_configs.py
@@ -60,6 +60,38 @@ def _add_to_model_dictionary(model_dictionary: dict[str, MaxTextModel], maxtext_
 )
 
 
+default_basic_1_pw = _add_to_model_dictionary(
+  trillium_model_dict,
+  MaxTextModel(
+    model_name="default-basic-1-pw",
+    model_type="default",
+    tuning_params={
+        "per_device_batch_size": 1,
+        "remat_policy": "full",
+        "global_parameter_scale": 1,
+        "attention": "flash",
+        "dataset_path": "gs://max-datasets-rogue",
+        "dataset_type": "synthetic",
+        "reuse_example_batch": 1,
+        "enable_checkpointing": False,
+        # "profiler": "xplane",
+
+        # Additional tuning params for pathways long running test.
+        "enable_checkpointing": True,
+        "async_checkpointing": True,
+        "checkpoint_period": 100,
+        "checkpoint_storage_use_ocdbt": False,
+        "checkpoint_storage_use_zarr3": False,
+        "metrics_file": "metrics.txt",
+        "goodput_upload_interval_seconds": 30,
+        # "enable_pathways_goodput": True,
+        "enable_checkpoint_cloud_logger": True,
+        "enable_single_controller": True,
+    },
+    xla_flags="",
+  )
+)
+
 default_32 = _add_to_model_dictionary(
   trillium_model_dict,
   MaxTextModel(
@@ -274,6 +306,48 @@ def _add_to_model_dictionary(model_dictionary: dict[str, MaxTextModel], maxtext_
   )
 )
 
+llama2_7b_4096_pw = _add_to_model_dictionary(
+  trillium_model_dict,
+  MaxTextModel(
+    model_name="llama2-7b-4096-pw",
+    model_type="llama2-7b",
+    tuning_params={
+        "per_device_batch_size": 4,
+        "ici_fsdp_parallelism": -1,
+        "remat_policy": "full",
+        "max_target_length": 4096,
+        "attention": "flash",
+        "gcs_metrics": True,
+        "use_iota_embed": True,
+        "dataset_path": "gs://max-datasets-rogue",
+        "dataset_type": "synthetic",
+        "reuse_example_batch": 1,
+        "enable_checkpointing": False,
+        "profiler": "xplane",
+        "sa_block_q": 1024,
+        "sa_block_q_dkv": 2048,
+        "sa_block_q_dq": 2048,
+        "steps": 1000000,
+
+        # Additional tuning params for pathways long running test.
+        "enable_checkpointing": True,
+        "async_checkpointing": True,
+        "checkpoint_period": 100,
+        "checkpoint_storage_use_ocdbt": False,
+        "checkpoint_storage_use_zarr3": False,
+        "metrics_file": "metrics.txt",
+        "goodput_upload_interval_seconds": 30,
+        # "enable_pathways_goodput": True,
+        "enable_checkpoint_cloud_logger": True,
+        "enable_single_controller": True,
+    },
+    xla_flags=(
+        xla_flags_library.DENSE_VMEM_LIMIT_FLAG
+        + xla_flags_library.CF_FOR_ALL_GATHER
+    ),
+  )
+)
+
 
 llama2_70b_4096 = _add_to_model_dictionary(
   trillium_model_dict,
diff --git a/benchmarks/maxtext_xpk_runner.py b/benchmarks/maxtext_xpk_runner.py
@@ -53,6 +53,7 @@ class PathwaysConfig:
   server_image: str
   proxy_image: str
   runner_image: str
+  remote_python_sidecar_image: str
 
 
 # TODO(@vbarr): Split out parameters related to XPK workload and a General workload
@@ -375,6 +376,8 @@ def generate_xpk_workload_cmd(
         '--use-pathways'
         f' --server-image={pw_config.server_image}'
         f' --proxy-server-image={pw_config.proxy_image}'
+        f' --remote-python-sidecar-image={pw_config.remote_python_sidecar_image}'
+        if pw_config.remote_python_sidecar_image is not None else ''
         ' --termination-grace-period-seconds=300'
         f' --pathways-gcs-location={wl_config.base_output_directory}'
         f' --restart-on-user-code-failure'
diff --git a/benchmarks/recipes/__init__.py b/benchmarks/recipes/__init__.py
diff --git a/benchmarks/recipes/args_helper.py b/benchmarks/recipes/args_helper.py
@@ -0,0 +1,94 @@
+"""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 argparse
+import os
+import sys
+
+# Needed to import files from the parent directory
+parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.append(parent_dir)
+
+import maxtext_xpk_runner as mxr
+
+# Constants for defining supported actions
+DELETE = "delete"
+
+
+def _handle_delete(
+    cluster_config: mxr.XpkClusterConfig, user: str, **kwargs
+) -> int:
+  """Handles the deletion of workloads.
+
+  Args:
+      cluster_config: mxr.XpkClusterConfig object
+      user: User string
+      **kwargs: Optional keyword arguments, such as xpk_path
+  """
+  xpk_path = kwargs.get("xpk_path", "xpk")  # Default to "xpk" if not provided
+  first_five_chars = user[:5]
+  delete_command = (
+      f"python3 {xpk_path}/xpk.py workload delete "
+      f"--project={cluster_config.project} --cluster={cluster_config.cluster_name}"
+      f" --filter-by-job={first_five_chars} --zone={cluster_config.zone}"
+  )
+  print(
+      f"Deleting workloads starting with: {first_five_chars} using command:"
+      f" {delete_command}"
+  )
+  os.system(delete_command)
+
+
+def handle_cmd_args(
+    cluster_config: mxr.XpkClusterConfig, *actions: str, **kwargs
+) -> bool:
+  """Parses command-line arguments and executes the specified actions.
+
+  Args:
+      cluster_config: Contains Cluster configuration information that's helpful
+        for running the actions.
+      *actions: Variable number of string arguments representing the actions to
+        be performed.
+      **kwargs: Optional keyword arguments to be passed to action handlers.
+
+  Raises:
+    ValueError: If an unsupported action is provided or if unknown arguments are
+    passed.
+  """
+
+  parser = argparse.ArgumentParser()
+
+  if DELETE in actions:
+    parser.add_argument(
+        "--delete",
+        action="store_true",
+        help="Delete workloads starting with the user's first five characters.",
+    )
+
+  known_args, unknown_args = parser.parse_known_args()
+
+  if unknown_args:
+    raise ValueError(f"Unrecognized arguments: {unknown_args}")
+
+  # Get user
+  user = os.environ["USER"]
+
+  # Handle actions
+  should_continue = True
+  if DELETE in actions and known_args.delete:
+    _handle_delete(cluster_config, user, **kwargs)
+    should_continue = False
+
+  return should_continue
diff --git a/benchmarks/recipes/pw_remote_python_recipe.py b/benchmarks/recipes/pw_remote_python_recipe.py
@@ -0,0 +1,116 @@
+"""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 args_helper as helper
+
+parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+sys.path.append(parent_dir)
+
+import maxtext_trillium_model_configs as model_configs
+import maxtext_xpk_runner as mxr
+
+
+def main() -> int:
+  # V6e cluster config
+  cluster_config = mxr.XpkClusterConfig(
+      cluster_name="v6e-256-cluster",
+      project="tpu-project",
+      zone="us-east5-b",
+      device_type="v6e-256",
+  )
+
+  xpk_path = "xpk"
+
+  # 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
+
+  # Configure test images
+  user = os.environ["USER"]
+  region = "-".join(cluster_config.zone.split("-")[:-1])
+  proxy_image = (
+      f"us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/{user}/"
+      "proxy_server:latest"
+  )
+  server_image = (
+      f"us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/gke/{user}/"
+      "server:latest"
+  )
+  remote_python_image = f"gcr.io/{cluster_config.project}/{user}/remote_python_sidecar_latest:latest"
+  runner = f"gcr.io/{cluster_config.project}/{user}_latest:latest"
+  base_output_directory = f"gs://{user}-{region}/{user}"
+
+  list_of_models = [
+      model_configs.default_basic_1_pw,
+  ]
+  pathways_config = mxr.PathwaysConfig(
+      server_image=server_image,
+      proxy_image=proxy_image,
+      runner_image=runner,
+      remote_python_sidecar_image=remote_python_image,
+  )
+  num_slices_list = [1]
+
+  xpk_workload_cmds = []
+  xpk_workload_names = []
+
+  for model in list_of_models:
+    # Run workloads on the below clusters
+    for cluster_config in [
+        cluster_config,
+    ]:
+      # Run workloads in the following slice configurations
+      for num_slices in num_slices_list:
+        wl_config = mxr.WorkloadConfig(
+            model=model,
+            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,
+            num_steps=1000000,
+        )
+        command, name = mxr.generate_xpk_workload_cmd(
+            cluster_config=cluster_config, wl_config=wl_config
+        )
+
+        print(f"Name of the workload is: {name} \n")
+        xpk_workload_names.append(name)
+
+        print(f"XPK command to be used is: {command} \n")
+        xpk_workload_cmds.append(command)
+
+  for xpk_workload_name, xpk_workload_cmd in zip(
+      xpk_workload_names, xpk_workload_cmds
+  ):
+    return_code = mxr.run_command_with_updates(
+        xpk_workload_cmd, xpk_workload_name
+    )
+    if return_code != 0:
+      print(f"Unable to run xpk workload: {xpk_workload_name}")
+
+
+if __name__ == "__main__":
+  main()
