xpk 0.14.3__py3-none-any.whl → 0.15.0__py3-none-any.whl

xpk/core/gcloud_context.py

@@ -19,6 +19,7 @@ import sys
 from dataclasses import dataclass
 
 from ..utils.console import xpk_print, xpk_exit
+from ..utils.versions import ReleaseChannel
 from .commands import run_command_for_value
 from functools import lru_cache
 
@@ -117,15 +118,18 @@ def get_cluster_location(project: str, name: str, zone: str) -> str:
 class GkeServerConfig:
   """Stores the valid gke versions based on gcloud recommendations."""
 
-  default_rapid_gke_version: str
+  default_gke_version: str
   valid_versions: set[str]
 
 
-def get_gke_server_config(args) -> tuple[int, GkeServerConfig | None]:
+def get_gke_server_config(
+    args, release_channel: ReleaseChannel
+) -> tuple[int, GkeServerConfig | None]:
   """Determine the GKE versions supported by gcloud currently.
 
   Args:
     args: user provided arguments for running the command.
+    release_channel: the release channel to use.
 
   Returns:
     Tuple of
@@ -136,22 +140,24 @@ def get_gke_server_config(args) -> tuple[int, GkeServerConfig | None]:
       'gcloud container get-server-config'
       f' --project={args.project} --region={zone_to_region(args.zone)}'
   )
-  default_rapid_gke_version_cmd = (
+  default_gke_version_cmd = (
       base_command
-      + ' --flatten="channels" --filter="channels.channel=RAPID"'
+      + ' --flatten="channels"'
+      f' --filter="channels.channel={release_channel.value}"'
       ' --format="value(channels.defaultVersion)"'
   )
   valid_versions_cmd = (
       base_command
-      + ' --flatten="channels" --filter="channels.channel=RAPID"'
+      + ' --flatten="channels"'
+      f' --filter="channels.channel={release_channel.value}"'
       ' --format="value(channels.validVersions)"'
   )
   base_command_description = 'Determine server supported GKE versions for '
 
   server_config_commands_and_descriptions = [
       (
-          default_rapid_gke_version_cmd,
-          base_command_description + 'default rapid gke version',
+          default_gke_version_cmd,
+          base_command_description + 'default gke version',
       ),
       (
           valid_versions_cmd,
@@ -172,8 +178,8 @@ def get_gke_server_config(args) -> tuple[int, GkeServerConfig | None]:
     command_outputs.append(cmd_output)
 
   return 0, GkeServerConfig(
-      default_rapid_gke_version=command_outputs[0].strip(),
-      valid_versions=set(command_outputs[1].split(';')),
+      default_gke_version=command_outputs[0].strip(),
+      valid_versions=set([s.strip() for s in command_outputs[1].split(';')]),
   )
 
 
@@ -196,7 +202,7 @@ def get_gke_control_plane_version(
   if args.gke_version is not None:
     master_gke_version = args.gke_version
   else:
-    master_gke_version = gke_server_config.default_rapid_gke_version
+    master_gke_version = gke_server_config.default_gke_version
 
   is_valid_version = master_gke_version in gke_server_config.valid_versions
 
@@ -204,7 +210,7 @@ def get_gke_control_plane_version(
     xpk_print(
         f'Planned GKE Version: {master_gke_version}\n Valid Versions:'
         f'\n{gke_server_config.valid_versions}\nRecommended / Default GKE'
-        f' Version: {gke_server_config.default_rapid_gke_version}'
+        f' Version: {gke_server_config.default_gke_version}'
     )
     xpk_print(
         f'Error: Planned GKE Version {master_gke_version} is not valid.'
@@ -213,7 +219,7 @@ def get_gke_control_plane_version(
     xpk_print(
         'Please select a gke version from the above list using --gke-version=x'
         ' argument or rely on the default gke version:'
-        f' {gke_server_config.default_rapid_gke_version}'
+        f' {gke_server_config.default_gke_version}'
     )
     return 1, None
 

xpk/core/gcloud_context_test.py

@@ -15,7 +15,20 @@ limitations under the License.
 """
 
 import pytest
-from .gcloud_context import get_cluster_location, zone_to_region
+from unittest.mock import MagicMock
+from .gcloud_context import (
+    get_cluster_location,
+    get_gke_control_plane_version,
+    get_gke_server_config,
+    GkeServerConfig,
+    zone_to_region,
+)
+from ..utils.versions import ReleaseChannel
+
+
+@pytest.fixture(autouse=True)
+def xpk_print(mocker):
+  return mocker.patch("xpk.core.gcloud_context.xpk_print")
 
 
 def test_zone_to_region_raises_when_zone_is_invalid():
@@ -94,3 +107,100 @@ def test_get_cluster_location_invokes_command_for_different_input_args(mocker):
   get_cluster_location(project="project6", name="name6", zone="us-central1-a")
 
   assert mock.call_count == 2
+
+
+def test_get_gke_server_config_success(mocker):
+  mock_run_command = mocker.patch(
+      "xpk.core.gcloud_context.run_command_for_value",
+      side_effect=[
+          (0, "1.2.3"),
+          (0, "1.2.3;1.2.4;1.3.0"),
+      ],
+  )
+  args = mocker.Mock(project="test-project", zone="us-central1")
+
+  return_code, config = get_gke_server_config(args, ReleaseChannel.STABLE)
+
+  assert return_code == 0
+  assert isinstance(config, GkeServerConfig)
+  assert config.default_gke_version == "1.2.3"
+  assert config.valid_versions == {"1.2.3", "1.2.4", "1.3.0"}
+  assert mock_run_command.call_count == 2
+
+
+def test_get_gke_server_config_fails_on_default_version_command(mocker):
+  mocker.patch(
+      "xpk.core.gcloud_context.run_command_for_value",
+      return_value=(1, "error"),
+  )
+  args = mocker.Mock(project="test-project", zone="us-central1")
+
+  return_code, config = get_gke_server_config(args, ReleaseChannel.STABLE)
+
+  assert return_code == 1
+  assert config is None
+
+
+def test_get_gke_server_config_fails_on_valid_versions_command(mocker):
+  mocker.patch(
+      "xpk.core.gcloud_context.run_command_for_value",
+      side_effect=[(0, "1.2.3"), (1, "error")],
+  )
+  args = mocker.Mock(project="test-project", zone="us-central1")
+
+  return_code, config = get_gke_server_config(args, ReleaseChannel.STABLE)
+
+  assert return_code == 1
+  assert config is None
+
+
+def test_get_gke_control_plane_version_uses_default_when_not_specified(mocker):
+  args = mocker.Mock(gke_version=None)
+  gke_server_config = GkeServerConfig(
+      default_gke_version="1.2.3", valid_versions={"1.2.3", "1.2.4"}
+  )
+
+  return_code, version = get_gke_control_plane_version(args, gke_server_config)
+
+  assert return_code == 0
+  assert version == "1.2.3"
+
+
+def test_get_gke_control_plane_version_uses_user_version_when_valid(mocker):
+  args = mocker.Mock(gke_version="1.2.4")
+  gke_server_config = GkeServerConfig(
+      default_gke_version="1.2.3", valid_versions={"1.2.3", "1.2.4"}
+  )
+
+  return_code, version = get_gke_control_plane_version(args, gke_server_config)
+
+  assert return_code == 0
+  assert version == "1.2.4"
+
+
+def test_get_gke_control_plane_version_fails_for_invalid_user_version(
+    mocker, xpk_print: MagicMock
+):
+  args = mocker.Mock(gke_version="1.2.5")
+  gke_server_config = GkeServerConfig(
+      default_gke_version="1.2.3", valid_versions={"1.2.3", "1.2.4"}
+  )
+
+  return_code, version = get_gke_control_plane_version(args, gke_server_config)
+
+  assert return_code == 1
+  assert version is None
+  assert "Planned GKE Version: 1.2.5" in xpk_print.mock_calls[0].args[0]
+  assert (
+      "Recommended / Default GKE Version: 1.2.3"
+      in xpk_print.mock_calls[0].args[0]
+  )
+  assert (
+      "Error: Planned GKE Version 1.2.5 is not valid."
+      in xpk_print.mock_calls[1].args[0]
+  )
+  assert (
+      "Please select a gke version from the above list using --gke-version=x"
+      " argument or rely on the default gke version: 1.2.3"
+      in xpk_print.mock_calls[2].args[0]
+  )

xpk/core/kjob.py

@@ -38,7 +38,7 @@ from .config import (
     KJOB_SHELL_IMAGE,
     KJOB_SHELL_INTERACTIVE_COMMAND,
     KJOB_SHELL_WORKING_DIRECTORY,
-    XpkConfig,
+    xpk_config,
 )
 from .network import get_cluster_subnetworks
 from .system_characteristics import AcceleratorType, SystemCharacteristics
@@ -234,8 +234,7 @@ def get_pod_template_interactive_command() -> str:
   Returns:
     str - PodTemplate's interactive command
   """
-  config = XpkConfig()
-  pod_command = config.get(KJOB_SHELL_INTERACTIVE_COMMAND)
+  pod_command = xpk_config.get(KJOB_SHELL_INTERACTIVE_COMMAND)
   if pod_command is None or len(pod_command) == 0:
     pod_command = PodTemplateDefaults.INTERACTIVE_COMMAND.value
 
@@ -287,11 +286,10 @@ def create_job_template_instance(
   Returns:
     exit_code > 0 if creating JobTemplate fails, 0 otherwise
   """
-  config = XpkConfig()
-  job_image = config.get(KJOB_BATCH_IMAGE)
+  job_image = xpk_config.get(KJOB_BATCH_IMAGE)
   if job_image is None or len(job_image) == 0:
     job_image = JobTemplateDefaults.IMAGE.value
-  working_directory = config.get(KJOB_BATCH_WORKING_DIRECTORY)
+  working_directory = xpk_config.get(KJOB_BATCH_WORKING_DIRECTORY)
   if working_directory is None or len(working_directory) == 0:
     working_directory = JobTemplateDefaults.WORKING_DIRECTORY.value
   resources = (
@@ -332,11 +330,10 @@ def create_pod_template_instance(service_account: str) -> int:
   Returns:
     exit_code > 0 if creating PodTemplate fails, 0 otherwise
   """
-  config = XpkConfig()
-  pod_image = config.get(KJOB_SHELL_IMAGE)
+  pod_image = xpk_config.get(KJOB_SHELL_IMAGE)
   if pod_image is None or len(pod_image) == 0:
     pod_image = PodTemplateDefaults.IMAGE.value
-  working_directory = config.get(KJOB_SHELL_WORKING_DIRECTORY)
+  working_directory = xpk_config.get(KJOB_SHELL_WORKING_DIRECTORY)
   if working_directory is None or len(working_directory) == 0:
     working_directory = PodTemplateDefaults.WORKING_DIRECTORY.value
 

xpk/core/kueue_manager.py

@@ -20,15 +20,17 @@ from dataclasses import dataclass
 from typing import Optional, List, Dict, Any
 import json
 from jinja2 import Environment, FileSystemLoader
-from ..utils.execution_context import is_dry_run
-from ..utils.kueue import is_queued_cluster
 
+from ..utils.topology import get_slice_topology_level, get_topology_product, is_topology_contained
+from ..utils.kueue import is_queued_cluster
+from kubernetes.utils import parse_quantity
 from .capacity import B200_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
 from .scheduling import (
     create_accelerator_label,
     create_machine_label,
 )
 from .system_characteristics import (
+    SUB_SLICING_TOPOLOGIES,
     AcceleratorTypeToAcceleratorCharacteristics,
     SystemCharacteristics,
 )
@@ -38,10 +40,12 @@ from ..core.commands import (
     run_command_with_updates_retry,
 )
 from ..utils.file import write_tmp_file
-from ..utils.console import xpk_print, xpk_exit
+from ..utils.console import xpk_print, xpk_exit, ask_for_user_consent
 from ..utils.templates import TEMPLATE_PATH, get_templates_absolute_path
 from packaging.version import Version
 
+KUEUE_VERSION = Version("v0.14.3")
+LATEST_BREAKING_VERSION = Version("v0.14.0")
 WAIT_FOR_KUEUE_TIMEOUT = "10m"
 CLUSTER_QUEUE_NAME = "cluster-queue"
 LOCAL_QUEUE_NAME = "multislice-queue"
@@ -52,10 +56,9 @@ KUEUE_CONTROLLER_MANAGER_JINJA_FILE = "kueue_controller_manager.yaml.j2"
 KUEUE_SUB_SLICING_TOPOLOGY_JINJA_FILE = "kueue_sub_slicing_topology.yaml.j2"
 MEMORY_SIZE_PER_VM = 1.2
 MIN_MEMORY_LIMIT_SIZE = 4096
-KUEUE_VERSION = Version("v0.12.2")
 
 
-@dataclass
+@dataclass(frozen=True)
 class KueueConfig:
   system: SystemCharacteristics
   total_chips: int
@@ -68,7 +71,7 @@ class KueueConfig:
   num_slices: int = 1
 
 
-@dataclass
+@dataclass(frozen=True)
 class _NameAndYaml:
   name: str
   yaml: str
@@ -79,9 +82,13 @@ class KueueManager:
 
   def __init__(
       self,
+      project: str,
+      zone: str,
       kueue_version: Version = KUEUE_VERSION,
       template_path=TEMPLATE_PATH,
   ):
+    self.project = project
+    self.zone = zone
     self.kueue_version = kueue_version
 
     self.template_env = Environment(
@@ -102,10 +109,10 @@ class KueueManager:
     Args:
         tolerations: An optional list of tolerations to apply to the kueue-controller-manager.
     """
-    return_code, installed_version = self.get_installed_kueue_version()
+    return_code, installed_version = get_installed_kueue_version()
 
-    if return_code == 0:
-      if installed_version and installed_version > self.kueue_version:
+    if return_code == 0 and installed_version:
+      if installed_version > self.kueue_version:
         xpk_print(
             f"Cluster has a newer Kueue version, {installed_version}. Skipping"
             " installation."
@@ -113,6 +120,10 @@ class KueueManager:
         return 0
       else:
         xpk_print(f"Upgrading Kueue to version v{self.kueue_version}...")
+        assert installed_version
+        prepare_code = self.__prepare_for_upgrade(installed_version)
+        if prepare_code != 0:
+          return prepare_code
     else:
       xpk_print(f"Installing Kueue version v{self.kueue_version}...")
 
@@ -122,24 +133,6 @@ class KueueManager:
 
     return self.__configure(kueue_config)
 
-  def get_installed_kueue_version(self) -> tuple[int, Version | None]:
-    command = (
-        "kubectl get deployment kueue-controller-manager -n kueue-system -o"
-        " jsonpath='{.spec.template.spec.containers[0].image}'"
-    )
-    task = "Get kueue version on server"
-    return_code, val = run_command_for_value(
-        command,
-        task,
-        dry_run_return_val="",
-    )
-    if return_code != 0:
-      return return_code, None
-    version_tag = val.split(":")
-    if len(version_tag) == 1:
-      return 1, None
-    return return_code, Version(version_tag[-1])
-
   def __install(
       self,
       tolerations: Optional[List[Dict[str, Any]]] = None,
@@ -161,6 +154,60 @@ class KueueManager:
 
     return self.__wait_for_kueue_available()
 
+  def __prepare_for_upgrade(self, installed_version: Version) -> int:
+    if installed_version >= LATEST_BREAKING_VERSION:
+      return 0
+
+    xpk_print(
+        f"Currently installed Kueue version v{installed_version} is"
+        f" incompatible with the newer v{self.kueue_version}."
+    )
+
+    changelog_link = f"https://github.com/kubernetes-sigs/kueue/blob/main/CHANGELOG/CHANGELOG-{self.kueue_version.major}.{self.kueue_version.minor}.md"
+    agreed = ask_for_user_consent(
+        "Do you want to allow XPK to update Kueue automatically? This will"
+        " delete all existing Kueue resources and create new ones. If you"
+        " decline, you will need to upgrade the Kueue manually (see"
+        f" {changelog_link} for help)."
+    )
+    if not agreed:
+      return 1
+
+    return self.__delete_all_kueue_resources()
+
+  def __delete_all_kueue_resources(self) -> int:
+    return_code, kueue_crds_string = run_command_for_value(
+        "kubectl get crd -o name | grep .kueue.x-k8s.io", "Get Kueue CRDs"
+    )
+    if return_code != 0:
+      return return_code
+
+    kueue_crds = [
+        line.strip().removeprefix(
+            "customresourcedefinition.apiextensions.k8s.io/"
+        )
+        for line in kueue_crds_string.strip().split("\n")
+    ]
+
+    for crd in kueue_crds:
+      return_code = run_command_with_updates(
+          f"kubectl delete {crd} --all", f"Delete all resources of type {crd}"
+      )
+      if return_code != 0:
+        return return_code
+
+    for crd in kueue_crds:
+      return_code = run_command_with_updates(
+          f"kubectl delete crd {crd}", f"Delete CRD {crd}"
+      )
+      if return_code != 0:
+        return return_code
+
+    return run_command_with_updates(
+        "kubectl delete deployment kueue-controller-manager -n kueue-system",
+        "Delete Kueue Controller Manager deployment",
+    )
+
   def __install_kueue_crs(self) -> int:
     manifest_url = f"https://github.com/kubernetes-sigs/kueue/releases/download/v{self.kueue_version}/manifests.yaml"
     install_command = (
@@ -228,6 +275,7 @@ class KueueManager:
     topology_name = (
         topology_name_and_yaml.name if topology_name_and_yaml else None
     )
+    cpu_limit, memory_limit = self.__autocorrect_resource_limits(kueue_config)
 
     # The manager builds the context internally based on its opinionated logic
     context = self.__build_template_context(
@@ -237,8 +285,8 @@ class KueueManager:
         autoprovisioning=kueue_config.autoprovisioning_enabled,
         flex=kueue_config.flex,
         num_slices=kueue_config.num_slices,
-        cpu_limit=kueue_config.cpu_limit,
-        memory_limit=kueue_config.memory_limit,
+        cpu_limit=cpu_limit,
+        memory_limit=memory_limit,
         topology_name=topology_name,
     )
 
@@ -364,12 +412,25 @@ class KueueManager:
           ).render(),
       )
     elif configure_sub_slicing:
+      sorted_topologies = sorted(
+          SUB_SLICING_TOPOLOGIES, key=get_topology_product, reverse=True
+      )
+      levels = [
+          get_slice_topology_level(topology)
+          for topology in sorted_topologies
+          if is_topology_contained(
+              contained=topology, container=system.topology
+          )
+      ]
+      levels.append("kubernetes.io/hostname")
+
       return _NameAndYaml(
           name=SUB_SLICE_TOPOLOGY_NAME,
           yaml=self.template_env.get_template(
               KUEUE_SUB_SLICING_TOPOLOGY_JINJA_FILE
           ).render({
               "sub_slice_topology_name": SUB_SLICE_TOPOLOGY_NAME,
+              "levels": levels,
           }),
       )
     else:
@@ -377,8 +438,6 @@ class KueueManager:
 
   def __apply_manifest(self, manifest: str) -> int:
     task = "Applying Kueue Custom Resources"
-    if is_dry_run():
-      xpk_print(f"Applying following Kueue resources:{manifest}")
     tmp_file = write_tmp_file(manifest)
     command = f"kubectl apply -f {tmp_file}"
     return run_command_with_updates(command, task)
@@ -422,13 +481,114 @@ class KueueManager:
       xpk_print(f"{task} returned ERROR {return_code}")
     return return_code
 
+  def __autocorrect_resource_limits(
+      self, kueue_config: KueueConfig
+  ) -> tuple[int, str]:
+    """Verify specified CPU and memory limits against machine type."""
+
+    cpu_limit = kueue_config.cpu_limit
+    memory_limit_str = kueue_config.memory_limit
+    if not cpu_limit and not memory_limit_str:
+      return cpu_limit, memory_limit_str
+
+    # Get CPU and memory capacity from machine type
+    command = (
+        "gcloud compute machine-types describe"
+        f" {kueue_config.system.gce_machine_type} "
+        f" --project={self.project} --zone={self.zone}"
+        " --format='value(guestCpus,memoryMb)'"
+    )
+    return_code, out = run_command_for_value(
+        command,
+        "Get vCPU and memory capacity for machine type",
+        dry_run_return_val="10 10",
+    )
+    if return_code != 0:
+      xpk_print(
+          "Unable to verify vCPU and memory capacity for machine type."
+          " XPK will proceed with using user-defined  limits."
+      )
+      return cpu_limit, memory_limit_str
+
+    cpu_capacity_str, memory_capacity_MB_str = out.split()
+    if cpu_limit:
+      cpu_limit = _autocorrect_cpu_limit(cpu_limit, int(cpu_capacity_str))
+    if memory_limit_str:
+      memory_limit_str = _autocorrect_memory_limit(
+          memory_limit_str, memory_capacity_MB_str
+      )
+    return cpu_limit, memory_limit_str
+
+
+def get_installed_kueue_version(
+    dry_run_version: Version | None = None,
+) -> tuple[int, Version | None]:
+  command = (
+      "kubectl get deployment kueue-controller-manager -n kueue-system -o"
+      " jsonpath='{.spec.template.spec.containers[0].image}'"
+  )
+  task = "Get kueue version on server"
+  return_code, val = run_command_for_value(
+      command,
+      task,
+      dry_run_return_val=(
+          f"registry.k8s.io/kueue/kueue:v{dry_run_version}"
+          if dry_run_version
+          else ""
+      ),
+  )
+  if return_code != 0:
+    return return_code, None
+  version_tag = val.split(":")
+  if len(version_tag) == 1:
+    return 1, None
+  return return_code, Version(version_tag[-1])
+
 
 def has_sub_slicing_enabled() -> tuple[int, bool | None]:
   return_code, value = run_command_for_value(
-      command="kubectl get topology", task="Get defined topologies"
+      command="kubectl get topology",
+      task="Get defined topologies",
+      dry_run_return_val=SUB_SLICE_TOPOLOGY_NAME,
   )
 
   if return_code != 0:
     return return_code, None
 
   return return_code, SUB_SLICE_TOPOLOGY_NAME in value
+
+
+def _autocorrect_cpu_limit(cpu_limit: int, cpu_capacity: int) -> int:
+  if cpu_limit > cpu_capacity:
+    xpk_print(
+        "The CPU limit is above the available capacity."
+        f" We will set CPU limit to {cpu_capacity}."
+    )
+  elif cpu_limit < cpu_capacity:
+    xpk_print(
+        "The CPU limit is below the available capacity, which would lead"
+        f" to underutilization. We will set CPU limit to {cpu_capacity}."
+    )
+  return cpu_capacity
+
+
+def _autocorrect_memory_limit(
+    memory_limit_str: str, memory_capacity_MB_str: str
+) -> str:
+  memory_limit_bytes = parse_quantity(memory_limit_str)
+  memory_capacity_bytes = int(memory_capacity_MB_str) << 20
+  if memory_limit_bytes == memory_capacity_bytes:
+    return memory_limit_str
+  memory_limit_str = memory_capacity_MB_str + "Mi"
+  if memory_limit_bytes > memory_capacity_bytes:
+    xpk_print(
+        "The memory limit is above the available capacity. We will set"
+        f" memory limit to {memory_limit_str}."
+    )
+  else:
+    xpk_print(
+        "The memory limit is below the available capacity, which would"
+        " lead to underutilization. We will set the memory limit to"
+        f" {memory_limit_str}."
+    )
+  return memory_limit_str