xpk 0.12.0__py3-none-any.whl → 0.14.0__py3-none-any.whl

xpk/core/system_characteristics.py

@@ -15,8 +15,8 @@ limitations under the License.
 """
 
 from dataclasses import dataclass
-from functools import reduce
-from operator import mul
+from ..utils.topology import get_topology_product
+
 
 AcceleratorType = {'TPU': 1, 'GPU': 2, 'CPU': 3}
 
@@ -29,27 +29,52 @@ class AcceleratorCharacteristics:
 
 
 AcceleratorTypeToAcceleratorCharacteristics = {
-    # TPU
     AcceleratorType['TPU']: AcceleratorCharacteristics(
-        'google.com/tpu',
-        'cloud.google.com/gke-tpu-accelerator',
-        'cloud.google.com/gke-tpu-topology',
+        resource_type='google.com/tpu',
+        accelerator_label='cloud.google.com/gke-tpu-accelerator',
+        machine_label='cloud.google.com/gke-tpu-topology',
     ),
-    # GPU
     AcceleratorType['GPU']: AcceleratorCharacteristics(
-        'nvidia.com/gpu',
-        'cloud.google.com/gke-accelerator',
-        'cloud.google.com/gce-machine-type',
+        resource_type='nvidia.com/gpu',
+        accelerator_label='cloud.google.com/gke-accelerator',
+        machine_label='cloud.google.com/gce-machine-type',
     ),
-    # CPU
     AcceleratorType['CPU']: AcceleratorCharacteristics(
-        'cpu', '', 'cloud.google.com/gke-nodepool'
+        resource_type='cpu',
+        accelerator_label='',
+        machine_label='cloud.google.com/gke-nodepool',
     ),
 }
 
 
 @dataclass
 class SystemCharacteristics:
+  """Contains the defining characteristics of a specific accelerator system.
+
+  This dataclass holds the hardware and configuration details for a given
+  accelerator type, such as its topology, machine type, and chip count. It
+  provides a standardized way to access system-specific information throughout
+  the application.
+
+  Attributes:
+    topology: The physical or logical layout of the accelerator chips (e.g.,
+      '2x2x1' for TPUs, 'N/A' for single-VM GPUs).
+    vms_per_slice: The number of Virtual Machines that constitute a single
+      accelerator slice.
+    gke_accelerator: The name of the accelerator as recognized by GKE (e.g.,
+      'nvidia-l4', 'tpu7x').
+    gce_machine_type: The GCE machine type that hosts the accelerator (e.g.,
+      'g2-standard-12').
+    chips_per_vm: The number of accelerator chips attached to a single VM.
+    accelerator_type: The category of the accelerator (e.g., TPU, GPU, CPU)
+      from the AcceleratorType enum.
+    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.
+    requires_workload_policy: A boolean indicating if a GCE resource
+      workload policy is required. This is automatically set to True for GPUs.
+  """
+
   topology: str
   vms_per_slice: int
   gke_accelerator: str
@@ -57,6 +82,12 @@ class SystemCharacteristics:
   chips_per_vm: int
   accelerator_type: int  # TODO: use enums
   device_type: str
+  supports_sub_slicing: bool
+  requires_workload_policy: bool = False
+
+  def __post_init__(self):
+    if self.accelerator_type == AcceleratorType['GPU']:
+      self.requires_workload_policy = True
 
 
 def get_system_characteristics(
@@ -99,21 +130,25 @@ def get_tpu_system_characteristics_map(
     gke_accelerator: str,
     machine_type: str,
     supported_topologies: list[str],
+    supports_sub_slicing: bool,
+    requires_workload_policy: bool = False,
 ) -> dict[str, SystemCharacteristics]:
   system_characteristics_map = {}
   for topology in supported_topologies:
-    total_chips = reduce(mul, (int(x) for x in topology.split('x')), 1)
+    total_chips = get_topology_product(topology)
     num_tensorcores = total_chips * tensorcores_per_chip
     chips_per_vm = 1 if total_chips == 1 else 4
     vms_per_slice = total_chips // chips_per_vm
     system = SystemCharacteristics(
-        topology,
-        vms_per_slice,
-        gke_accelerator,
-        machine_type,
-        chips_per_vm,
-        AcceleratorType['TPU'],
-        f'{prefix}-{num_tensorcores}',
+        topology=topology,
+        vms_per_slice=vms_per_slice,
+        gke_accelerator=gke_accelerator,
+        gce_machine_type=machine_type,
+        chips_per_vm=chips_per_vm,
+        accelerator_type=AcceleratorType['TPU'],
+        device_type=f'{prefix}-{num_tensorcores}',
+        requires_workload_policy=requires_workload_policy,
+        supports_sub_slicing=supports_sub_slicing,
     )
     system_characteristics_map[f'{prefix}-{topology}'] = system
     system_characteristics_map[f'{prefix}-{num_tensorcores}'] = system
@@ -131,126 +166,166 @@ UserFacingNameToSystemCharacteristics = {
     # GPU system characteristics
     # l4-$CHIPSc
     'l4-1': SystemCharacteristics(
-        'N/A',
-        1,
-        'nvidia-l4',
-        'g2-standard-12',
-        1,
-        AcceleratorType['GPU'],
-        'l4-1',
+        topology='N/A',
+        vms_per_slice=1,
+        gke_accelerator='nvidia-l4',
+        gce_machine_type='g2-standard-12',
+        chips_per_vm=1,
+        accelerator_type=AcceleratorType['GPU'],
+        device_type='l4-1',
+        supports_sub_slicing=False,
     ),
     'l4-2': SystemCharacteristics(
-        'N/A',
-        1,
-        'nvidia-l4',
-        'g2-standard-24',
-        2,
-        AcceleratorType['GPU'],
-        'l4-2',
+        topology='N/A',
+        vms_per_slice=1,
+        gke_accelerator='nvidia-l4',
+        gce_machine_type='g2-standard-24',
+        chips_per_vm=2,
+        accelerator_type=AcceleratorType['GPU'],
+        device_type='l4-2',
+        supports_sub_slicing=False,
     ),
     'l4-4': SystemCharacteristics(
-        'N/A',
-        1,
-        'nvidia-l4',
-        'g2-standard-48',
-        4,
-        AcceleratorType['GPU'],
-        'l4-4',
+        topology='N/A',
+        vms_per_slice=1,
+        gke_accelerator='nvidia-l4',
+        gce_machine_type='g2-standard-48',
+        chips_per_vm=4,
+        accelerator_type=AcceleratorType['GPU'],
+        device_type='l4-4',
+        supports_sub_slicing=False,
     ),
     'l4-8': SystemCharacteristics(
-        'N/A',
-        1,
-        'nvidia-l4',
-        'g2-standard-96',
-        8,
-        AcceleratorType['GPU'],
-        'l4-8',
+        topology='N/A',
+        vms_per_slice=1,
+        gke_accelerator='nvidia-l4',
+        gce_machine_type='g2-standard-96',
+        chips_per_vm=8,
+        accelerator_type=AcceleratorType['GPU'],
+        device_type='l4-8',
+        supports_sub_slicing=False,
     ),
     # A100-40gb-$CHIPSc
     'a100-40gb-1': SystemCharacteristics(
-        'N/A',
-        1,
-        'nvidia-tesla-a100',
-        'a2-highgpu-1g',
-        1,
-        AcceleratorType['GPU'],
-        'a100-40gb-1',
+        topology='N/A',
+        vms_per_slice=1,
+        gke_accelerator='nvidia-tesla-a100',
+        gce_machine_type='a2-highgpu-1g',
+        chips_per_vm=1,
+        accelerator_type=AcceleratorType['GPU'],
+        device_type='a100-40gb-1',
+        supports_sub_slicing=False,
     ),
     'a100-40gb-2': SystemCharacteristics(
-        'N/A',
-        1,
-        'nvidia-tesla-a100',
-        'a2-highgpu-2g',
-        2,
-        AcceleratorType['GPU'],
-        'a100-40gb-2',
+        topology='N/A',
+        vms_per_slice=1,
+        gke_accelerator='nvidia-tesla-a100',
+        gce_machine_type='a2-highgpu-2g',
+        chips_per_vm=2,
+        accelerator_type=AcceleratorType['GPU'],
+        device_type='a100-40gb-2',
+        supports_sub_slicing=False,
     ),
     'a100-40gb-4': SystemCharacteristics(
-        'N/A',
-        1,
-        'nvidia-tesla-a100',
-        'a2-highgpu-4g',
-        4,
-        AcceleratorType['GPU'],
-        'a100-40gb-4',
+        topology='N/A',
+        vms_per_slice=1,
+        gke_accelerator='nvidia-tesla-a100',
+        gce_machine_type='a2-highgpu-4g',
+        chips_per_vm=4,
+        accelerator_type=AcceleratorType['GPU'],
+        device_type='a100-40gb-4',
+        supports_sub_slicing=False,
     ),
     'a100-40gb-8': SystemCharacteristics(
-        'N/A',
-        1,
-        'nvidia-tesla-a100',
-        'a2-highgpu-8g',
-        8,
-        AcceleratorType['GPU'],
-        'a100-40gb-8',
+        topology='N/A',
+        vms_per_slice=1,
+        gke_accelerator='nvidia-tesla-a100',
+        gce_machine_type='a2-highgpu-8g',
+        chips_per_vm=8,
+        accelerator_type=AcceleratorType['GPU'],
+        device_type='a100-40gb-8',
+        supports_sub_slicing=False,
+    ),
+    'gb200-4': SystemCharacteristics(
+        topology='1x72',
+        vms_per_slice=1,
+        gke_accelerator='nvidia-gb200',
+        gce_machine_type='a4x-highgpu-4g',
+        chips_per_vm=4,
+        accelerator_type=AcceleratorType['GPU'],
+        device_type='gb200-4',
+        supports_sub_slicing=False,
+    ),
+    'gb200-4-nolssd': SystemCharacteristics(
+        topology='1x72',
+        vms_per_slice=1,
+        gke_accelerator='nvidia-gb200',
+        gce_machine_type='a4x-highgpu-4g-nolssd',
+        chips_per_vm=4,
+        accelerator_type=AcceleratorType['GPU'],
+        device_type='gb200-4',
+        supports_sub_slicing=False,
     ),
     'b200-8': SystemCharacteristics(
-        'N/A',
-        1,
-        'nvidia-b200',
-        'a4-highgpu-8g',
-        8,
-        AcceleratorType['GPU'],
-        'b200-8',
+        topology='N/A',
+        vms_per_slice=1,
+        gke_accelerator='nvidia-b200',
+        gce_machine_type='a4-highgpu-8g',
+        chips_per_vm=8,
+        accelerator_type=AcceleratorType['GPU'],
+        device_type='b200-8',
+        supports_sub_slicing=False,
     ),
     'h200-141gb-8': SystemCharacteristics(
-        'N/A',
-        1,
-        'nvidia-h200-141gb',
-        'a3-ultragpu-8g',
-        8,
-        AcceleratorType['GPU'],
-        'h200-141gb-8',
+        topology='N/A',
+        vms_per_slice=1,
+        gke_accelerator='nvidia-h200-141gb',
+        gce_machine_type='a3-ultragpu-8g',
+        chips_per_vm=8,
+        accelerator_type=AcceleratorType['GPU'],
+        device_type='h200-141gb-8',
+        supports_sub_slicing=False,
     ),
     # H100-80gb-$CHIPS
     'h100-80gb-8': SystemCharacteristics(
-        'N/A',
-        1,
-        'nvidia-h100-80gb',
-        'a3-highgpu-8g',
-        8,
-        AcceleratorType['GPU'],
-        'h100-80gb-8',
+        topology='N/A',
+        vms_per_slice=1,
+        gke_accelerator='nvidia-h100-80gb',
+        gce_machine_type='a3-highgpu-8g',
+        chips_per_vm=8,
+        accelerator_type=AcceleratorType['GPU'],
+        device_type='h100-80gb-8',
+        supports_sub_slicing=False,
     ),
     # H100-mega-80gb-$CHIPS
     'h100-mega-80gb-8': SystemCharacteristics(
-        'N/A',
-        1,
-        'nvidia-h100-mega-80gb',
-        'a3-megagpu-8g',
-        8,
-        AcceleratorType['GPU'],
-        'h100-mega-80gb-8',
+        topology='N/A',
+        vms_per_slice=1,
+        gke_accelerator='nvidia-h100-mega-80gb',
+        gce_machine_type='a3-megagpu-8g',
+        chips_per_vm=8,
+        accelerator_type=AcceleratorType['GPU'],
+        device_type='h100-mega-80gb-8',
+        supports_sub_slicing=False,
     ),
     # TPU system characteristics
     **get_tpu_system_characteristics_map(
-        'tpu7x', 2, 'tpu7x', 'tpu7x-standard-1t', ['1x1x1']
+        prefix='tpu7x',
+        tensorcores_per_chip=2,
+        gke_accelerator='tpu7x',
+        machine_type='tpu7x-standard-1t',
+        supported_topologies=['1x1x1'],
+        requires_workload_policy=True,
+        supports_sub_slicing=False,
     ),
     **get_tpu_system_characteristics_map(
-        'tpu7x',
-        2,
-        'tpu7x',
-        'tpu7x-standard-4t',
-        [
+        prefix='tpu7x',
+        tensorcores_per_chip=2,
+        gke_accelerator='tpu7x',
+        machine_type='tpu7x-standard-4t',
+        requires_workload_policy=True,
+        supports_sub_slicing=False,
+        supported_topologies=[
             '12x12x12',
             '12x12x16',
             '12x12x20',
@@ -352,21 +427,36 @@ UserFacingNameToSystemCharacteristics = {
         ],
     ),
     **get_tpu_system_characteristics_map(
-        'v6e', 1, 'tpu-v6e-slice', 'ct6e-standard-1t', ['1x1']
+        prefix='v6e',
+        tensorcores_per_chip=1,
+        gke_accelerator='tpu-v6e-slice',
+        machine_type='ct6e-standard-1t',
+        supports_sub_slicing=False,
+        supported_topologies=['1x1'],
     ),
     **get_tpu_system_characteristics_map(
-        'v6e',
-        1,
-        'tpu-v6e-slice',
-        'ct6e-standard-4t',
-        ['2x2', '2x4', '4x4', '4x8', '8x8', '8x16', '16x16'],
+        prefix='v6e',
+        tensorcores_per_chip=1,
+        gke_accelerator='tpu-v6e-slice',
+        machine_type='ct6e-standard-4t',
+        supports_sub_slicing=True,
+        supported_topologies=[
+            '2x2',
+            '2x4',
+            '4x4',
+            '4x8',
+            '8x8',
+            '8x16',
+            '16x16',
+        ],
     ),
     **get_tpu_system_characteristics_map(
-        'v5p',
-        2,
-        'tpu-v5p-slice',
-        'ct5p-hightpu-4t',
-        [
+        prefix='v5p',
+        tensorcores_per_chip=2,
+        gke_accelerator='tpu-v5p-slice',
+        machine_type='ct5p-hightpu-4t',
+        supports_sub_slicing=False,
+        supported_topologies=[
             '2x2x1',
             '2x2x2',
             '2x2x4',
@@ -466,18 +556,20 @@ UserFacingNameToSystemCharacteristics = {
         ],
     ),
     **get_tpu_system_characteristics_map(
-        'v5litepod',
-        1,
-        'tpu-v5-lite-podslice',
-        'ct5lp-hightpu-4t',
-        ['2x4', '4x4', '4x8', '8x8', '8x16', '16x16'],
+        prefix='v5litepod',
+        tensorcores_per_chip=1,
+        gke_accelerator='tpu-v5-lite-podslice',
+        machine_type='ct5lp-hightpu-4t',
+        supports_sub_slicing=False,
+        supported_topologies=['2x4', '4x4', '4x8', '8x8', '8x16', '16x16'],
     ),
     **get_tpu_system_characteristics_map(
-        'v4',
-        2,
-        'tpu-v4-podslice',
-        'ct4p-hightpu-4t',
-        [
+        prefix='v4',
+        tensorcores_per_chip=2,
+        gke_accelerator='tpu-v4-podslice',
+        machine_type='ct4p-hightpu-4t',
+        supports_sub_slicing=False,
+        supported_topologies=[
             '2x2x1',
             '2x2x2',
             '2x2x4',
@@ -496,131 +588,145 @@ UserFacingNameToSystemCharacteristics = {
     # There are no chips in CPUs.
     # m1-megamem-#vCPUs-#VMs
     'm1-megamem-96-1': SystemCharacteristics(
-        'N/A',
-        1,
-        'N/A',
-        'm1-megamem-96',
-        96,
-        AcceleratorType['CPU'],
-        'm1-megamem-96-1',
+        topology='N/A',
+        vms_per_slice=1,
+        gke_accelerator='N/A',
+        gce_machine_type='m1-megamem-96',
+        chips_per_vm=96,
+        accelerator_type=AcceleratorType['CPU'],
+        device_type='m1-megamem-96-1',
+        supports_sub_slicing=False,
     ),
     # n2-standard-#vCPUs-#VMs
     'n2-standard-64-1': SystemCharacteristics(
-        'N/A',
-        1,
-        'N/A',
-        'n2-standard-64',
-        64,
-        AcceleratorType['CPU'],
-        'n2-standard-64-1',
+        topology='N/A',
+        vms_per_slice=1,
+        gke_accelerator='N/A',
+        gce_machine_type='n2-standard-64',
+        chips_per_vm=64,
+        accelerator_type=AcceleratorType['CPU'],
+        device_type='n2-standard-64-1',
+        supports_sub_slicing=False,
     ),
     'n2-standard-32-1': SystemCharacteristics(
-        'N/A',
-        1,
-        'N/A',
-        'n2-standard-32',
-        32,
-        AcceleratorType['CPU'],
-        'n2-standard-32-1',
+        topology='N/A',
+        vms_per_slice=1,
+        gke_accelerator='N/A',
+        gce_machine_type='n2-standard-32',
+        chips_per_vm=32,
+        accelerator_type=AcceleratorType['CPU'],
+        device_type='n2-standard-32-1',
+        supports_sub_slicing=False,
     ),
     'n2-standard-32-2': SystemCharacteristics(
-        'N/A',
-        2,
-        'N/A',
-        'n2-standard-32',
-        32,
-        AcceleratorType['CPU'],
-        'n2-standard-32-2',
+        topology='N/A',
+        vms_per_slice=2,
+        gke_accelerator='N/A',
+        gce_machine_type='n2-standard-32',
+        chips_per_vm=32,
+        accelerator_type=AcceleratorType['CPU'],
+        device_type='n2-standard-32-2',
+        supports_sub_slicing=False,
     ),
     'n2-standard-32-4': SystemCharacteristics(
-        'N/A',
-        4,
-        'N/A',
-        'n2-standard-32',
-        32,
-        AcceleratorType['CPU'],
-        'n2-standard-32-4',
+        topology='N/A',
+        vms_per_slice=4,
+        gke_accelerator='N/A',
+        gce_machine_type='n2-standard-32',
+        chips_per_vm=32,
+        accelerator_type=AcceleratorType['CPU'],
+        device_type='n2-standard-32-4',
+        supports_sub_slicing=False,
     ),
     'n2-standard-32-8': SystemCharacteristics(
-        'N/A',
-        8,
-        'N/A',
-        'n2-standard-32',
-        32,
-        AcceleratorType['CPU'],
-        'n2-standard-32-8',
+        topology='N/A',
+        vms_per_slice=8,
+        gke_accelerator='N/A',
+        gce_machine_type='n2-standard-32',
+        chips_per_vm=32,
+        accelerator_type=AcceleratorType['CPU'],
+        device_type='n2-standard-32-8',
+        supports_sub_slicing=False,
     ),
     'n2-standard-32-16': SystemCharacteristics(
-        'N/A',
-        16,
-        'N/A',
-        'n2-standard-32',
-        32,
-        AcceleratorType['CPU'],
-        'n2-standard-32-16',
+        topology='N/A',
+        vms_per_slice=16,
+        gke_accelerator='N/A',
+        gce_machine_type='n2-standard-32',
+        chips_per_vm=32,
+        accelerator_type=AcceleratorType['CPU'],
+        device_type='n2-standard-32-16',
+        supports_sub_slicing=False,
     ),
     'n2-standard-32-32': SystemCharacteristics(
-        'N/A',
-        32,
-        'N/A',
-        'n2-standard-32',
-        32,
-        AcceleratorType['CPU'],
-        'n2-standard-32-32',
+        topology='N/A',
+        vms_per_slice=32,
+        gke_accelerator='N/A',
+        gce_machine_type='n2-standard-32',
+        chips_per_vm=32,
+        accelerator_type=AcceleratorType['CPU'],
+        device_type='n2-standard-32-32',
+        supports_sub_slicing=False,
     ),
     'n2-standard-32-64': SystemCharacteristics(
-        'N/A',
-        64,
-        'N/A',
-        'n2-standard-32',
-        32,
-        AcceleratorType['CPU'],
-        'n2-standard-32-64',
+        topology='N/A',
+        vms_per_slice=64,
+        gke_accelerator='N/A',
+        gce_machine_type='n2-standard-32',
+        chips_per_vm=32,
+        accelerator_type=AcceleratorType['CPU'],
+        device_type='n2-standard-32-64',
+        supports_sub_slicing=False,
     ),
     'n2-standard-32-128': SystemCharacteristics(
-        'N/A',
-        128,
-        'N/A',
-        'n2-standard-32',
-        32,
-        AcceleratorType['CPU'],
-        'n2-standard-32-128',
+        topology='N/A',
+        vms_per_slice=128,
+        gke_accelerator='N/A',
+        gce_machine_type='n2-standard-32',
+        chips_per_vm=32,
+        accelerator_type=AcceleratorType['CPU'],
+        device_type='n2-standard-32-128',
+        supports_sub_slicing=False,
     ),
     'n2-standard-32-256': SystemCharacteristics(
-        'N/A',
-        256,
-        'N/A',
-        'n2-standard-32',
-        32,
-        AcceleratorType['CPU'],
-        'n2-standard-32-256',
+        topology='N/A',
+        vms_per_slice=256,
+        gke_accelerator='N/A',
+        gce_machine_type='n2-standard-32',
+        chips_per_vm=32,
+        accelerator_type=AcceleratorType['CPU'],
+        device_type='n2-standard-32-256',
+        supports_sub_slicing=False,
     ),
     'n2-standard-32-512': SystemCharacteristics(
-        'N/A',
-        512,
-        'N/A',
-        'n2-standard-32',
-        32,
-        AcceleratorType['CPU'],
-        'n2-standard-32-512',
+        topology='N/A',
+        vms_per_slice=512,
+        gke_accelerator='N/A',
+        gce_machine_type='n2-standard-32',
+        chips_per_vm=32,
+        accelerator_type=AcceleratorType['CPU'],
+        device_type='n2-standard-32-512',
+        supports_sub_slicing=False,
     ),
     'n2-standard-32-1024': SystemCharacteristics(
-        'N/A',
-        1024,
-        'N/A',
-        'n2-standard-32',
-        32,
-        AcceleratorType['CPU'],
-        'n2-standard-32-1024',
+        topology='N/A',
+        vms_per_slice=1024,
+        gke_accelerator='N/A',
+        gce_machine_type='n2-standard-32',
+        chips_per_vm=32,
+        accelerator_type=AcceleratorType['CPU'],
+        device_type='n2-standard-32-1024',
+        supports_sub_slicing=False,
     ),
     'n2-standard-32-2048': SystemCharacteristics(
-        'N/A',
-        2048,
-        'N/A',
-        'n2-standard-32',
-        32,
-        AcceleratorType['CPU'],
-        'n2-standard-32-2048',
+        topology='N/A',
+        vms_per_slice=2048,
+        gke_accelerator='N/A',
+        gce_machine_type='n2-standard-32',
+        chips_per_vm=32,
+        accelerator_type=AcceleratorType['CPU'],
+        device_type='n2-standard-32-2048',
+        supports_sub_slicing=False,
     ),
 }
 """ If you modify UserFacingNameToSystemCharacteristics you should also modify