Validate kueue subslice constraints in workload create (#731)

feat: validate kueue subslice constraints in workload create

Author: scaliby

src/xpk/commands/workload.py

@@ -26,7 +26,8 @@
     get_cluster_credentials,
     setup_k8s_env,
 )
-from ..core.commands import run_command_with_updates, run_commands
+from ..core.commands import run_command_with_updates, run_commands, run_command_for_value
+from ..core.kueue_manager import KueueManager, SUB_SLICE_TOPOLOGY_NAME
 from ..core.config import (VERTEX_TENSORBOARD_FEATURE_FLAG, XPK_CURRENT_VERSION)
 from ..core.docker_container import (
     get_main_container_docker_image,
@@ -95,6 +96,7 @@
     tcpxo_decorator,
 )
 from ..utils.console import get_user_input, xpk_exit, xpk_print
+from packaging.version import Version
 from ..utils.file import write_tmp_file
 from ..utils.execution_context import is_dry_run
 from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
@@ -283,6 +285,7 @@
 """
 
 SUB_SLICING_TOPOLOGIES = ['2x2', '2x4', '4x4', '4x8', '8x8', '8x16', '16x16']
+SUB_SLICING_MINIMUM_KUEUE_VERSION = Version('0.13.0')
 
 
 def workload_create_pathways(args) -> None:
@@ -340,6 +343,7 @@ def workload_create(args) -> None:
     xpk_exit(return_code)
 
   if FeatureFlags.SUB_SLICING_ENABLED and args.sub_slicing_topology is not None:
+    _validate_sub_slicing_availability()
     _validate_sub_slicing_topology(system, args.sub_slicing_topology)
 
   if not check_if_workload_can_schedule(args, system):
@@ -678,6 +682,43 @@ def workload_create(args) -> None:
   xpk_exit(0)
 
 
+def _validate_sub_slicing_availability():
+  return_code, value = run_command_for_value(
+      command='kubectl get topology', task='Get defined topologies'
+  )
+
+  if return_code != 0:
+    xpk_print(
+        'Error: Unable to validate sub-slicing support on a given cluster.'
+    )
+    xpk_exit(1)
+
+  if SUB_SLICE_TOPOLOGY_NAME not in value:
+    xpk_print(
+        'Error: Cluster has not been not set up for Sub-slicing. Please enable'
+        ' --sub-slicing in "cluster create" command first.'
+    )
+    xpk_exit(1)
+
+  kueue_manager = KueueManager()
+  return_code, current_version = kueue_manager.get_installed_kueue_version()
+  if return_code != 0:
+    xpk_print(
+        'Error: Unable to validate sub-slicing support on a given cluster.'
+    )
+    xpk_exit(1)
+
+  if current_version < SUB_SLICING_MINIMUM_KUEUE_VERSION:
+    xpk_print(
+        "Error: Current Kueue version ({current_version}) doesn't support"
+        ' Sub-slicing. The minimal required version is'
+        ' v{SUB_SLICING_MINIMUM_KUEUE_VERSION}. Please either update Kueue'
+        ' manually, or run "cluster create --sub-slicing" on the existing'
+        ' cluster.'
+    )
+    xpk_exit(1)
+
+
 def _validate_sub_slicing_topology(
     system_characteristics: SystemCharacteristics, sub_slicing_topology: str
 ) -> None:

src/xpk/commands/workload_test.py

@@ -18,7 +18,8 @@
 from unittest.mock import MagicMock, patch
 import pytest
 from ..core.system_characteristics import SystemCharacteristics
-from .workload import _validate_sub_slicing_topology
+from .workload import _validate_sub_slicing_topology, _validate_sub_slicing_availability
+from packaging.version import Version
 
 
 SYSTEM_CHARACTERISTICS = SystemCharacteristics(
@@ -40,7 +41,7 @@ def xpk_print(mocker):
 
 
 def test_validate_sub_slicing_topology_exits_for_unsupported_topology(
-    xpk_print,
+    xpk_print: MagicMock,
 ):
   with pytest.raises(SystemExit):
     _validate_sub_slicing_topology(SYSTEM_CHARACTERISTICS, '2x1')
@@ -50,7 +51,9 @@ def test_validate_sub_slicing_topology_exits_for_unsupported_topology(
   )
 
 
-def test_validate_sub_slicing_topology_exits_for_too_large_topology(xpk_print):
+def test_validate_sub_slicing_topology_exits_for_too_large_topology(
+    xpk_print: MagicMock,
+):
   with pytest.raises(SystemExit):
     _validate_sub_slicing_topology(SYSTEM_CHARACTERISTICS, '16x16')
 
@@ -64,6 +67,86 @@ def test_validate_sub_slicing_topology_does_nothing_for_supported_topology():
   _validate_sub_slicing_topology(SYSTEM_CHARACTERISTICS, '4x4')
 
 
+def test_validate_sub_slicing_availability_exits_when_getting_topologies_fails(
+    xpk_print: MagicMock, mocker
+):
+  mocker.patch(
+      'xpk.commands.workload.run_command_for_value',
+      return_value=(1, ''),
+  )
+  with pytest.raises(SystemExit):
+    _validate_sub_slicing_availability()
+
+  assert (
+      'Unable to validate sub-slicing support'
+      in xpk_print.mock_calls[0].args[0]
+  )
+
+
+def test_validate_sub_slicing_availability_exits_when_subslicing_topology_is_not_defined(
+    xpk_print: MagicMock, mocker
+):
+  mocker.patch(
+      'xpk.commands.workload.run_command_for_value',
+      return_value=(0, ''),
+  )
+  with pytest.raises(SystemExit):
+    _validate_sub_slicing_availability()
+
+  assert (
+      'Cluster has not been not set up for Sub-slicing.'
+      in xpk_print.mock_calls[0].args[0]
+  )
+
+
+def test_validate_sub_slicing_availability_exits_when_kueue_version_cannot_be_determined(
+    xpk_print: MagicMock, mocker
+):
+  mocker.patch(
+      'xpk.commands.workload.run_command_for_value',
+      return_value=(0, 'sub-slice-topology'),
+  )
+  mocker.patch(
+      'xpk.commands.workload.KueueManager.get_installed_kueue_version',
+      return_value=(1, None),
+  )
+  with pytest.raises(SystemExit):
+    _validate_sub_slicing_availability()
+
+  assert 'Unable to validate sub-slicing' in xpk_print.mock_calls[0].args[0]
+
+
+def test_validate_sub_slicing_availability_exits_when_kueue_version_does_not_meet_minimum_requirements(
+    xpk_print: MagicMock, mocker
+):
+  mocker.patch(
+      'xpk.commands.workload.run_command_for_value',
+      return_value=(0, 'sub-slice-topology'),
+  )
+  mocker.patch(
+      'xpk.commands.workload.KueueManager.get_installed_kueue_version',
+      return_value=(0, Version('0.0.0')),
+  )
+  with pytest.raises(SystemExit):
+    _validate_sub_slicing_availability()
+
+  assert 'The minimal required version is' in xpk_print.mock_calls[0].args[0]
+
+
+def test_validate_sub_slicing_availability_does_nothing_when_cluster_is_correctly_configured_for_subslicing(
+    mocker,
+):
+  mocker.patch(
+      'xpk.commands.workload.run_command_for_value',
+      return_value=(0, 'sub-slice-topology'),
+  )
+  mocker.patch(
+      'xpk.commands.workload.KueueManager.get_installed_kueue_version',
+      return_value=(0, Version('0.13.0')),
+  )
+  _validate_sub_slicing_availability()
+
+
 @patch('xpk.commands.common.xpk_print')
 def test_validate_sub_slicing_topology_fails_for_unsupported_system(
     common_xpk_print: MagicMock,