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

xpk/core/kueue.py

@@ -15,18 +15,24 @@ limitations under the License.
 """
 
 from argparse import Namespace
-from packaging.version import Version
+
 import packaging
+from packaging.version import Version
+
+from ..utils.console import xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
-from ..utils.console import xpk_print, xpk_exit
-from .commands import run_command_with_updates, run_command_with_updates_retry, run_command_for_value
-from .core import (
-    AutoprovisioningConfig,
+from .commands import (
+    run_command_for_value,
+    run_command_with_updates,
+    run_command_with_updates_retry,
+)
+from .pathways import add_pw_resource_flavors, add_pw_resources_to_kueue
+from .resources import AutoprovisioningConfig
+from .scheduling import (
     create_accelerator_label,
     create_machine_label,
     get_total_chips_requested_from_args,
 )
-from .pathways import add_pw_resource_flavors, add_pw_resources_to_kueue
 from .system_characteristics import (
     AcceleratorTypeToAcceleratorCharacteristics,
     SystemCharacteristics,

xpk/core/monitoring.py

@@ -0,0 +1,134 @@
+"""
+Copyright 2025 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.
+"""
+
+from ..utils.console import xpk_print
+from .commands import run_command_for_value
+
+
+def get_gke_dashboard(args, dashboard_filter) -> tuple[bool, str | None]:
+  """Get the identifier of GKE dashboard deployed in the project.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    bool:
+      True if 'gcloud monitoring dashboards list' returned an error or
+      multiple dashboards with same filter exist in the project,
+      False otherwise.
+    str:
+      identifier of dashboard if deployed in project,
+      None otherwise.
+  """
+  command = (
+      'gcloud monitoring dashboards list'
+      f' --project={args.project} --filter="{dashboard_filter}"'
+      ' --format="value(name)" --verbosity=error'
+  )
+
+  return_code, return_value = run_command_for_value(
+      command, 'GKE Dashboard List', args
+  )
+
+  if return_code != 0:
+    xpk_print(
+        f'GKE Dashboard List request returned ERROR {return_code}. If there is'
+        ' a permissions error, please check'
+        ' https://github.com/google/xpk/blob/main/README.md#roles-needed-based-on-permission-errors'
+        ' for possible solutions.'
+    )
+    return True, None
+
+  if not return_value:
+    xpk_print(
+        f'No dashboard with {dashboard_filter} found in the'
+        f' project:{args.project}.'
+    )
+    return False, return_value
+
+  dashboards = return_value.strip().split('\n')
+  if len(dashboards) > 1:
+    xpk_print(
+        f'Multiple dashboards with same {dashboard_filter} exist in the'
+        f' project:{args.project}. Delete all but one dashboard deployed using'
+        ' https://github.com/google/cloud-tpu-monitoring-debugging.'
+    )
+    return True, None
+
+  if dashboards[0]:
+    return False, dashboards[0].strip().split('/')[-1]
+
+  return True, None
+
+
+def get_gke_outlier_dashboard(args) -> str | None:
+  """Get the identifier of GKE outlier dashboard deployed in the project.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    str:
+      identifier of outlier dashboard if deployed in project,
+      None otherwise.
+  """
+  outlier_dashboard_filter = "displayName:'GKE - TPU Monitoring Dashboard'"
+  is_error, dashboard_id = get_gke_dashboard(args, outlier_dashboard_filter)
+
+  # 'gcloud monitoring dashboards list' returned an error or multiple dashboards with same filter exist in the project
+  if is_error:
+    return None
+
+  # 'gcloud monitoring dashboards list' succeeded but no dashboard for the filter exist in the project
+  if not is_error and not dashboard_id:
+    xpk_print(
+        'Follow https://github.com/google/cloud-tpu-monitoring-debugging to'
+        ' deploy monitoring dashboard to view statistics and outlier mode of'
+        ' GKE metrics.'
+    )
+    return None
+
+  return str(dashboard_id)
+
+
+def get_gke_debugging_dashboard(args) -> str | None:
+  """Get the identifier of GKE debugging dashboard deployed in the project.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    str:
+      identifier of debugging dashboard if deployed in project,
+      None otherwise.
+  """
+  debugging_dashboard_filter = "displayName:'GKE - TPU Logging Dashboard'"
+  is_error, dashboard_id = get_gke_dashboard(args, debugging_dashboard_filter)
+
+  # 'gcloud monitoring dashboards list' returned an error or multiple dashboards with same filter exist in the project
+  if is_error:
+    return None
+
+  # 'gcloud monitoring dashboards list' succeeded but no dashboard for the filter exist in the project
+  if not is_error and not dashboard_id:
+    xpk_print(
+        'Follow https://github.com/google/cloud-tpu-monitoring-debugging to'
+        ' deploy debugging dashboard to view stack traces collected in Cloud'
+        ' Logging.'
+    )
+    return None
+
+  return str(dashboard_id)

xpk/core/nap.py

@@ -14,29 +14,31 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from ..core.core import (
+from ..utils.console import xpk_print
+from ..utils.file import write_tmp_file
+from ..utils.objects import get_value_from_map
+from .capacity import (
     AUTOPROVISIONING_CONFIG_VALUE,
     CAPACITY_TYPE_CONFIG_KEY,
-    CLUSTER_METADATA_CONFIGMAP,
-    CLUSTER_RESOURCES_CONFIGMAP,
     RESERVATION_CONFIG_KEY,
-    AutoprovisioningConfig,
     CapacityType,
-    get_all_nodepools_programmatic,
     get_capacity_node_selectors_from_capacity_type,
     get_capacity_type,
-    get_cluster_configmap,
-    get_total_chips_requested_from_args,
     verify_reservation_exists,
-    zone_to_region,
 )
-from ..utils.objects import get_value_from_map
-from ..utils.file import write_tmp_file
-from ..utils.console import xpk_print
 from .commands import run_command_with_updates, run_commands
+from .gcloud_context import zone_to_region
+from .nodepool import get_all_nodepools_programmatic
+from .resources import (
+    CLUSTER_METADATA_CONFIGMAP,
+    CLUSTER_RESOURCES_CONFIGMAP,
+    AutoprovisioningConfig,
+    get_cluster_configmap,
+)
+from .scheduling import get_total_chips_requested_from_args
 from .system_characteristics import AcceleratorType, SystemCharacteristics
 
-autoprovisioning_config_file = """
+AUTOPROVISIONING_CONFIG_FILE = """
 management:
   autoRepair: true
   autoUpgrade: true
@@ -44,8 +46,7 @@ autoprovisioningLocations:
   {zones}
 {resource_limits}
 """
-
-autoprovisioning_resource_limits = """
+AUTOPROVISIONING_RESOURCE_LIMITS = """
 resourceLimits:
 - resourceType: 'cpu'
   {cpu_limits}
@@ -53,8 +54,7 @@ resourceLimits:
   {memory_limits}
 {custom_resource_type}
 """
-
-autoprovisioning_custom_resource_type = """
+AUTOPROVISIONING_CUSTOM_RESOURCE_TYPE = """
 - resourceType: {resource_type}
   minimum: {minimum}
   maximum: {maximum}
@@ -218,19 +218,19 @@ def create_autoprovisioning_config(
       ' small, rescaling will not work well.'
   )
 
-  custom_resource_string = autoprovisioning_custom_resource_type.format(
+  custom_resource_string = AUTOPROVISIONING_CUSTOM_RESOURCE_TYPE.format(
       resource_type=system.gke_accelerator,
       minimum=minimum,
       maximum=maximum,
   )
 
-  resource_limits = autoprovisioning_resource_limits.format(
+  resource_limits = AUTOPROVISIONING_RESOURCE_LIMITS.format(
       cpu_limits=cpu_limits,
       memory_limits=memory_limits,
       custom_resource_type=custom_resource_string,
   )
 
-  yml_string = autoprovisioning_config_file.format(
+  yml_string = AUTOPROVISIONING_CONFIG_FILE.format(
       resource_limits=resource_limits,
       zones=f'- {args.zone}',
   )
@@ -266,7 +266,7 @@ def is_autoprovisioning_enabled(
     return False, 0
 
   return_code, autoprovisioning_value = get_value_from_map(
-      system.gke_accelerator, cluster_config_map
+      system.gke_accelerator, cluster_config_map, verbose=False
   )
   if return_code != 0:
     xpk_print(
@@ -347,3 +347,15 @@ def get_autoprovisioning_node_selector_args(args) -> tuple[str, int]:
     return node_selector_args, return_code
 
   return node_selector_args, return_code
+
+
+def get_cluster_provisioner(args) -> str:
+  metadata_configmap_name = f'{args.cluster}-{CLUSTER_METADATA_CONFIGMAP}'
+  cluster_config_map = get_cluster_configmap(args, metadata_configmap_name)
+  cluster_provisioner = 'gcloud'
+  if not cluster_config_map is None:
+    provisioner = cluster_config_map.get('provisioner')
+    if not provisioner is None:
+      cluster_provisioner = provisioner
+  xpk_print(f'Cluster provisioner: {cluster_provisioner}')
+  return cluster_provisioner

xpk/core/network.py

@@ -0,0 +1,377 @@
+"""
+Copyright 2025 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.
+"""
+
+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 = """
+apiVersion: networking.gke.io/v1
+kind: Network
+metadata:
+  name: vpc1
+spec:
+  parametersRef:
+    group: networking.gke.io
+    kind: GKENetworkParamSet
+    name: vpc1
+  type: Device
+---
+apiVersion: networking.gke.io/v1
+kind: Network
+metadata:
+  name: vpc2
+spec:
+  parametersRef:
+    group: networking.gke.io
+    kind: GKENetworkParamSet
+    name: vpc2
+  type: Device
+---
+apiVersion: networking.gke.io/v1
+kind: Network
+metadata:
+  name: vpc3
+spec:
+  parametersRef:
+    group: networking.gke.io
+    kind: GKENetworkParamSet
+    name: vpc3
+  type: Device
+---
+apiVersion: networking.gke.io/v1
+kind: Network
+metadata:
+  name: vpc4
+spec:
+  parametersRef:
+    group: networking.gke.io
+    kind: GKENetworkParamSet
+    name: vpc4
+  type: Device
+---
+apiVersion: networking.gke.io/v1
+kind: GKENetworkParamSet
+metadata:
+  name: vpc1
+spec:
+  vpc: {cluster_name}-net-1
+  vpcSubnet: {cluster_name}-sub-1
+  deviceMode: NetDevice
+---
+apiVersion: networking.gke.io/v1
+kind: GKENetworkParamSet
+metadata:
+  name: vpc2
+spec:
+  vpc: {cluster_name}-net-2
+  vpcSubnet: {cluster_name}-sub-2
+  deviceMode: NetDevice
+---
+apiVersion: networking.gke.io/v1
+kind: GKENetworkParamSet
+metadata:
+  name: vpc3
+spec:
+  vpc: {cluster_name}-net-3
+  vpcSubnet: {cluster_name}-sub-3
+  deviceMode: NetDevice
+---
+apiVersion: networking.gke.io/v1
+kind: GKENetworkParamSet
+metadata:
+  name: vpc4
+spec:
+  vpc: {cluster_name}-net-4
+  vpcSubnet: {cluster_name}-sub-4
+  deviceMode: NetDevice
+"""
+
+
+def create_cluster_network(args, index) -> int:
+  """Create one GKE Cluster network.
+
+  Args:
+    args: user provided arguments for running the command.
+    index: index number for the network to be created.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  existing_network_names, return_code = get_all_networks_programmatic(args)
+  if return_code > 0:
+    xpk_print('Listing all networks failed!')
+    return return_code
+
+  network_name = f'{args.cluster}-net-{index}'
+  if network_name not in existing_network_names:
+    command = (
+        f'gcloud compute --project={args.project}'
+        f' networks create {network_name}'
+        ' --subnet-mode=custom --mtu=8244'
+    )
+    return_code = run_command_with_updates(
+        command, 'Create Cluster Network', args, verbose=False
+    )
+
+    if return_code != 0:
+      xpk_print(f'Create Cluster Network request returned ERROR {return_code}')
+      return 1
+  else:
+    xpk_print(f'Reusing existing network {network_name}')
+
+  return 0
+
+
+def create_cluster_subnet(args, index) -> int:
+  """Create one GKE Cluster subnet.
+
+  Args:
+    args: user provided arguments for running the command.
+    index: index number for the subnet to be created.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  existing_subnet_names, return_code = get_all_subnets_programmatic(args)
+  if return_code > 0:
+    xpk_print('Listing all subnets failed!')
+    return return_code
+  subnet_name = f'{args.cluster}-{zone_to_region(args.zone)}-sub-{index}'
+  if subnet_name not in existing_subnet_names:
+    command = (
+        f'gcloud compute --project={args.project}'
+        f' networks subnets create {subnet_name}'
+        f' --network={args.cluster}-net-{index}'
+        f' --region={zone_to_region(args.zone)} --range=192.168.{index}.0/24'
+    )
+    return_code = run_command_with_updates(
+        command, 'Create Cluster Subnet', args, verbose=False
+    )
+
+    if return_code != 0:
+      xpk_print(f'Create Cluster Subnet request returned ERROR {return_code}')
+      return 1
+  else:
+    xpk_print(f'Reusing existing subnet {subnet_name}')
+
+  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.
+
+  Args:
+    args: user provided arguments for running the command.
+    index: index number for the firewall rule to be created.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  existing_firewall_rules_names, return_code = (
+      get_all_firewall_rules_programmatic(args)
+  )
+  if return_code > 0:
+    xpk_print('Listing all firewall rules failed!')
+    return return_code
+  firewall_rule_name = f'{args.cluster}-internal-{index}'
+  if firewall_rule_name not in existing_firewall_rules_names:
+    command = (
+        f'gcloud compute --project={args.project} firewall-rules create'
+        f' {firewall_rule_name} --network={args.cluster}-net-{index} --action=ALLOW'
+        ' --rules=tcp:0-65535,udp:0-65535,icmp --source-ranges=192.168.0.0/16'
+    )
+    return_code = run_command_with_updates(
+        command, 'Create Cluster Firewall Rule', args, verbose=False
+    )
+
+    if return_code != 0:
+      xpk_print(
+          f'Create Cluster Firewall Rule request returned ERROR {return_code}'
+      )
+      return 1
+  else:
+    xpk_print(f'Reusing existing firewall rule {firewall_rule_name}')
+  return 0
+
+
+def create_cluster_network_config(args) -> int:
+  """Run the Create GKE Cluster Network Config request.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  yml_string = CLUSTER_NETWORK_YAML.format(cluster_name=args.cluster)
+  tmp = write_tmp_file(yml_string)
+  command = f'kubectl apply -f {str(tmp.file.name)}'
+
+  return_code = run_command_with_updates(
+      command, 'GKE Cluster Create Network Config', args
+  )
+  if return_code != 0:
+    xpk_print(
+        f'GKE Cluster Create ConfigMap request returned ERROR {return_code}'
+    )
+    return 1
+
+  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.
+
+  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):
+    return_code = create_cluster_network(args, i)
+    if return_code != 0:
+      return 1
+    return_code = create_cluster_subnet(args, i)
+    if return_code != 0:
+      return 1
+    return_code = create_cluster_firewall_rule(args, i)
+    if return_code != 0:
+      return 1
+  return 0
+
+
+def delete_cluster_subnets(args) -> int:
+  """Delete GKE Cluster subnets.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  existing_subnet_names, return_code = get_all_subnets_programmatic(args)
+  if return_code > 0:
+    xpk_print('Listing all subnets failed!')
+    return return_code
+
+  for subnet_name in existing_subnet_names:
+    command = (
+        f'gcloud compute networks subnets delete {subnet_name}'
+        f' --region={zone_to_region(args.zone)} --project={args.project} --quiet'
+    )
+
+    return_code = run_command_with_updates(
+        command, 'Delete Cluster Subnet', args, verbose=False
+    )
+
+    if return_code != 0:
+      xpk_print(f'Delete Cluster Subnet request returned ERROR {return_code}')
+      return 1
+    else:
+      xpk_print(f'Deleted existing subnet {subnet_name}')
+
+  return 0
+
+
+def get_all_networks_programmatic(args) -> tuple[list[str], int]:
+  """Gets all the networks associated with project .
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    List of networks and 0 if successful and 1 otherwise.
+  """
+  command = 'gcloud compute networks list --format="csv[no-heading](name)"'
+  return_code, raw_network_output = run_command_for_value(
+      command, 'Get All Networks', args
+  )
+  if return_code != 0:
+    xpk_print(f'Get All Networks returned ERROR {return_code}')
+    return [], 1
+
+  return raw_network_output.splitlines(), 0
+
+
+def get_all_subnets_programmatic(args) -> tuple[list[str], int]:
+  """Gets all the subnets associated with the project.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    List of subnets and 0 if successful and 1 otherwise.
+  """
+  subnet_name_filter = f'{args.cluster}-{zone_to_region(args.zone)}-sub-*'
+
+  command = (
+      'gcloud compute networks subnets list'
+      f' --filter=name~"{subnet_name_filter}" --project={args.project}'
+  )
+  return_code, raw_subnets_output = run_command_for_value(
+      command, 'Get All Subnets', args
+  )
+  if return_code != 0:
+    xpk_print(f'Get All Subnets returned ERROR {return_code}')
+    return [], 1
+
+  all_outputs = raw_subnets_output.splitlines()
+  all_networks = [
+      all_outputs[i].split(' ')[0] for i in range(1, len(all_outputs))
+  ]
+  return all_networks, 0
+
+
+def get_all_firewall_rules_programmatic(args) -> tuple[list[str], int]:
+  """Gets all the firewall rules associated with the project.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    List of firewall rules and 0 if successful and 1 otherwise.
+  """
+  command = (
+      'gcloud compute firewall-rules list --format="csv[no-heading](name)"'
+  )
+  return_code, raw_subnets_output = run_command_for_value(
+      command, 'Get All Firewall Rules', args
+  )
+  if return_code != 0:
+    xpk_print(f'Get All Firewall Rules returned ERROR {return_code}')
+    return [], 1
+
+  return raw_subnets_output.splitlines(), 0