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

xpk/core/ray.py

@@ -102,16 +102,16 @@ def install_ray_cluster(args, system) -> int:
     0 if successful and 1 otherwise.
   """
 
-  delete_ray_cluster(args)
+  delete_ray_cluster()
 
   label = 'cloud.google.com/gke-nodepool=default-pool'
   available_head_cpu, available_head_mem = generate_available_resources(
-      label, args, HEAD_CPU
+      label, HEAD_CPU
   )
 
   label = f'cloud.google.com/gke-tpu-accelerator={system.gke_accelerator}'
   available_worker_cpu, available_worker_mem = generate_available_resources(
-      label, args, WORKER_CPU
+      label, WORKER_CPU
   )
 
   yml_string = ray_cluster_crd_yaml.format(
@@ -136,7 +136,7 @@ def install_ray_cluster(args, system) -> int:
   task = 'Applying RayCluster'
   retry_attempts = 1
   return_code = run_command_with_updates_retry(
-      command, task, args, num_retry_attempts=retry_attempts
+      command, task, num_retry_attempts=retry_attempts
   )
   if return_code != 0:
     xpk_print(f'{task} not successful.')
@@ -144,12 +144,9 @@ def install_ray_cluster(args, system) -> int:
   return return_code
 
 
-def delete_ray_cluster(args) -> None:
+def delete_ray_cluster() -> None:
   """Delete all RayClusters on the cluster
 
-  Args:
-    args: user provided arguments for running the command.
-
   Returns:
     None
   """
@@ -158,7 +155,7 @@ def delete_ray_cluster(args) -> None:
   task = 'Deleting old RayCluster'
   retry_attempts = 1
   return_code = run_command_with_updates_retry(
-      command, task, args, num_retry_attempts=retry_attempts
+      command, task, num_retry_attempts=retry_attempts
   )
 
   if return_code != 0:
@@ -168,12 +165,11 @@ def delete_ray_cluster(args) -> None:
   return
 
 
-def generate_available_resources(label, args, percent) -> tuple:
+def generate_available_resources(label, percent) -> tuple:
   """Generate the available resources for the nodes that match the given label
 
   Args:
     label: the label used to match the appropriate nodes
-    args: user provided arguments for running the command
     percent: the percent of the available resources to use
 
   Returns:
@@ -184,13 +180,13 @@ def generate_available_resources(label, args, percent) -> tuple:
       f"kubectl get nodes -l {label} -o jsonpath='{{.items[0].metadata.name}}'"
   )
   task = f'Getting nodes with label {label}'
-  _, node_name = run_command_for_value(command, task, args)
+  _, node_name = run_command_for_value(command, task)
 
   command = (
       f"kubectl get node {node_name} -o jsonpath='{{.status.allocatable.cpu}}'"
   )
   task = 'Fetching available CPU on node'
-  _, available_cpu = run_command_for_value(command, task, args)
+  _, available_cpu = run_command_for_value(command, task)
   match = re.match(r'(\d+)([a-zA-Z]+)', available_cpu)
   if not match:
     xpk_print(
@@ -207,7 +203,7 @@ def generate_available_resources(label, args, percent) -> tuple:
       " jsonpath='{.status.allocatable.memory}'"
   )
   task = 'Fetching available memory on node'
-  _, available_memory = run_command_for_value(command, task, args)
+  _, available_memory = run_command_for_value(command, task)
   match = re.match(r'(\d+)([a-zA-Z]+)', available_memory)
   if not match:
     xpk_print(

xpk/core/resources.py

@@ -50,11 +50,10 @@ class AutoprovisioningConfig:
   maximum_chips: int
 
 
-def get_cluster_configmap(args, configmap_name) -> dict[str, str] | None:
+def get_cluster_configmap(configmap_name) -> dict[str, str] | None:
   """Run the Get GKE Cluster ConfigMap request.
 
   Args:
-    args: user provided arguments for running the command.
     configmap_name: name of the configmap.
 
   Returns:
@@ -68,7 +67,6 @@ def get_cluster_configmap(args, configmap_name) -> dict[str, str] | None:
   return_code, return_value = run_command_for_value(
       command,
       'GKE Cluster Get ConfigMap',
-      args,
       dry_run_return_val='map[]',
   )
   if return_code != 0:
@@ -155,12 +153,10 @@ def create_cluster_configmaps(
       args=args, name=metadata_configmap_name, data=metadata
   )
   configmap_yml[metadata_configmap_name] = metadata_yml
-  return create_or_update_cluster_configmap(configmap_yml, args.dry_run)
+  return create_or_update_cluster_configmap(configmap_yml)
 
 
-def create_or_update_cluster_configmap(
-    configmap_yml: dict, dry_run: bool
-) -> int:
+def create_or_update_cluster_configmap(configmap_yml: dict) -> int:
   """
   Args:
     configmap_yml: dict containing ConfigMap name and yml string.
@@ -181,7 +177,6 @@ def create_or_update_cluster_configmap(
       commands,
       'GKE Cluster CreateOrUpdate ConfigMap(s)',
       task_names,
-      dry_run=dry_run,
   )
   if return_code != 0:
     xpk_print(
@@ -206,7 +201,7 @@ def check_cluster_resources(args, system) -> tuple[bool, bool]:
     True if device_type/gke_accelerator exists in the cluster, False otherwise.
   """
   resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
-  resources_config_map = get_cluster_configmap(args, resources_configmap_name)
+  resources_config_map = get_cluster_configmap(resources_configmap_name)
   if resources_config_map is None:
     xpk_print(
         f'No ConfigMap exist for cluster with the name {resources_config_map}.'
@@ -229,7 +224,7 @@ def get_cluster_system_characteristics(args) -> SystemCharacteristics | None:
     returns system characteristics
   """
   resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
-  cluster_config_map = get_cluster_configmap(args, resources_configmap_name)
+  cluster_config_map = get_cluster_configmap(resources_configmap_name)
 
   if cluster_config_map is None:
     return None
@@ -251,7 +246,7 @@ def get_cluster_capacity_type(args) -> CapacityType | None:
     returns system characteristics
   """
   metadata_configmap_name = f'{args.cluster}-{CLUSTER_METADATA_CONFIGMAP}'
-  cluster_config_map = get_cluster_configmap(args, metadata_configmap_name)
+  cluster_config_map = get_cluster_configmap(metadata_configmap_name)
 
   if cluster_config_map is None:
     return None

xpk/core/scheduling.py

@@ -36,7 +36,7 @@ def check_if_workload_can_schedule(args, system: SystemCharacteristics) -> bool:
     returns true if workload can schedule, otherwise returns false.
   """
   resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
-  cluster_config_map = get_cluster_configmap(args, resources_configmap_name)
+  cluster_config_map = get_cluster_configmap(resources_configmap_name)
 
   # Prevents workload creation failure for existing clusters with no ConfigMap
   if cluster_config_map is None:
@@ -291,3 +291,21 @@ def create_tpu_topology(
   ):
     return f'{system.topology}'
   return ''
+
+
+def create_sub_slicing_annotations(sub_slicing_topology: str) -> list[str]:
+  """Generates subslicing annotations.
+
+  Args:
+    sub_slicing_topology: subslice topology.
+
+  Returns:
+    Annotations to be rendered in deployment yaml.
+  """
+  return [
+      (
+          'kueue.x-k8s.io/podset-required-topology:'
+          f' "google.com/gke-tpu-slice-{sub_slicing_topology}-id"'
+      ),
+      f'cloud.google.com/gke-tpu-slice-topology: {sub_slicing_topology}',
+  ]

xpk/core/scheduling_test.py

@@ -0,0 +1,31 @@
+"""
+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.
+"""
+
+from .scheduling import create_sub_slicing_annotations
+
+
+def test_create_sub_slicing_annotations_returns_valid_annotations():
+  subslicing_topology = '2x2'
+
+  result = create_sub_slicing_annotations(subslicing_topology)
+
+  assert result == [
+      (
+          'kueue.x-k8s.io/podset-required-topology:'
+          ' "google.com/gke-tpu-slice-2x2-id"'
+      ),
+      'cloud.google.com/gke-tpu-slice-topology: 2x2',
+  ]