xpk 0.16.1__py3-none-any.whl → 0.17.1__py3-none-any.whl

xpk/core/scheduling.py

@@ -16,7 +16,7 @@ limitations under the License.
 
 from enum import Enum
 
-from .kueue_manager import get_installed_kueue_version, has_sub_slicing_enabled
+from .kueue_manager import get_installed_kueue_version, has_sub_slicing_enabled, has_super_slicing_enabled
 from ..utils.feature_flags import FeatureFlags
 from ..utils.topology import get_slice_topology_level
 from ..utils.console import xpk_print
@@ -33,12 +33,14 @@ from .system_characteristics import (
 from packaging.version import Version
 
 _SUB_SLICING_MINIMUM_KUEUE_VERSION = Version('0.13.0')
+_SUPER_SLICING_MINIMUM_KUEUE_VERSION = Version('0.14.0')
 
 
 class WorkloadScheduling(Enum):
   UNAVAILABLE = 0
   AVAILABLE = 1
   SUB_SLICING_AVAILABLE = 2
+  SUPER_SLICING_AVAILABLE = 3
 
 
 def check_if_workload_can_schedule(
@@ -94,10 +96,9 @@ def check_if_workload_can_schedule(
     else:
       return WorkloadScheduling.UNAVAILABLE
 
-  if _check_sub_slicing_availability(
+  if cluster_system and _check_sub_slicing_availability(
       workload_system=workload_system, cluster_system=cluster_system
   ):
-    assert cluster_system
     if _check_workload_size_fits(
         args,
         workload_system,
@@ -107,6 +108,18 @@ def check_if_workload_can_schedule(
     else:
       return WorkloadScheduling.UNAVAILABLE
 
+  if cluster_system and _check_super_slicing_availability(
+      workload_system=workload_system, cluster_system=cluster_system
+  ):
+    if _check_workload_size_fits(
+        args,
+        workload_system,
+        max_vm_in_cluster=int(resources_config_map[cluster_system.device_type]),
+    ):
+      return WorkloadScheduling.SUPER_SLICING_AVAILABLE
+    else:
+      return WorkloadScheduling.UNAVAILABLE
+
   xpk_print(
       'Workload scheduling validation failed. XPK will not create the workload'
       f' {args.workload}.'
@@ -147,11 +160,10 @@ def _check_workload_size_fits(
 
 def _check_sub_slicing_availability(
     workload_system: SystemCharacteristics,
-    cluster_system: SystemCharacteristics | None,
+    cluster_system: SystemCharacteristics,
 ) -> bool:
   if (
       (not FeatureFlags.SUB_SLICING_ENABLED)
-      or (not cluster_system)
       or (workload_system.gke_accelerator != cluster_system.gke_accelerator)
       or (not cluster_system.supports_sub_slicing)
       or (workload_system.topology not in SUB_SLICING_TOPOLOGIES)
@@ -163,7 +175,7 @@ def _check_sub_slicing_availability(
     return False
 
   return_code, current_version = get_installed_kueue_version(
-      dry_run_version=Version('0.13')
+      dry_run_version=_SUB_SLICING_MINIMUM_KUEUE_VERSION
   )
 
   return (
@@ -173,6 +185,33 @@ def _check_sub_slicing_availability(
   )
 
 
+def _check_super_slicing_availability(
+    workload_system: SystemCharacteristics,
+    cluster_system: SystemCharacteristics,
+) -> bool:
+  # TODO: b/465447813 - Add super-slicing workload topology validation.
+  if (
+      (not FeatureFlags.SUPER_SLICING_ENABLED)
+      or (workload_system.gke_accelerator != cluster_system.gke_accelerator)
+      or (not cluster_system.supports_super_slicing)
+  ):
+    return False
+
+  return_code, sub_slicing_enabled = has_super_slicing_enabled()
+  if return_code != 0 or not sub_slicing_enabled:
+    return False
+
+  return_code, current_version = get_installed_kueue_version(
+      dry_run_version=_SUPER_SLICING_MINIMUM_KUEUE_VERSION
+  )
+
+  return (
+      return_code == 0
+      and current_version is not None
+      and current_version >= _SUPER_SLICING_MINIMUM_KUEUE_VERSION
+  )
+
+
 def get_total_chips_requested_from_args(
     args, system: SystemCharacteristics
 ) -> int:
@@ -303,13 +342,18 @@ def create_sub_slicing_annotations(sub_slicing_topology: str) -> list[str]:
   ]
 
 
-def create_placement_policy_label(system: SystemCharacteristics) -> str:
-  name = get_placement_policy_name(system)
+def create_placement_policy_label(
+    system: SystemCharacteristics, super_slicing: bool
+) -> str:
+  name = get_placement_policy_name(system, super_slicing)
   return f'cloud.google.com/placement-policy-name: {name}'
 
 
-def get_placement_policy_name(system: SystemCharacteristics) -> str:
-  return f'{system.device_type}-{system.topology}-placement-policy'
+def get_placement_policy_name(
+    system: SystemCharacteristics, super_slicing: bool
+) -> str:
+  super_slicing_part = '-ss' if super_slicing else ''
+  return f'{system.device_type}-{system.topology}{super_slicing_part}-placement-policy'
 
 
 def is_placement_policy_supported(system: SystemCharacteristics) -> bool:

xpk/core/scheduling_test.py

@@ -65,9 +65,12 @@ def test_create_placement_policy_label_returns_valid_label():
       device_type='tpu7x',
       accelerator_type=AcceleratorType.TPU,
       supports_sub_slicing=False,
+      supports_super_slicing=False,
       docker_platform=DockerPlatform.ARM,
   )
-  label = create_placement_policy_label(system_characteristics)
+  label = create_placement_policy_label(
+      system_characteristics, super_slicing=False
+  )
   assert (
       label
       == 'cloud.google.com/placement-policy-name: tpu7x-1x1x1-placement-policy'
@@ -85,12 +88,31 @@ def test_get_placement_policy_name_returns_valid_name():
       device_type='tpu7x',
       accelerator_type=AcceleratorType.TPU,
       supports_sub_slicing=False,
+      supports_super_slicing=False,
       docker_platform=DockerPlatform.ARM,
   )
-  name = get_placement_policy_name(system_characteristics)
+  name = get_placement_policy_name(system_characteristics, super_slicing=False)
   assert name == 'tpu7x-1x1x1-placement-policy'
 
 
+def test_get_placement_policy_name_super_slicing_returns_valid_name():
+  system_characteristics = SystemCharacteristics(
+      chips_per_vm=1,
+      gce_machine_type='tpu7x-standard-1t',
+      gke_accelerator='tpu7x',
+      requires_workload_policy=False,
+      topology='1x1x1',
+      vms_per_slice=1,
+      device_type='tpu7x',
+      accelerator_type=AcceleratorType.TPU,
+      supports_sub_slicing=False,
+      supports_super_slicing=False,
+      docker_platform=DockerPlatform.ARM,
+  )
+  name = get_placement_policy_name(system_characteristics, super_slicing=True)
+  assert name == 'tpu7x-1x1x1-ss-placement-policy'
+
+
 def test_is_placement_policy_supported_returns_true_for_system_characteristics_supporting_workload_policy_and_having_valid_topology():
   system_characteristics = SystemCharacteristics(
       chips_per_vm=1,
@@ -102,6 +124,7 @@ def test_is_placement_policy_supported_returns_true_for_system_characteristics_s
       device_type='tpu7x',
       accelerator_type=AcceleratorType.TPU,
       supports_sub_slicing=False,
+      supports_super_slicing=False,
       docker_platform=DockerPlatform.ARM,
   )
   assert is_placement_policy_supported(system_characteristics) is True
@@ -118,6 +141,7 @@ def test_is_placement_policy_supported_returns_false_for_system_characteristics_
       device_type='tpu7x',
       accelerator_type=AcceleratorType.TPU,
       supports_sub_slicing=False,
+      supports_super_slicing=False,
       docker_platform=DockerPlatform.ARM,
   )
   assert is_placement_policy_supported(system_characteristics) is False
@@ -134,6 +158,7 @@ def test_is_placement_policy_supported_returns_false_for_system_characteristics_
       device_type='tpu7x',
       accelerator_type=AcceleratorType.TPU,
       supports_sub_slicing=False,
+      supports_super_slicing=False,
       docker_platform=DockerPlatform.ARM,
   )
   assert is_placement_policy_supported(system_characteristics) is False
@@ -145,28 +170,42 @@ class SchedulingTestCase:
   num_slices: int = 1
   cluster_system: SystemCharacteristics | None = None
   resources_config_map: dict[str, str] | None = None
-  sub_slicing_feature_enabled: bool = False
   kueue_version: str | None = None
+  sub_slicing_feature_enabled: bool = False
   sub_slicing_topology_set: bool = False
+  super_slicing_feature_enabled: bool = False
+  super_slicing_topology_set: bool = False
 
 
+NAP_CASE = SchedulingTestCase(
+    workload_system=_get_system_characteristics_or_die('v6e-8'),
+    cluster_system=None,
+    resources_config_map={
+        'tpu-v6e-slice': AUTOPROVISIONING_CONFIG_VALUE,
+        AUTOPROVISIONING_CONFIG_MAXIMUM_KEY: '10',
+    },
+)
+
 SUB_SLICING_CASE = SchedulingTestCase(
     workload_system=_get_system_characteristics_or_die('v6e-8'),
     cluster_system=_get_system_characteristics_or_die('v6e-16'),
-    resources_config_map={'v6e-16': '8'},
-    sub_slicing_feature_enabled=True,
+    # 2 slices:
+    resources_config_map={'v6e-16': str(8 // 4 * 2)},
     kueue_version='0.13.0',
+    sub_slicing_feature_enabled=True,
     sub_slicing_topology_set=True,
     num_slices=1,
 )
 
-NAP_CASE = SchedulingTestCase(
-    workload_system=_get_system_characteristics_or_die('v6e-8'),
-    cluster_system=None,
-    resources_config_map={
-        'tpu-v6e-slice': AUTOPROVISIONING_CONFIG_VALUE,
-        AUTOPROVISIONING_CONFIG_MAXIMUM_KEY: '10',
-    },
+SUPER_SLICING_CASE = SchedulingTestCase(
+    workload_system=_get_system_characteristics_or_die('tpu7x-4x4x16'),
+    cluster_system=_get_system_characteristics_or_die('tpu7x-4x4x4'),
+    # 5 4x4x4 cubes:
+    resources_config_map={'tpu7x-128': str(64 // 4 * 5)},
+    kueue_version='0.14.0',
+    super_slicing_feature_enabled=True,
+    super_slicing_topology_set=True,
+    num_slices=1,
 )
 
 
@@ -283,6 +322,66 @@ NAP_CASE = SchedulingTestCase(
             ),
             WorkloadScheduling.AVAILABLE,
         ),
+        (
+            'Correct Super-slicing',
+            SUPER_SLICING_CASE,
+            WorkloadScheduling.SUPER_SLICING_AVAILABLE,
+        ),
+        (
+            'Super-slicing, but disabled flag',
+            dataclasses.replace(
+                SUPER_SLICING_CASE, super_slicing_feature_enabled=False
+            ),
+            WorkloadScheduling.UNAVAILABLE,
+        ),
+        (
+            'Super-slicing, but low Kueue version',
+            dataclasses.replace(SUPER_SLICING_CASE, kueue_version='0.13.0'),
+            WorkloadScheduling.UNAVAILABLE,
+        ),
+        (
+            'Super-slicing, but no super-slicing-topology',
+            dataclasses.replace(
+                SUPER_SLICING_CASE, super_slicing_topology_set=False
+            ),
+            WorkloadScheduling.UNAVAILABLE,
+        ),
+        (
+            'Super-slicing, but workload too big',
+            dataclasses.replace(SUPER_SLICING_CASE, num_slices=100),
+            WorkloadScheduling.UNAVAILABLE,
+        ),
+        (
+            'Super-slicing, but cluster system is incorrect',
+            dataclasses.replace(
+                SUPER_SLICING_CASE,
+                cluster_system=_get_system_characteristics_or_die(
+                    'tpu7x-4x4x8'
+                ),
+            ),
+            WorkloadScheduling.UNAVAILABLE,
+        ),
+        (
+            'Super-slicing, but workload system is incorrect',
+            dataclasses.replace(
+                SUPER_SLICING_CASE,
+                workload_system=_get_system_characteristics_or_die('v6e-8'),
+            ),
+            WorkloadScheduling.UNAVAILABLE,
+        ),
+        (
+            (
+                'Super-slicing should be ignored when a given device is already'
+                ' present in the cluster'
+            ),
+            dataclasses.replace(
+                SUPER_SLICING_CASE,
+                workload_system=_get_system_characteristics_or_die('tpu7x-64'),
+                cluster_system=_get_system_characteristics_or_die('tpu7x-64'),
+                resources_config_map={'tpu7x-64': '16'},
+            ),
+            WorkloadScheduling.AVAILABLE,
+        ),
     ],
 )
 def test_check_if_workload_can_schedule(
@@ -292,6 +391,7 @@ def test_check_if_workload_can_schedule(
     expected: WorkloadScheduling,
 ):
   FeatureFlags.SUB_SLICING_ENABLED = case.sub_slicing_feature_enabled
+  FeatureFlags.SUPER_SLICING_ENABLED = case.super_slicing_feature_enabled
   commands_tester.set_result_for_command(
       (
           0,
@@ -302,8 +402,13 @@ def test_check_if_workload_can_schedule(
       'kubectl get deployment',
       'image',
   )
+  topology_response = ''
+  if case.sub_slicing_topology_set:
+    topology_response = 'sub-slice-topology'
+  elif case.super_slicing_topology_set:
+    topology_response = 'super-slice-topology'
   commands_tester.set_result_for_command(
-      (0, 'sub-slice-topology' if case.sub_slicing_topology_set else ''),
+      (0, topology_response),
       'kubectl get topology',
   )
   args = Namespace(

xpk/core/system_characteristics.py

@@ -137,6 +137,7 @@ class SystemCharacteristics:
     device_type: A user-facing name for the specific hardware configuration
       (e.g., 'l4-1', 'h100-80gb-8').
     supports_sub_slicing: Whether the Sub-slicing feature is supported.
+    supports_super_slicing: Whether the Super-slicing feature is supported.
     requires_workload_policy: A boolean indicating if a GCE resource
       workload policy is required. This is automatically set to True for GPUs.
   """
@@ -149,6 +150,7 @@ class SystemCharacteristics:
   accelerator_type: AcceleratorType
   device_type: str
   supports_sub_slicing: bool
+  supports_super_slicing: bool
   docker_platform: DockerPlatform
   requires_workload_policy: bool = False
   gpu_config: Optional[GpuConfig] = None
@@ -239,14 +241,16 @@ def get_tpu_system_characteristics_map(
     gke_accelerator: str,
     machine_type: str,
     supported_topologies: list[str],
-    supports_sub_slicing: bool,
     docker_platform: DockerPlatform,
     tpu_type_requires_workload_policy: bool = False,
     default_topologies: set[str] | None = None,
+    sub_slicing_topologies: set[str] | None = None,
+    super_slicing_topologies: set[str] | None = None,
 ) -> dict[str, SystemCharacteristics]:
   system_characteristics_map = {}
-  if default_topologies is None:
-    default_topologies = set()
+  default_topologies = default_topologies or set()
+  sub_slicing_topologies = sub_slicing_topologies or set()
+  super_slicing_topologies = super_slicing_topologies or set()
   for topology in supported_topologies:
     chips_per_vm = compute_chips_per_vm(topology)
     vms_per_slice = compute_vms_per_slice(topology)
@@ -262,7 +266,8 @@ def get_tpu_system_characteristics_map(
         device_type=device_type,
         requires_workload_policy=tpu_type_requires_workload_policy
         and vms_per_slice > 1,
-        supports_sub_slicing=supports_sub_slicing,
+        supports_sub_slicing=topology in sub_slicing_topologies,
+        supports_super_slicing=topology in super_slicing_topologies,
         docker_platform=docker_platform,
     )
     system_characteristics_map[f'{prefix}-{topology}'] = system
@@ -306,6 +311,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.GPU,
         device_type='l4-1',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         gpu_config=GpuConfig(requires_topology=False),
         docker_platform=AMD_PLATFORM,
     ),
@@ -318,6 +324,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.GPU,
         device_type='l4-2',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         gpu_config=GpuConfig(requires_topology=False),
         docker_platform=AMD_PLATFORM,
     ),
@@ -330,6 +337,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.GPU,
         device_type='l4-4',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         gpu_config=GpuConfig(requires_topology=False),
         docker_platform=AMD_PLATFORM,
     ),
@@ -342,6 +350,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.GPU,
         device_type='l4-8',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         gpu_config=GpuConfig(requires_topology=False),
         docker_platform=AMD_PLATFORM,
     ),
@@ -355,6 +364,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.GPU,
         device_type='a100-40gb-1',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         gpu_config=GpuConfig(requires_topology=False),
         docker_platform=AMD_PLATFORM,
     ),
@@ -367,6 +377,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.GPU,
         device_type='a100-40gb-2',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         gpu_config=GpuConfig(requires_topology=False),
         docker_platform=AMD_PLATFORM,
     ),
@@ -379,6 +390,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.GPU,
         device_type='a100-40gb-4',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         gpu_config=GpuConfig(requires_topology=False),
         docker_platform=AMD_PLATFORM,
     ),
@@ -391,6 +403,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.GPU,
         device_type='a100-40gb-8',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         gpu_config=GpuConfig(requires_topology=False),
         docker_platform=AMD_PLATFORM,
     ),
@@ -403,6 +416,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.GPU,
         device_type='gb200-4',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         gpu_config=GpuConfig(
             requires_topology=True,
             nccl_installer=INSTALLER_NCCL_RDMA_A4X,
@@ -421,6 +435,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.GPU,
         device_type='gb200-4',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         gpu_config=GpuConfig(
             requires_topology=True,
             nccl_installer=INSTALLER_NCCL_RDMA_A4X,
@@ -439,6 +454,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.GPU,
         device_type='b200-8',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         gpu_config=GpuConfig(
             requires_topology=True,
             nccl_installer=INSTALLER_NCCL_RDMA,
@@ -457,6 +473,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.GPU,
         device_type='h200-141gb-8',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         gpu_config=GpuConfig(
             requires_topology=True,
             nccl_installer=INSTALLER_NCCL_RDMA,
@@ -476,6 +493,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.GPU,
         device_type='h100-80gb-8',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         gpu_config=GpuConfig(
             requires_topology=True,
             nccl_installer=INSTALLER_NCCL_TCPX,
@@ -495,6 +513,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.GPU,
         device_type='h100-mega-80gb-8',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         gpu_config=GpuConfig(
             requires_topology=True,
             nccl_installer=INSTALLER_NCCL_TCPXO,
@@ -512,7 +531,6 @@ UserFacingNameToSystemCharacteristics = {
         machine_type='tpu7x-standard-1t',
         supported_topologies=['1x1x1'],
         tpu_type_requires_workload_policy=True,
-        supports_sub_slicing=False,
         docker_platform=AMD_PLATFORM,
     ),
     **get_tpu_system_characteristics_map(
@@ -521,9 +539,9 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='tpu7x',
         machine_type='tpu7x-standard-4t',
         tpu_type_requires_workload_policy=True,
-        supports_sub_slicing=False,
         docker_platform=AMD_PLATFORM,
         supported_topologies=generate_tpu_topologies(max_cubes=144),
+        super_slicing_topologies=set(['4x4x4']),
         default_topologies=set([
             '12x12x12',
             '12x12x16',
@@ -630,7 +648,6 @@ UserFacingNameToSystemCharacteristics = {
         tensorcores_per_chip=1,
         gke_accelerator='tpu-v6e-slice',
         machine_type='ct6e-standard-1t',
-        supports_sub_slicing=False,
         supported_topologies=['1x1'],
         docker_platform=AMD_PLATFORM,
     ),
@@ -639,19 +656,8 @@ UserFacingNameToSystemCharacteristics = {
         tensorcores_per_chip=1,
         gke_accelerator='tpu-v6e-slice',
         machine_type='ct6e-standard-4t',
-        supports_sub_slicing=False,
-        supported_topologies=[
-            '2x2',
-        ],
-        docker_platform=AMD_PLATFORM,
-    ),
-    **get_tpu_system_characteristics_map(
-        prefix='v6e',
-        tensorcores_per_chip=1,
-        gke_accelerator='tpu-v6e-slice',
-        machine_type='ct6e-standard-4t',
-        supports_sub_slicing=True,
-        supported_topologies=SUB_SLICING_TOPOLOGIES,
+        supported_topologies=['2x2'] + SUB_SLICING_TOPOLOGIES,
+        sub_slicing_topologies=set(SUB_SLICING_TOPOLOGIES),
         docker_platform=AMD_PLATFORM,
     ),
     **get_tpu_system_characteristics_map(
@@ -659,7 +665,6 @@ UserFacingNameToSystemCharacteristics = {
         tensorcores_per_chip=2,
         gke_accelerator='tpu-v5p-slice',
         machine_type='ct5p-hightpu-4t',
-        supports_sub_slicing=False,
         docker_platform=AMD_PLATFORM,
         supported_topologies=generate_tpu_topologies(max_cubes=140),
         default_topologies=set([
@@ -767,7 +772,6 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='tpu-v5-lite-podslice',
         machine_type='ct5lp-hightpu-4t',
         docker_platform=AMD_PLATFORM,
-        supports_sub_slicing=False,
         supported_topologies=['2x4', '4x4', '4x8', '8x8', '8x16', '16x16'],
     ),
     **get_tpu_system_characteristics_map(
@@ -776,7 +780,6 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='tpu-v4-podslice',
         machine_type='ct4p-hightpu-4t',
         docker_platform=AMD_PLATFORM,
-        supports_sub_slicing=False,
         supported_topologies=generate_tpu_topologies(
             max_cubes=64, enforce_nondecreasing=False
         ),
@@ -807,6 +810,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.CPU,
         device_type='m1-megamem-96-1',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         docker_platform=AMD_PLATFORM,
     ),
     # n2-standard-#vCPUs-#VMs
@@ -819,6 +823,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-64-1',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-1': SystemCharacteristics(
@@ -830,6 +835,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-1',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-2': SystemCharacteristics(
@@ -841,6 +847,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-2',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-4': SystemCharacteristics(
@@ -852,6 +859,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-4',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-8': SystemCharacteristics(
@@ -863,6 +871,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-8',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-16': SystemCharacteristics(
@@ -874,6 +883,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-16',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-32': SystemCharacteristics(
@@ -885,6 +895,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-32',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-64': SystemCharacteristics(
@@ -896,6 +907,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-64',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-128': SystemCharacteristics(
@@ -907,6 +919,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-128',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-256': SystemCharacteristics(
@@ -918,6 +931,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-256',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-512': SystemCharacteristics(
@@ -929,6 +943,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-512',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-1024': SystemCharacteristics(
@@ -940,6 +955,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-1024',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         docker_platform=AMD_PLATFORM,
     ),
     'n2-standard-32-2048': SystemCharacteristics(
@@ -951,6 +967,7 @@ UserFacingNameToSystemCharacteristics = {
         accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-2048',
         supports_sub_slicing=False,
+        supports_super_slicing=False,
         docker_platform=AMD_PLATFORM,
     ),
 }
@@ -983,7 +1000,7 @@ def create_accelerator_label(system: SystemCharacteristics) -> str:
 def create_machine_label(system: SystemCharacteristics) -> str:
   if system.accelerator_type == AcceleratorType.TPU:
     return (
-        f'{AcceleratorTypeToAcceleratorCharacteristics[system.accelerator_type].machine_label}:'
+        f'{AcceleratorTypeToAcceleratorCharacteristics[AcceleratorType.TPU].machine_label}:'
         f' {system.topology}'
     )
   return ''