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

xpk/core/mtc.py

@@ -0,0 +1,195 @@
+"""
+Copyright 2024 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 requests
+import yaml
+
+from ..core.cluster import JOBSET_VERSION
+from ..core.cluster import setup_k8s_env
+from ..utils import templates
+from ..utils.console import xpk_exit
+from ..utils.console import xpk_print
+from ..utils.kubectl import apply_kubectl_manifest
+
+
+MTC_CPC_PATH = "/../templates/mtc-cpc.yaml"
+
+
+def create_mtc_cpc(
+    mtc_gcs_bucket: str,
+    mtc_machine_type: str,
+    mtc_toleration_key: str,
+    mtc_ramdisk_size: str,
+) -> dict:
+  """Create MTC Checkpoint Configuration.
+
+  Args:
+    mtc_gcs_bucket: GCS bucket for MTC
+    mtc_machine_type: Machine type for MTC
+    mtc_toleration_key: Toleration key for MTC
+    mtc_ramdisk_size: Ramdisk size for MTC
+
+  Returns:
+    MTC Checkpoint Configuration
+  """
+  data = templates.load(MTC_CPC_PATH)
+
+  data["spec"]["cloudStorageBucketName"] = mtc_gcs_bucket
+  data["spec"]["nodeSelector"][
+      "node.kubernetes.io/instance-type"
+  ] = mtc_machine_type
+  data["spec"]["tolerations"][0]["key"] = mtc_toleration_key
+  data["spec"]["inMemoryVolumeSize"] = mtc_ramdisk_size
+
+  return data
+
+
+def install_mtc_on_cluster(args, system) -> int:
+  """Install MTC on the cluster.
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system related information.
+
+  Returns:
+    return code of the command.
+  """
+  if args.mtc_gcs_bucket is None:
+    xpk_print("MTC GCS bucket is required.")
+    xpk_exit(1)
+  if args.mtc_gcs_bucket.startswith("gs://"):
+    args.mtc_gcs_bucket = args.mtc_gcs_bucket.replace("gs://", "")
+
+  if args.mtc_ramdisk_size is None:
+    xpk_print("MTC ramdisk size is required.")
+    xpk_exit(1)
+
+  if args.mtc_toleration_key is None:
+    args.mtc_toleration_key = "google.com/tpu"
+
+  k8s_api_client = setup_k8s_env(args)
+  jobset_manifest = update_jobset_manifest()
+  if jobset_manifest is None:
+    xpk_print(
+        "Updated jobset manifest is empty, not updating the jobset controller."
+    )
+
+  xpk_print("Applying Jobset with MTC Configuration")
+  return_code = apply_kubectl_manifest(k8s_api_client, [jobset_manifest])
+  if return_code != 0:
+    return return_code
+
+  mtc_checkpoint_configuration_crd_data = create_mtc_cpc(
+      args.mtc_gcs_bucket,
+      system.gce_machine_type,
+      args.mtc_toleration_key,
+      args.mtc_ramdisk_size,
+  )
+  xpk_print("Applying MTC Checkpoint Configuration")
+  return_code = apply_kubectl_manifest(
+      k8s_api_client, [mtc_checkpoint_configuration_crd_data]
+  )
+
+  return return_code
+
+
+def update_jobset_manifest():
+  """Update the jobset manifest to increase the resources for the jobset controller manager.
+
+  Returns:
+    The updated jobset manifest.
+  """
+  manifest_url = f"https://github.com/kubernetes-sigs/jobset/releases/download/{JOBSET_VERSION}/manifests.yaml"
+  manifest_content = None
+  # Fetch the manifest content
+  try:
+    response = requests.get(manifest_url, timeout=10)
+    response.raise_for_status()  # Raise an exception for HTTP errors
+    manifest_content = response.text
+  except requests.exceptions.Timeout as e:
+    xpk_print(f"Error: Request to {manifest_url} after 10 seconds: {e}")
+    xpk_exit(1)
+  except requests.exceptions.RequestException as e:
+    xpk_print(f"Error fetching manifest from {manifest_url}: {e}")
+    xpk_exit(1)
+
+  if manifest_content is None:
+    xpk_print("Manifest content not found.")
+    xpk_exit(1)
+
+  # Load all YAML documents from the manifest
+  yaml_data_list = list(yaml.safe_load_all(manifest_content))
+  # Iterate through the yaml_data to find the Deployment for
+  # jobset-controller-manager
+  update_manifest = False
+  for yaml_data in yaml_data_list:
+    if (
+        yaml_data
+        and yaml_data.get("apiVersion") == "apps/v1"
+        and yaml_data.get("kind") == "Deployment"
+        and yaml_data.get("metadata", {}).get("name")
+        == "jobset-controller-manager"
+    ):
+      # Found the Deployment, now modify the resources
+      containers = yaml_data["spec"]["template"]["spec"]["containers"]
+      for container in containers:
+        if container["name"] == "manager":
+          # Update resource limits and requests
+          current_cpu_request = (
+              container["resources"].get("requests", {}).get("cpu", "0m")
+          )
+          current_memory_request = (
+              container["resources"].get("requests", {}).get("memory", "0Mi")
+          )
+          current_memory_limit = (
+              container["resources"].get("limits", {}).get("memory", "0Mi")
+          )
+
+          # Define new values for comparison
+          new_cpu_request = "1000m"
+          new_memory_request = "1Gi"
+          new_memory_limit = "2Gi"
+
+          if parse_resource_value(current_cpu_request) < parse_resource_value(
+              new_cpu_request
+          ):
+            container["resources"]["requests"]["cpu"] = new_cpu_request
+            update_manifest = True
+          if parse_resource_value(
+              current_memory_request
+          ) < parse_resource_value(new_memory_request):
+            container["resources"]["requests"]["memory"] = new_memory_request
+            update_manifest = True
+          if parse_resource_value(current_memory_limit) < parse_resource_value(
+              new_memory_limit
+          ):
+            container["resources"]["limits"]["memory"] = new_memory_limit
+            update_manifest = True
+          break
+      if update_manifest:
+        xpk_print("Jobset controller updation required.")
+        return yaml_data
+  xpk_print("Jobset controller no updation required.")
+
+
+def parse_resource_value(value) -> int:
+  if value.endswith("m"):
+    return int(value[:-1])
+  if value.endswith("Mi"):
+    return int(value[:-2])
+  if value.endswith("Gi"):
+    return int(value[:-2]) * 1024
+  return int(value)

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

@@ -14,12 +14,10 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from ..utils.console import xpk_print
+from ..utils.console import xpk_exit, 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,40 @@ 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 get_cluster_subnetworks(args) -> list[str]:
+  """Gets the list of cluster networks.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    list[str]: list of cluster networks
+  """
+  command = 'kubectl get GKENetworkParamSet'
+  return_code, stdout = run_command_for_value(
+      command, 'Get Cluster Networks', args
+  )
+  if return_code != 0:
+    xpk_print('GKE Cluster Get NetworkParamSet failed')
+    xpk_exit(return_code)
+
+  networks = [line.split()[0] for line in stdout.splitlines()][1:]
+
+  return networks
+
+
+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 +323,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 +376,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
     ):