feat: Add --sub-slicing flag to cluster create

pylintrc

@@ -32,7 +32,8 @@ unsafe-load-any-extension=no
 [MESSAGES CONTROL]
 
 per-file-ignores =
-  .*_test\.py:protected-access # ignore "protected-access" errors in test files
+  .*_test\.py:protected-access,redefined-outer-name,unused-argument # ignore errors in test files. outer-name and unused-arguments are common patterns when using fixtures.
+
 
 # Only show warnings with the listed confidence levels. Leave empty to show
 # all. Valid levels: HIGH, INFERENCE, INFERENCE_FAILURE, UNDEFINED

src/xpk/parser/cluster.py

@@ -31,6 +31,7 @@
 from ..core.vertex import DEFAULT_VERTEX_TENSORBOARD_NAME
 from .common import add_shared_arguments, ParserOrArgumentGroup
 from .validators import name_type
+from ..utils.feature_flags import SUB_SLICING_ENABLED
 
 
 def set_cluster_parser(cluster_parser: ArgumentParser):
@@ -142,6 +143,12 @@ def set_cluster_create_parser(cluster_create_parser: ArgumentParser):
           ' enable cluster to accept Pathways workloads.'
       ),
   )
+  if SUB_SLICING_ENABLED:
+    cluster_create_optional_arguments.add_argument(
+        '--sub-slicing',
+        action='store_true',
+        help='Whether to set up cluster to support sub-slicing',
+    )
 
   autoprovisioning_arguments = cluster_create_parser.add_argument_group(
       'Autoprovisioning Arguments',

src/xpk/parser/cluster_test.py

@@ -0,0 +1,65 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import argparse
+from xpk.parser.cluster import set_cluster_create_parser
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def with_sub_slicing_enabled(mocker):
+  mocker.patch("xpk.parser.cluster.SUB_SLICING_ENABLED", True)
+
+
+def test_cluster_create_sub_slicing_is_hidden_with_flag_off(mocker):
+  mocker.patch("xpk.parser.cluster.SUB_SLICING_ENABLED", False)
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  help_str = parser.format_help()
+
+  assert "--sub-slicing" not in help_str
+
+
+def test_cluster_create_sub_slicing_is_shown_with_flag_on():
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  help_str = parser.format_help()
+
+  assert "--sub-slicing" in help_str
+
+
+def test_cluster_create_sub_slicing_is_false_by_default():
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  args = parser.parse_args(
+      ["--cluster", "test-cluster", "--tpu-type", "test-tpu"]
+  )
+
+  assert args.sub_slicing is False
+
+
+def test_cluster_create_sub_slicing_can_be_set():
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  args = parser.parse_args(
+      ["--cluster", "test-cluster", "--tpu-type", "test-tpu", "--sub-slicing"]
+  )
+
+  assert args.sub_slicing is True

src/xpk/utils/feature_flags.py

@@ -0,0 +1,24 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import os
+
+
+def __getBooleanFlag(flag: str, default: bool) -> bool:
+  return os.getenv(flag, str(default)).lower() == "true"
+
+
+SUB_SLICING_ENABLED = __getBooleanFlag("SUB_SLICING_ENABLED", default=False)