xpk 0.16.0__py3-none-any.whl → 0.17.0__py3-none-any.whl

xpk/core/system_characteristics_test.py

@@ -32,7 +32,6 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_1x1_topol
       gke_accelerator="test",
       machine_type="test",
       supported_topologies=["1x1"],
-      supports_sub_slicing=False,
       docker_platform=DockerPlatform.AMD,
       tpu_type_requires_workload_policy=False,
   )
@@ -46,6 +45,7 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_1x1_topol
       accelerator_type=AcceleratorType.TPU,
       device_type="test-1",
       supports_sub_slicing=False,
+      supports_super_slicing=False,
       docker_platform=DockerPlatform.AMD,
       requires_workload_policy=False,
   )
@@ -62,7 +62,6 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2_topol
       gke_accelerator="test",
       machine_type="test",
       supported_topologies=["2x2"],
-      supports_sub_slicing=False,
       docker_platform=DockerPlatform.AMD,
       tpu_type_requires_workload_policy=True,
   )
@@ -76,6 +75,7 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2_topol
       accelerator_type=AcceleratorType.TPU,
       device_type="test-8",
       supports_sub_slicing=False,
+      supports_super_slicing=False,
       docker_platform=DockerPlatform.AMD,
       requires_workload_policy=False,
   )
@@ -92,7 +92,6 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2x2_top
       gke_accelerator="test",
       machine_type="test",
       supported_topologies=["2x2x2"],
-      supports_sub_slicing=False,
       docker_platform=DockerPlatform.AMD,
       tpu_type_requires_workload_policy=True,
   )
@@ -106,6 +105,7 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2x2_top
       accelerator_type=AcceleratorType.TPU,
       device_type="test-16",
       supports_sub_slicing=False,
+      supports_super_slicing=False,
       docker_platform=DockerPlatform.AMD,
       requires_workload_policy=True,
   )
@@ -115,6 +115,38 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2x2_top
   }
 
 
+def test_get_tpu_system_characteristics_map_sets_sub_slicing_support():
+  result = get_tpu_system_characteristics_map(
+      prefix="test",
+      tensorcores_per_chip=2,
+      gke_accelerator="test",
+      machine_type="test",
+      supported_topologies=["4x4x4", "4x4x8", "4x4x16"],
+      docker_platform=DockerPlatform.AMD,
+      sub_slicing_topologies=set(["4x4x8", "4x4x16"]),
+  )
+
+  assert result["test-4x4x4"].supports_sub_slicing is False
+  assert result["test-4x4x8"].supports_sub_slicing is True
+  assert result["test-4x4x16"].supports_sub_slicing is True
+
+
+def test_get_tpu_system_characteristics_map_sets_super_slicing_support():
+  result = get_tpu_system_characteristics_map(
+      prefix="test",
+      tensorcores_per_chip=2,
+      gke_accelerator="test",
+      machine_type="test",
+      supported_topologies=["4x4x4", "4x4x8", "4x4x16"],
+      docker_platform=DockerPlatform.AMD,
+      super_slicing_topologies=set(["4x4x8", "4x4x16"]),
+  )
+
+  assert result["test-4x4x4"].supports_super_slicing is False
+  assert result["test-4x4x8"].supports_super_slicing is True
+  assert result["test-4x4x16"].supports_super_slicing is True
+
+
 def test_get_tpu_system_characteristics_map_prefers_default_topologies():
   result = get_tpu_system_characteristics_map(
       prefix="test",
@@ -122,7 +154,6 @@ def test_get_tpu_system_characteristics_map_prefers_default_topologies():
       gke_accelerator="test",
       machine_type="test",
       supported_topologies=["4x4x4", "4x4x32", "4x8x16", "8x8x8"],
-      supports_sub_slicing=False,
       docker_platform=DockerPlatform.AMD,
       default_topologies=set(["4x8x16"]),
   )
@@ -174,6 +205,7 @@ def test_system_characteristics_post_init_sets_workload_policy_for_gpu():
       accelerator_type=AcceleratorType.GPU,
       device_type="l4-1",
       supports_sub_slicing=False,
+      supports_super_slicing=False,
       docker_platform=DockerPlatform.AMD,
       gpu_config=GpuConfig(requires_topology=False),
   )
@@ -192,5 +224,6 @@ def test_system_characteristics_post_init_throws_for_gpu_without_config():
         accelerator_type=AcceleratorType.GPU,
         device_type="l4-1",
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         docker_platform=DockerPlatform.AMD,
     )

xpk/core/telemetry.py

@@ -124,6 +124,7 @@ class MetricsEventMetadataKey(Enum):
   EXIT_CODE = "XPK_EXIT_CODE"
   RUNNING_AS_PIP = "XPK_RUNNING_AS_PIP"
   RUNNING_FROM_SOURCE = "XPK_RUNNING_FROM_SOURCE"
+  LATENCY_SECONDS = "XPK_LATENCY_SECONDS"
 
 
 @dataclass
@@ -190,10 +191,14 @@ def _generate_payload(events: list[_MetricsEvent]) -> str:
   base_concord_event = _get_base_concord_event()
   base_event_metadata = _get_base_event_metadata()
   serialized_events = []
+  first_time = events[0].time if len(events) > 0 else 0
   for event in events:
     metadata = {
         **base_event_metadata,
         **event.metadata,
+        MetricsEventMetadataKey.LATENCY_SECONDS: str(
+            int(event.time - first_time)
+        ),
     }
     serialized_events.append({
         "event_time_ms": int(event.time * 1000),

xpk/core/telemetry_test.py

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import itertools
 import pytest
 import json
 from .config import get_config, CLIENT_ID_KEY, SEND_TELEMETRY_KEY
@@ -26,7 +27,7 @@ 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('time.time', side_effect=itertools.count())
   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')
@@ -76,6 +77,7 @@ def test_metrics_collector_logs_start_event_correctly():
           {'key': 'XPK_RUNNING_AS_PIP', 'value': 'false'},
           {'key': 'XPK_RUNNING_FROM_SOURCE', 'value': 'true'},
           {'key': 'XPK_COMMAND', 'value': 'test'},
+          {'key': 'XPK_LATENCY_SECONDS', 'value': '0'},
       ],
       'event_name': 'start',
       'event_type': 'commands',
@@ -106,6 +108,7 @@ def test_metrics_collector_logs_complete_event_correctly():
           {'key': 'XPK_RUNNING_AS_PIP', 'value': 'false'},
           {'key': 'XPK_RUNNING_FROM_SOURCE', 'value': 'true'},
           {'key': 'XPK_EXIT_CODE', 'value': '2'},
+          {'key': 'XPK_LATENCY_SECONDS', 'value': '0'},
       ],
       'event_name': 'complete',
       'event_type': 'commands',
@@ -129,6 +132,7 @@ def test_metrics_collector_logs_custom_event_correctly():
           {'key': 'XPK_RUNNING_AS_PIP', 'value': 'false'},
           {'key': 'XPK_RUNNING_FROM_SOURCE', 'value': 'true'},
           {'key': 'XPK_PROVISIONING_MODE', 'value': 'flex'},
+          {'key': 'XPK_LATENCY_SECONDS', 'value': '0'},
       ],
       'event_name': 'test',
       'event_type': 'custom',
@@ -136,6 +140,19 @@ def test_metrics_collector_logs_custom_event_correctly():
   }
 
 
+def test_metrics_collector_computest_latency_correctly():
+  MetricsCollector.log_start(command='test')
+  MetricsCollector.log_complete(exit_code=0)
+  payload = json.loads(MetricsCollector.flush())
+  extension_json = json.loads(payload['log_event'][1]['source_extension_json'])
+  latency = (
+      el['value']
+      for el in extension_json['event_metadata']
+      if el['key'] == 'XPK_LATENCY_SECONDS'
+  )
+  assert next(latency, None) == '1'
+
+
 def test_metrics_collector_logs_correct_envelope():
   MetricsCollector.log_start(command='test')
   MetricsCollector.log_custom(
@@ -145,7 +162,7 @@ def test_metrics_collector_logs_correct_envelope():
   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 payload['request_time_ms'] == 3000
   assert len(payload['log_event']) == 3
 
 

xpk/core/updates.py

@@ -28,7 +28,7 @@ def get_latest_xpk_version() -> tuple[int, Version | None]:
     return 0, Version(__version__)
 
   return_code, result = run_command_for_value(
-      command="pip index versions xpk --json",
+      command="pip index versions xpk --json --no-input",
       task="Retrieve latest XPK version",
       quiet=True,
   )

xpk/main.py

@@ -36,6 +36,7 @@ import argcomplete
 import sys
 
 from .parser.core import set_parser
+from .parser.common import extract_command_path
 from .core.updates import print_xpk_hello
 from .core.config import set_config, FileSystemConfig
 from .core.telemetry import MetricsCollector, send_clearcut_payload, should_send_telemetry
@@ -78,7 +79,7 @@ def main() -> None:
             or ('force' in main_args and main_args.force)
         ),
     )
-    MetricsCollector.log_start(main_args.xpk_subcommands)
+    MetricsCollector.log_start(command=extract_command_path(parser, main_args))
     print_xpk_hello()
     main_args.func(main_args)
     xpk_print('XPK Done.', flush=True)

xpk/parser/cluster.py

@@ -132,6 +132,10 @@ def set_cluster_create_parser(cluster_create_parser: ArgumentParser):
 
   if FeatureFlags.SUB_SLICING_ENABLED:
     add_cluster_create_sub_slicing_arguments(cluster_create_optional_arguments)
+  if FeatureFlags.SUPER_SLICING_ENABLED:
+    add_cluster_create_super_slicing_arguments(
+        cluster_create_optional_arguments
+    )
 
   autoprovisioning_arguments = cluster_create_parser.add_argument_group(
       'Autoprovisioning Arguments',
@@ -205,6 +209,10 @@ def set_cluster_create_pathways_parser(
     add_cluster_create_sub_slicing_arguments(
         cluster_create_pathways_optional_arguments
     )
+  if FeatureFlags.SUPER_SLICING_ENABLED:
+    add_cluster_create_super_slicing_arguments(
+        cluster_create_pathways_optional_arguments
+    )
 
   autoprovisioning_arguments = (
       cluster_create_pathways_parser.add_argument_group(
@@ -330,7 +338,7 @@ 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, sub_slicing=False
+      func=cluster_create_ray_cluster, sub_slicing=False, super_slicing=False
   )
 
 
@@ -596,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,
   )
@@ -910,3 +921,24 @@ def add_cluster_create_sub_slicing_arguments(
       action='store_true',
       help='Whether to set up cluster to support sub-slicing',
   )
+
+
+def add_cluster_create_super_slicing_arguments(
+    parser_or_group: ParserOrArgumentGroup,
+):
+  parser_or_group.add_argument(
+      '--super-slicing',
+      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,
+  )

xpk/parser/cluster_test.py

@@ -144,3 +144,120 @@ def test_cluster_create_enable_lustre_legacy_port_can_be_set():
   ])
 
   assert args.enable_legacy_lustre_port is True
+
+
+def test_cluster_create_super_slicing_is_hidden_with_flag_off():
+  FeatureFlags.SUPER_SLICING_ENABLED = False
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  help_str = parser.format_help()
+
+  assert "--super-slicing" not in help_str
+
+
+def test_cluster_create_super_slicing_is_shown_with_flag_on():
+  FeatureFlags.SUPER_SLICING_ENABLED = True
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  help_str = parser.format_help()
+
+  assert "--super-slicing" in help_str
+
+
+def test_cluster_create_super_slicing_is_false_by_default():
+  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.super_slicing is False
+
+
+def test_cluster_create_super_slicing_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", "--super-slicing"],
+  )
+
+  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

xpk/parser/common.py

@@ -369,3 +369,35 @@ def add_tpu_and_device_type_arguments(
 ) -> None:
   add_tpu_type_argument(custom_parser_or_group)
   add_device_type_argument(custom_parser_or_group)
+
+
+def extract_command_path(parser: argparse.ArgumentParser, args):
+  """
+  Reconstructs the command path (e.g. 'cluster create').
+  """
+
+  def _get_path_segments(current_parser):
+    subparser_action = next(
+        (
+            action
+            for action in current_parser._actions  # pylint: disable=protected-access
+            if isinstance(action, argparse._SubParsersAction)  # pylint: disable=protected-access
+        ),
+        None,
+    )
+
+    if subparser_action is None:
+      return []
+
+    chosen_command = getattr(args, subparser_action.dest, None)
+
+    if chosen_command is None:
+      return []
+
+    if chosen_command in subparser_action.choices:
+      next_parser = subparser_action.choices[chosen_command]
+      return [chosen_command] + _get_path_segments(next_parser)
+
+    return [chosen_command]
+
+  return ' '.join(_get_path_segments(parser))

xpk/parser/common_test.py

@@ -0,0 +1,49 @@
+"""
+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 .common import extract_command_path
+from .core import set_parser
+
+
+def test_extract_zero_level_nested_command():
+  parser = argparse.ArgumentParser()
+  set_parser(parser=parser)
+  args = parser.parse_args([])
+  assert extract_command_path(parser, args) == ""
+
+
+def test_extract_one_level_nested_command():
+  parser = argparse.ArgumentParser()
+  set_parser(parser=parser)
+  args = parser.parse_args(["version"])
+  assert extract_command_path(parser, args) == "version"
+
+
+def test_extract_two_level_nested_command():
+  parser = argparse.ArgumentParser()
+  set_parser(parser=parser)
+  args = parser.parse_args(["cluster", "list"])
+  assert extract_command_path(parser, args) == "cluster list"
+
+
+def test_extract_two_level_nested_command_with_flags():
+  parser = argparse.ArgumentParser()
+  set_parser(parser=parser)
+  args = parser.parse_args(
+      ["cluster", "list", "--project=abc", "--zone=us-central1-a"]
+  )
+  assert extract_command_path(parser, args) == "cluster list"

xpk/templates/kueue_config.yaml.j2

@@ -1,15 +1,16 @@
-{% for flavor in flavors %}
+{%- for flavor in flavors %}
 apiVersion: kueue.x-k8s.io/v1beta1
 kind: ResourceFlavor
 metadata:
   name: "{{ flavor.name }}"
 spec:
   nodeLabels: {{ flavor.nodeLabels | tojson }}
-  {% if flavor.topologyLabel %}
+  {%- if flavor.topologyLabel %}
   {{ flavor.topologyLabel }}
-  {% endif %}
+  {%- endif %}
 ---
-{% endfor %}
+{%- endfor %}
+{%- if 'dws-prov' in admission_checks %}
 apiVersion: kueue.x-k8s.io/v1beta1
 kind: AdmissionCheck
 metadata:
@@ -21,6 +22,16 @@ spec:
     kind: ProvisioningRequestConfig
     name: dws-config
 ---
+{%- endif %}
+{%- if 'ss-kueue-operator' in admission_checks %}
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: AdmissionCheck
+metadata:
+  name: ss-kueue-operator
+spec:
+  controllerName: accelerator.gke.io/slice
+---
+{%- endif %}
 apiVersion: kueue.x-k8s.io/v1beta1
 kind: ProvisioningRequestConfig
 metadata:
@@ -44,7 +55,12 @@ spec:
     withinClusterQueue: LowerPriority
   namespaceSelector: {} # match all.
   resourceGroups: {{ resource_groups }}
-  {{ admission_checks | indent(2) }}
+  {%- if admission_checks %}
+  admissionChecks:
+  {%- for check in admission_checks %}
+  - {{ check }}
+  {%- endfor %}
+  {%- endif %}
 ---
 apiVersion: kueue.x-k8s.io/v1beta1
 kind: LocalQueue

xpk/templates/kueue_super_slicing_topology.yaml.j2

@@ -0,0 +1,9 @@
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: Topology
+metadata:
+  name: {{ super_slice_topology_name }}
+spec:
+  levels:
+  - nodeLabel: cloud.google.com/gce-topology-block
+  - nodeLabel: cloud.google.com/gke-tpu-partition-4x4x4-id
+  - nodeLabel: kubernetes.io/hostname

xpk/utils/kueue.py

@@ -14,7 +14,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from ..core.system_characteristics import AcceleratorType
 
-def is_queued_cluster(num_slices: int) -> bool:
+
+def is_queued_cluster(
+    num_slices: int, accelerator_type: AcceleratorType
+) -> bool:
   """Determines if admission checks should be enabled and cluster queued."""
-  return num_slices <= 1
+  return num_slices <= 1 and accelerator_type == AcceleratorType.GPU

{xpk-0.16.0.dist-info → xpk-0.17.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xpk
-Version: 0.16.0
+Version: 0.17.0
 Summary: xpk helps Cloud developers to orchestrate training jobs on accelerators on GKE.
 Author-email: XPK team <xpk-code-reviewers@google.com>
 License: Apache-2.0
@@ -25,6 +25,7 @@ Requires-Dist: packaging==24.2
 Requires-Dist: google-cloud-filestore==1.12.0
 Requires-Dist: google-cloud-storage
 Requires-Dist: Jinja2==3.1.6
+Requires-Dist: urllib3<2.6.0
 Provides-Extra: dev
 Requires-Dist: pyink==24.3.0; extra == "dev"
 Requires-Dist: pylint>=2.6.0; extra == "dev"