xpk 0.12.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(
@@ -132,11 +132,11 @@ def install_ray_cluster(args, system) -> int:
   )
 
   tmp = write_tmp_file(yml_string)
-  command = f'kubectl apply -f {str(tmp.file.name)}'
+  command = f'kubectl apply -f {str(tmp)}'
   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:
@@ -66,7 +65,9 @@ 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
+      command,
+      'GKE Cluster Get ConfigMap',
+      dry_run_return_val='map[]',
   )
   if return_code != 0:
     xpk_print(f'GKE Cluster Get ConfigMap request returned ERROR {return_code}')
@@ -81,8 +82,10 @@ def get_cluster_configmap(args, configmap_name) -> dict[str, str] | None:
     configs = return_value[4:-1].split(' ')
 
     for config in configs:
-      key, value = config.strip().split(':')
-      config_map[key] = value
+      parts = config.strip().split(':')
+      if len(parts) != 2:
+        continue
+      config_map[parts[0]] = parts[1]
   return config_map
 
 
@@ -165,13 +168,15 @@ def create_or_update_cluster_configmap(configmap_yml: dict) -> int:
   task_names = []
   for configmap_name, yml_string in configmap_yml.items():
     tmp = write_tmp_file(yml_string)
-    command = f'kubectl apply -f {str(tmp.file.name)}'
+    command = f'kubectl apply -f {str(tmp)}'
     commands.append(command)
     task_name = f'ConfigMap CreateOrUpdate-{configmap_name}'
     task_names.append(task_name)
 
   return_code = run_commands(
-      commands, 'GKE Cluster CreateOrUpdate ConfigMap(s)', task_names
+      commands,
+      'GKE Cluster CreateOrUpdate ConfigMap(s)',
+      task_names,
   )
   if return_code != 0:
     xpk_print(
@@ -196,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}.'
@@ -219,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
@@ -241,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

@@ -15,6 +15,7 @@ limitations under the License.
 """
 
 from ..utils.console import xpk_print
+from ..utils.execution_context import is_dry_run
 from .capacity import AUTOPROVISIONING_CONFIG_MAXIMUM_KEY, AUTOPROVISIONING_CONFIG_VALUE
 from .resources import CLUSTER_RESOURCES_CONFIGMAP, get_cluster_configmap
 from .system_characteristics import (
@@ -35,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:
@@ -45,6 +46,9 @@ def check_if_workload_can_schedule(args, system: SystemCharacteristics) -> bool:
     )
     return True
 
+  if is_dry_run():
+    return True
+
   # Check for gke accelerator type:
   missing_gke_accelerator_type = False
   if not cluster_config_map.get(system.gke_accelerator):
@@ -287,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',
+  ]