Merge branch 'develop'

Author: scaliby

src/xpk/commands/cluster.py

@@ -17,7 +17,7 @@
 from tabulate import tabulate
 
 from ..utils.feature_flags import FeatureFlags
-from ..core.capacity import H100_DEVICE_TYPE, H200_DEVICE_TYPE, B200_DEVICE_TYPE
+from ..core.capacity import H100_DEVICE_TYPE, H200_DEVICE_TYPE, B200_DEVICE_TYPE, get_reservation_deployment_type
 from ..core.cluster import (
     get_all_clusters_programmatic,
     get_cluster_credentials,
@@ -204,6 +204,38 @@ def cluster_adapt(args) -> None:
 def _validate_cluster_create_args(args, system: SystemCharacteristics):
   if FeatureFlags.SUB_SLICING_ENABLED and args.sub_slicing:
     validate_sub_slicing_system(system)
+    _validate_sub_slicing_reservation(args)
+
+
+def _validate_sub_slicing_reservation(args):
+  if args.reservation is None:
+    xpk_print(
+        'Error: Validation failed: Sub-slicing cluster creation requires'
+        ' Cluster Director reservation to be specified.'
+    )
+    xpk_exit(1)
+
+  deployment_type = get_reservation_deployment_type(
+      reservation=args.reservation, project=args.project, zone=args.zone
+  )
+  if deployment_type != 'DENSE':
+    xpk_print(
+        'Error: Validation failed: The specified reservation'
+        f' "{args.reservation}" is not a Cluster Director reservation.'
+    )
+    xpk_print(
+        'Please provide a reservation created for Cluster Director to proceed.'
+    )
+    xpk_print('To list valid Cluster Director reservations, run:')
+    xpk_print(
+        '  gcloud compute reservations list --filter="deploymentType=DENSE"'
+    )
+    xpk_print(
+        'Refer to the documentation for more information on creating Cluster'
+        ' Director reservations:'
+        ' https://cloud.google.com/cluster-director/docs/reserve-capacity'
+    )
+    xpk_exit(1)
 
 
 def cluster_create(args) -> None:

src/xpk/commands/cluster_test.py

@@ -27,7 +27,8 @@
 @dataclass
 class _Mocks:
   common_print_mock: MagicMock
-  common_exit_mock: MagicMock
+  commands_print_mock: MagicMock
+  commands_get_reservation_deployment_type: MagicMock
 
 
 @pytest.fixture
@@ -36,12 +37,17 @@ def mock_common_print_and_exit(mocker):
       'xpk.commands.common.xpk_print',
       return_value=None,
   )
-  common_exit_mock = mocker.patch(
-      'xpk.commands.common.xpk_exit',
-      return_value=None,
+  commands_print_mock = mocker.patch(
+      'xpk.commands.cluster.xpk_print', return_value=None
+  )
+  commands_get_reservation_deployment_type = mocker.patch(
+      'xpk.commands.cluster.get_reservation_deployment_type',
+      return_value='DENSE',
   )
   return _Mocks(
-      common_print_mock=common_print_mock, common_exit_mock=common_exit_mock
+      common_print_mock=common_print_mock,
+      commands_get_reservation_deployment_type=commands_get_reservation_deployment_type,
+      commands_print_mock=commands_print_mock,
   )
 
 
@@ -61,32 +67,82 @@ def test_validate_cluster_create_args_for_correct_args_pass(
   _validate_cluster_create_args(args, DEFAULT_TEST_SYSTEM)
 
   assert mock_common_print_and_exit.common_print_mock.call_count == 0
-  assert mock_common_print_and_exit.common_exit_mock.call_count == 0
 
 
 def test_validate_cluster_create_args_for_correct_sub_slicing_args_pass(
     mock_common_print_and_exit: _Mocks,
 ):
   FeatureFlags.SUB_SLICING_ENABLED = True
-  args = Namespace(sub_slicing=True)
+  args = Namespace(
+      sub_slicing=True,
+      reservation='test-reservation',
+      project='project',
+      zone='zone',
+  )
 
   _validate_cluster_create_args(args, SUB_SLICING_SYSTEM)
 
   assert mock_common_print_and_exit.common_print_mock.call_count == 0
-  assert mock_common_print_and_exit.common_exit_mock.call_count == 0
 
 
 def test_validate_cluster_create_args_for_not_supported_system_throws(
     mock_common_print_and_exit: _Mocks,
 ):
   FeatureFlags.SUB_SLICING_ENABLED = True
-  args = Namespace(sub_slicing=True)
+  args = Namespace(
+      sub_slicing=True,
+      reservation='test-reservation',
+      project='project',
+      zone='zone',
+  )
 
-  _validate_cluster_create_args(args, DEFAULT_TEST_SYSTEM)
+  with pytest.raises(SystemExit):
+    _validate_cluster_create_args(args, DEFAULT_TEST_SYSTEM)
 
   assert mock_common_print_and_exit.common_print_mock.call_count == 1
   assert (
       mock_common_print_and_exit.common_print_mock.call_args[0][0]
       == 'Error: l4-1 does not support Sub-slicing.'
   )
-  assert mock_common_print_and_exit.common_exit_mock.call_count == 1
+
+
+def test_validate_cluster_create_args_for_missing_reservation(
+    mock_common_print_and_exit: _Mocks,
+):
+  FeatureFlags.SUB_SLICING_ENABLED = True
+  args = Namespace(
+      sub_slicing=True, project='project', zone='zone', reservation=None
+  )
+
+  with pytest.raises(SystemExit):
+    _validate_cluster_create_args(args, SUB_SLICING_SYSTEM)
+
+  assert mock_common_print_and_exit.commands_print_mock.call_count == 1
+  assert (
+      'Validation failed: Sub-slicing cluster creation requires'
+      in mock_common_print_and_exit.commands_print_mock.call_args[0][0]
+  )
+
+
+def test_validate_cluster_create_args_for_invalid_reservation(
+    mock_common_print_and_exit: _Mocks,
+):
+  FeatureFlags.SUB_SLICING_ENABLED = True
+  args = Namespace(
+      sub_slicing=True,
+      project='project',
+      zone='zone',
+      reservation='test-reservation',
+  )
+  mock_common_print_and_exit.commands_get_reservation_deployment_type.return_value = (
+      'SPARSE'
+  )
+
+  with pytest.raises(SystemExit):
+    _validate_cluster_create_args(args, SUB_SLICING_SYSTEM)
+
+  assert mock_common_print_and_exit.commands_print_mock.call_count == 5
+  assert (
+      'Refer to the documentation for more information on creating Cluster'
+      in mock_common_print_and_exit.commands_print_mock.call_args[0][0]
+  )

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,

src/xpk/core/capacity.py

@@ -152,6 +152,23 @@ def get_reservation_placement_policy(
   return output.strip()
 
 
+def get_reservation_deployment_type(
+    reservation: str, zone: str, project: str
+) -> str:
+  """Get reservation deployment type."""
+  command = (
+      f'gcloud beta compute reservations describe {reservation}'
+      f' --project={project} --zone={zone} --format="value(deploymentType)"'
+  )
+  return_code, output = run_command_for_value(
+      command, 'Get reservation deployment type', dry_run_return_val='DENSE'
+  )
+  if return_code != 0:
+    xpk_print(f'Get reservation deployment type ERROR {return_code}')
+    xpk_exit(1)
+  return output.strip()
+
+
 def verify_reservation_exists(args) -> int:
   """Verify the reservation exists.