xpk 0.7.2__py3-none-any.whl → 0.9.0__py3-none-any.whl

xpk/core/cluster.py

@@ -14,27 +14,37 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import yaml
 from google.api_core.exceptions import PermissionDenied
 from google.cloud import resourcemanager_v3
 from kubernetes import client as k8s_client
 from kubernetes import config
 from kubernetes.client.exceptions import ApiException
-from .resources import get_cluster_system_characteristics
 
 from ..utils.console import xpk_exit, xpk_print
-from .capacity import H100_DEVICE_TYPE
+from .capacity import B200_DEVICE_TYPE, H100_DEVICE_TYPE, H200_DEVICE_TYPE
 from .commands import (
     run_command_for_value,
     run_command_with_updates,
     run_command_with_updates_retry,
 )
-from .gcloud_context import add_zone_and_project, get_gke_server_config, zone_to_region
+from .gcloud_context import (
+    add_zone_and_project,
+    get_gke_server_config,
+    zone_to_region,
+)
 from .nodepool import upgrade_gke_nodepools_version
+from .resources import get_cluster_system_characteristics
 from .system_characteristics import SystemCharacteristics
 
-JOBSET_VERSION = 'v0.7.2'
-INSTALLER_NCC_TCPX = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpx/nccl-tcpx-installer.yaml'
-INSTALLER_NCC_TCPXO = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/nccl-tcpxo-installer.yaml'
+JOBSET_VERSION = 'v0.8.0'
+PATHWAYS_JOB_VERSION = 'v0.1.1'
+INSTALLER_NCCL_TCPX = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpx/nccl-tcpx-installer.yaml'
+INSTALLER_NCCL_TCPXO = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/nccl-tcpxo-installer.yaml'
+INSTALLER_NCCL_RDMA = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer.yaml'
+CONFIG_NCCL_TCPX = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpx/nccl-config.yaml'
+NRI_DEVICE_INJECTOR = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/nri_device_injector/nri-device-injector.yaml'
+MGLRU_DISABLE = 'https://raw.githubusercontent.com/GoogleCloudPlatform/cluster-toolkit/refs/heads/main/examples/gke-a3-ultragpu/mglru-disable.yaml'
 
 DEFAULT_NAMESPACE = 'default'
 XPK_SA = 'xpk-sa'
@@ -71,6 +81,35 @@ def set_jobset_on_cluster(args) -> int:
   return return_code
 
 
+def set_pathways_job_on_cluster(args) -> int:
+  """Add PathwaysJob command on server side and ask user to verify it is created.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = (
+      'kubectl apply --server-side -f'
+      f' https://github.com/google/pathways-job/releases/download/{PATHWAYS_JOB_VERSION}/install.yaml'
+  )
+  task = f'Install PathwaysJob on {args.cluster}'
+  return_code = run_command_with_updates_retry(command, task, args)
+
+  if return_code != 0:
+    xpk_print(f'{task} returned with ERROR {return_code}.\n')
+    xpk_print(
+        "This LIKELY means you're missing Kubernetes Permissions, you can"
+        ' validate this by checking if the error references permission problems'
+        ' such as `requires one of ["container.*"] permission(s)`. Follow our'
+        ' readme:'
+        ' https://github.com/google/xpk/blob/main/README.md#troubleshooting for'
+        ' instructions on how to fix these permissions.'
+    )
+  return return_code
+
+
 def install_nccl_on_cluster(args, system: SystemCharacteristics) -> int:
   """Install NCCL plugin on the cluster.
 
@@ -82,9 +121,11 @@ def install_nccl_on_cluster(args, system: SystemCharacteristics) -> int:
     0 if successful and 1 otherwise.
   """
   if system.device_type == H100_DEVICE_TYPE:
-    command = f'kubectl apply -f {INSTALLER_NCC_TCPX}'
+    command = f'kubectl apply -f {INSTALLER_NCCL_TCPX}'
+  elif system.device_type in [H200_DEVICE_TYPE, B200_DEVICE_TYPE]:
+    command = f'kubectl apply -f {INSTALLER_NCCL_RDMA}'
   else:
-    command = f'kubectl apply -f {INSTALLER_NCC_TCPXO}'
+    command = f'kubectl apply -f {INSTALLER_NCCL_TCPXO}'
 
   return_code = run_command_with_updates(
       command, 'Install NCCL Plugin On Cluster', args
@@ -96,9 +137,108 @@ def install_nccl_on_cluster(args, system: SystemCharacteristics) -> int:
     )
     return 1
 
+  if system.device_type == H100_DEVICE_TYPE:
+    command = f'kubectl apply -f {CONFIG_NCCL_TCPX}'
+
+    return_code = run_command_with_updates(
+        command, 'Install NCCL Config On Cluster', args
+    )
+
+    if return_code != 0:
+      xpk_print(
+          f'Install NCCL Config On Cluster request returned ERROR {return_code}'
+      )
+      return 1
+
+  return 0
+
+
+def disable_mglru_on_cluster(args) -> int:
+  """Disable MGLRU on the cluster.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = f'kubectl apply -f {MGLRU_DISABLE}'
+  return_code = run_command_with_updates(
+      command, 'Disable MGLRU On Cluster', args
+  )
+
+  if return_code != 0:
+    xpk_print('Disablig MGLRU On Cluster request returned ERROR')
+    return 1
+
+  return 0
+
+
+def install_nri_on_cluster(args) -> int:
+  """Install NRI Device Injector on the cluster.
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system characteristics.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = f'kubectl apply -f {NRI_DEVICE_INJECTOR}'
+  return_code = run_command_with_updates(
+      command, 'Install NRI Device Injector On Cluster', args
+  )
+
+  if return_code != 0:
+    xpk_print(
+        'Install NRI Device Injector On Cluster request returned ERROR'
+        f' {return_code}'
+    )
+    return 1
+
   return 0
 
 
+def get_cluster_nodes_info(args) -> list[dict]:
+  """Get list of cluster's nodes descrition in yaml format
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    List of nodes info yaml objects.
+  """
+  xpk_print("Getting cluster's info...")
+  command = 'kubectl get nodes -o yaml'
+  err_code, val = run_command_for_value(
+      command=command,
+      task='Get cluster nodes info',
+      global_args=args,
+  )
+  if err_code != 0:
+    xpk_exit(err_code)
+  data = yaml.safe_load(val)
+  return data['items']  # pytype: disable=bad-return-type
+
+
+def count_nodes_on_cluster(args, system: SystemCharacteristics) -> int:
+  """Count cluster nodes by accelerator type"""
+  nodes_info = get_cluster_nodes_info(args)
+  accelerators = [
+      node['metadata']['labels']['cloud.google.com/gke-accelerator']
+      for node in nodes_info
+      if 'cloud.google.com/gke-accelerator' in node['metadata']['labels']
+  ]
+  if system.device_type != H200_DEVICE_TYPE:
+    xpk_print(
+        'Automatic node detection is not supported for device type:'
+        f' {system.device_type}'
+    )
+    xpk_exit(1)
+  num_nodes: int = sum(acc == system.gke_accelerator for acc in accelerators)
+  return num_nodes
+
+
 def get_cluster_network(args) -> str:
   xpk_print("Getting cluster's VPC network...")
   cluster_network_cmd = (
@@ -135,8 +275,48 @@ def update_cluster_with_gcpfilestore_driver_if_necessary(args) -> int:
   return 0
 
 
+def update_cluster_with_parallelstore_driver_if_necessary(args) -> int:
+  """Updates a GKE cluster to enable Parallelstore CSI driver, if not enabled already.
+  Args:
+    args: user provided arguments for running the command.
+  Returns:
+    0 if successful and error code otherwise.
+  """
+  if is_driver_enabled_on_cluster(args, driver='parallelstoreCsiDriver'):
+    return 0
+  cluster_update_return_code = update_gke_cluster_with_addon(
+      args, 'ParallelstoreCsiDriver'
+  )
+  if cluster_update_return_code > 0:
+    xpk_print('Updating GKE cluster to enable Parallelstore CSI driver failed!')
+    return cluster_update_return_code
+
+  return 0
+
+
+def update_cluster_with_pd_driver_if_necessary(args) -> int:
+  """Updates a GKE cluster to enable PersistentDisk CSI driver, if not enabled already.
+  Args:
+    args: user provided arguments for running the command.
+  Returns:
+    0 if successful and error code otherwise.
+  """
+  if is_driver_enabled_on_cluster(args, driver='gcePersistentDiskCsiDriver'):
+    return 0
+  cluster_update_return_code = update_gke_cluster_with_addon(
+      args, 'GcePersistentDiskCsiDriver'
+  )
+  if cluster_update_return_code > 0:
+    xpk_print(
+        'Updating GKE cluster to enable PersistentDisk CSI driver failed!'
+    )
+    return cluster_update_return_code
+
+  return 0
+
+
 def is_driver_enabled_on_cluster(args, driver: str) -> bool:
-  """Checks if GCSFuse CSI driver is enabled on the cluster.
+  """Checks if the CSI driver is enabled on the cluster.
   Args:
     args: user provided arguments for running the command.
     driver (str) : name of the driver
@@ -148,14 +328,14 @@ def is_driver_enabled_on_cluster(args, driver: str) -> bool:
       f' --project={args.project} --region={zone_to_region(args.zone)}'
       f' --format="value(addonsConfig.{driver}Config.enabled)"'
   )
-  return_code, gcsfuse_driver_enabled = run_command_for_value(
+  return_code, driver_enabled = run_command_for_value(
       command,
       f'Checks if {driver} driver is enabled in cluster describe.',
       args,
   )
   if return_code != 0:
     xpk_exit(return_code)
-  if gcsfuse_driver_enabled.lower() == 'true':
+  if driver_enabled.strip().lower() == 'true':
     xpk_print(f'{driver} driver is enabled on the cluster, no update needed.')
     return True
   return False
@@ -446,7 +626,7 @@ def is_gcsfuse_driver_enabled_on_cluster(args) -> bool:
   )
   if return_code != 0:
     xpk_exit(return_code)
-  if gcsfuse_driver_enabled.lower() == 'true':
+  if gcsfuse_driver_enabled.strip().lower() == 'true':
     xpk_print('GCSFuse CSI driver is enabled on the cluster, no update needed.')
     return True
   return False
@@ -551,6 +731,7 @@ def get_cluster_credentials(args) -> None:
   command = (
       'gcloud container clusters get-credentials'
       f' {args.cluster} --region={zone_to_region(args.zone)}'
+      ' --dns-endpoint'
       f' --project={args.project} &&'
       ' kubectl config view && kubectl config set-context --current'
       ' --namespace=default'

xpk/core/config.py

@@ -24,7 +24,7 @@ from ..utils.console import xpk_print
 from .system_characteristics import AcceleratorType, SystemCharacteristics
 
 # This is the version for XPK PyPI package
-__version__ = 'v0.7.2'
+__version__ = 'v0.9.0'
 XPK_CURRENT_VERSION = __version__
 XPK_CONFIG_FILE = os.path.expanduser('~/.config/xpk/config.yaml')
 
@@ -39,6 +39,7 @@ KJOB_SHELL_IMAGE = 'shell-image'
 KJOB_SHELL_INTERACTIVE_COMMAND = 'shell-interactive-command'
 KJOB_SHELL_WORKING_DIRECTORY = 'shell-working-directory'
 CONFIGS_KEY = 'configs'
+GKE_ENDPOINT_KEY = 'gke-endpoint'
 DEPENDENCIES_KEY = 'deps-verified-version'
 XPK_CONFIG_FILE = os.path.expanduser('~/.config/xpk/config.yaml')
 
@@ -47,6 +48,7 @@ DEFAULT_KEYS = [
     CLUSTER_NAME_KEY,
     PROJECT_KEY,
     ZONE_KEY,
+    GKE_ENDPOINT_KEY,
     DEPENDENCIES_KEY,
     KJOB_BATCH_IMAGE,
     KJOB_BATCH_WORKING_DIRECTORY,

xpk/core/docker_manager.py

@@ -30,7 +30,7 @@ import time
 DockerRunCommandExitCode = 135
 dockerBuildErrorCode = 134
 ctk_dockerfile_path = "Dockerfile"
-ctk_build_ref = "v1.45.1"
+ctk_build_ref = "v1.48.0"
 ctk_docker_image = "xpk-ctk"
 ctk_container_name = "xpk-ctk-container"
 gcloud_cfg_mount_path = "/root/.config/gcloud"

xpk/core/docker_resources.py

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from .capacity import H100_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
+from .capacity import H100_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE, B200_DEVICE_TYPE
 from .cluster import setup_k8s_env
 from .storage import GCS_FUSE_TYPE, GCP_FILESTORE_TYPE, Storage, get_storages_to_mount
 from .system_characteristics import AcceleratorType, SystemCharacteristics
@@ -64,22 +64,6 @@ def get_env_container(args, system: SystemCharacteristics) -> str:
     str:
       YAML with the env config for the main container, as a YAML string.
   """
-  pw_env_yaml = """
-                - name: XCLOUD_ENVIRONMENT
-                  value: GCP
-                - name: JAX_PLATFORMS
-                  value: proxy
-                - name: JAX_BACKEND_TARGET
-                  value: {proxy_address}
-                - name: JOBSET_NAME
-                  valueFrom:
-                    fieldRef:
-                      fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name']"""
-  if args.use_pathways:
-    return pw_env_yaml.format(
-        args=args, proxy_address=args.pathways_proxy_address
-    )
-
   gpu_env_yaml = """
                   - name: REPLICATED_JOB_NAME
                     valueFrom:
@@ -182,11 +166,14 @@ def get_volumes(args, system: SystemCharacteristics) -> str:
                 name: dshm-2
               """
 
-  if args.ramdisk_directory != '':
-    volumes += """
+  if hasattr(args, 'ramdisk_directory') and args.ramdisk_directory != '':
+    driver = 'phase1-checkpoint.csi.storage.gke.io'
+    if hasattr(args, 'mtc_enabled') and args.mtc_enabled:
+      driver = 'multitier-checkpoint.csi.storage.gke.io'
+    volumes += f"""
               - name: cache
                 csi:
-                  driver: phase1-checkpoint.csi.storage.gke.io"""
+                  driver: {driver}"""
 
   if (
       system.accelerator_type == AcceleratorType['TPU']
@@ -229,7 +216,7 @@ def get_volume_mounts(args, system: SystemCharacteristics) -> str:
                   name: dshm-2
                 """
 
-  if args.ramdisk_directory != '':
+  if hasattr(args, 'ramdisk_directory') and args.ramdisk_directory != '':
     volume_mount_yaml += f"""
                 - mountPath: /{args.ramdisk_directory}
                   name: cache"""
@@ -262,6 +249,7 @@ def get_volume_mounts(args, system: SystemCharacteristics) -> str:
     elif (
         system.device_type == H100_MEGA_DEVICE_TYPE
         or system.device_type == H200_DEVICE_TYPE
+        or system.device_type == B200_DEVICE_TYPE
     ):
       volume_mount_yaml = ''
 

xpk/core/filestore.py

@@ -230,7 +230,11 @@ class FilestoreClient:
     return data
 
   def manifest(
-      self, name: str, vol: str, access_mode: str, network: str
+      self,
+      name: str,
+      vol: str,
+      access_mode: str,
+      network: str,
   ) -> list[dict]:
     self.load_instance()
     pv = self.create_pv(name, vol, access_mode)

xpk/core/gcsfuse.py

@@ -20,11 +20,22 @@ FUSE_PV_PATH = "/../templates/fuse-pv.yaml"
 FUSE_PVC_PATH = "/../templates/fuse-pvc.yaml"
 
 
-def create_pv(name: str, size: int, bucket: str) -> dict:
+def create_pv(
+    name: str,
+    size: int,
+    bucket: str,
+    mount_options: str,
+    prefetch_metadata: bool,
+) -> dict:
   data = templates.load(FUSE_PV_PATH)
   data["metadata"]["name"] = f"{name}-pv"
   data["spec"]["capacity"]["storage"] = f"{size}Gi"
   data["spec"]["csi"]["volumeHandle"] = bucket
+  if prefetch_metadata:
+    data["spec"]["csi"]["volumeAttributes"][
+        "gcsfuseMetadataPrefetchOnMount"
+    ] = "true"
+  data["spec"]["mountOptions"] = mount_options.split(",")
   return data
 
 
@@ -36,15 +47,25 @@ def create_pvc(name: str, size: int) -> dict:
   return data
 
 
-def manifest(name: str, bucket: str, size: int) -> list[dict]:
-  """Creates GCS FUSE manifest file.
+def manifest(
+    name: str,
+    bucket: str,
+    size: int,
+    mount_options: str,
+    prefetch_metadata: bool,
+) -> list[dict]:
+  """Creates GCS FUSE storage manifest file.
 
   Args:
-      path (str): path to the file where the manifest will be created
       name (str): base name of the volumes
       bucket (str): name of the storage bucket
-      size (str): size of the storage
+      size (str): size of the storage (in GB)
+      prefetch_metadata (bool): if set, then enables metadata pre-population when mounting the volume
+      mount_options (str): comma-separated list of mountOptions for PersistentVolume
+
+  Returns:
+      list[dict]: list of manifests
   """
-  pv = create_pv(name, size, bucket)
+  pv = create_pv(name, size, bucket, mount_options, prefetch_metadata)
   pvc = create_pvc(name, size)
   return [pv, pvc]

xpk/core/kjob.py

@@ -14,27 +14,50 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from ..core.blueprint.blueprint_generator import get_subnetworks_for_a3mega, get_subnetworks_for_a3ultra
-from ..core.capacity import H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
 from argparse import Namespace
-import yaml
-from .workload_decorators.tcpxo_decorator import get_tcpxo_deamon_entry
-from ..utils.console import xpk_print, xpk_exit
+from enum import Enum
 
-from ..utils import templates
+import yaml
 from kubernetes import client as k8s_client
 from kubernetes.client import ApiClient
 from kubernetes.client.rest import ApiException
-from .cluster import setup_k8s_env, XPK_SA, DEFAULT_NAMESPACE
-from .storage import get_auto_mount_storages, get_auto_mount_gcsfuse_storages
-from .commands import run_command_for_value, run_kubectl_apply, run_command_with_updates
-from .config import XpkConfig, KJOB_SHELL_IMAGE, KJOB_SHELL_INTERACTIVE_COMMAND, KJOB_SHELL_WORKING_DIRECTORY, KJOB_BATCH_IMAGE, KJOB_BATCH_WORKING_DIRECTORY
-from .resources import get_cluster_system_characteristics, SystemCharacteristics, AcceleratorType
-from enum import Enum
 
-from ..core.workload_decorators import tcpxo_decorator
-
-from ..core.workload_decorators import rdma_decorator
+from ..utils import templates
+from ..utils.console import xpk_exit, xpk_print
+from .capacity import H100_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
+from .cluster import DEFAULT_NAMESPACE, XPK_SA, setup_k8s_env
+from .commands import (
+    run_command_for_value,
+    run_command_with_updates,
+    run_kubectl_apply,
+)
+from .config import (
+    KJOB_BATCH_IMAGE,
+    KJOB_BATCH_WORKING_DIRECTORY,
+    KJOB_SHELL_IMAGE,
+    KJOB_SHELL_INTERACTIVE_COMMAND,
+    KJOB_SHELL_WORKING_DIRECTORY,
+    XpkConfig,
+)
+from .network import get_cluster_subnetworks
+from .resources import (
+    AcceleratorType,
+    SystemCharacteristics,
+    get_cluster_system_characteristics,
+)
+from .storage import (
+    GCS_FUSE_ANNOTATIONS,
+    PARALLELSTORE_ANNOTATIONS,
+    get_auto_mount_gcsfuse_storages,
+    get_auto_mount_parallelstore_storages,
+    get_auto_mount_storages,
+)
+from .workload_decorators import (
+    rdma_decorator,
+    tcpx_decorator,
+    tcpxo_decorator,
+)
+from .workload_decorators.tcpxo_decorator import get_tcpxo_deamon_entry
 
 KJOB_API_GROUP_NAME = "kjobctl.x-k8s.io"
 KJOB_API_GROUP_VERSION = "v1alpha1"
@@ -146,8 +169,20 @@ Kueue_TAS_annotation = "kueue.x-k8s.io/podset-preferred-topology=cloud.google.co
 default_interface_annotation = "networking.gke.io/default-interface=eth0"
 
 
+def get_a4_pod_template_annotations(args) -> tuple[str, str]:
+  sub_networks = get_cluster_subnetworks(args)
+  interfaces_key, interfaces_value = rdma_decorator.get_interfaces_entry(
+      sub_networks
+  )
+
+  return (
+      default_interface_annotation,
+      f"{interfaces_key}=$'{interfaces_value}'",
+  )
+
+
 def get_a3ultra_pod_template_annotations(args: Namespace) -> tuple[str, str]:
-  sub_networks = get_subnetworks_for_a3ultra(args.cluster)
+  sub_networks = get_cluster_subnetworks(args)
   interfaces_key, interfaces_value = rdma_decorator.get_interfaces_entry(
       sub_networks
   )
@@ -162,7 +197,7 @@ def get_a3mega_pod_template_annotations(
     args: Namespace,
 ) -> tuple[str, str, str]:
   """Adds or updates annotations in the Pod template."""
-  sub_networks = get_subnetworks_for_a3mega(args.cluster)
+  sub_networks = get_cluster_subnetworks(args)
   tcpxo_deamon_key, tcpxo_deamon_paths = get_tcpxo_deamon_entry()
   interfaces_key, interfaces_value = tcpxo_decorator.get_interfaces_entry(
       sub_networks
@@ -237,6 +272,8 @@ def create_app_profile_instance(
 
 def decorate_job_template_with_gpu(yml_string: str, gpu_type: str) -> str:
   job_spec = yaml.safe_load(yml_string)["template"]
+  if gpu_type == H100_DEVICE_TYPE:
+    job_spec = tcpx_decorator.decorate_kjob_template(job_spec)
   if gpu_type == H100_MEGA_DEVICE_TYPE:
     job_spec = tcpxo_decorator.decorate_kjob_template(job_spec)
   if gpu_type == H200_DEVICE_TYPE:
@@ -436,9 +473,18 @@ def create_volume_bundle_instance(
       xpk_exit(1)
 
 
-def get_gcsfuse_annotation(args: Namespace) -> str | None:
+def get_storage_annotations(args: Namespace) -> list[str]:
+  annotations = []
   k8s_api_client = setup_k8s_env(args)
+
   gcsfuse_storages = get_auto_mount_gcsfuse_storages(k8s_api_client)
   if len(gcsfuse_storages) > 0:
-    return "gke-gcsfuse/volumes=true"
-  return None
+    for key, value in GCS_FUSE_ANNOTATIONS.items():
+      annotations.append(f"{key}={value}")
+
+  parallelstore_storages = get_auto_mount_parallelstore_storages(k8s_api_client)
+  if len(parallelstore_storages) > 0:
+    for key, value in PARALLELSTORE_ANNOTATIONS.items():
+      annotations.append(f"{key}={value}")
+
+  return annotations

xpk/core/kueue.py

@@ -21,6 +21,7 @@ from packaging.version import Version
 
 from ..utils.console import xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
+from .capacity import B200_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
 from .commands import (
     run_command_for_value,
     run_command_with_updates,
@@ -45,6 +46,19 @@ WAIT_FOR_KUEUE_TIMEOUT = '5m'
 
 packaging.version.VERSION_PATTERN = r'^v\d+\.\d+\.\d+$'
 
+topology_yaml = """apiVersion: kueue.x-k8s.io/v1alpha1
+kind: Topology
+metadata:
+  name: "gke-default"
+spec:
+  levels:
+  - nodeLabel: "cloud.google.com/gce-topology-block"
+  - nodeLabel: "cloud.google.com/gce-topology-subblock"
+  - nodeLabel: "cloud.google.com/gce-topology-host"
+  - nodeLabel: "kubernetes.io/hostname"
+---
+"""
+
 cluster_set_crd_yaml = """apiVersion: kueue.x-k8s.io/v1beta1
 kind: ResourceFlavor
 metadata:
@@ -53,6 +67,7 @@ spec:
   nodeLabels:
     {accelerator_label}
     {machine_label}
+  {topology_label}
 ---
 {pw_resource_flavors}
 apiVersion: kueue.x-k8s.io/v1beta1
@@ -300,6 +315,14 @@ def install_kueue_crs(
       resource_type=resource_type,
       total_chips=total_chips,
   )
+  topology_label = ''
+  if system.device_type in [
+      H100_MEGA_DEVICE_TYPE,
+      H200_DEVICE_TYPE,
+      B200_DEVICE_TYPE,
+  ]:
+    topology_label = 'topologyName: "gke-default"'
+
   yml_string = cluster_set_crd_yaml.format(
       system=system,
       cluster_hardware_name=cluster_hardware_name,
@@ -309,6 +332,7 @@ def install_kueue_crs(
       machine_label=create_machine_label(
           system.accelerator_type, system, autoprovisioning_enabled
       ),
+      topology_label=topology_label,
       covered_resources_config=covered_resources_config,
       resource_type=AcceleratorTypeToAcceleratorCharacteristics[
           system.accelerator_type
@@ -318,6 +342,12 @@ def install_kueue_crs(
       cluster_queue_name=CLUSTER_QUEUE_NAME,
       local_queue_name=LOCAL_QUEUE_NAME,
   )
+  if system.device_type in [
+      H100_MEGA_DEVICE_TYPE,
+      H200_DEVICE_TYPE,
+      B200_DEVICE_TYPE,
+  ]:
+    yml_string = topology_yaml + yml_string
 
   tmp = write_tmp_file(yml_string)
   command = f'kubectl apply -f {str(tmp.file.name)}'