xpk 0.16.1__py3-none-any.whl → 0.17.1__py3-none-any.whl

xpk/commands/cluster.py

@@ -19,7 +19,7 @@ from tabulate import tabulate
 from ..utils.feature_flags import FeatureFlags
 from ..utils.versions import ReleaseChannel
 from ..core.pathways import get_pathways_machine_types
-from ..core.capacity import H100_DEVICE_TYPE, get_reservation_deployment_type
+from ..core.capacity import H100_DEVICE_TYPE, get_reservation_deployment_type, parse_reservation
 from ..core.cluster import (
     get_all_clusters_programmatic,
     get_cluster_credentials,
@@ -79,7 +79,7 @@ from ..utils.file import write_tmp_file
 from ..utils.execution_context import is_dry_run, is_quiet
 from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
 from . import cluster_gcluster
-from .common import set_cluster_command, validate_sub_slicing_system
+from .common import set_cluster_command, validate_sub_slicing_system, validate_super_slicing_system
 from jinja2 import Environment, FileSystemLoader
 from ..utils.templates import get_templates_absolute_path
 import shutil
@@ -211,6 +211,11 @@ 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)
+  if FeatureFlags.SUPER_SLICING_ENABLED:
+    _validate_num_slices_and_set_default(args)
+    if args.super_slicing:
+      validate_super_slicing_system(system)
+      _validate_super_slicing_reservation(args)
   if args.enable_pathways:
     _validate_pathways_machine(args)
 
@@ -233,15 +238,30 @@ def _validate_pathways_machine(args):
 
 
 def _validate_sub_slicing_reservation(args):
+  _validate_gsc_reservation(args, 'Sub-slicing')
+
+
+def _validate_super_slicing_reservation(args):
+  _validate_gsc_reservation(args, 'Super-slicing')
+  reservation = parse_reservation(args.reservation, args.project)
+  if reservation.block_name is None:
+    xpk_print(
+        'Error: Validation failed: Super-slicing cluster creation'
+        ' requires a block or sub-block reservation.'
+    )
+    xpk_exit(1)
+
+
+def _validate_gsc_reservation(args, creation_description: str):
   if args.reservation is None:
     xpk_print(
-        'Error: Validation failed: Sub-slicing cluster creation requires'
-        ' Cluster Director reservation to be specified.'
+        f'Error: Validation failed: {creation_description} 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
+      reservation_path=args.reservation, project=args.project, zone=args.zone
   )
   if deployment_type != 'DENSE':
     xpk_print(
@@ -263,6 +283,22 @@ def _validate_sub_slicing_reservation(args):
     xpk_exit(1)
 
 
+def _validate_num_slices_and_set_default(args):
+  if args.num_cubes is not None and not args.super_slicing:
+    xpk_print('--num-cubes can only be used with --super-slicing')
+    xpk_exit(1)
+
+  if (
+      args.num_cubes is not None
+      and args.num_slices is not None
+      and args.num_cubes != args.num_slices
+  ):
+    xpk_print('--num-cubes must not be different from --num-slices')
+    xpk_exit(1)
+
+  args.num_slices = args.num_slices or args.num_cubes or 1
+
+
 def cluster_create(args) -> None:
   """Function around cluster creation.
 
@@ -374,6 +410,7 @@ def cluster_create(args) -> None:
   )
   if return_code != 0:
     xpk_exit(return_code)
+  assert gke_node_pool_version
 
   run_gke_node_pool_create_command_code = run_gke_node_pool_create_command(
       args, system, gke_node_pool_version
@@ -1256,6 +1293,9 @@ def run_gke_cluster_create_command(
     addons_str = ','.join(addons)
     command += f' --addons={addons_str}'
 
+  if FeatureFlags.SUPER_SLICING_ENABLED and args.super_slicing:
+    command += ' --enable-slice-controller'
+
   if args.custom_cluster_arguments:
     command += f' {args.custom_cluster_arguments}'
 
@@ -1347,6 +1387,9 @@ def _install_kueue(
           configure_sub_slicing=(
               FeatureFlags.SUB_SLICING_ENABLED and args.sub_slicing
           ),
+          configure_super_slicing=(
+              FeatureFlags.SUPER_SLICING_ENABLED and args.super_slicing
+          ),
       )
   )
 

xpk/commands/cluster_gcluster.py

@@ -173,6 +173,9 @@ def __install_kueue(args) -> int:
           configure_sub_slicing=(
               FeatureFlags.SUB_SLICING_ENABLED and args.sub_slicing
           ),
+          configure_super_slicing=(
+              FeatureFlags.SUPER_SLICING_ENABLED and args.super_slicing
+          ),
       ),
       tolerations=tolerations,
   )

xpk/commands/cluster_gcluster_test.py

@@ -97,6 +97,7 @@ def test_install_kueue_standard(
       accelerator_type=AcceleratorType.GPU,
       device_type="h100-mega-80gb-8",
       supports_sub_slicing=False,
+      supports_super_slicing=False,
       docker_platform=DockerPlatform.ARM,
       gpu_config=GpuConfig(requires_topology=True),
   )
@@ -150,6 +151,7 @@ def test_install_kueue_with_autoprovisioning(
       accelerator_type=AcceleratorType.GPU,
       device_type="h100-mega-80gb-8",
       supports_sub_slicing=False,
+      supports_super_slicing=False,
       docker_platform=DockerPlatform.ARM,
       gpu_config=GpuConfig(requires_topology=True),
   )

xpk/commands/cluster_test.py

@@ -110,6 +110,7 @@ def construct_args(**kwargs: Any) -> Namespace:
       cluster='test-cluster',
       default_pool_cpu_num_nodes='100',
       sub_slicing=False,
+      super_slicing=False,
       gke_version='',
       private=False,
       authorized_networks=None,
@@ -226,6 +227,9 @@ GPU_TEST_SYSTEM: SystemCharacteristics = UserFacingNameToSystemCharacteristics[
 SUB_SLICING_SYSTEM: SystemCharacteristics = (
     UserFacingNameToSystemCharacteristics['v6e-4x4']
 )
+SUPER_SLICING_SYSTEM: SystemCharacteristics = (
+    UserFacingNameToSystemCharacteristics['tpu7x-4x4x4']
+)
 TPU_TEST_SYSTEM: SystemCharacteristics = UserFacingNameToSystemCharacteristics[
     'v6e-4x4'
 ]
@@ -608,3 +612,202 @@ def test_cluster_create_calls_run_command_with_correct_channel_and_version(
   ]
 
   mocks.commands_tester.assert_command_run(*expected_command_parts)
+
+
+def test_run_gke_cluster_create_command_with_super_slicing_enables_slice_controller(
+    mocks: _Mocks,
+):
+  FeatureFlags.SUPER_SLICING_ENABLED = True
+  result = run_gke_cluster_create_command(
+      args=construct_args(gke_version='1.2.3', super_slicing=True),
+      gke_control_plane_version='1.2.3',
+      system=SUPER_SLICING_SYSTEM,
+      release_channel=ReleaseChannel.REGULAR,
+  )
+
+  assert result == 0
+  mocks.commands_tester.assert_command_run(
+      'clusters create', '--enable-slice-controller'
+  )
+
+
+def test_validate_cluster_create_args_for_correct_super_slicing_args_pass(
+    mocks: _Mocks,
+):
+  FeatureFlags.SUPER_SLICING_ENABLED = True
+  args = construct_args(
+      super_slicing=True,
+      reservation='test-reservation/reservationBlocks/block',
+      num_cubes=None,
+      num_slices=None,
+  )
+
+  _validate_cluster_create_args(args, SUPER_SLICING_SYSTEM)
+  args = construct_args(
+      super_slicing=True,
+      reservation='test-reservation/reservationBlocks/block/reservationSubBlocks/subblock',
+      num_cubes=None,
+      num_slices=None,
+  )
+  _validate_cluster_create_args(
+      args, UserFacingNameToSystemCharacteristics['tpu7x-128']
+  )
+
+  assert mocks.common_print_mock.call_count == 0
+
+
+def test_validate_cluster_create_args_for_super_slicing_system_not_supported_throws(
+    mocks: _Mocks,
+):
+  FeatureFlags.SUPER_SLICING_ENABLED = True
+  args = construct_args(
+      super_slicing=True,
+      reservation='test-reservation/reservationBlocks/block',
+      num_cubes=None,
+      num_slices=None,
+  )
+
+  with pytest.raises(SystemExit):
+    _validate_cluster_create_args(
+        args, UserFacingNameToSystemCharacteristics['tpu7x-4x4x8']
+    )
+
+  assert mocks.common_print_mock.call_count == 1
+  assert (
+      mocks.common_print_mock.call_args[0][0]
+      == 'Error: tpu7x-256 does not support Super-slicing.'
+  )
+
+
+def test_validate_cluster_create_args_for_super_slicing_missing_reservation(
+    mocks: _Mocks,
+):
+  FeatureFlags.SUPER_SLICING_ENABLED = True
+  args = construct_args(
+      super_slicing=True,
+      reservation=None,
+      num_cubes=None,
+      num_slices=None,
+  )
+
+  with pytest.raises(SystemExit):
+    _validate_cluster_create_args(args, SUPER_SLICING_SYSTEM)
+
+  assert mocks.commands_print_mock.call_count == 1
+  assert (
+      'Validation failed: Super-slicing cluster creation requires'
+      in mocks.commands_print_mock.call_args[0][0]
+  )
+
+
+def test_validate_cluster_create_args_for_super_slicing_reservation_no_blocks(
+    mocks: _Mocks,
+):
+  FeatureFlags.SUPER_SLICING_ENABLED = True
+  args = construct_args(
+      super_slicing=True,
+      reservation='reservation',
+      num_cubes=None,
+      num_slices=None,
+  )
+
+  with pytest.raises(SystemExit):
+    _validate_cluster_create_args(args, SUPER_SLICING_SYSTEM)
+
+  assert mocks.commands_print_mock.call_count == 1
+  assert (
+      'requires a block or sub-block reservation'
+      in mocks.commands_print_mock.call_args[0][0]
+  )
+
+
+def test_validate_cluster_create_args_for_super_slicing_sparse_deployment_type_reservation(
+    mocks: _Mocks,
+):
+  FeatureFlags.SUPER_SLICING_ENABLED = True
+  args = construct_args(
+      super_slicing=True,
+      reservation='test-reservation/reservationBlocks/block',
+      num_cubes=None,
+      num_slices=None,
+  )
+  mocks.commands_get_reservation_deployment_type.return_value = 'SPARSE'
+
+  with pytest.raises(SystemExit):
+    _validate_cluster_create_args(args, SUPER_SLICING_SYSTEM)
+
+  assert mocks.commands_print_mock.call_count == 5
+  assert (
+      'Refer to the documentation for more information on creating Cluster'
+      in mocks.commands_print_mock.call_args[0][0]
+  )
+
+
+def test_validate_cluster_create_args_forbids_num_cubes_without_superslicing(
+    mocks: _Mocks,
+):
+  FeatureFlags.SUPER_SLICING_ENABLED = True  # enable the feature
+  args = construct_args(
+      super_slicing=False,  # but disable the flag
+      reservation='test-reservation/reservationBlocks/block',
+      num_cubes=1,
+      num_slices=None,
+  )
+
+  with pytest.raises(SystemExit):
+    _validate_cluster_create_args(args, SUPER_SLICING_SYSTEM)
+
+  assert mocks.commands_print_mock.call_count == 1
+  assert (
+      '--num-cubes can only be used with --super-slicing'
+      in mocks.commands_print_mock.call_args[0][0]
+  )
+
+
+def test_validate_cluster_create_args_forbids_num_cubes_different_from_num_slices(
+    mocks: _Mocks,
+):
+  FeatureFlags.SUPER_SLICING_ENABLED = True
+  args = construct_args(
+      super_slicing=True,
+      reservation='test-reservation/reservationBlocks/block',
+      num_cubes=1,
+      num_slices=2,
+  )
+
+  with pytest.raises(SystemExit):
+    _validate_cluster_create_args(args, SUPER_SLICING_SYSTEM)
+
+  assert mocks.commands_print_mock.call_count == 1
+  assert (
+      '--num-cubes must not be different from --num-slices'
+      in mocks.commands_print_mock.call_args[0][0]
+  )
+
+
+@pytest.mark.parametrize(
+    'num_cubes, num_slices, expected',
+    [
+        (None, None, 1),
+        (3, None, 3),
+        (None, 3, 3),
+        (3, 3, 3),
+    ],
+)
+def test_validate_cluster_create_args_sets_correct_num_slices(
+    mocks: _Mocks,
+    num_cubes: int | None,
+    num_slices: int | None,
+    expected: int,
+):
+  FeatureFlags.SUPER_SLICING_ENABLED = True
+  args = construct_args(
+      super_slicing=True,
+      reservation='test-reservation/reservationBlocks/block',
+      num_cubes=num_cubes,
+      num_slices=num_slices,
+  )
+
+  _validate_cluster_create_args(args, SUPER_SLICING_SYSTEM)
+
+  assert args.num_slices == expected

xpk/commands/common.py

@@ -73,3 +73,9 @@ def validate_sub_slicing_system(system: SystemCharacteristics):
   if not system.supports_sub_slicing:
     xpk_print(f'Error: {system.device_type} does not support Sub-slicing.')
     xpk_exit(1)
+
+
+def validate_super_slicing_system(system: SystemCharacteristics):
+  if not system.supports_super_slicing:
+    xpk_print(f'Error: {system.device_type} does not support Super-slicing.')
+    xpk_exit(1)

xpk/commands/kind.py

@@ -98,6 +98,7 @@ def cluster_create(args) -> None:
       AcceleratorType.CPU,
       'kind',
       supports_sub_slicing=False,
+      supports_super_slicing=False,
       docker_platform=DockerPlatform.ARM,
   )
 
@@ -113,6 +114,7 @@ def cluster_create(args) -> None:
           is_pathways_cluster=False,
           flex=False,
           configure_sub_slicing=False,
+          configure_super_slicing=False,
       ),
   )
 

xpk/commands/workload.py

@@ -132,6 +132,7 @@ spec:
               annotations:
                 {storage_annotations}
                 {sub_slicing_annotations}
+                {annotations_machine_label}
             spec:
               schedulerName: {args.scheduler}
               imagePullSecrets:
@@ -140,7 +141,7 @@ spec:
               {affinity}
               nodeSelector:
                 {accelerator_label}
-                {machine_label}
+                {node_selector_machine_label}
                 {placement_policy_label}
                 {autoprovisioning_args}
               priorityClassName: {args.priority}
@@ -498,19 +499,24 @@ def workload_create(args) -> None:
                 operator: NotIn
                 values: [{restart_on_exit_codes}]"""
 
-  if is_placement_policy_supported(workload_system):
+  placement_policy_label = ''
+  if (
+      # Don't bother with placement for sub/super-slicing workloads:
+      workload_scheduling == WorkloadScheduling.AVAILABLE
+      and is_placement_policy_supported(workload_system)
+  ):
     ensure_resource_policy_exists(
-        resource_policy_name=get_placement_policy_name(workload_system),
+        resource_policy_name=get_placement_policy_name(
+            workload_system, super_slicing=False
+        ),
         project=args.project,
         zone=args.zone,
         topology=workload_system.topology,
+        super_slicing=False,
+    )
+    placement_policy_label = create_placement_policy_label(
+        workload_system, super_slicing=False
     )
-
-  placement_policy_label = (
-      create_placement_policy_label(workload_system)
-      if is_placement_policy_supported(workload_system)
-      else ''
-  )
 
   # Create the workload file based on accelerator type or workload type.
   if workload_system.accelerator_type == AcceleratorType.GPU:
@@ -616,12 +622,26 @@ def workload_create(args) -> None:
     use_sub_slicing = (
         workload_scheduling == WorkloadScheduling.SUB_SLICING_AVAILABLE
     )
+    use_super_slicing = (
+        workload_scheduling == WorkloadScheduling.SUPER_SLICING_AVAILABLE
+    )
     if use_sub_slicing:
       xpk_print('Workload will be scheduled using the Sub-slicing feature.')
+    if use_super_slicing:
+      xpk_print('Workload will be scheduled using the Super-slicing feature.')
 
     container, debugging_dashboard_id = get_user_workload_container(
         args, workload_system
     )
+
+    machine_label = (
+        create_machine_label(cluster_system)
+        if use_sub_slicing and cluster_system
+        else create_machine_label(workload_system)
+    )
+    node_selector_machine_label = machine_label if not use_super_slicing else ''
+    annotations_machine_label = machine_label if use_super_slicing else ''
+
     yml_string = WORKLOAD_CREATE_YAML.format(
         args=args,
         container=container,
@@ -636,11 +656,8 @@ def workload_create(args) -> None:
             else ''
         ),
         placement_policy_label=placement_policy_label,
-        machine_label=(
-            create_machine_label(cluster_system)
-            if use_sub_slicing and cluster_system
-            else create_machine_label(workload_system)
-        ),
+        node_selector_machine_label=node_selector_machine_label,
+        annotations_machine_label=annotations_machine_label,
         local_queue_name=LOCAL_QUEUE_NAME,
         autoprovisioning_args=autoprovisioning_args,
         volumes=get_volumes(args, workload_system),
@@ -822,12 +839,15 @@ def workload_delete(args) -> None:
     if len(workloads) == 1:
       return_code = run_command_with_updates(commands[0], 'Delete Workload')
     else:
-      return_code = run_commands(
+      maybe_failure = run_commands(
           commands,
           'Delete Workload',
           task_names,
           batch=100,
       )
+      return_code = (
+          maybe_failure.return_code if maybe_failure is not None else 0
+      )
 
     if return_code != 0:
       xpk_print(f'Delete Workload request returned ERROR {return_code}')

xpk/commands/workload_test.py

@@ -34,6 +34,7 @@ SYSTEM_CHARACTERISTICS = SystemCharacteristics(
     accelerator_type=AcceleratorType.TPU,
     device_type='l4-1',
     supports_sub_slicing=True,
+    supports_super_slicing=False,
     requires_workload_policy=False,
     docker_platform=DockerPlatform.AMD,
 )