Move sub-slicing validation to kueue-manager (#756)

* feat: move sub-slicing validation to kueue-manager

* style: remove unused import

* style: remove unused import

* feat: refactor to use free function

Author: scaliby

src/xpk/commands/workload.py

@@ -26,8 +26,8 @@
     get_cluster_credentials,
     setup_k8s_env,
 )
-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.commands import run_command_with_updates, run_commands
+from ..core.kueue_manager import KueueManager, has_sub_slicing_enabled
 from ..core.config import (VERTEX_TENSORBOARD_FEATURE_FLAG, XPK_CURRENT_VERSION)
 from ..core.docker_container import (
     get_main_container_docker_image,
@@ -683,17 +683,14 @@ def workload_create(args) -> None:
 
 
 def _validate_sub_slicing_availability():
-  return_code, value = run_command_for_value(
-      command='kubectl get topology', task='Get defined topologies'
-  )
-
+  return_code, sub_slicing_enabled = has_sub_slicing_enabled()
   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:
+  if not sub_slicing_enabled:
     xpk_print(
         'Error: Cluster has not been not set up for Sub-slicing. Please enable'
         ' --sub-slicing in "cluster create" command first.'

src/xpk/commands/workload_test.py

@@ -71,8 +71,8 @@ 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, ''),
+      'xpk.commands.workload.has_sub_slicing_enabled',
+      return_value=(1, None),
   )
   with pytest.raises(SystemExit):
     _validate_sub_slicing_availability()
@@ -87,8 +87,8 @@ def test_validate_sub_slicing_availability_exits_when_subslicing_topology_is_not
     xpk_print: MagicMock, mocker
 ):
   mocker.patch(
-      'xpk.commands.workload.run_command_for_value',
-      return_value=(0, ''),
+      'xpk.commands.workload.has_sub_slicing_enabled',
+      return_value=(0, False),
   )
   with pytest.raises(SystemExit):
     _validate_sub_slicing_availability()
@@ -103,8 +103,8 @@ def test_validate_sub_slicing_availability_exits_when_kueue_version_cannot_be_de
     xpk_print: MagicMock, mocker
 ):
   mocker.patch(
-      'xpk.commands.workload.run_command_for_value',
-      return_value=(0, 'sub-slice-topology'),
+      'xpk.commands.workload.has_sub_slicing_enabled',
+      return_value=(0, True),
   )
   mocker.patch(
       'xpk.commands.workload.KueueManager.get_installed_kueue_version',
@@ -120,8 +120,8 @@ def test_validate_sub_slicing_availability_exits_when_kueue_version_does_not_mee
     xpk_print: MagicMock, mocker
 ):
   mocker.patch(
-      'xpk.commands.workload.run_command_for_value',
-      return_value=(0, 'sub-slice-topology'),
+      'xpk.commands.workload.has_sub_slicing_enabled',
+      return_value=(0, True),
   )
   mocker.patch(
       'xpk.commands.workload.KueueManager.get_installed_kueue_version',
@@ -137,8 +137,8 @@ def test_validate_sub_slicing_availability_does_nothing_when_cluster_is_correctl
     mocker,
 ):
   mocker.patch(
-      'xpk.commands.workload.run_command_for_value',
-      return_value=(0, 'sub-slice-topology'),
+      'xpk.commands.workload.has_sub_slicing_enabled',
+      return_value=(0, True),
   )
   mocker.patch(
       'xpk.commands.workload.KueueManager.get_installed_kueue_version',

src/xpk/core/kueue_manager.py

@@ -421,3 +421,14 @@ def __update_kueue_resources_if_necessary(self) -> int:
     if return_code != 0:
       xpk_print(f"{task} returned ERROR {return_code}")
     return return_code
+
+
+def has_sub_slicing_enabled() -> tuple[int, bool | None]:
+  return_code, value = run_command_for_value(
+      command="kubectl get topology", task="Get defined topologies"
+  )
+
+  if return_code != 0:
+    return return_code, None
+
+  return return_code, SUB_SLICE_TOPOLOGY_NAME in value

src/xpk/core/kueue_manager_test.py

@@ -19,7 +19,7 @@
 import yaml
 from unittest.mock import MagicMock, patch
 
-from xpk.core.kueue_manager import KueueConfig, KueueManager
+from xpk.core.kueue_manager import KueueConfig, KueueManager, has_sub_slicing_enabled
 from xpk.core.system_characteristics import AcceleratorType, SystemCharacteristics
 from packaging.version import Version
 
@@ -562,6 +562,36 @@ def _trigger_installation(self, kueue_config: KueueConfig) -> str:
     return manifest
 
 
+@patch("xpk.core.kueue_manager.run_command_for_value")
+def test_has_sub_slicing_enabled_returns_exit_code_when_command_fails(
+    mock_run_for_value,
+):
+  mock_run_for_value.return_value = (1, None)
+  return_code, result = has_sub_slicing_enabled()
+  assert return_code == 1
+  assert result is None
+
+
+@patch("xpk.core.kueue_manager.run_command_for_value")
+def test_has_sub_slicing_enabled_returns_false_when_sub_slicing_topology_is_not_present(
+    mock_run_for_value,
+):
+  mock_run_for_value.return_value = (0, "")
+  return_code, result = has_sub_slicing_enabled()
+  assert return_code == 0
+  assert result is False
+
+
+@patch("xpk.core.kueue_manager.run_command_for_value")
+def test_has_sub_slicing_enabled_returns_true_when_sub_slicing_topology_is_not_present(
+    mock_run_for_value,
+):
+  mock_run_for_value.return_value = (0, "sub-slice-topology")
+  return_code, result = has_sub_slicing_enabled()
+  assert return_code == 0
+  assert result is True
+
+
 T = TypeVar("T")