diff --git a/pylintrc b/pylintrc index 393b508d4..61262ff37 100644 --- a/pylintrc +++ b/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 diff --git a/src/xpk/parser/cluster.py b/src/xpk/parser/cluster.py index 18dd2d04f..97f6b0408 100644 --- a/src/xpk/parser/cluster.py +++ b/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', diff --git a/src/xpk/parser/cluster_test.py b/src/xpk/parser/cluster_test.py new file mode 100644 index 000000000..86ea15f48 --- /dev/null +++ b/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 diff --git a/src/xpk/utils/feature_flags.py b/src/xpk/utils/feature_flags.py new file mode 100644 index 000000000..089394571 --- /dev/null +++ b/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)