xpk 0.13.0__py3-none-any.whl → 0.14.1__py3-none-any.whl

xpk/core/cluster.py

@@ -30,15 +30,14 @@ from .commands import (
 )
 from .gcloud_context import (
     add_zone_and_project,
-    get_gke_server_config,
+    get_cluster_location,
     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.8.0'
-PATHWAYS_JOB_VERSION = 'v0.1.2'
+PATHWAYS_JOB_VERSION = 'v0.1.3'
 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'
@@ -66,7 +65,7 @@ def set_jobset_on_cluster(args) -> int:
       f' -f https://github.com/kubernetes-sigs/jobset/releases/download/{JOBSET_VERSION}/manifests.yaml'
   )
   task = f'Install Jobset on {args.cluster}'
-  return_code = run_command_with_updates_retry(command, task, args)
+  return_code = run_command_with_updates_retry(command, task)
 
   if return_code != 0:
     xpk_print(f'{task} returned with ERROR {return_code}.\n')
@@ -95,7 +94,7 @@ def set_pathways_job_on_cluster(args) -> int:
       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)
+  return_code = run_command_with_updates_retry(command, task)
 
   if return_code != 0:
     xpk_print(f'{task} returned with ERROR {return_code}.\n')
@@ -110,11 +109,10 @@ def set_pathways_job_on_cluster(args) -> int:
   return return_code
 
 
-def install_nccl_on_cluster(args, system: SystemCharacteristics) -> int:
+def install_nccl_on_cluster(system: SystemCharacteristics) -> int:
   """Install NCCL plugin on the cluster.
 
   Args:
-    args: user provided arguments for running the command.
     system: system characteristics.
 
   Returns:
@@ -128,7 +126,7 @@ def install_nccl_on_cluster(args, system: SystemCharacteristics) -> int:
     command = f'kubectl apply -f {INSTALLER_NCCL_TCPXO}'
 
   return_code = run_command_with_updates(
-      command, 'Install NCCL Plugin On Cluster', args
+      command, 'Install NCCL Plugin On Cluster'
   )
 
   if return_code != 0:
@@ -141,7 +139,7 @@ def install_nccl_on_cluster(args, system: SystemCharacteristics) -> int:
     command = f'kubectl apply -f {CONFIG_NCCL_TCPX}'
 
     return_code = run_command_with_updates(
-        command, 'Install NCCL Config On Cluster', args
+        command, 'Install NCCL Config On Cluster'
     )
 
     if return_code != 0:
@@ -153,19 +151,14 @@ def install_nccl_on_cluster(args, system: SystemCharacteristics) -> int:
   return 0
 
 
-def disable_mglru_on_cluster(args) -> int:
+def disable_mglru_on_cluster() -> 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
-  )
+  return_code = run_command_with_updates(command, 'Disable MGLRU On Cluster')
 
   if return_code != 0:
     xpk_print('Disablig MGLRU On Cluster request returned ERROR')
@@ -174,11 +167,10 @@ def disable_mglru_on_cluster(args) -> int:
   return 0
 
 
-def install_nri_on_cluster(args) -> int:
+def install_nri_on_cluster() -> int:
   """Install NRI Device Injector on the cluster.
 
   Args:
-    args: user provided arguments for running the command.
     system: system characteristics.
 
   Returns:
@@ -186,7 +178,7 @@ def install_nri_on_cluster(args) -> int:
   """
   command = f'kubectl apply -f {NRI_DEVICE_INJECTOR}'
   return_code = run_command_with_updates(
-      command, 'Install NRI Device Injector On Cluster', args
+      command, 'Install NRI Device Injector On Cluster'
   )
 
   if return_code != 0:
@@ -199,12 +191,9 @@ def install_nri_on_cluster(args) -> int:
   return 0
 
 
-def get_cluster_nodes_info(args) -> list[dict]:
+def get_cluster_nodes_info() -> 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.
   """
@@ -213,7 +202,6 @@ def get_cluster_nodes_info(args) -> list[dict]:
   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)
@@ -221,9 +209,9 @@ def get_cluster_nodes_info(args) -> list[dict]:
   return data['items']
 
 
-def count_nodes_on_cluster(args, system: SystemCharacteristics) -> int:
+def count_nodes_on_cluster(system: SystemCharacteristics) -> int:
   """Count cluster nodes by accelerator type"""
-  nodes_info = get_cluster_nodes_info(args)
+  nodes_info = get_cluster_nodes_info()
   accelerators = [
       node['metadata']['labels']['cloud.google.com/gke-accelerator']
       for node in nodes_info
@@ -243,12 +231,11 @@ def get_cluster_network(args) -> str:
   xpk_print("Getting cluster's VPC network...")
   cluster_network_cmd = (
       'gcloud container clusters describe'
-      f' {args.cluster} --zone={zone_to_region(args.zone)} --project={args.project} --format="value(network)"'
+      f' {args.cluster} --location={get_cluster_location(args.project, args.cluster, args.zone)} --project={args.project} --format="value(network)"'
   )
   err_code, val = run_command_for_value(
       command=cluster_network_cmd,
       task='Get network cluster is in',
-      global_args=args,
   )
   if err_code != 0:
     xpk_exit(err_code)
@@ -354,14 +341,13 @@ def is_driver_enabled_on_cluster(
   """
   command = (
       f'gcloud container clusters describe {args.cluster}'
-      f' --project={args.project} --region={zone_to_region(args.zone)}'
+      f' --project={args.project} --location={get_cluster_location(args.project, args.cluster, args.zone)}'
       f' --format="value(addonsConfig.{driver}Config.{config_key})"'
   )
   return_code, driver_enabled = run_command_for_value(
       command,
       f"Checks if {driver} driver's {config_key} is enabled in cluster"
       ' describe.',
-      args,
   )
   if return_code != 0:
     xpk_exit(return_code)
@@ -382,14 +368,12 @@ def update_gke_cluster_with_addon(args, addon: str) -> int:
   """
   command = (
       'gcloud container clusters update'
-      f' {args.cluster} --project={args.project}'
-      f' --region={zone_to_region(args.zone)}'
-      f' --update-addons {addon}=ENABLED'
-      ' --quiet'
+      f' {args.cluster} --project={args.project} --location={get_cluster_location(args.project, args.cluster, args.zone)} --update-addons'
+      f' {addon}=ENABLED --quiet'
   )
   xpk_print(f'Updating GKE cluster to enable {addon}, may take a while!')
   return_code = run_command_with_updates(
-      command, f'GKE Cluster Update to enable {addon}', args
+      command, f'GKE Cluster Update to enable {addon}'
   )
   if return_code != 0:
     xpk_print(f'GKE Cluster Update request returned ERROR {return_code}')
@@ -408,11 +392,12 @@ def get_all_clusters_programmatic(args) -> tuple[list[str], int]:
   """
   command = (
       'gcloud container clusters list'
-      f' --project={args.project} --region={zone_to_region(args.zone)}'
+      f' --project={args.project} '
+      f'--filter=location~"{zone_to_region(args.zone)}.*"'
       ' --format="csv[no-heading](name)"'
   )
   return_code, raw_cluster_output = run_command_for_value(
-      command, 'Find if Cluster Exists', args
+      command, 'Find if Cluster Exists'
   )
   if return_code != 0:
     xpk_print(f'Find if Cluster Exists returned ERROR {return_code}')
@@ -578,34 +563,6 @@ def create_role_binding(sa: str, role_name: str) -> None:
       xpk_exit(1)
 
 
-def update_gke_cluster_with_clouddns(args) -> int:
-  """Run the GKE cluster update command for existing clusters and enable CloudDNS.
-
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    0 if successful and 1 otherwise.
-  """
-  command = (
-      'gcloud container clusters update'
-      f' {args.cluster} --project={args.project}'
-      f' --region={zone_to_region(args.zone)}'
-      ' --cluster-dns=clouddns'
-      ' --cluster-dns-scope=vpc'
-      f' --cluster-dns-domain={args.cluster}-domain'
-      ' --quiet'
-  )
-  xpk_print('Updating GKE cluster to use Cloud DNS, may take a while!')
-  return_code = run_command_with_updates(
-      command, 'GKE Cluster Update to enable Cloud DNS', args
-  )
-  if return_code != 0:
-    xpk_print(f'GKE Cluster Update request returned ERROR {return_code}')
-    return 1
-  return 0
-
-
 def update_gke_cluster_with_workload_identity_enabled(args) -> int:
   """Run the GKE cluster update command for existing cluster and enable Workload Identity Federation.
   Args:
@@ -615,9 +572,7 @@ def update_gke_cluster_with_workload_identity_enabled(args) -> int:
   """
   command = (
       'gcloud container clusters update'
-      f' {args.cluster} --project={args.project}'
-      f' --region={zone_to_region(args.zone)}'
-      f' --workload-pool={args.project}.svc.id.goog'
+      f' {args.cluster} --project={args.project} --location={get_cluster_location(args.project, args.cluster, args.zone)} --workload-pool={args.project}.svc.id.goog'
       ' --quiet'
   )
   xpk_print(
@@ -625,7 +580,7 @@ def update_gke_cluster_with_workload_identity_enabled(args) -> int:
       ' while!'
   )
   return_code = run_command_with_updates(
-      command, 'GKE Cluster Update to enable Workload Identity Federation', args
+      command, 'GKE Cluster Update to enable Workload Identity Federation'
   )
   if return_code != 0:
     xpk_print(f'GKE Cluster Update request returned ERROR {return_code}')
@@ -642,16 +597,14 @@ def update_gke_cluster_with_gcsfuse_driver_enabled(args) -> int:
   """
   command = (
       'gcloud container clusters update'
-      f' {args.cluster} --project={args.project}'
-      f' --region={zone_to_region(args.zone)}'
-      ' --update-addons GcsFuseCsiDriver=ENABLED'
-      ' --quiet'
+      f' {args.cluster} --project={args.project} --location={get_cluster_location(args.project, args.cluster, args.zone)} --update-addons'
+      ' GcsFuseCsiDriver=ENABLED --quiet'
   )
   xpk_print(
       'Updating GKE cluster to enable GCSFuse CSI driver, may take a while!'
   )
   return_code = run_command_with_updates(
-      command, 'GKE Cluster Update to enable GCSFuse CSI driver', args
+      command, 'GKE Cluster Update to enable GCSFuse CSI driver'
   )
   if return_code != 0:
     xpk_print(f'GKE Cluster Update request returned ERROR {return_code}')
@@ -668,16 +621,14 @@ def update_gke_cluster_with_lustre_driver_enabled(args) -> int:
   """
   command = (
       'gcloud container clusters update'
-      f' {args.cluster} --project={args.project}'
-      f' --region={zone_to_region(args.zone)}'
-      ' --enable-legacy-lustre-port'
+      f' {args.cluster} --project={args.project} --location={get_cluster_location(args.project, args.cluster, args.zone)} --enable-legacy-lustre-port'
       ' --quiet'
   )
   xpk_print(
       'Updating GKE cluster to enable Lustre CSI driver, may take a while!'
   )
   return_code = run_command_with_updates(
-      command, 'GKE Cluster Update to enable Lustre CSI driver', args
+      command, 'GKE Cluster Update to enable Lustre CSI driver'
   )
   if return_code != 0:
     xpk_print(f'GKE Cluster Update request returned ERROR {return_code}')
@@ -685,63 +636,6 @@ def update_gke_cluster_with_lustre_driver_enabled(args) -> int:
   return 0
 
 
-def upgrade_gke_control_plane_version(args, default_rapid_gke_version) -> int:
-  """Upgrade GKE cluster's control plane version before updating nodepools to use CloudDNS.
-
-  Args:
-    args: user provided arguments for running the command.
-    default_rapid_gke_version: Rapid default version for the upgrade.
-
-  Returns:
-    0 if successful and 1 otherwise.
-  """
-  command = (
-      'gcloud container clusters upgrade'
-      f' {args.cluster} --project={args.project}'
-      f' --region={zone_to_region(args.zone)}'
-      f' --cluster-version={default_rapid_gke_version}'
-      ' --master'
-      ' --quiet'
-  )
-  xpk_print("Updating GKE cluster's control plane version, may take a while!")
-  return_code = run_command_with_updates(
-      command,
-      'GKE Cluster control plane version update to enable Cloud DNS',
-      args,
-  )
-  if return_code != 0:
-    xpk_print(
-        "GKE cluster's control plane version update request returned"
-        f' ERROR {return_code}'
-    )
-    return 1
-  return 0
-
-
-def is_cluster_using_clouddns(args) -> bool:
-  """Checks if cluster is using CloudDNS.
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    True if cluster is using CloudDNS and False otherwise.
-  """
-  command = (
-      f'gcloud container clusters describe {args.cluster}'
-      f' --project={args.project} --region={zone_to_region(args.zone)}'
-      ' 2> /dev/null | grep "clusterDns: CLOUD_DNS"'
-  )
-  return_code, _ = run_command_for_value(
-      command,
-      'Check if Cloud DNS is enabled in cluster describe.',
-      args,
-  )
-  if return_code == 0:
-    xpk_print('Cloud DNS is enabled on the cluster, no update needed.')
-    return True
-  return False
-
-
 def is_workload_identity_enabled_on_cluster(args) -> bool:
   """Checks if Workload Identity Federation is enabled on the cluster.
   Args:
@@ -751,13 +645,12 @@ def is_workload_identity_enabled_on_cluster(args) -> bool:
   """
   command = (
       f'gcloud container clusters describe {args.cluster}'
-      f' --project={args.project} --region={zone_to_region(args.zone)}'
+      f' --project={args.project} --location={get_cluster_location(args.project, args.cluster, args.zone)}'
       ' --format="value(workloadIdentityConfig.workloadPool)"'
   )
   return_code, workload_pool = run_command_for_value(
       command,
       'Checks if Workload Identity Federation is enabled in cluster describe.',
-      args,
   )
   if return_code != 0:
     xpk_exit(return_code)
@@ -779,13 +672,12 @@ def is_gcsfuse_driver_enabled_on_cluster(args) -> bool:
   """
   command = (
       f'gcloud container clusters describe {args.cluster}'
-      f' --project={args.project} --region={zone_to_region(args.zone)}'
+      f' --project={args.project} --location={get_cluster_location(args.project, args.cluster, args.zone)}'
       ' --format="value(addonsConfig.gcsFuseCsiDriverConfig.enabled)"'
   )
   return_code, gcsfuse_driver_enabled = run_command_for_value(
       command,
       'Checks if GCSFuse CSI driver is enabled in cluster describe.',
-      args,
   )
   if return_code != 0:
     xpk_exit(return_code)
@@ -795,53 +687,6 @@ def is_gcsfuse_driver_enabled_on_cluster(args) -> bool:
   return False
 
 
-def update_cluster_with_clouddns_if_necessary(args) -> int:
-  """Updates a GKE cluster to use CloudDNS, if not enabled already.
-
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    0 if successful and error code otherwise.
-  """
-  all_clusters, return_code = get_all_clusters_programmatic(args)
-  if return_code > 0:
-    xpk_print('Listing all clusters failed!')
-    return 1
-  if args.cluster in all_clusters:
-    # If cluster is already using clouddns, no update necessary!
-    if is_cluster_using_clouddns(args):
-      return 0
-    cluster_update_return_code = update_gke_cluster_with_clouddns(args)
-    if cluster_update_return_code > 0:
-      xpk_print('Updating GKE cluster to use CloudDNS failed!')
-      return cluster_update_return_code
-
-    # Find default rapid control plane version and update the control plane to the same.
-    server_config_return_code, gke_server_config = get_gke_server_config(args)
-    if server_config_return_code != 0:
-      xpk_exit(server_config_return_code)
-    assert gke_server_config
-
-    upgrade_master_return_code = upgrade_gke_control_plane_version(
-        args,
-        gke_server_config.default_rapid_gke_version,
-    )
-    if upgrade_master_return_code > 0:
-      xpk_print("Updating GKE cluster's control plane upgrade failed!")
-      return upgrade_master_return_code
-
-    # Upgrade nodepools version after the master upgrade.
-    node_pool_update_code = upgrade_gke_nodepools_version(
-        args,
-        gke_server_config.default_rapid_gke_version,
-    )
-    if node_pool_update_code > 0:
-      xpk_print('Upgrading nodepools version failed!')
-      return node_pool_update_code
-  return 0
-
-
 def update_cluster_with_workload_identity_if_necessary(args) -> int:
   """Updates a GKE cluster to enable Workload Identity Federation, if not enabled already.
   Args:
@@ -884,26 +729,78 @@ def update_cluster_with_gcsfuse_driver_if_necessary(args) -> int:
   return 0
 
 
-def get_cluster_credentials(args) -> None:
-  """Run cluster configuration command to set the kubectl config.
+def test_and_retry_credentials_with_dns_logic(args) -> int:
+  """Tests kubectl credentials and retries with default settings if a DNS error is found.
 
   Args:
     args: user provided arguments for running the command.
 
   Returns:
-    0 if successful and 1 otherwise.
+    0 if credentials are valid after retrying, 1 otherwise.
   """
-  command = (
+
+  xpk_print('Testing credentials with kubectl...')
+  kubectl_command = 'kubectl get pods'
+  kubectl_return_code, kubectl_output = run_command_for_value(
+      kubectl_command, 'kubectl get pods'
+  )
+  if kubectl_return_code == 0:
+    xpk_print('Credentials test succeeded.')
+    return 0
+
+  dns_endpoint_error = (
+      'control_plane_endpoints_config.dns_endpoint_config.allow_external_traffic'
+      ' is disabled'
+  )
+  if dns_endpoint_error not in kubectl_output:
+    xpk_print(f'kubectl failed with an unhandled error: {kubectl_output}')
+    xpk_exit(kubectl_return_code)
+  xpk_print(
+      'Detected DNS endpoint-related error. Retrying without --dns-endpoint'
+      ' flag...'
+  )
+
+  location = get_cluster_location(args.project, args.cluster, args.zone)
+  without_dns_command = (
       'gcloud container clusters get-credentials'
-      f' {args.cluster} --region={zone_to_region(args.zone)}'
+      f' {args.cluster} --location={location}'
       f' --project={args.project} &&'
       ' kubectl config view && kubectl config set-context --current'
       ' --namespace=default'
   )
-  task = f'get-credentials to cluster {args.cluster}'
-  return_code = run_command_with_updates_retry(
-      command, task, args, verbose=False
+  return_code = run_command_with_updates(
+      without_dns_command, 'get-credentials to cluster', verbose=False
   )
+  if return_code != 0:
+    xpk_print('Failed to get credentials even without --dns-endpoint. Exiting.')
+    xpk_exit(return_code)
+  return 0
+
+
+def get_cluster_credentials(args) -> int:
+  """Run cluster configuration command to set the kubectl config.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  location = get_cluster_location(args.project, args.cluster, args.zone)
+  command = (
+      'gcloud container clusters get-credentials'
+      f' {args.cluster} --location={location} --dns-endpoint'
+      f' --project={args.project} && kubectl config view && kubectl config'
+      ' set-context --current --namespace=default'
+  )
+  task = f'get-credentials-dns-endpoint to cluster {args.cluster}'
+  return_code = run_command_with_updates_retry(command, task, verbose=False)
+
   if return_code != 0:
     xpk_print(f'{task} returned ERROR {return_code}')
     xpk_exit(return_code)
+
+  return_code = test_and_retry_credentials_with_dns_logic(args)
+  xpk_print('Finished get-credentials and kubectl setup.')
+
+  return return_code

xpk/core/cluster_private.py

@@ -22,7 +22,7 @@ from ..utils.network import (
 from ..utils.execution_context import is_dry_run
 from ..utils.objects import is_text_true
 from .commands import run_command_for_value, run_command_with_updates
-from .gcloud_context import zone_to_region
+from .gcloud_context import get_cluster_location
 
 
 def authorize_private_cluster_access_if_necessary(args) -> int:
@@ -127,13 +127,12 @@ def is_cluster_private(args) -> bool:
   """
   command = (
       f'gcloud container clusters describe {args.cluster}'
-      f' --project={args.project} --region={zone_to_region(args.zone)}'
+      f' --project={args.project} --location={get_cluster_location(args.project, args.cluster, args.zone)}'
       ' --format="value(privateClusterConfig.enablePrivateNodes)"'
   )
   return_code, private_nodes_enabled = run_command_for_value(
       command,
       'Check if Private Nodes is enabled in cluster.',
-      args,
   )
 
   if return_code != 0:
@@ -158,13 +157,12 @@ def get_cluster_authorized_networks(args) -> list[str]:
   """
   command = (
       f'gcloud container clusters describe {args.cluster}'
-      f' --project={args.project} --region={zone_to_region(args.zone)}'
+      f' --project={args.project} --location={get_cluster_location(args.project, args.cluster, args.zone)}'
       ' --format="value(masterAuthorizedNetworksConfig.cidrBlocks[].cidrBlock)"'
   )
   return_code, authorized_networks = run_command_for_value(
       command,
       'Fetching the list of authorized network from cluster describe.',
-      args,
       dry_run_return_val='127.0.0.1/32',
   )
 
@@ -189,15 +187,12 @@ def update_cluster_authorized_networks(args, authorized_networks) -> int:
   """
   command = (
       'gcloud container clusters update'
-      f' {args.cluster} --project={args.project}'
-      f' --region={zone_to_region(args.zone)}'
-      ' --enable-master-authorized-networks'
-      f' --master-authorized-networks={",".join(authorized_networks)}'
-      ' --quiet'
+      f' {args.cluster} --project={args.project} --location={get_cluster_location(args.project, args.cluster, args.zone)} --enable-master-authorized-networks'
+      f' --master-authorized-networks={",".join(authorized_networks)} --quiet'
   )
 
   return_code = run_command_with_updates(
-      command, 'GKE Cluster Update master authorized networks', args
+      command, 'GKE Cluster Update master authorized networks'
   )
 
   if return_code != 0: