Super-slicing cluster create arguments validation (#911)

* Super-slicing cluster create arguments validation

* lint and goldens fixes

* fix f-string-without-interpolation

Author: SikaGrr

goldens.yaml

@@ -18,7 +18,7 @@ goldens:
   "Cluster create sub-slicing":
     command: SUB_SLICING_ENABLED=true xpk cluster create --project=golden-project --zone=us-central1-a --cluster=golden-cluster --tpu-type=v6e-4x4 --reservation=golden-reservation --sub-slicing --dry-run
   "Cluster create super-slicing":
-    command: SUPER_SLICING_ENABLED=true xpk cluster create --project=golden-project --zone=us-central1-a --cluster=golden-cluster --tpu-type=tpu7x-4x4x4 --reservation=golden-reservation --super-slicing --num-slices=5 --dry-run
+    command: SUPER_SLICING_ENABLED=true xpk cluster create --project=golden-project --zone=us-central1-a --cluster=golden-cluster --tpu-type=tpu7x-4x4x4 --reservation=golden-reservation/reservationBlocks/block/reservationSubBlocks/subblock --super-slicing --num-slices=5 --dry-run
   "Cluster create private":
     command: xpk cluster create-pathways --project=golden-project --zone=us-central1-a --cluster=golden-cluster-private --private --tpu-type=v5p-8  --num-slices=1 --default-pool-cpu-machine-type=n1-standard-16 --default-pool-cpu-num-nodes=4 --reservation=golden-reservation --dry-run
   "Cluster create with Managed Lustre driver":

goldens/Cluster_create_super-slicing.txt

@@ -1,7 +1,9 @@
-$ SUPER_SLICING_ENABLED=true xpk cluster create --project=golden-project --zone=us-central1-a --cluster=golden-cluster --tpu-type=tpu7x-4x4x4 --reservation=golden-reservation --super-slicing --num-slices=5 --dry-run
+$ SUPER_SLICING_ENABLED=true xpk cluster create --project=golden-project --zone=us-central1-a --cluster=golden-cluster --tpu-type=tpu7x-4x4x4 --reservation=golden-reservation/reservationBlocks/block/reservationSubBlocks/subblock --super-slicing --num-slices=5 --dry-run
 [XPK] Starting xpk v0.0.0
 [XPK] Starting cluster create for cluster golden-cluster:
 [XPK] Working on golden-project and us-central1-a
+[XPK] Task: `Get reservation deployment type` is implemented by the following command not running since it is a dry run. 
+gcloud beta compute reservations describe golden-reservation --project=golden-project --zone=us-central1-a --format="value(deploymentType)"
 [XPK] Task: `Determine server supported GKE versions for default gke version` is implemented by the following command not running since it is a dry run. 
 gcloud container get-server-config --project=golden-project --region=us-central1 --flatten="channels" --filter="channels.channel=RAPID" --format="value(channels.defaultVersion)"
 [XPK] Task: `Determine server supported GKE versions for valid versions` is implemented by the following command not running since it is a dry run. 
@@ -51,11 +53,11 @@ kubectl get configmap golden-cluster-resources-configmap -o=custom-columns="Conf
 [XPK] Existing node pool names  ['0']
 [XPK] Task: `Retrieve resource policy` is implemented by the following command not running since it is a dry run. 
 gcloud compute resource-policies describe tpu7x-128-4x4x4-ss-placement-policy --project=golden-project --region=us-central1
-[XPK] To complete NodepoolCreate-golden-cluster-np-0 we are executing gcloud beta container node-pools create golden-cluster-np-0 --location=us-central1 --cluster=golden-cluster --project=golden-project --node-locations=us-central1-a --machine-type=tpu7x-standard-4t --host-maintenance-interval=AS_NEEDED --reservation-affinity=specific --reservation=golden-reservation --placement-policy=tpu7x-128-4x4x4-ss-placement-policy --enable-gvnic --node-version=0 --num-nodes=16 --scopes=storage-full,gke-default,"https://www.googleapis.com/auth/cloud-platform" --max-pods-per-node 15  
-[XPK] To complete NodepoolCreate-golden-cluster-np-1 we are executing gcloud beta container node-pools create golden-cluster-np-1 --location=us-central1 --cluster=golden-cluster --project=golden-project --node-locations=us-central1-a --machine-type=tpu7x-standard-4t --host-maintenance-interval=AS_NEEDED --reservation-affinity=specific --reservation=golden-reservation --placement-policy=tpu7x-128-4x4x4-ss-placement-policy --enable-gvnic --node-version=0 --num-nodes=16 --scopes=storage-full,gke-default,"https://www.googleapis.com/auth/cloud-platform" --max-pods-per-node 15  
-[XPK] To complete NodepoolCreate-golden-cluster-np-2 we are executing gcloud beta container node-pools create golden-cluster-np-2 --location=us-central1 --cluster=golden-cluster --project=golden-project --node-locations=us-central1-a --machine-type=tpu7x-standard-4t --host-maintenance-interval=AS_NEEDED --reservation-affinity=specific --reservation=golden-reservation --placement-policy=tpu7x-128-4x4x4-ss-placement-policy --enable-gvnic --node-version=0 --num-nodes=16 --scopes=storage-full,gke-default,"https://www.googleapis.com/auth/cloud-platform" --max-pods-per-node 15  
-[XPK] To complete NodepoolCreate-golden-cluster-np-3 we are executing gcloud beta container node-pools create golden-cluster-np-3 --location=us-central1 --cluster=golden-cluster --project=golden-project --node-locations=us-central1-a --machine-type=tpu7x-standard-4t --host-maintenance-interval=AS_NEEDED --reservation-affinity=specific --reservation=golden-reservation --placement-policy=tpu7x-128-4x4x4-ss-placement-policy --enable-gvnic --node-version=0 --num-nodes=16 --scopes=storage-full,gke-default,"https://www.googleapis.com/auth/cloud-platform" --max-pods-per-node 15  
-[XPK] To complete NodepoolCreate-golden-cluster-np-4 we are executing gcloud beta container node-pools create golden-cluster-np-4 --location=us-central1 --cluster=golden-cluster --project=golden-project --node-locations=us-central1-a --machine-type=tpu7x-standard-4t --host-maintenance-interval=AS_NEEDED --reservation-affinity=specific --reservation=golden-reservation --placement-policy=tpu7x-128-4x4x4-ss-placement-policy --enable-gvnic --node-version=0 --num-nodes=16 --scopes=storage-full,gke-default,"https://www.googleapis.com/auth/cloud-platform" --max-pods-per-node 15  
+[XPK] To complete NodepoolCreate-golden-cluster-np-0 we are executing gcloud beta container node-pools create golden-cluster-np-0 --location=us-central1 --cluster=golden-cluster --project=golden-project --node-locations=us-central1-a --machine-type=tpu7x-standard-4t --host-maintenance-interval=AS_NEEDED --reservation-affinity=specific --reservation=golden-reservation/reservationBlocks/block/reservationSubBlocks/subblock --placement-policy=tpu7x-128-4x4x4-ss-placement-policy --enable-gvnic --node-version=0 --num-nodes=16 --scopes=storage-full,gke-default,"https://www.googleapis.com/auth/cloud-platform" --max-pods-per-node 15  
+[XPK] To complete NodepoolCreate-golden-cluster-np-1 we are executing gcloud beta container node-pools create golden-cluster-np-1 --location=us-central1 --cluster=golden-cluster --project=golden-project --node-locations=us-central1-a --machine-type=tpu7x-standard-4t --host-maintenance-interval=AS_NEEDED --reservation-affinity=specific --reservation=golden-reservation/reservationBlocks/block/reservationSubBlocks/subblock --placement-policy=tpu7x-128-4x4x4-ss-placement-policy --enable-gvnic --node-version=0 --num-nodes=16 --scopes=storage-full,gke-default,"https://www.googleapis.com/auth/cloud-platform" --max-pods-per-node 15  
+[XPK] To complete NodepoolCreate-golden-cluster-np-2 we are executing gcloud beta container node-pools create golden-cluster-np-2 --location=us-central1 --cluster=golden-cluster --project=golden-project --node-locations=us-central1-a --machine-type=tpu7x-standard-4t --host-maintenance-interval=AS_NEEDED --reservation-affinity=specific --reservation=golden-reservation/reservationBlocks/block/reservationSubBlocks/subblock --placement-policy=tpu7x-128-4x4x4-ss-placement-policy --enable-gvnic --node-version=0 --num-nodes=16 --scopes=storage-full,gke-default,"https://www.googleapis.com/auth/cloud-platform" --max-pods-per-node 15  
+[XPK] To complete NodepoolCreate-golden-cluster-np-3 we are executing gcloud beta container node-pools create golden-cluster-np-3 --location=us-central1 --cluster=golden-cluster --project=golden-project --node-locations=us-central1-a --machine-type=tpu7x-standard-4t --host-maintenance-interval=AS_NEEDED --reservation-affinity=specific --reservation=golden-reservation/reservationBlocks/block/reservationSubBlocks/subblock --placement-policy=tpu7x-128-4x4x4-ss-placement-policy --enable-gvnic --node-version=0 --num-nodes=16 --scopes=storage-full,gke-default,"https://www.googleapis.com/auth/cloud-platform" --max-pods-per-node 15  
+[XPK] To complete NodepoolCreate-golden-cluster-np-4 we are executing gcloud beta container node-pools create golden-cluster-np-4 --location=us-central1 --cluster=golden-cluster --project=golden-project --node-locations=us-central1-a --machine-type=tpu7x-standard-4t --host-maintenance-interval=AS_NEEDED --reservation-affinity=specific --reservation=golden-reservation/reservationBlocks/block/reservationSubBlocks/subblock --placement-policy=tpu7x-128-4x4x4-ss-placement-policy --enable-gvnic --node-version=0 --num-nodes=16 --scopes=storage-full,gke-default,"https://www.googleapis.com/auth/cloud-platform" --max-pods-per-node 15  
 [XPK] Breaking up a total of 5 commands into 1 batches
 [XPK] Pretending all the jobs succeeded
 [XPK] Create or delete node pool request complete.
@@ -70,15 +72,15 @@ metadata:
 data:
   tpu7x-128: "80"
 
-[XPK] Temp file (a3dd06b296e1eb6792c99ad309e6eb714888c53e8b8fb8adc3beb8f250ef163c) content: 
+[XPK] Temp file (cd1d35a150fd66bdee3f462567a1cf4ed56b8af9ab774e15dcdac83b6b82e48e) content: 
 kind: ConfigMap
 apiVersion: v1
 metadata:
   name: golden-cluster-metadata-configmap
 data:
   xpk_version: v0.0.0
   capacity_type: RESERVATION
-  reservation_id: golden-reservation
+  reservation_id: golden-reservation/reservationBlocks/block/reservationSubBlocks/subblock
 
 [XPK] Breaking up a total of 2 commands into 1 batches
 [XPK] Pretending all the jobs succeeded

src/xpk/commands/cluster.py

@@ -19,7 +19,7 @@
 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.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,9 @@ 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 and args.super_slicing:
+    validate_super_slicing_system(system)
+    _validate_super_slicing_reservation(args)
   if args.enable_pathways:
     _validate_pathways_machine(args)
 
@@ -233,15 +236,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(

src/xpk/commands/cluster_test.py

@@ -629,3 +629,103 @@ def test_run_gke_cluster_create_command_with_super_slicing_enables_slice_control
   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',
+  )
+
+  _validate_cluster_create_args(args, SUPER_SLICING_SYSTEM)
+  args = construct_args(
+      super_slicing=True,
+      reservation='test-reservation/reservationBlocks/block/reservationSubBlocks/subblock',
+  )
+  _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',
+  )
+
+  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,
+  )
+
+  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',
+  )
+
+  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',
+  )
+  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]
+  )

src/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)