xpk 0.7.1__py3-none-any.whl → 0.8.0__py3-none-any.whl

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

@@ -200,7 +200,9 @@ class FilestoreClient:
     ] = f"projects/{self.project}/global/networks/{network}"
     return data
 
-  def create_pv(self, name: str, vol: str, access_mode: str) -> dict:
+  def create_pv(
+      self, name: str, vol: str, access_mode: str, mount_options: str
+  ) -> dict:
     """Create a yaml representing filestore PersistentVolume."""
     data = templates.load(FS_PV_PATH)
     data["metadata"]["name"] = get_pv_name(name)
@@ -215,6 +217,7 @@ class FilestoreClient:
         0
     ].ip_addresses[0]
     data["spec"]["csi"]["volumeAttributes"]["volume"] = vol
+    data["spec"]["mountOptions"] = mount_options.split(",")
     return data
 
   def create_pvc(self, name: str, access_mode: str) -> dict:
@@ -230,10 +233,15 @@ 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,
+      mount_options: str,
   ) -> list[dict]:
     self.load_instance()
-    pv = self.create_pv(name, vol, access_mode)
+    pv = self.create_pv(name, vol, access_mode, mount_options)
     pvc = self.create_pvc(name, access_mode)
     sc = self.create_sc(name, network)
     return [pv, pvc, sc]

xpk/core/gcsfuse.py

@@ -20,11 +20,12 @@ 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) -> 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
+  data["spec"]["mountOptions"] = mount_options.split(",")
   return data
 
 
@@ -36,15 +37,17 @@ def create_pvc(name: str, size: int) -> dict:
   return data
 
 
-def manifest(name: str, bucket: str, size: int) -> list[dict]:
+def manifest(
+    name: str, bucket: str, size: int, mount_options: str
+) -> list[dict]:
   """Creates GCS FUSE 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)
+      mount_options (str): comma-separated list of mountOptions for PersistentVolume
   """
-  pv = create_pv(name, size, bucket)
+  pv = create_pv(name, size, bucket, mount_options)
   pvc = create_pvc(name, size)
   return [pv, pvc]

xpk/core/kjob.py

@@ -14,27 +14,45 @@ 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 ..core.blueprint.blueprint_generator import (
+    get_subnetworks_for_a3mega,
+    get_subnetworks_for_a3ultra,
+    get_subnetworks_for_a4,
+)
+from ..core.capacity import H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
+from ..core.storage import GCS_FUSE_ANNOTATIONS, PARALLELSTORE_ANNOTATIONS
+from ..core.workload_decorators import rdma_decorator, tcpxo_decorator
+from ..utils import templates
+from ..utils.console import xpk_exit, xpk_print
+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 .resources import (
+    AcceleratorType,
+    SystemCharacteristics,
+    get_cluster_system_characteristics,
+)
+from .storage import get_auto_mount_gcsfuse_storages, get_auto_mount_storages, get_auto_mount_parallelstore_storages
+from .workload_decorators.tcpxo_decorator import get_tcpxo_deamon_entry
 
 KJOB_API_GROUP_NAME = "kjobctl.x-k8s.io"
 KJOB_API_GROUP_VERSION = "v1alpha1"
@@ -146,6 +164,18 @@ 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() -> tuple[str, str]:
+  sub_networks = get_subnetworks_for_a4()
+  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)
   interfaces_key, interfaces_value = rdma_decorator.get_interfaces_entry(
@@ -436,9 +466,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/nap.py

@@ -255,6 +255,10 @@ def is_autoprovisioning_enabled(
     bool is true if autoprovisioning is enabled, false otherwise.
     int of 0 if successful and 1 otherwise.
   """
+  # Currently autoprovisioning is not enabled for Pathways workloads. b/360898087
+  if args.use_pathways:
+    return False, 0
+
   resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
   cluster_config_map = get_cluster_configmap(args, resources_configmap_name)
 

xpk/core/network.py

@@ -16,10 +16,8 @@ limitations under the License.
 
 from ..utils.console import xpk_print
 from ..utils.file import write_tmp_file
-from .capacity import H100_DEVICE_TYPE
 from .commands import run_command_for_value, run_command_with_updates
 from .gcloud_context import zone_to_region
-from .system_characteristics import SystemCharacteristics
 
 # cluster_network_yaml: the config when creating the network for a3 cluster
 CLUSTER_NETWORK_YAML = """
@@ -175,16 +173,6 @@ def create_cluster_subnet(args, index) -> int:
   return 0
 
 
-def get_subnetworks_for_a3mega(cluster_name: str) -> list[str]:
-  return [f'{cluster_name}-gpunet-{i}-subnet' for i in range(8)]
-
-
-def get_subnetworks_for_a3ultra(cluster_name: str) -> list[str]:
-  return [f'{cluster_name}-sub-1'] + [
-      f'{cluster_name}-rdma-sub-{i}' for i in range(8)
-  ]
-
-
 def create_cluster_firewall_rule(args, index) -> int:
   """Create one GKE Cluster firewall rule.
 
@@ -247,20 +235,18 @@ def create_cluster_network_config(args) -> int:
   return 0
 
 
-def set_up_cluster_network_for_gpu(args, system: SystemCharacteristics) -> int:
-  """Set up GKE Cluster networks, subnets and firewall rules for A3/A3+.
-  Note: there are 4 NICs for GPU-GPU bw and 1 NIC for host in an A3 node,
-  and there are 8 NICs for GPU-GPU bw and 1 NIC for host in an A3+ node.
+def set_up_cluster_network_for_a3(args) -> int:
+  """Set up GKE Cluster networks, subnets and firewall rules for A3.
+  Note: there are 4 NICs for GPU-GPU bw and 1 NIC for host in an A3 node.
 
   Args:
     args: user provided arguments for running the command.
-    system: system characteristics.
 
   Returns:
     0 if successful and 1 otherwise.
   """
-  num_networks = 5 if system.device_type == H100_DEVICE_TYPE else 9
-  for i in range(1, num_networks):
+  num_networks = 4
+  for i in range(1, num_networks + 1):
     return_code = create_cluster_network(args, i)
     if return_code != 0:
       return 1
@@ -315,7 +301,10 @@ def get_all_networks_programmatic(args) -> tuple[list[str], int]:
   Returns:
     List of networks and 0 if successful and 1 otherwise.
   """
-  command = 'gcloud compute networks list --format="csv[no-heading](name)"'
+  command = (
+      'gcloud compute networks list --format="csv[no-heading](name)" '
+      f' --project={args.project}'
+  )
   return_code, raw_network_output = run_command_for_value(
       command, 'Get All Networks', args
   )
@@ -365,7 +354,8 @@ def get_all_firewall_rules_programmatic(args) -> tuple[list[str], int]:
     List of firewall rules and 0 if successful and 1 otherwise.
   """
   command = (
-      'gcloud compute firewall-rules list --format="csv[no-heading](name)"'
+      'gcloud compute firewall-rules list --format="csv[no-heading](name)" '
+      f' --project={args.project}'
   )
   return_code, raw_subnets_output = run_command_for_value(
       command, 'Get All Firewall Rules', args

xpk/core/nodepool.py

@@ -37,6 +37,8 @@ CLOUD_PLATFORM_AUTH_SCOPE_URL = (
     '"https://www.googleapis.com/auth/cloud-platform"'
 )
 
+OLDER_PATHWAYS_CPU_NP_TO_DELETE = ['cpu-rm-np', 'cpu-proxy-np', 'cpu-user-np']
+
 
 def run_gke_node_pool_create_command(
     args, system, gke_node_pool_version
@@ -122,7 +124,10 @@ def run_gke_node_pool_create_command(
         args, system, existing_node_pool_names, desired_node_pool_names
     )
     for node_pool_name in existing_node_pool_names:
-      if node_pool_name.find(f'{args.cluster}-np-') != 0:
+      if (
+          node_pool_name.find(f'{args.cluster}-np-') != 0
+          and node_pool_name not in OLDER_PATHWAYS_CPU_NP_TO_DELETE
+      ):
         continue
 
       if node_pool_name in node_pools_to_delete:
@@ -283,28 +288,15 @@ def run_gke_node_pool_create_command(
       command += (
           ' --accelerator'
           f' type={system.gke_accelerator},count={str(system.chips_per_vm)},gpu-driver-version=latest'
-          ' --no-enable-autoupgrade '
-          f' --scopes={CLOUD_PLATFORM_AUTH_SCOPE_URL} --additional-node-network'
-          f' network={args.cluster}-net-1,subnetwork={subnet_prefix}-sub-1'
-          ' --additional-node-network'
-          f' network={args.cluster}-net-2,subnetwork={subnet_prefix}-sub-2'
-          ' --additional-node-network'
-          f' network={args.cluster}-net-3,subnetwork={subnet_prefix}-sub-3'
-          ' --additional-node-network'
-          f' network={args.cluster}-net-4,subnetwork={subnet_prefix}-sub-4'
+          f' --no-enable-autoupgrade --scopes={CLOUD_PLATFORM_AUTH_SCOPE_URL}'
       )
       if device_type == H100_MEGA_DEVICE_TYPE:
-        command += (
-            ' --additional-node-network'
-            f' network={args.cluster}-net-5,subnetwork={subnet_prefix}-sub-5'
-            ' --additional-node-network'
-            f' network={args.cluster}-net-6,subnetwork={subnet_prefix}-sub-6'
-            ' --additional-node-network'
-            f' network={args.cluster}-net-7,subnetwork={subnet_prefix}-sub-7'
-            ' --additional-node-network'
-            f' network={args.cluster}-net-8,subnetwork={subnet_prefix}-sub-8'
-            ' --max-pods-per-node=32'
-        )
+        for i in range(1, 9):
+          command += (
+              ' --additional-node-network'
+              f' network={args.cluster}-net-{i},subnetwork={subnet_prefix}-sub-{i}'
+          )
+        command += ' --max-pods-per-node=32'
     elif system.accelerator_type == AcceleratorType['CPU']:
       command += f' --num-nodes={system.vms_per_slice}'
       command += (
@@ -318,7 +310,7 @@ def run_gke_node_pool_create_command(
     create_commands.append(command)
     create_task_names.append(task)
 
-  desired_pw_cpu_node_pools = ['cpu-user-np', 'cpu-rm-np', 'cpu-proxy-np']
+  desired_pw_cpu_node_pools = ['cpu-np']
   if args.enable_pathways:
     # Pathways needs CPU nodepools in addition to TPU nodepools
     for node_pool_name in desired_pw_cpu_node_pools:
@@ -368,11 +360,9 @@ def get_node_pools_to_delete(
   check_resource, is_requested_resource_in_cluster = check_cluster_resources(
       args, system
   )
-  for existing_node_pool_name in existing_node_pool_names:
-    # Deletion logic would leave behind any Pathways CPU nodepools.
-    if existing_node_pool_name.find(f'{args.cluster}-np-') != 0:
-      continue
+  xpk_print('Existing node pool names ', existing_node_pool_names)
 
+  for existing_node_pool_name in existing_node_pool_names:
     # Nodepools will be deleted in two scenarios:
     # Scenario 1: Cluster exists with 3 nodepools of 'x' device_type/gke_accelerator and now we are updating
     # the cluster to 2 nodepools of 'x' device_type/gke_accelerator. In this case, we will delete
@@ -380,6 +370,18 @@ def get_node_pools_to_delete(
     # Scenario 2: Cluster exists with 2 nodepools of 'x' device_type/gke_accelerator and now we are updating
     # the cluster to 2 nodepools of 'y' device_type/gke_accelerator. In this case, we will delete
     # '{args.cluster}-np-0' and '{args.cluster}-np-1' from the cluster.
+    # Scenario 3: Deletes older Pathways CPU nodepools named cpu-rm-np, cpu-proxy-np and cpu-user-np
+
+    if existing_node_pool_name in OLDER_PATHWAYS_CPU_NP_TO_DELETE:
+      node_pools_to_delete.append(existing_node_pool_name)
+      xpk_print(
+          'Upgrading Pathways version on the cluster. Deleting older pathways'
+          ' nodepool ',
+          existing_node_pool_name,
+      )
+
+    if existing_node_pool_name.find(f'{args.cluster}-np-') != 0:
+      continue
     if existing_node_pool_name not in desired_node_pool_names or (
         check_resource and not is_requested_resource_in_cluster
     ):