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

xpk/commands/batch.py

@@ -14,18 +14,25 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import re
 from argparse import Namespace
 
-from ..core.cluster import create_xpk_k8s_service_account
+from ..core.cluster import (
+    create_xpk_k8s_service_account,
+    get_cluster_credentials,
+)
 from ..core.commands import run_command_for_value
 from ..core.gcloud_context import add_zone_and_project
+from ..core.kjob import (
+    AppProfileDefaults,
+    JobTemplateDefaults,
+    get_storage_annotations,
+    prepare_kjob,
+)
 from ..core.kueue import LOCAL_QUEUE_NAME
 from ..utils.console import xpk_exit, xpk_print
-from .common import set_cluster_command
-from ..core.kjob import AppProfileDefaults, JobTemplateDefaults, prepare_kjob, Kueue_TAS_annotation, get_gcsfuse_annotation
-from .kjob_common import add_gpu_networking_annotations_to_command
 from .kind import set_local_cluster_command
-import re
+from .kjob_common import add_gpu_networking_annotations_to_command, add_TAS_annotations_to_command
 
 
 def batch(args: Namespace) -> None:
@@ -38,12 +45,11 @@ def batch(args: Namespace) -> None:
   """
   if not args.kind_cluster:
     add_zone_and_project(args)
-    set_cluster_command_code = set_cluster_command(args)
+    get_cluster_credentials(args)
   else:
     set_cluster_command_code = set_local_cluster_command(args)
-
-  if set_cluster_command_code != 0:
-    xpk_exit(set_cluster_command_code)
+    if set_cluster_command_code != 0:
+      xpk_exit(set_cluster_command_code)
 
   err_code = prepare_kjob(args)
   if err_code > 0:
@@ -61,14 +67,14 @@ def submit_job(args: Namespace) -> None:
       'kubectl kjob create slurm'
       f' --profile {AppProfileDefaults.NAME.value}'
       f' --localqueue {LOCAL_QUEUE_NAME}'
-      f' --pod-template-annotation {Kueue_TAS_annotation}'
       f' --worker-container {JobTemplateDefaults.CONTAINER_NAME.value}'
       ' --first-node-ip'
   )
   cmd = add_gpu_networking_annotations_to_command(args, cmd)
-  gcsfuse_annotation = get_gcsfuse_annotation(args)
-  if gcsfuse_annotation is not None:
-    cmd += f' --pod-template-annotation {gcsfuse_annotation}'
+  cmd = add_TAS_annotations_to_command(args, cmd)
+
+  for annotation in get_storage_annotations(args):
+    cmd += f' --pod-template-annotation {annotation}'
 
   if args.ignore_unknown_flags:
     cmd += ' --ignore-unknown-flags'

xpk/commands/cluster.py

@@ -16,14 +16,21 @@ limitations under the License.
 
 from tabulate import tabulate
 
-from ..core.capacity import H100_DEVICE_TYPE
+from ..core.capacity import H100_DEVICE_TYPE, H200_DEVICE_TYPE, B200_DEVICE_TYPE
 from ..core.cluster import (
     get_all_clusters_programmatic,
     get_cluster_credentials,
     install_nccl_on_cluster,
+    install_nri_on_cluster,
     set_jobset_on_cluster,
+    set_pathways_job_on_cluster,
     setup_k8s_env,
+    disable_mglru_on_cluster,
+    count_nodes_on_cluster,
+    update_cluster_with_gcpfilestore_driver_if_necessary,
     update_cluster_with_gcsfuse_driver_if_necessary,
+    update_cluster_with_parallelstore_driver_if_necessary,
+    update_cluster_with_pd_driver_if_necessary,
     update_cluster_with_workload_identity_if_necessary,
 )
 from ..core.cluster_private import authorize_private_cluster_access_if_necessary
@@ -46,10 +53,14 @@ from ..core.nap import enable_autoprovisioning_on_cluster
 from ..core.network import (
     create_cluster_network_config,
     delete_cluster_subnets,
-    set_up_cluster_network_for_gpu,
+    set_up_cluster_network_for_a3,
+)
+from ..core.nodepool import (
+    get_gke_node_pool_version,
+    run_gke_node_pool_create_command,
 )
-from ..core.nodepool import get_gke_node_pool_version, run_gke_node_pool_create_command
 from ..core.ray import install_ray_cluster
+from ..core.mtc import install_mtc_on_cluster
 from ..core.resources import create_cluster_configmaps
 from ..core.storage import install_storage_crd
 from ..core.system_characteristics import (
@@ -64,7 +75,118 @@ from ..utils.console import get_user_input, xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
 from . import cluster_gcluster
 from .common import set_cluster_command
-from ..core.cluster import update_cluster_with_gcpfilestore_driver_if_necessary
+
+
+def cluster_adapt(args) -> None:
+  """Function that performs cluster adaptation.
+
+  Args:
+    args: user provided arguments for running the command.
+  """
+  args.enable_pathways = False
+
+  system, return_code = get_system_characteristics(args)
+
+  if return_code > 0:
+    xpk_print('Fetching system characteristics failed!')
+    xpk_exit(return_code)
+
+  xpk_print(
+      f'Starting cluster adaptation for cluster {args.cluster}:', flush=True
+  )
+  add_zone_and_project(args)
+
+  if system.accelerator_type == AcceleratorType['GPU'] and not getattr(
+      args, 'num_nodes'
+  ):
+    xpk_print(
+        'Argument --num-nodes was not provided, trying to determine number of'
+        ' nodes based on the available nodes in the cluster...'
+    )
+    args.num_nodes = count_nodes_on_cluster(args, system)
+    if args.num_nodes == 0:
+      xpk_print(
+          'Found unexpected number of nodes. Is the --device-type correct?'
+      )
+      xpk_exit(1)
+    else:
+      xpk_print(f'Using {args.num_nodes} nodes.')
+
+  # ToDo(roshanin@) - Re-enable CloudDNS on Pathways clusters conditionally.
+  # Enable WorkloadIdentity if not enabled already.
+  if args.enable_workload_identity or args.enable_gcsfuse_csi_driver:
+    update_cluster_command_code = (
+        update_cluster_with_workload_identity_if_necessary(args)
+    )
+    if update_cluster_command_code != 0:
+      xpk_exit(update_cluster_command_code)
+
+  get_cluster_credentials(args)
+
+  k8s_client = setup_k8s_env(args)
+
+  install_storage_crd(k8s_client)
+  install_storage_csis(args)
+
+  # create Vertex Tensorboard for new and existing clusters if create-vertex-tensorboard is set
+  tensorboard_config = {}
+  if VERTEX_TENSORBOARD_FEATURE_FLAG and args.create_vertex_tensorboard:
+    tensorboard_config = create_vertex_tensorboard(args)
+    # exit if failed to create Tensorboard in Vertex AI
+    if not tensorboard_config:
+      xpk_exit(1)
+
+  # Provision node pools dynamically based on incoming workloads:
+  # Currently autoprovisioning is not supported with Pathways.
+  autoprovisioning_config = None
+  if args.enable_autoprovisioning:
+    xpk_print('Enabling Autoprovisioning')
+    autoprovisioning_config, return_code = enable_autoprovisioning_on_cluster(
+        args, system
+    )
+    if return_code != 0:
+      xpk_exit(return_code)
+
+  xpk_print('Creating ConfigMap for cluster')
+  create_cluster_configmaps_code = create_cluster_configmaps(
+      args, system, tensorboard_config, autoprovisioning_config
+  )
+  if create_cluster_configmaps_code != 0:
+    xpk_exit(create_cluster_configmaps_code)
+
+  xpk_print(
+      'Enabling the jobset API on our cluster, to be deprecated when Jobset is'
+      ' globally available'
+  )
+  set_jobset_on_cluster_code = set_jobset_on_cluster(args)
+  if set_jobset_on_cluster_code != 0:
+    xpk_exit(set_jobset_on_cluster_code)
+
+  # TODO: Uncomment when cluster_adapt will support TPU cluters
+  # set_pathways_job_on_cluster_code = set_pathways_job_on_cluster(args)
+  # if set_pathways_job_on_cluster_code != 0:
+  #   xpk_exit(set_pathways_job_on_cluster_code)
+
+  install_kueue(args, system, autoprovisioning_config)
+
+  install_kjob(args)
+
+  if system.accelerator_type == AcceleratorType['GPU']:
+    prepare_gpus(args, system)
+
+  if args.enable_ray_cluster:
+    return_code = install_ray_cluster(args, system)
+    if return_code != 0:
+      xpk_print('Installation of RayCluster failed.')
+      xpk_exit(return_code)
+
+  xpk_print('GKE commands done! Resources are created.')
+  xpk_print(
+      'See your GKE Cluster here:'
+      # pylint: disable=line-too-long
+      f' https://console.cloud.google.com/kubernetes/clusters/details/{zone_to_region(args.zone)}/{args.cluster}/details?project={args.project}'
+  )
+  xpk_exit(0)
 
 
 def cluster_create(args) -> None:
@@ -72,9 +194,6 @@ def cluster_create(args) -> None:
 
   Args:
     args: user provided arguments for running the command.
-
-  Returns:
-    0 if successful and 1 otherwise.
   """
   system, return_code = get_system_characteristics(args)
 
@@ -117,35 +236,19 @@ def cluster_create(args) -> None:
 
   # ToDo(roshanin@) - Re-enable CloudDNS on Pathways clusters conditionally.
   # Enable WorkloadIdentity if not enabled already.
-  if (
-      args.enable_workload_identity
-      or args.enable_gcsfuse_csi_driver
-      or args.enable_gcpfilestore_csi_driver
-  ):
+  if args.enable_workload_identity or args.enable_gcsfuse_csi_driver:
     update_cluster_command_code = (
         update_cluster_with_workload_identity_if_necessary(args)
     )
     if update_cluster_command_code != 0:
       xpk_exit(update_cluster_command_code)
 
-  # Enable GCSFuse CSI Driver if not enabled already.
-  if args.enable_gcsfuse_csi_driver:
-    update_cluster_command_code = (
-        update_cluster_with_gcsfuse_driver_if_necessary(args)
-    )
-    if update_cluster_command_code != 0:
-      xpk_exit(update_cluster_command_code)
-
-  if args.enable_gcpfilestore_csi_driver:
-    update_cluster_command_code = (
-        update_cluster_with_gcpfilestore_driver_if_necessary(args)
-    )
-    if update_cluster_command_code != 0:
-      xpk_exit(update_cluster_command_code)
+  get_cluster_credentials(args)
 
-  # Update Pathways clusters with CloudDNS if not enabled already.
+  k8s_client = setup_k8s_env(args)
 
-  get_cluster_credentials(args)
+  install_storage_crd(k8s_client)
+  install_storage_csis(args)
 
   # create Vertex Tensorboard for new and existing clusters if create-vertex-tensorboard is set
   tensorboard_config = {}
@@ -155,13 +258,12 @@ def cluster_create(args) -> None:
     if not tensorboard_config:
       xpk_exit(1)
 
-  if system.accelerator_type == AcceleratorType['GPU']:
+  if system.device_type == H100_DEVICE_TYPE:
     xpk_print('Setting up Network for cluster')
-    set_up_cluster_network_code = set_up_cluster_network_for_gpu(args, system)
+    set_up_cluster_network_code = set_up_cluster_network_for_a3(args)
     if set_up_cluster_network_code != 0:
       xpk_exit(set_up_cluster_network_code)
 
-  if system.device_type == H100_DEVICE_TYPE:
     xpk_print('Creating Network Config for cluster')
     create_cluster_network_config_code = create_cluster_network_config(args)
     if create_cluster_network_config_code != 0:
@@ -207,45 +309,16 @@ def cluster_create(args) -> None:
   if set_jobset_on_cluster_code != 0:
     xpk_exit(set_jobset_on_cluster_code)
 
-  xpk_print('Enabling Kueue on the cluster')
-  install_kueue_on_cluster_code = install_kueue_on_cluster(args)
-  if install_kueue_on_cluster_code != 0:
-    xpk_exit(install_kueue_on_cluster_code)
+  set_pathways_job_on_cluster_code = set_pathways_job_on_cluster(args)
+  if set_pathways_job_on_cluster_code != 0:
+    xpk_exit(set_pathways_job_on_cluster_code)
 
-  xpk_print('Verifying kjob installation')
-  err_code = verify_kjob_installed(args)
-  if err_code > 0:
-    xpk_exit(err_code)
+  install_kueue(args, system, autoprovisioning_config)
 
-  xpk_print('Applying kjob CDRs')
-  err_code = apply_kjob_crds(args)
-  if err_code > 0:
-    xpk_exit(err_code)
-
-  err_code = prepare_kjob(args)
-  if err_code > 0:
-    xpk_exit(err_code)
-
-  k8s_client = setup_k8s_env(args)
-  install_storage_crd(k8s_client)
-
-  xpk_print('Wait for Kueue to be fully available')
-  wait_for_kueue_available_code = wait_for_kueue_available(args)
-  if wait_for_kueue_available_code != 0:
-    xpk_exit(wait_for_kueue_available_code)
-
-  xpk_print('Install Kueue Custom Resources')
-  enable_kueue_credentials_code = install_kueue_crs(
-      args, system, autoprovisioning_config
-  )
-  if enable_kueue_credentials_code != 0:
-    xpk_exit(enable_kueue_credentials_code)
+  install_kjob(args)
 
   if system.accelerator_type == AcceleratorType['GPU']:
-    xpk_print('Installing NCCL Plugin for cluster')
-    install_nccl_code = install_nccl_on_cluster(args, system)
-    if install_nccl_code != 0:
-      xpk_exit(install_nccl_code)
+    prepare_gpus(args, system)
 
   if args.enable_ray_cluster:
     return_code = install_ray_cluster(args, system)
@@ -253,6 +326,12 @@ def cluster_create(args) -> None:
       xpk_print('Installation of RayCluster failed.')
       xpk_exit(return_code)
 
+  if hasattr(args, 'enable_mtc') and args.enable_mtc:
+    return_code = install_mtc_on_cluster(args, system)
+    if return_code != 0:
+      xpk_print('Installation of MTC failed.')
+      xpk_exit(return_code)
+
   xpk_print('GKE commands done! Resources are created.')
   xpk_print(
       'See your GKE Cluster here:'
@@ -757,6 +836,7 @@ def run_gke_cluster_create_command(
       f' --num-nodes {args.default_pool_cpu_num_nodes}'
       f' {args.custom_cluster_arguments}'
       f' {rapid_release_cmd}'
+      ' --enable-dns-access'
   )
 
   enable_ip_alias = False
@@ -783,11 +863,7 @@ def run_gke_cluster_create_command(
   if args.enable_ray_cluster:
     command += ' --addons RayOperator'
 
-  if (
-      args.enable_workload_identity
-      or args.enable_gcsfuse_csi_driver
-      or args.enable_gcpfilestore_csi_driver
-  ):
+  if args.enable_workload_identity or args.enable_gcsfuse_csi_driver:
     command += f' --workload-pool={args.project}.svc.id.goog'
 
   addons = []
@@ -797,6 +873,15 @@ def run_gke_cluster_create_command(
   if args.enable_gcpfilestore_csi_driver:
     addons.append('GcpFilestoreCsiDriver')
 
+  if args.enable_parallelstore_csi_driver:
+    addons.append('ParallelstoreCsiDriver')
+
+  if args.enable_pd_csi_driver:
+    addons.append('GcePersistentDiskCsiDriver')
+
+  if hasattr(args, 'enable_mtc') and args.enable_mtc:
+    addons.append('HighScaleCheckpointing')
+
   if len(addons) > 0:
     addons_str = ','.join(addons)
     command += f' --addons={addons_str}'
@@ -806,3 +891,87 @@ def run_gke_cluster_create_command(
     xpk_print(f'GKE Cluster Create request returned ERROR {return_code}')
     return 1
   return 0
+
+
+def install_storage_csis(args):
+  if args.enable_gcsfuse_csi_driver:
+    update_cluster_command_code = (
+        update_cluster_with_gcsfuse_driver_if_necessary(args)
+    )
+    if update_cluster_command_code != 0:
+      xpk_exit(update_cluster_command_code)
+
+  if args.enable_gcpfilestore_csi_driver:
+    update_cluster_command_code = (
+        update_cluster_with_gcpfilestore_driver_if_necessary(args)
+    )
+    if update_cluster_command_code != 0:
+      xpk_exit(update_cluster_command_code)
+
+  if args.enable_parallelstore_csi_driver:
+    update_cluster_command_code = (
+        update_cluster_with_parallelstore_driver_if_necessary(args)
+    )
+    if update_cluster_command_code != 0:
+      xpk_exit(update_cluster_command_code)
+
+  if args.enable_pd_csi_driver:
+    update_cluster_command_code = update_cluster_with_pd_driver_if_necessary(
+        args
+    )
+    if update_cluster_command_code != 0:
+      xpk_exit(update_cluster_command_code)
+
+
+def install_kjob(args):
+  xpk_print('Verifying kjob installation')
+  err_code = verify_kjob_installed(args)
+  if err_code > 0:
+    xpk_exit(err_code)
+
+  xpk_print('Applying kjob CDRs')
+  err_code = apply_kjob_crds(args)
+  if err_code > 0:
+    xpk_exit(err_code)
+
+  err_code = prepare_kjob(args)
+  if err_code > 0:
+    xpk_exit(err_code)
+
+
+def install_kueue(args, system: SystemCharacteristics, autoprovisioning_config):
+  xpk_print('Enabling Kueue on the cluster')
+  install_kueue_on_cluster_code = install_kueue_on_cluster(args)
+  if install_kueue_on_cluster_code != 0:
+    xpk_exit(install_kueue_on_cluster_code)
+
+  xpk_print('Wait for Kueue to be fully available')
+  wait_for_kueue_available_code = wait_for_kueue_available(args)
+  if wait_for_kueue_available_code != 0:
+    xpk_exit(wait_for_kueue_available_code)
+
+  xpk_print('Install Kueue Custom Resources')
+  enable_kueue_credentials_code = install_kueue_crs(
+      args, system, autoprovisioning_config
+  )
+  if enable_kueue_credentials_code != 0:
+    xpk_exit(enable_kueue_credentials_code)
+
+
+def prepare_gpus(args, system: SystemCharacteristics):
+  xpk_print('Installing NCCL Plugin for cluster')
+  install_nccl_code = install_nccl_on_cluster(args, system)
+  if install_nccl_code != 0:
+    xpk_exit(install_nccl_code)
+
+  if system.device_type == H100_DEVICE_TYPE:
+    xpk_print('Installing NRI device injector for cluster')
+    install_nri_code = install_nri_on_cluster(args)
+    if install_nri_code != 0:
+      xpk_exit(install_nri_code)
+
+  if system.device_type in [H200_DEVICE_TYPE, B200_DEVICE_TYPE]:
+    xpk_print('Disabling MGLRU')
+    err_code = disable_mglru_on_cluster(args)
+    if err_code > 0:
+      xpk_exit(err_code)

xpk/commands/cluster_gcluster.py

@@ -16,26 +16,27 @@ limitations under the License.
 
 import os
 
-from ..core.remote_state.remote_state_client import RemoteStateClient
-from ..core.remote_state.fuse_remote_state import FuseStateClient
 from ..core.blueprint.blueprint_generator import (
     BlueprintGenerator,
     BlueprintGeneratorOutput,
     a3mega_device_type,
     a3ultra_device_type,
+    a4_device_type,
     supported_device_types,
 )
-from ..core.commands import run_command_for_value
 from ..core.capacity import get_capacity_type
+from ..core.cluster import get_cluster_credentials
+from ..core.commands import run_command_for_value
 from ..core.docker_manager import DockerManager
 from ..core.gcloud_context import zone_to_region
 from ..core.gcluster_manager import GclusterManager
+from ..core.kjob import apply_kjob_crds, prepare_kjob
+from ..core.remote_state.fuse_remote_state import FuseStateClient
+from ..core.remote_state.remote_state_client import RemoteStateClient
 from ..utils.console import xpk_exit, xpk_print
 from ..utils.file import ensure_directory_exists
 from ..utils.network import all_IPs_cidr
 from ..utils.objects import hash_string
-from ..core.cluster import get_cluster_credentials
-from ..core.kjob import apply_kjob_crds, prepare_kjob
 
 blueprints_path = os.path.abspath('xpkclusters/blueprints')
 gcluster_working_dir = os.path.abspath('xpkclusters/gcluster-out')
@@ -266,4 +267,20 @@ def generate_blueprint(
           system_node_pool_min_node_count=args.default_pool_cpu_num_nodes,
           gcs_bucket=args.cluster_state_gcs_bucket,
       )
+    if args.device_type == a4_device_type:
+      num_nodes = args.num_nodes if not args.num_nodes is None else 2
+      return bpg.generate_a4_blueprint(
+          blueprint_name=blueprint_name,
+          prefix=prefix,
+          cluster_name=args.cluster,
+          region=zone_to_region(args.zone),
+          project_id=args.project,
+          zone=args.zone,
+          auth_cidr=all_IPs_cidr,
+          num_nodes=num_nodes,
+          reservation=args.reservation if args.reservation else None,
+          capacity_type=capacity_type,
+          system_node_pool_machine_type=args.default_pool_cpu_machine_type,
+          system_node_pool_min_node_count=args.default_pool_cpu_num_nodes,
+      )
   return None

xpk/commands/common.py

@@ -15,8 +15,10 @@ limitations under the License.
 """
 
 from ..core.commands import run_command_with_updates_retry
+from ..core.resources import get_cluster_capacity_type, get_cluster_system_characteristics
+from ..core.capacity import H100_MEGA_DEVICE_TYPE, CapacityType
 from ..core.gcloud_context import zone_to_region
-from ..utils.console import xpk_print
+from ..utils.console import xpk_print, xpk_exit
 
 
 def set_cluster_command(args) -> int:
@@ -31,6 +33,7 @@ def set_cluster_command(args) -> int:
   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'
@@ -42,3 +45,32 @@ def set_cluster_command(args) -> int:
   if return_code != 0:
     xpk_print(f'{task} returned ERROR {return_code}')
   return return_code
+
+
+def is_TAS_possible(args) -> bool:
+  """Check cluster's machine_type and capacity type to determine if Kueue TAS is possible
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    True if possible and False otherwise.
+  """
+  system_characteristics = get_cluster_system_characteristics(args)
+  capacity_type = get_cluster_capacity_type(args)
+
+  if system_characteristics is None:
+    xpk_print('system_characteristics data was not found in configmaps.')
+    xpk_exit(1)
+
+  if capacity_type is None:
+    xpk_print('capacity_type data was not found in configmaps.')
+    xpk_exit(1)
+
+  if (
+      system_characteristics.device_type == H100_MEGA_DEVICE_TYPE
+      and capacity_type == CapacityType.SPOT
+  ):
+    return False
+
+  return True

xpk/commands/info.py

@@ -20,10 +20,10 @@ from argparse import Namespace
 from tabulate import tabulate
 
 from ..core.commands import run_command_for_value
+from ..core.cluster import get_cluster_credentials
 from ..core.gcloud_context import add_zone_and_project
 from ..core.kueue import verify_kueuectl
 from ..utils.console import xpk_exit, xpk_print
-from .common import set_cluster_command
 
 table_fmt = 'plain'
 
@@ -37,9 +37,7 @@ def info(args: Namespace) -> None:
     None
   """
   add_zone_and_project(args)
-  set_cluster_command_code = set_cluster_command(args)
-  if set_cluster_command_code != 0:
-    xpk_exit(set_cluster_command_code)
+  get_cluster_credentials(args)
 
   verify_kueuectl(args)
   lq, cq = bool(args.localqueue), bool(args.clusterqueue)

xpk/commands/job.py

@@ -20,10 +20,10 @@ import sys
 from ruamel.yaml import YAML
 
 from ..core.commands import run_command_for_value, run_command_with_updates
+from ..core.cluster import get_cluster_credentials
 from ..core.gcloud_context import add_zone_and_project
 from ..core.kjob import AppProfileDefaults
 from ..utils.console import xpk_exit, xpk_print
-from .common import set_cluster_command
 from .kind import set_local_cluster_command
 
 
@@ -143,14 +143,14 @@ def job_list(args) -> None:
   """
   if not args.kind_cluster:
     add_zone_and_project(args)
-    set_cluster_command_code = set_cluster_command(args)
+    get_cluster_credentials(args)
     msg = f'Listing jobs for project {args.project} and zone {args.zone}:'
   else:
     set_cluster_command_code = set_local_cluster_command(args)
     msg = 'Listing jobs:'
+    if set_cluster_command_code != 0:
+      xpk_exit(set_cluster_command_code)
 
-  if set_cluster_command_code != 0:
-    xpk_exit(set_cluster_command_code)
   xpk_print(msg, flush=True)
 
   return_code = run_slurm_job_list_command(args)
@@ -178,12 +178,11 @@ def job_cancel(args) -> None:
   xpk_print(f'Starting job cancel for job: {args.name}', flush=True)
   if not args.kind_cluster:
     add_zone_and_project(args)
-    set_cluster_command_code = set_cluster_command(args)
+    get_cluster_credentials(args)
   else:
     set_cluster_command_code = set_local_cluster_command(args)
-
-  if set_cluster_command_code != 0:
-    xpk_exit(set_cluster_command_code)
+    if set_cluster_command_code != 0:
+      xpk_exit(set_cluster_command_code)
 
   return_code = run_slurm_job_delete_command(args)
   xpk_exit(return_code)