Merge branch 'develop' into chzheng/docker_image_flag

Author: ycchenzheng

README.md

@@ -41,13 +41,15 @@ XPK supports the following TPU types:
 * v5e
 * v5p
 * Trillium (v6e)
+* Ironwood (tpu7x)
 
 and the following GPU types:
 * A100
 * A3-Highgpu (h100)
 * A3-Mega (h100-mega) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
 * A3-Ultra (h200) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
 * A4 (b200) - [Create cluster](#provisioning-a3-ultra-a3-mega-and-a4-clusters-gpu-machines), [Create workloads](#workloads-for-a3-ultra-a3-mega-and-a4-clusters-gpu-machines)
+* A4X (gb200)
 
 and the following CPU types:
 * n2-standard-32

golden_buddy.sh

@@ -83,6 +83,11 @@ GREEN='\033[0;32m'
 YELLOW='\033[0;33m'
 NC='\033[0m'
 
+if ! command -v yq &> /dev/null; then
+  echo -e "${RED}Error: 'yq' command not found. Please install yq to continue.${NC}" >&2
+  exit 1
+fi
+
 if [[ "$MODE" != "update" && "$MODE" != "verify" ]]; then
   echo "Error: Unsupported mode '$MODE'. Must be 'update' or 'verify'." >&2
   exit 1

src/xpk/commands/run.py

@@ -126,7 +126,7 @@ def submit_job(args: Namespace) -> None:
   if args.time is not None:
     cmd += f' --time {args.time}'
 
-  return_code = run_command_with_full_controls(cmd, 'run task', args)
+  return_code = run_command_with_full_controls(cmd, 'run task')
 
   if return_code != 0:
     xpk_print(f'Running task returned ERROR {return_code}')

src/xpk/commands/shell.py

@@ -38,9 +38,7 @@ def shell(args: Namespace):
   if exisitng_shell_pod_name is None:
     return_code = connect_to_new_interactive_shell(args)
   else:
-    return_code = connect_to_existing_interactive_shell(
-        exisitng_shell_pod_name, args
-    )
+    return_code = connect_to_existing_interactive_shell(exisitng_shell_pod_name)
 
   if return_code != 0:
     xpk_print(f'The command failed with code {return_code}.')
@@ -94,21 +92,17 @@ def connect_to_new_interactive_shell(args: Namespace) -> int:
   return run_command_with_full_controls(
       command=cmd,
       task='Creating new interactive shell and entering it',
-      global_args=args,
       instructions=exit_instructions,
   )
 
 
-def connect_to_existing_interactive_shell(
-    pod_name: str, args: Namespace
-) -> int:
+def connect_to_existing_interactive_shell(pod_name: str) -> int:
   return run_command_with_full_controls(
       command=(
           f'kubectl exec --stdin --tty {pod_name} --'
           f' {get_pod_template_interactive_command()}'
       ),
       task='Entering existing interactive shell',
-      global_args=args,
       instructions=exit_instructions,
   )
 

src/xpk/commands/workload.py

@@ -732,7 +732,6 @@ def workload_delete(args) -> None:
           'Delete Workload',
           task_names,
           batch=100,
-          dry_run=args.dry_run,
       )
 
     if return_code != 0:

src/xpk/core/commands.py

@@ -18,23 +18,21 @@
 import subprocess
 import sys
 import time
-from argparse import Namespace
 
 from ..utils.objects import chunks
 from ..utils.file import make_tmp_files, write_tmp_file
 from ..utils.console import xpk_print
 from ..utils.execution_context import is_dry_run
 
 
-def run_commands(commands, jobname, per_command_name, batch=10, dry_run=False):
+def run_commands(commands, jobname, per_command_name, batch=10):
   """Run commands in groups of `batch`.
 
   Args:
     commands: list of command.
     jobname: the name of the job.
     per_command_name: list of command names.
     batch: number of commands to run in parallel.
-    dry_run: enables dry_run if set to true.
 
   Returns:
     0 if successful and 1 otherwise.
@@ -47,7 +45,7 @@ def run_commands(commands, jobname, per_command_name, batch=10, dry_run=False):
       f'Breaking up a total of {len(commands)} commands into'
       f' {len(commands_batched)} batches'
   )
-  if dry_run:
+  if is_dry_run():
     xpk_print('Pretending all the jobs succeeded')
     return 0
 
@@ -302,7 +300,6 @@ def run_command_for_value(
 def run_command_with_full_controls(
     command: str,
     task: str,
-    global_args: Namespace,
     instructions: str | None = None,
 ) -> int:
   """Run command in current shell with system out, in and error handles. Wait
@@ -311,13 +308,12 @@ def run_command_with_full_controls(
   Args:
     command: command to execute
     task: user-facing name of the task
-    global_args: user provided arguments for running the command.
     verbose: shows stdout and stderr if set to true. Set to True by default.
 
   Returns:
     0 if successful and 1 otherwise.
   """
-  if global_args.dry_run:
+  if is_dry_run():
     xpk_print(
         f'Task: `{task}` is implemented by the following command'
         ' not running since it is a dry run.'

src/xpk/core/docker_image.py

@@ -17,7 +17,6 @@
 import datetime
 import os
 import random
-import re
 import string
 import subprocess
 
@@ -195,7 +194,10 @@ def setup_docker_image(args) -> tuple[int, str]:
     docker_image = DEFAULT_DOCKER_IMAGE
 
   result = subprocess.run(
-      ["docker", "images", "-q", docker_image], capture_output=True, text=True
+      ['docker', 'images', '-q', docker_image],
+      capture_output=True,
+      text=True,
+      check=True,
   )
   is_local = bool(result.stdout.strip())
 

src/xpk/core/nap.py

@@ -156,7 +156,6 @@ def enable_autoprovisioning_on_cluster(
       commands,
       'Update node pools with autoprovisioning support',
       task_names,
-      dry_run=args.dry_run,
   )
   if max_return_code != 0:
     xpk_print(

src/xpk/core/nodepool.py

@@ -212,7 +212,6 @@ def run_gke_node_pool_create_command(
         delete_commands,
         'Delete Nodepools',
         delete_task_names,
-        dry_run=args.dry_run,
     )
     if max_return_code != 0:
       xpk_print(f'Delete Nodepools returned ERROR {max_return_code}')
@@ -240,7 +239,6 @@ def run_gke_node_pool_create_command(
           update_WI_commands,
           'Enable Workload Identity on existing Nodepools',
           update_WI_task_names,
-          dry_run=args.dry_run,
       )
       if max_return_code != 0:
         xpk_print(
@@ -265,9 +263,7 @@ def run_gke_node_pool_create_command(
       )
       configmap_yml = {}
       configmap_yml[resources_configmap_name] = resources_yml
-      return_code = create_or_update_cluster_configmap(
-          configmap_yml, args.dry_run
-      )
+      return_code = create_or_update_cluster_configmap(configmap_yml)
       if return_code != 0:
         return 1
 
@@ -369,7 +365,6 @@ def run_gke_node_pool_create_command(
       create_commands,
       'Create Nodepools',
       create_task_names,
-      dry_run=args.dry_run,
   )
   if max_return_code != 0:
     xpk_print(f'Create Nodepools returned ERROR {max_return_code}')
@@ -582,7 +577,6 @@ def upgrade_gke_nodepools_version(args, default_rapid_gke_version) -> int:
       commands,
       'Update GKE node pools to default RAPID GKE version',
       task_names,
-      dry_run=args.dry_run,
   )
   if max_return_code != 0:
     xpk_print(

src/xpk/core/pathways.py

@@ -326,7 +326,7 @@ def try_to_delete_pathwaysjob_first(args, workloads) -> bool:
     return_code = run_command_with_updates(commands[0], 'Delete Workload')
   else:
     return_code = run_commands(
-        commands, 'Delete Workload', task_names, batch=100, dry_run=args.dry_run
+        commands, 'Delete Workload', task_names, batch=100
     )
 
   if return_code != 0: