fix: relax Kueue version check for sub and super-slicing workloads (#1111)

fix: relax Kueue version check for custom image SHAs

Author: jamOne-

src/xpk/core/kueue_manager.py

@@ -40,7 +40,7 @@
 from ..utils.file import write_tmp_file
 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
+from packaging.version import Version, InvalidVersion
 
 KUEUE_VERSION = Version("v0.15.2")
 LATEST_BREAKING_VERSION = Version("v0.15.0")
@@ -113,7 +113,13 @@ def install_or_upgrade(
     return_code, installed_version = get_installed_kueue_version()
 
     if return_code == 0 and installed_version:
-      if installed_version > self.kueue_version:
+      if isinstance(installed_version, str):
+        xpk_print(
+            "Cluster has an unknown or custom Kueue version installed. Skipping"
+            " installation."
+        )
+        return 0
+      elif installed_version > self.kueue_version:
         xpk_print(
             f"Cluster has a newer Kueue version, {installed_version}. Skipping"
             " installation."
@@ -532,7 +538,16 @@ def __autocorrect_resource_limits(
 
 def get_installed_kueue_version(
     dry_run_version: Version | None = None,
-) -> tuple[int, Version | None]:
+) -> tuple[int, Version | str | None]:
+  """Gets the currently installed Kueue version from the cluster.
+
+  Returns a tuple containing:
+  - The return code of the kubectl command.
+  - The version information:
+    - None if the command fails or the value is empty (e.g. Kueue not installed).
+    - A Version object if the image tag can be parsed as a valid version.
+    - A string if the image tag contains a custom SHA or is unparseable.
+  """
   command = (
       "kubectl get deployment kueue-controller-manager -n kueue-system -o"
       " jsonpath='{.spec.template.spec.containers[0].image}'"
@@ -547,12 +562,17 @@ def get_installed_kueue_version(
           else ""
       ),
   )
-  if return_code != 0:
+  if return_code != 0 or not val:
     return return_code, None
-  version_tag = val.split(":")
-  if len(version_tag) == 1:
-    return 1, None
-  return return_code, Version(version_tag[-1])
+
+  if "@sha256:" in val or ":" not in val:
+    return return_code, val
+
+  version_tag = val.split(":")[-1]
+  try:
+    return return_code, Version(version_tag)
+  except InvalidVersion:
+    return return_code, val
 
 
 def has_sub_slicing_enabled() -> tuple[int, bool | None]:

src/xpk/core/kueue_manager_test.py

@@ -66,6 +66,16 @@ def set_installed_kueue_version(
   )
 
 
+def set_custom_installed_kueue_version(
+    commands_tester: CommandsTester, image: str
+):
+  commands_tester.set_result_for_command(
+      (0, image),
+      "kubectl get deployment kueue-controller-manager",
+      "containers[0].image",
+  )
+
+
 @pytest.fixture(autouse=True)
 def mock_ask_for_user_consent(mocker: MockerFixture) -> MagicMock:
   return mocker.patch(
@@ -104,6 +114,21 @@ def test_install_or_upgrade_when_newer_version_already_installed(
   mock_commands.assert_command_not_run("kubectl apply")
 
 
+def test_install_or_upgrade_when_custom_version_already_installed(
+    mock_commands: CommandsTester, kueue_manager: KueueManager
+):
+  """Test install_or_upgrade when a custom Kueue build is installed."""
+  set_custom_installed_kueue_version(
+      mock_commands,
+      "us-central1-docker.pkg.dev/dummy-project/kueue/kueue@sha256:1234567890abcdef1234567890abcdef1234567890abcdef1234567890abcdef",
+  )
+
+  result = kueue_manager.install_or_upgrade(KUEUE_CONFIG)
+
+  assert result == 0
+  mock_commands.assert_command_not_run("kubectl apply")
+
+
 def test_install_or_upgrade_when_outdated(
     mock_commands: CommandsTester, kueue_manager: KueueManager
 ):

src/xpk/core/scheduling.py

@@ -35,6 +35,26 @@
 _SUB_SLICING_MINIMUM_KUEUE_VERSION = Version('0.13.0')
 _SUPER_SLICING_MINIMUM_KUEUE_VERSION = Version('0.15.2')
 _SUPER_SLICING_MAX_CUBES = 144
+
+
+def _is_kueue_version_sufficient(
+    return_code: int,
+    current_version: Version | str | None,
+    minimum_version: Version,
+) -> bool:
+  if (
+      return_code != 0
+      or current_version is None
+      or isinstance(current_version, str)
+  ):
+    xpk_print(
+        'Warning: Could not determine Kueue version. Proceeding with workload'
+        ' submission, but scheduling might fail if Kueue is outdated.'
+    )
+    return True
+  return current_version >= minimum_version
+
+
 ONE_TO_ONE_REPLICA_NODE_POOL_ASSIGNMENT_ANNOTATION = (
     'alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool'
 )
@@ -181,11 +201,8 @@ def _check_sub_slicing_availability(
   return_code, current_version = get_installed_kueue_version(
       dry_run_version=_SUB_SLICING_MINIMUM_KUEUE_VERSION
   )
-
-  return (
-      return_code == 0
-      and current_version is not None
-      and current_version >= _SUB_SLICING_MINIMUM_KUEUE_VERSION
+  return _is_kueue_version_sufficient(
+      return_code, current_version, _SUB_SLICING_MINIMUM_KUEUE_VERSION
   )
 
 
@@ -205,11 +222,8 @@ def _check_super_slicing_availability(
   return_code, current_version = get_installed_kueue_version(
       dry_run_version=_SUPER_SLICING_MINIMUM_KUEUE_VERSION
   )
-
-  return (
-      return_code == 0
-      and current_version is not None
-      and current_version >= _SUPER_SLICING_MINIMUM_KUEUE_VERSION
+  return _is_kueue_version_sufficient(
+      return_code, current_version, _SUPER_SLICING_MINIMUM_KUEUE_VERSION
   )
 
 

src/xpk/core/scheduling_test.py

@@ -174,6 +174,7 @@ class SchedulingTestCase:
   cluster_system: SystemCharacteristics | None = None
   resources_config_map: dict[str, str] | None = None
   kueue_version: str | None = None
+  kueue_version_return_code: int = 0
   sub_slicing_feature_enabled: bool = False
   sub_slicing_topology_set: bool = False
   super_slicing_topology_set: bool = False
@@ -269,6 +270,16 @@ class SchedulingTestCase:
             ),
             WorkloadScheduling.UNAVAILABLE,
         ),
+        (
+            'Sub-slicing, with unparseable Kueue version (SHA)',
+            dataclasses.replace(SUB_SLICING_CASE, kueue_version='sha256-12345'),
+            WorkloadScheduling.SUB_SLICING_AVAILABLE,
+        ),
+        (
+            'Sub-slicing, with failing Kueue version check',
+            dataclasses.replace(SUB_SLICING_CASE, kueue_version_return_code=1),
+            WorkloadScheduling.SUB_SLICING_AVAILABLE,
+        ),
         (
             'Sub-slicing, but low Kueue version',
             dataclasses.replace(SUB_SLICING_CASE, kueue_version='0.12.0'),
@@ -328,6 +339,20 @@ class SchedulingTestCase:
             SUPER_SLICING_CASE,
             WorkloadScheduling.SUPER_SLICING_AVAILABLE,
         ),
+        (
+            'Super-slicing, with unparseable Kueue version (SHA)',
+            dataclasses.replace(
+                SUPER_SLICING_CASE, kueue_version='sha256-12345'
+            ),
+            WorkloadScheduling.SUPER_SLICING_AVAILABLE,
+        ),
+        (
+            'Super-slicing, with failing Kueue version check',
+            dataclasses.replace(
+                SUPER_SLICING_CASE, kueue_version_return_code=1
+            ),
+            WorkloadScheduling.SUPER_SLICING_AVAILABLE,
+        ),
         (
             'Super-slicing, but low Kueue version',
             dataclasses.replace(SUPER_SLICING_CASE, kueue_version='0.13.0'),
@@ -418,7 +443,7 @@ def test_check_if_workload_can_schedule(
   FeatureFlags.SUB_SLICING_ENABLED = case.sub_slicing_feature_enabled
   commands_tester.set_result_for_command(
       (
-          0,
+          case.kueue_version_return_code,
           f'registry.k8s.io/kueue/kueue:v{case.kueue_version}'
           if case.kueue_version
           else '',