xpk 0.14.4__py3-none-any.whl → 0.15.0__py3-none-any.whl

xpk/core/telemetry_test.py

@@ -0,0 +1,211 @@
+"""
+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 pytest
+import json
+from .config import xpk_config, CLIENT_ID_KEY, SEND_TELEMETRY_KEY
+from .telemetry import MetricsCollector, MetricsEventMetadataKey, should_send_telemetry
+from ..utils.execution_context import set_dry_run
+from ..utils.feature_flags import FeatureFlags
+from pytest_mock import MockerFixture
+
+
+@pytest.fixture(autouse=True)
+def setup_mocks(mocker: MockerFixture):
+  mocker.patch('xpk.core.telemetry._get_session_id', return_value='321231')
+  mocker.patch('time.time', return_value=0)
+  mocker.patch('platform.python_version', return_value='99.99.99')
+  mocker.patch('os.path.basename', return_value='xpk.py')
+  mocker.patch('os.path.abspath', return_value='/home/xpk_user')
+  set_dry_run(False)
+  xpk_config.set(CLIENT_ID_KEY, 'client_id')
+  yield
+  xpk_config.set(CLIENT_ID_KEY, None)
+
+
+@pytest.mark.parametrize(
+    argnames='feature_flag,config_value,expected',
+    argvalues=[
+        (True, 'true', True),
+        (False, 'true', False),
+        (True, None, True),
+        (True, 'false', False),
+    ],
+)
+def test_should_send_telemetry_returns_correct_value(
+    feature_flag: bool, config_value: str, expected: bool
+):
+  xpk_config.set(SEND_TELEMETRY_KEY, config_value)
+  FeatureFlags.TELEMETRY_ENABLED = feature_flag
+  assert should_send_telemetry() is expected
+
+
+def test_metrics_collector_generates_client_id_if_not_present():
+  xpk_config.set(CLIENT_ID_KEY, None)
+  MetricsCollector.log_start(command='test')
+  payload = json.loads(MetricsCollector.flush())
+  extension_json = json.loads(payload['log_event'][0]['source_extension_json'])
+  assert extension_json['client_install_id'] is not None
+  assert len(extension_json['client_install_id']) > 0
+
+
+def test_metrics_collector_logs_start_event_correctly():
+  MetricsCollector.log_start(command='test')
+  payload = json.loads(MetricsCollector.flush())
+  extension_json = json.loads(payload['log_event'][0]['source_extension_json'])
+  assert extension_json == {
+      'client_install_id': 'client_id',
+      'console_type': 'XPK',
+      'event_metadata': [
+          {'key': 'XPK_SESSION_ID', 'value': '321231'},
+          {'key': 'XPK_DRY_RUN', 'value': 'false'},
+          {'key': 'XPK_PYTHON_VERSION', 'value': '99.99.99'},
+          {'key': 'XPK_RUNNING_AS_PIP', 'value': 'false'},
+          {'key': 'XPK_RUNNING_FROM_SOURCE', 'value': 'true'},
+          {'key': 'XPK_COMMAND', 'value': 'test'},
+      ],
+      'event_name': 'start',
+      'event_type': 'commands',
+      'release_version': 'v0.15.0',
+  }
+
+
+def test_metrics_collector_generates_client_id_when_not_present():
+  xpk_config.set(CLIENT_ID_KEY, None)
+  MetricsCollector.log_start(command='test')
+  payload = json.loads(MetricsCollector.flush())
+  extension_json = json.loads(payload['log_event'][0]['source_extension_json'])
+  assert extension_json['client_install_id'] is not None
+  assert len(extension_json['client_install_id']) > 0
+
+
+def test_metrics_collector_logs_complete_event_correctly():
+  MetricsCollector.log_complete(exit_code=2)
+  payload = json.loads(MetricsCollector.flush())
+  extension_json = json.loads(payload['log_event'][0]['source_extension_json'])
+  assert extension_json == {
+      'client_install_id': 'client_id',
+      'console_type': 'XPK',
+      'event_metadata': [
+          {'key': 'XPK_SESSION_ID', 'value': '321231'},
+          {'key': 'XPK_DRY_RUN', 'value': 'false'},
+          {'key': 'XPK_PYTHON_VERSION', 'value': '99.99.99'},
+          {'key': 'XPK_RUNNING_AS_PIP', 'value': 'false'},
+          {'key': 'XPK_RUNNING_FROM_SOURCE', 'value': 'true'},
+          {'key': 'XPK_EXIT_CODE', 'value': '2'},
+      ],
+      'event_name': 'complete',
+      'event_type': 'commands',
+      'release_version': 'v0.15.0',
+  }
+
+
+def test_metrics_collector_logs_custom_event_correctly():
+  MetricsCollector.log_custom(
+      name='test', metadata={MetricsEventMetadataKey.PROVISIONING_MODE: 'flex'}
+  )
+  payload = json.loads(MetricsCollector.flush())
+  extension_json = json.loads(payload['log_event'][0]['source_extension_json'])
+  assert extension_json == {
+      'client_install_id': 'client_id',
+      'console_type': 'XPK',
+      'event_metadata': [
+          {'key': 'XPK_SESSION_ID', 'value': '321231'},
+          {'key': 'XPK_DRY_RUN', 'value': 'false'},
+          {'key': 'XPK_PYTHON_VERSION', 'value': '99.99.99'},
+          {'key': 'XPK_RUNNING_AS_PIP', 'value': 'false'},
+          {'key': 'XPK_RUNNING_FROM_SOURCE', 'value': 'true'},
+          {'key': 'XPK_PROVISIONING_MODE', 'value': 'flex'},
+      ],
+      'event_name': 'test',
+      'event_type': 'custom',
+      'release_version': 'v0.15.0',
+  }
+
+
+def test_metrics_collector_logs_correct_envelope():
+  MetricsCollector.log_start(command='test')
+  MetricsCollector.log_custom(
+      name='test', metadata={MetricsEventMetadataKey.PROVISIONING_MODE: 'flex'}
+  )
+  MetricsCollector.log_complete(exit_code=2)
+  payload = json.loads(MetricsCollector.flush())
+  assert payload['client_info'] == {'client_type': 'XPK'}
+  assert payload['log_source_name'] == 'CONCORD'
+  assert payload['request_time_ms'] == 0
+  assert len(payload['log_event']) == 3
+
+
+def test_metrics_collector_does_not_flush_event_twice():
+  MetricsCollector.log_start(command='test')
+  MetricsCollector.flush()
+  MetricsCollector.log_start(command='version')
+  payload = json.loads(MetricsCollector.flush())
+  assert len(payload['log_event']) == 1
+
+
+@pytest.mark.parametrize(
+    argnames='dry_run,expected', argvalues=[(False, 'false'), (True, 'true')]
+)
+def test_metrics_collector_logs_correct_dry_run_value(
+    dry_run: bool, expected: str
+):
+  set_dry_run(dry_run)
+  MetricsCollector.log_start(command='test')
+  payload = MetricsCollector.flush()
+  assert _get_metadata_value(payload, 'XPK_DRY_RUN') == expected
+
+
+@pytest.mark.parametrize(
+    argnames='basename,expected',
+    argvalues=[
+        ('xpk', 'true'),
+        ('xpk.py', 'false'),
+    ],
+)
+def test_metrics_collectors_logs_correct_running_as_pip_value(
+    basename: str, expected: str, mocker: MockerFixture
+):
+  mocker.patch('os.path.basename', return_value=basename)
+  MetricsCollector.log_start(command='test')
+  payload = MetricsCollector.flush()
+  assert _get_metadata_value(payload, 'XPK_RUNNING_AS_PIP') == expected
+
+
+@pytest.mark.parametrize(
+    argnames='abspath,expected',
+    argvalues=[
+        ('/site-packages/', 'false'),
+        ('/dist-packages/', 'false'),
+        ('/home/xpk_user', 'true'),
+    ],
+)
+def test_metrics_collectors_logs_correct_running_from_source_value(
+    abspath: str, expected: str, mocker: MockerFixture
+):
+  mocker.patch('os.path.abspath', return_value=abspath)
+  MetricsCollector.log_start(command='test')
+  payload = MetricsCollector.flush()
+  assert _get_metadata_value(payload, 'XPK_RUNNING_FROM_SOURCE') == expected
+
+
+def _get_metadata_value(payload_str: str, key: str) -> str | None:
+  payload = json.loads(payload_str)
+  metadata = json.loads(payload['log_event'][0]['source_extension_json'])[
+      'event_metadata'
+  ]
+  matching = (item['value'] for item in metadata if item['key'] == key)
+  return next(matching, None)

xpk/main.py

@@ -32,11 +32,13 @@ Next Steps:
 """
 
 import argparse
+import argcomplete
 import sys
 
 from .parser.core import set_parser
 from .core.updates import print_xpk_hello
-from .utils.console import xpk_print
+from .core.telemetry import MetricsCollector, send_clearcut_payload, should_send_telemetry
+from .utils.console import xpk_print, exit_code_to_int
 from .utils.execution_context import set_context
 ################### Compatibility Check ###################
 # Check that the user runs the below version or greater.
@@ -59,19 +61,35 @@ if (
 
 
 def main() -> None:
-  # Create top level parser for xpk command.
-  parser = argparse.ArgumentParser(description='xpk command', prog='xpk')
-  set_parser(parser=parser)
+  try:
+    # Create top level parser for xpk command.
+    parser = argparse.ArgumentParser(description='xpk command', prog='xpk')
+    set_parser(parser=parser)
+    argcomplete.autocomplete(parser)
 
-  main_args = parser.parse_args()
-  main_args.enable_ray_cluster = False
-  set_context(
-      dry_run_value='dry_run' in main_args and main_args.dry_run,
-      quiet_value='quiet' in main_args and main_args.quiet,
-  )
-  print_xpk_hello()
-  main_args.func(main_args)
-  xpk_print('XPK Done.', flush=True)
+    main_args = parser.parse_args()
+    main_args.enable_ray_cluster = False
+    set_context(
+        dry_run_value='dry_run' in main_args and main_args.dry_run,
+        quiet_value=(
+            ('quiet' in main_args and main_args.quiet)
+            or ('force' in main_args and main_args.force)
+        ),
+    )
+    MetricsCollector.log_start(main_args.xpk_subcommands)
+    print_xpk_hello()
+    main_args.func(main_args)
+    xpk_print('XPK Done.', flush=True)
+    MetricsCollector.log_complete(0)
+  except SystemExit as e:
+    MetricsCollector.log_complete(exit_code_to_int(e.code))
+    raise
+  except:
+    MetricsCollector.log_complete(-1)
+    raise
+  finally:
+    if should_send_telemetry():
+      send_clearcut_payload(MetricsCollector.flush())
 
 
 if __name__ == '__main__':

xpk/parser/cluster.py

@@ -26,7 +26,8 @@ from ..commands.cluster import (
     cluster_describe,
     cluster_list,
 )
-from ..commands.config import xpk_cfg
+from ..core.config import xpk_config
+from ..core.system_characteristics import get_system_characteristics_keys_by_accelerator_type, AcceleratorType
 from ..core.config import CFG_BUCKET_KEY
 from ..core.vertex import DEFAULT_VERTEX_TENSORBOARD_NAME
 from .common import add_shared_arguments, ParserOrArgumentGroup
@@ -103,6 +104,10 @@ def set_cluster_create_parser(cluster_create_parser: ArgumentParser):
       type=str,
       default=None,
       help='The tpu type to use, v5litepod-16, etc.',
+      metavar='TPU_TYPE',
+      choices=get_system_characteristics_keys_by_accelerator_type(
+          [AcceleratorType.TPU]
+      ),
   )
   cluster_device_group.add_argument(
       '--device-type',
@@ -112,6 +117,8 @@ def set_cluster_create_parser(cluster_create_parser: ArgumentParser):
           'The device type to use (can be tpu or gpu or cpu), v5litepod-16,'
           ' h100-80gb-8, n2-standard-32-4 etc.'
       ),
+      metavar='DEVICE_TYPE',
+      choices=get_system_characteristics_keys_by_accelerator_type(),
   )
 
   ### Optional arguments specific to "cluster create"
@@ -124,7 +131,7 @@ def set_cluster_create_parser(cluster_create_parser: ArgumentParser):
   cluster_create_optional_arguments.add_argument(
       '--cluster-state-gcs-bucket',
       type=str,
-      default=xpk_cfg.get(CFG_BUCKET_KEY),
+      default=xpk_config.get(CFG_BUCKET_KEY),
       help='The name of the bucket to store cluster state.',
       required=False,
   )
@@ -144,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',
@@ -209,6 +212,10 @@ def set_cluster_create_pathways_parser(
       type=str,
       default=None,
       help='The tpu type to use, v5litepod-16, etc.',
+      metavar='TPU_TYPE',
+      choices=get_system_characteristics_keys_by_accelerator_type(
+          [AcceleratorType.TPU]
+      ),
   )
 
   ### Optional arguments specific to "cluster create-pathways"
@@ -221,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(
@@ -287,6 +298,10 @@ def set_cluster_create_ray_parser(cluster_create_ray_parser: ArgumentParser):
       default=None,
       help='The tpu type to use, v5litepod-16, etc.',
       required=True,
+      metavar='TPU_TYPE',
+      choices=get_system_characteristics_keys_by_accelerator_type(
+          [AcceleratorType.TPU]
+      ),
   )
   # TODO(bzmarke): Add --device-type to support GPU/CPU
   cluster_create_ray_required_arguments.add_argument(
@@ -350,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):
@@ -375,7 +392,7 @@ def set_cluster_delete_parser(cluster_delete_parser: ArgumentParser):
   cluster_delete_optional_arguments.add_argument(
       '--cluster-state-gcs-bucket',
       type=str,
-      default=xpk_cfg.get(CFG_BUCKET_KEY),
+      default=xpk_config.get(CFG_BUCKET_KEY),
       help='The name of the bucket to store cluster state.',
       required=False,
   )
@@ -409,6 +426,10 @@ def set_cluster_cacheimage_parser(cluster_cacheimage_parser: ArgumentParser):
       type=str,
       default=None,
       help='The tpu type to cache images on, v5litepod-16, etc.',
+      metavar='TPU_TYPE',
+      choices=get_system_characteristics_keys_by_accelerator_type(
+          [AcceleratorType.TPU]
+      ),
   )
   cluster_cacheimage_group.add_argument(
       '--device-type',
@@ -418,6 +439,8 @@ def set_cluster_cacheimage_parser(cluster_cacheimage_parser: ArgumentParser):
           'The device type to cache images on (can be tpu or gpu),'
           ' v5litepod-16, h100-80gb-8, etc.'
       ),
+      metavar='DEVICE_TYPE',
+      choices=get_system_characteristics_keys_by_accelerator_type(),
   )
 
   ### Required arguments
@@ -508,6 +531,10 @@ def set_cluster_adapt_parser(cluster_adapt_parser: ArgumentParser):
       type=str,
       default=None,
       help='The tpu type used on cluster, v5litepod-16, etc.',
+      metavar='TPU_TYPE',
+      choices=get_system_characteristics_keys_by_accelerator_type(
+          [AcceleratorType.TPU]
+      ),
   )
   cluster_adapt_device_group.add_argument(
       '--device-type',
@@ -517,6 +544,8 @@ def set_cluster_adapt_parser(cluster_adapt_parser: ArgumentParser):
           'The device type used on cluster (can be tpu or gpu or cpu), eg.'
           ' h100-80gb-8, n2-standard-32-4 etc.'
       ),
+      metavar='DEVICE_TYPE',
+      choices=get_system_characteristics_keys_by_accelerator_type(),
   )
 
   cluster_adapt_optional_arguments = cluster_adapt_parser.add_argument_group(
@@ -937,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',
+  )

xpk/parser/cluster_test.py

@@ -15,7 +15,7 @@ limitations under the License.
 """
 
 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
 
@@ -49,7 +49,7 @@ def test_cluster_create_sub_slicing_is_false_by_default():
 
   set_cluster_create_parser(parser)
   args = parser.parse_args(
-      ["--cluster", "test-cluster", "--tpu-type", "test-tpu"]
+      ["--cluster", "test-cluster", "--tpu-type", "tpu7x-2"]
   )
 
   assert args.sub_slicing is False
@@ -60,7 +60,46 @@ def test_cluster_create_sub_slicing_can_be_set():
 
   set_cluster_create_parser(parser)
   args = parser.parse_args(
-      ["--cluster", "test-cluster", "--tpu-type", "test-tpu", "--sub-slicing"]
+      ["--cluster", "test-cluster", "--tpu-type", "tpu7x-2", "--sub-slicing"]
   )
 
   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

xpk/parser/workload.py

@@ -25,6 +25,7 @@ from ..core.docker_image import DEFAULT_DOCKER_IMAGE, DEFAULT_SCRIPT_DIR
 from .common import add_shared_arguments
 from .validators import directory_path_type, name_type
 from ..utils.feature_flags import FeatureFlags
+from ..core.system_characteristics import get_system_characteristics_keys_by_accelerator_type, AcceleratorType, SUB_SLICING_TOPOLOGIES
 
 
 def set_workload_parsers(workload_parser: ArgumentParser):
@@ -123,6 +124,10 @@ def set_workload_create_parser(workload_create_parser: ArgumentParser):
       type=str,
       default=None,
       help='The tpu type to use, v5litepod-16, etc.',
+      metavar='TPU_TYPE',
+      choices=get_system_characteristics_keys_by_accelerator_type(
+          [AcceleratorType.TPU]
+      ),
   )
   workload_device_group.add_argument(
       '--device-type',
@@ -132,6 +137,8 @@ def set_workload_create_parser(workload_create_parser: ArgumentParser):
           'The device type to use (can be tpu or gpu or cpu), v5litepod-16,'
           ' h100-80gb-8, n2-standard-32-4 etc.'
       ),
+      metavar='DEVICE_TYPE',
+      choices=get_system_characteristics_keys_by_accelerator_type(),
   )
 
   workload_create_parser_optional_arguments.add_argument(
@@ -285,6 +292,10 @@ def set_workload_create_pathways_parser(
       type=str,
       default=None,
       help='The tpu type to use, v5litepod-16, etc.',
+      metavar='TPU_TYPE',
+      choices=get_system_characteristics_keys_by_accelerator_type(
+          [AcceleratorType.TPU]
+      ),
   )
 
   ### "workload create-pathways" Optional arguments, specific to Pathways
@@ -665,6 +676,7 @@ def add_shared_workload_create_optional_arguments(args_parsers):
           type=str,
           help='Sub-slicing topology to use.',
           required=False,
+          choices=SUB_SLICING_TOPOLOGIES,
       )
 
 

xpk/parser/workload_test.py

@@ -56,7 +56,7 @@ def test_workload_create_sub_slicing_topology_is_none_by_default():
       "--workload",
       "test",
       "--tpu-type",
-      "test-tpu",
+      "tpu7x-2",
   ])
 
   assert args.sub_slicing_topology is None
@@ -74,9 +74,9 @@ def test_workload_create_sub_slicing_topology_can_be_set():
       "--workload",
       "test",
       "--tpu-type",
-      "test-tpu",
+      "tpu7x-8",
       "--sub-slicing-topology",
-      "2x2",
+      "2x4",
   ])
 
-  assert args.sub_slicing_topology is "2x2"
+  assert args.sub_slicing_topology is "2x4"

xpk/telemetry_uploader.py

@@ -0,0 +1,29 @@
+"""
+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 sys
+import os
+import requests
+import json
+
+file_path = sys.argv[1]
+if os.path.exists(file_path):
+  with open(file_path, mode="r", encoding="utf-8") as file:
+    kwargs = json.load(file)
+    response = requests.request(**kwargs)
+    print(f"Telemetry upload finished with {response.status_code} status code")
+
+  os.remove(file_path)

xpk/templates/kueue_gke_default_topology.yaml.j2

@@ -1,4 +1,4 @@
-apiVersion: kueue.x-k8s.io/v1alpha1
+apiVersion: kueue.x-k8s.io/v1beta1
 kind: Topology
 metadata:
   name: "gke-default"

xpk/templates/kueue_sub_slicing_topology.yaml.j2

@@ -4,11 +4,6 @@ metadata:
   name: {{ sub_slice_topology_name }}
 spec:
   levels:
-  - nodeLabel: "cloud.google.com/gke-tpu-slice-16x16-id"
-  - nodeLabel: "cloud.google.com/gke-tpu-slice-8x16-id"
-  - nodeLabel: "cloud.google.com/gke-tpu-slice-8x8-id"
-  - nodeLabel: "cloud.google.com/gke-tpu-slice-4x8-id"
-  - nodeLabel: "cloud.google.com/gke-tpu-slice-4x4-id"
-  - nodeLabel: "cloud.google.com/gke-tpu-slice-2x4-id"
-  - nodeLabel: "cloud.google.com/gke-tpu-slice-2x2-id"
-  - nodeLabel: "kubernetes.io/hostname"
+  {% for level in levels %}
+  - nodeLabel: "{{level}}"
+  {% endfor %}