Add --sub-slicing to cluster create-pathways and set default for create-ray

src/xpk/parser/cluster.py

@@ -151,11 +151,7 @@ def set_cluster_create_parser(cluster_create_parser: ArgumentParser):
       ),
   )
   if FeatureFlags.SUB_SLICING_ENABLED:
-    cluster_create_optional_arguments.add_argument(
-        '--sub-slicing',
-        action='store_true',
-        help='Whether to set up cluster to support sub-slicing',
-    )
+    add_cluster_create_sub_slicing_arguments(cluster_create_optional_arguments)
 
   autoprovisioning_arguments = cluster_create_parser.add_argument_group(
       'Autoprovisioning Arguments',
@@ -232,6 +228,10 @@ def set_cluster_create_pathways_parser(
   add_shared_cluster_create_optional_arguments(
       cluster_create_pathways_optional_arguments
   )
+  if FeatureFlags.SUB_SLICING_ENABLED:
+    add_cluster_create_sub_slicing_arguments(
+        cluster_create_pathways_optional_arguments
+    )
 
   autoprovisioning_arguments = (
       cluster_create_pathways_parser.add_argument_group(
@@ -365,7 +365,9 @@ def set_cluster_create_ray_parser(cluster_create_ray_parser: ArgumentParser):
   )
   add_resource_limits(cluster_create_resource_limits)
 
-  cluster_create_ray_parser.set_defaults(func=cluster_create_ray_cluster)
+  cluster_create_ray_parser.set_defaults(
+      func=cluster_create_ray_cluster, sub_slicing=False
+  )
 
 
 def set_cluster_delete_parser(cluster_delete_parser: ArgumentParser):
@@ -964,3 +966,13 @@ def add_resource_limits(parser_or_group: ParserOrArgumentGroup):
       default=None,
       help='The CPU limit for the Kueue controller manager.',
   )
+
+
+def add_cluster_create_sub_slicing_arguments(
+    parser_or_group: ParserOrArgumentGroup,
+):
+  parser_or_group.add_argument(
+      '--sub-slicing',
+      action='store_true',
+      help='Whether to set up cluster to support sub-slicing',
+  )

src/xpk/parser/cluster_test.py

@@ -15,7 +15,7 @@
 """
 
 import argparse
-from xpk.parser.cluster import set_cluster_create_parser
+from xpk.parser.cluster import set_cluster_create_parser, set_cluster_create_pathways_parser, set_cluster_create_ray_parser
 import pytest
 from ..utils.feature_flags import FeatureFlags
 
@@ -64,3 +64,42 @@ def test_cluster_create_sub_slicing_can_be_set():
   )
 
   assert args.sub_slicing is True
+
+
+def test_cluster_create_pathways_sub_slicing_is_hidden_with_flag_off():
+  FeatureFlags.SUB_SLICING_ENABLED = False
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_pathways_parser(parser)
+  help_str = parser.format_help()
+
+  assert "--sub-slicing" not in help_str
+
+
+def test_cluster_create_pathways_sub_slicing_can_be_set():
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_pathways_parser(parser)
+  args = parser.parse_args(
+      ["--cluster", "test-cluster", "--tpu-type", "tpu7x-2", "--sub-slicing"]
+  )
+
+  assert args.sub_slicing is True
+
+
+def test_cluster_create_ray_sub_slicing_is_hidden_but_set_to_false():
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_ray_parser(parser)
+  args = parser.parse_args([
+      "--cluster",
+      "test-cluster",
+      "--tpu-type",
+      "tpu7x-2",
+      "--ray-version",
+      "19.32.0",
+  ])
+  help_str = parser.format_help()
+
+  assert args.sub_slicing is False
+  assert "--sub-slicing" not in help_str