Add --num-cubes for superslicing cluster create (#914)

* Super-slicing cluster create arguments validation

* lint and goldens fixes

* fix f-string-without-interpolation

* add --num-cubes for superslicing cluster create

* update help message

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/reservationBlocks/block/reservationSubBlocks/subblock --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-cubes=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,4 +1,4 @@
-$ 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
+$ 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-cubes=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

src/xpk/commands/cluster.py

@@ -211,9 +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 and args.super_slicing:
-    validate_super_slicing_system(system)
-    _validate_super_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)
 
@@ -281,6 +283,22 @@ def _validate_gsc_reservation(args, creation_description: str):
     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.
 

src/xpk/commands/cluster_test.py

@@ -638,12 +638,16 @@ def test_validate_cluster_create_args_for_correct_super_slicing_args_pass(
   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']
@@ -659,6 +663,8 @@ def test_validate_cluster_create_args_for_super_slicing_system_not_supported_thr
   args = construct_args(
       super_slicing=True,
       reservation='test-reservation/reservationBlocks/block',
+      num_cubes=None,
+      num_slices=None,
   )
 
   with pytest.raises(SystemExit):
@@ -680,6 +686,8 @@ def test_validate_cluster_create_args_for_super_slicing_missing_reservation(
   args = construct_args(
       super_slicing=True,
       reservation=None,
+      num_cubes=None,
+      num_slices=None,
   )
 
   with pytest.raises(SystemExit):
@@ -699,6 +707,8 @@ def test_validate_cluster_create_args_for_super_slicing_reservation_no_blocks(
   args = construct_args(
       super_slicing=True,
       reservation='reservation',
+      num_cubes=None,
+      num_slices=None,
   )
 
   with pytest.raises(SystemExit):
@@ -718,6 +728,8 @@ def test_validate_cluster_create_args_for_super_slicing_sparse_deployment_type_r
   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'
 
@@ -729,3 +741,73 @@ def test_validate_cluster_create_args_for_super_slicing_sparse_deployment_type_r
       '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

src/xpk/parser/cluster.py

@@ -604,7 +604,10 @@ def add_shared_cluster_create_optional_arguments(
   parser_or_group.add_argument(
       '--num-slices',
       type=int,
-      default=1,
+      # removing default in case of super slicing because
+      # --num-slices must be equal to --num-cubes if both are set
+      # it will default to 1 during validation
+      default=1 if not FeatureFlags.SUPER_SLICING_ENABLED else None,
       help='The number of slices to run the job on, defaults to 1.',
       required=False,
   )
@@ -928,3 +931,14 @@ def add_cluster_create_super_slicing_arguments(
       action='store_true',
       help='Whether to set up cluster to support super-slicing',
   )
+  parser_or_group.add_argument(
+      '--num-cubes',
+      type=int,
+      # default value is set during validation because it needs to be compared
+      # against --num-slices
+      help=(
+          'Total number of cubes to create within a cluster, defaults to 1. Can'
+          ' only be used with --super-slicing.'
+      ),
+      required=False,
+  )

src/xpk/parser/cluster_test.py

@@ -188,3 +188,76 @@ def test_cluster_create_super_slicing_can_be_set():
   )
 
   assert args.super_slicing is True
+
+
+def test_cluster_create_num_cubes_is_hidden_with_flag_off():
+  FeatureFlags.SUPER_SLICING_ENABLED = False
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  help_str = parser.format_help()
+
+  assert "--num-cubes" not in help_str
+
+
+def test_cluster_create_num_cubes_is_shown_with_flag_on():
+  FeatureFlags.SUPER_SLICING_ENABLED = True
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  help_str = parser.format_help()
+
+  assert "--num-cubes" in help_str
+
+
+def test_cluster_create_num_cubes_can_be_set():
+  FeatureFlags.SUPER_SLICING_ENABLED = True
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  args = parser.parse_args(
+      [
+          "--cluster",
+          "test-cluster",
+          "--tpu-type",
+          "tpu7x-2",
+          "--num-cubes",
+          "5",
+      ],
+  )
+
+  assert args.num_cubes == 5
+
+
+def test_cluster_create_num_slices_defaults_to_1_if_no_superslicing_feature():
+  FeatureFlags.SUPER_SLICING_ENABLED = False
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  args = parser.parse_args(
+      [
+          "--cluster",
+          "test-cluster",
+          "--tpu-type",
+          "tpu7x-2",
+      ],
+  )
+
+  assert args.num_slices == 1
+
+
+def test_cluster_create_num_slices_has_no_default_if_superslicing_feature():
+  FeatureFlags.SUPER_SLICING_ENABLED = True
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  args = parser.parse_args(
+      [
+          "--cluster",
+          "test-cluster",
+          "--tpu-type",
+          "tpu7x-2",
+      ],
+  )
+
+  assert args.num_slices is None