xpk 0.13.0__py3-none-any.whl → 0.14.0__py3-none-any.whl

xpk/core/commands.py

@@ -18,14 +18,14 @@ import datetime
 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:
@@ -33,7 +33,6 @@ def run_commands(commands, jobname, per_command_name, batch=10, dry_run=False):
     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.
@@ -46,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
 
@@ -133,14 +132,13 @@ def run_command_batch(commands, jobname, per_command_name, output_logs):
 
 
 def run_command_with_updates_retry(
-    command, task, args, verbose=True, num_retry_attempts=5, wait_seconds=10
+    command, task, verbose=True, num_retry_attempts=5, wait_seconds=10
 ) -> int:
   """Generic run commands function with updates and retry logic.
 
   Args:
     command: command to execute
     task: user-facing name of the task
-    args: user provided arguments for running the command.
     verbose: shows stdout and stderr if set to true. Set to True by default.
     num_retry_attempts: number of attempts to retry the command.
         This has a default value in the function arguments.
@@ -160,23 +158,22 @@ def run_command_with_updates_retry(
       time.sleep(wait_seconds)
     i += 1
     xpk_print(f'Try {i}: {task}')
-    return_code = run_command_with_updates(command, task, args, verbose=verbose)
+    return_code = run_command_with_updates(command, task, verbose=verbose)
   return return_code
 
 
-def run_command_with_updates(command, task, global_args, verbose=True) -> int:
+def run_command_with_updates(command, task, verbose=True) -> int:
   """Generic run commands function with updates.
 
   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.'
@@ -226,7 +223,6 @@ def run_command_with_updates(command, task, global_args, verbose=True) -> int:
 def run_command_for_value(
     command,
     task,
-    global_args,
     dry_run_return_val='0',
     print_timer=False,
     hide_error=False,
@@ -239,7 +235,6 @@ def run_command_for_value(
   Args:
     command: user provided command to run.
     task: user provided task name for running the command.
-    global_args: user provided arguments for running the command.
     dry_run_return_val: return value of this command for dry run.
     print_timer: print out the time the command is running.
     hide_error: hide the error from the command output upon success.
@@ -249,7 +244,7 @@ def run_command_for_value(
     int: return_code, default is 0
     str: return_val, default is '0'
   """
-  if global_args is not None and 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.'
@@ -305,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
@@ -314,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.'
@@ -352,8 +345,8 @@ def run_command_with_full_controls(
   return return_code
 
 
-def run_kubectl_apply(yml_string: str, task: str, args: Namespace) -> int:
+def run_kubectl_apply(yml_string: str, task: str) -> int:
   tmp = write_tmp_file(yml_string)
   command = f'kubectl apply -f {str(tmp)}'
-  err_code = run_command_with_updates(command, task, args)
+  err_code = run_command_with_updates(command, task)
   return err_code

xpk/core/config.py

@@ -22,7 +22,7 @@ from ..utils import file
 from ..utils.console import xpk_print
 
 # This is the version for XPK PyPI package
-__version__ = 'v0.13.0'
+__version__ = 'v0.14.0'
 XPK_CURRENT_VERSION = __version__
 XPK_CONFIG_FILE = os.path.expanduser('~/.config/xpk/config.yaml')
 

xpk/core/docker_image.py

@@ -49,7 +49,7 @@ def validate_docker_image(docker_image, args) -> int:
       f'gcloud container images describe {docker_image} --project {project}'
   )
   return_code = run_command_with_updates(
-      command, 'Validate Docker Image', args, verbose=False
+      command, 'Validate Docker Image', verbose=False
   )
   if return_code != 0:
     xpk_print(
@@ -104,7 +104,6 @@ def build_docker_image_from_base_image(args, verbose=True) -> tuple[int, str]:
   return_code = run_command_with_updates(
       docker_build_command,
       'Building script_dir into docker image',
-      args,
       verbose=verbose,
   )
   if return_code != 0:
@@ -134,7 +133,7 @@ def build_docker_image_from_base_image(args, verbose=True) -> tuple[int, str]:
   # Tag the docker image.
   tag_docker_image_command = f'docker tag {docker_name} {cloud_docker_image}'
   return_code = run_command_with_updates(
-      tag_docker_image_command, 'Tag Docker Image', args, verbose=verbose
+      tag_docker_image_command, 'Tag Docker Image', verbose=verbose
   )
   if return_code != 0:
     xpk_print(
@@ -147,7 +146,7 @@ def build_docker_image_from_base_image(args, verbose=True) -> tuple[int, str]:
   # Upload image to Artifact Registry.
   upload_docker_image_command = f'docker push {cloud_docker_image}'
   return_code = run_command_with_updates(
-      upload_docker_image_command, 'Upload Docker Image', args, verbose=verbose
+      upload_docker_image_command, 'Upload Docker Image', verbose=verbose
   )
   if return_code != 0:
     xpk_print(

xpk/core/gcloud_context.py

@@ -18,8 +18,9 @@ import subprocess
 import sys
 from dataclasses import dataclass
 
-from ..utils.console import xpk_print
+from ..utils.console import xpk_print, xpk_exit
 from .commands import run_command_for_value
+from functools import lru_cache
 
 
 def get_project():
@@ -85,9 +86,33 @@ def zone_to_region(zone: str) -> str:
      The region name.
   """
   zone_terms = zone.split('-')
+  if len(zone_terms) != 2 and len(zone_terms) != 3:
+    raise ValueError(f'Invalid zone name: {zone}')
   return zone_terms[0] + '-' + zone_terms[1]
 
 
+@lru_cache()
+def get_cluster_location(project: str, name: str, zone: str) -> str:
+  """Helper function to resolve location for a given cluster"""
+  return_code, result = run_command_for_value(
+      command=(
+          'gcloud container clusters list '
+          f'--project={project} '
+          f'--filter=name={name} '
+          '--format="value(location)"'
+      ),
+      task='Find cluster region or zone',
+      dry_run_return_val=zone_to_region(zone),
+  )
+
+  if return_code != 0:
+    xpk_print('Error: Unable to determine cluster region or zone')
+    xpk_exit(return_code)
+
+  regions = result.strip().splitlines()
+  return zone if zone in regions else zone_to_region(zone)
+
+
 @dataclass
 class GkeServerConfig:
   """Stores the valid gke versions based on gcloud recommendations."""
@@ -139,7 +164,6 @@ def get_gke_server_config(args) -> tuple[int, GkeServerConfig | None]:
     return_code, cmd_output = run_command_for_value(
         command,
         command_description,
-        args,
         hide_error=True,
     )
     if return_code != 0:

xpk/core/gcloud_context_test.py

@@ -0,0 +1,96 @@
+"""
+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 pytest
+from .gcloud_context import get_cluster_location, zone_to_region
+
+
+def test_zone_to_region_raises_when_zone_is_invalid():
+  with pytest.raises(ValueError):
+    zone_to_region("us")
+
+
+def test_zone_to_region_returns_region_when_region_given():
+  assert zone_to_region("us-central1") == "us-central1"
+
+
+def test_zone_to_region_returns_region_when_zone_is_valid():
+  assert zone_to_region("us-central1-a") == "us-central1"
+
+
+def test_get_cluster_location_returns_cluster_region_when_cluster_is_regional(
+    mocker,
+):
+  mocker.patch(
+      "xpk.core.gcloud_context.run_command_for_value",
+      return_value=(0, "us-central1"),
+  )
+
+  result = get_cluster_location(
+      project="project1", name="name1", zone="us-central1-a"
+  )
+
+  assert result == "us-central1"
+
+
+def test_get_cluster_location_returns_cluster_zone_when_both_regional_and_zonal_clusters_exist(
+    mocker,
+):
+  mocker.patch(
+      "xpk.core.gcloud_context.run_command_for_value",
+      return_value=(0, "us-central1\nus-central1-a"),
+  )
+
+  result = get_cluster_location(
+      project="project2", name="name2", zone="us-central1-a"
+  )
+
+  assert result == "us-central1-a"
+
+
+def test_get_cluster_location_returns_given_zone_converted_to_region_when_cluster_is_not_found(
+    mocker,
+):
+  mocker.patch(
+      "xpk.core.gcloud_context.run_command_for_value", return_value=(0, "")
+  )
+
+  result = get_cluster_location(
+      project="project3", name="name3", zone="us-central1-a"
+  )
+
+  assert result == "us-central1"
+
+
+def test_get_cluster_location_caches_previous_command_result(mocker):
+  mock = mocker.patch(
+      "xpk.core.gcloud_context.run_command_for_value", return_value=(0, "")
+  )
+
+  get_cluster_location(project="project4", name="name4", zone="us-central1-a")
+
+  assert mock.call_count == 1
+
+
+def test_get_cluster_location_invokes_command_for_different_input_args(mocker):
+  mock = mocker.patch(
+      "xpk.core.gcloud_context.run_command_for_value", return_value=(0, "")
+  )
+
+  get_cluster_location(project="project5", name="name5", zone="us-central1-a")
+  get_cluster_location(project="project6", name="name6", zone="us-central1-a")
+
+  assert mock.call_count == 2

xpk/core/gcluster_manager.py

@@ -27,9 +27,6 @@ blueprint_file_name = 'xpk_blueprint.yaml'
 deployment_module = '/out/xpk-deployment'
 a3_utils_dir_name = 'a3-mega-xpk'
 config_map_repo_path = 'src/xpk/blueprints/a3-mega-xpk/config-map.yaml.tftpl'
-kueue_config_repo_path = (
-    'src/xpk/blueprints/a3-mega-xpk/kueue-xpk-configuration.yaml.tftpl'
-)
 
 
 class GclusterManager:

xpk/core/jobset.py

@@ -18,7 +18,7 @@ import math
 
 from ..utils.console import xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
-from ..core.kueue import (
+from ..core.kueue_manager import (
     MEMORY_SIZE_PER_VM,
     MIN_MEMORY_LIMIT_SIZE,
 )
@@ -110,19 +110,16 @@ spec:
 """
 
 
-def update_jobset_resources_if_necessary(args):
+def update_jobset_resources_if_necessary():
   """Update the jobset manifest to increase the resources for the jobset controller manager.
 
-  Args:
-    args: user provided arguments for running the command.
-
   Returns:
     0 if successful and 1 otherwise.
   """
   # Get total number of nodes
   cmd_total_node_num = 'kubectl get node --no-headers | wc -l'
   return_code, out = run_command_for_value(
-      cmd_total_node_num, 'Count total nodes', args
+      cmd_total_node_num, 'Count total nodes'
   )
   if return_code != 0:
     xpk_exit(1)
@@ -137,7 +134,7 @@ def update_jobset_resources_if_necessary(args):
   command = f'kubectl apply -f {str(tmp)}'
 
   task = 'Updating jobset Controller Manager resources'
-  return_code = run_command_with_updates_retry(command, task, args)
+  return_code = run_command_with_updates_retry(command, task)
   if return_code != 0:
     xpk_print(f'{task} returned ERROR {return_code}')
   return return_code

xpk/core/kjob.py

@@ -167,8 +167,8 @@ Kueue_TAS_annotation = "kueue.x-k8s.io/podset-preferred-topology=cloud.google.co
 default_interface_annotation = "networking.gke.io/default-interface=eth0"
 
 
-def get_a4_pod_template_annotations(args) -> tuple[str, str]:
-  sub_networks = get_cluster_subnetworks(args)
+def get_a4_pod_template_annotations() -> tuple[str, str]:
+  sub_networks = get_cluster_subnetworks()
   interfaces_key, interfaces_value = rdma_decorator.get_interfaces_entry(
       sub_networks
   )
@@ -179,8 +179,8 @@ def get_a4_pod_template_annotations(args) -> tuple[str, str]:
   )
 
 
-def get_a3ultra_pod_template_annotations(args: Namespace) -> tuple[str, str]:
-  sub_networks = get_cluster_subnetworks(args)
+def get_a3ultra_pod_template_annotations() -> tuple[str, str]:
+  sub_networks = get_cluster_subnetworks()
   interfaces_key, interfaces_value = rdma_decorator.get_interfaces_entry(
       sub_networks
   )
@@ -191,11 +191,9 @@ def get_a3ultra_pod_template_annotations(args: Namespace) -> tuple[str, str]:
   )
 
 
-def get_a3mega_pod_template_annotations(
-    args: Namespace,
-) -> tuple[str, str, str]:
+def get_a3mega_pod_template_annotations() -> tuple[str, str, str]:
   """Adds or updates annotations in the Pod template."""
-  sub_networks = get_cluster_subnetworks(args)
+  sub_networks = get_cluster_subnetworks()
   tcpxo_deamon_key, tcpxo_deamon_paths = get_tcpxo_deamon_entry()
   interfaces_key, interfaces_value = tcpxo_decorator.get_interfaces_entry(
       sub_networks
@@ -205,16 +203,14 @@ def get_a3mega_pod_template_annotations(
   return tcpxo, interfaces, default_interface_annotation
 
 
-def verify_kjob_installed(args: Namespace) -> int:
+def verify_kjob_installed() -> int:
   """Check if kjob is installed. If not provide user with proper communicate and exit.
-  Args:
-    args - user provided arguments.
   Returns:
     error code > if kjob not installed, otherwise 0
   """
   command = "kubectl-kjob help"
   task = "Verify kjob installation "
-  verify_kjob_installed_code, _ = run_command_for_value(command, task, args)
+  verify_kjob_installed_code, _ = run_command_for_value(command, task)
 
   if verify_kjob_installed_code == 0:
     xpk_print("kjob found")
@@ -246,9 +242,7 @@ def get_pod_template_interactive_command() -> str:
   return pod_command
 
 
-def create_app_profile_instance(
-    args: Namespace, volume_bundles: list[str]
-) -> int:
+def create_app_profile_instance(volume_bundles: list[str]) -> int:
   """Create new AppProfile instance on cluster with default settings.
 
   Args:
@@ -264,7 +258,6 @@ def create_app_profile_instance(
           volume_bundles=volume_bundles,
       ),
       task="Creating AppProfile",
-      args=args,
   )
 
 
@@ -332,15 +325,12 @@ def create_job_template_instance(
   return run_kubectl_apply(
       yml_string,
       task="Creating JobTemplate",
-      args=args,
   )
 
 
-def create_pod_template_instance(args: Namespace, service_account: str) -> int:
+def create_pod_template_instance(service_account: str) -> int:
   """Create new PodTemplate instance on cluster with default settings.
 
-  Args:
-    args - user provided arguments
   Returns:
     exit_code > 0 if creating PodTemplate fails, 0 otherwise
   """
@@ -362,7 +352,6 @@ def create_pod_template_instance(args: Namespace, service_account: str) -> int:
           service_account=service_account,
       ),
       task="Creating PodTemplate",
-      args=args,
   )
 
 
@@ -381,29 +370,27 @@ def prepare_kjob(args: Namespace) -> int:
   job_err_code = create_job_template_instance(args, system, service_account)
   if job_err_code > 0:
     return job_err_code
-  pod_err_code = create_pod_template_instance(args, service_account)
+  pod_err_code = create_pod_template_instance(service_account)
   if pod_err_code > 0:
     return pod_err_code
 
   volume_bundles = [item.name for item in storages]
 
-  return create_app_profile_instance(args, volume_bundles)
+  return create_app_profile_instance(volume_bundles)
 
 
-def apply_kjob_crds(args: Namespace) -> int:
+def apply_kjob_crds() -> int:
   """Apply kjob CRDs on cluster.
 
   This function install kjob CRDs files from kjobctl printcrds.
   It creates all neccessary kjob CRDs.
 
-  Args:
-    args - user provided arguments
   Returns:
     None
   """
   command = "kubectl kjob printcrds | kubectl apply --server-side -f -"
   task = "Create kjob CRDs on cluster"
-  return_code = run_command_with_updates(command, task, args)
+  return_code = run_command_with_updates(command, task)
   if return_code != 0:
     xpk_print(f"{task} returned ERROR {return_code}")
     return return_code