feat: remove args from run_command_with_updates_retry (#652)

Author: scaliby

src/xpk/commands/cluster.py

@@ -176,7 +176,7 @@ def cluster_adapt(args) -> None:
 
   install_kjob(args)
   if system.accelerator_type == AcceleratorType['GPU']:
-    prepare_gpus(args, system)
+    prepare_gpus(system)
 
   if args.enable_ray_cluster:
     return_code = install_ray_cluster(args, system)
@@ -249,7 +249,7 @@ def cluster_create(args) -> None:
 
   get_cluster_credentials(args)
 
-  update_coredns_command_code = update_coredns_if_necessary(args)
+  update_coredns_command_code = update_coredns_if_necessary()
   if update_coredns_command_code != 0:
     xpk_exit(update_cluster_command_code)
 
@@ -317,7 +317,7 @@ def cluster_create(args) -> None:
   set_jobset_on_cluster_code = set_jobset_on_cluster(args)
   if set_jobset_on_cluster_code != 0:
     xpk_exit(set_jobset_on_cluster_code)
-  update_jobset_resources_code = update_jobset_resources_if_necessary(args)
+  update_jobset_resources_code = update_jobset_resources_if_necessary()
   if update_jobset_resources_code != 0:
     xpk_exit(update_jobset_resources_code)
 
@@ -330,7 +330,7 @@ def cluster_create(args) -> None:
   install_kjob(args)
 
   if system.accelerator_type == AcceleratorType['GPU']:
-    prepare_gpus(args, system)
+    prepare_gpus(system)
 
   if args.enable_ray_cluster:
     return_code = install_ray_cluster(args, system)
@@ -416,15 +416,13 @@ def cluster_cacheimage(args) -> None:
   command_delete = f'kubectl delete -f {str(tmp)} --ignore-not-found=true'
 
   return_code = run_command_with_updates(
-      command_delete, 'Deleting Cached Image', args
+      command_delete, 'Deleting Cached Image'
   )
   if return_code != 0:
     xpk_print(f'Delete Cached Image returned ERROR {return_code}')
     xpk_exit(return_code)
 
-  return_code = run_command_with_updates(
-      command_apply, 'Creating Cached Image', args
-  )
+  return_code = run_command_with_updates(command_apply, 'Creating Cached Image')
   if return_code != 0:
     xpk_print(f'Create Cached Image returned ERROR {return_code}')
     xpk_exit(return_code)
@@ -704,20 +702,20 @@ def cluster_create_ray_cluster(args) -> None:
   cluster_create(args)
 
 
-def install_jq(args):
+def install_jq():
   """Installs 'jq' utility."""
   if shutil.which('jq'):
     xpk_print("Task: 'Install jq' skipped, jq already installed.")
     return
   command_jq_install = 'sudo apt install jq -y'
   xpk_print("Task: 'Install jq' in progress.")
-  return_code = run_command_with_updates(command_jq_install, 'Install jq', args)
+  return_code = run_command_with_updates(command_jq_install, 'Install jq')
   if return_code != 0:
     xpk_print(f'Install jq error {return_code}')
     xpk_exit(return_code)
 
 
-def clone_coredns_deployment_repo(args, coredns_repo_full_path: str):
+def clone_coredns_deployment_repo(coredns_repo_full_path: str):
   """Clones the CoreDNS deployment repository if it doesn't exist."""
   if os.path.exists(coredns_repo_full_path):
     xpk_print(
@@ -732,15 +730,13 @@ def clone_coredns_deployment_repo(args, coredns_repo_full_path: str):
       "Task: 'Clone deployment' in progress, Target"
       f' directory:{coredns_repo_full_path}.'
   )
-  return_code = run_command_with_updates(
-      command_git_clone, 'Clone deployment', args
-  )
+  return_code = run_command_with_updates(command_git_clone, 'Clone deployment')
   if return_code != 0:
     xpk_print(f'Clone deployment error {return_code}')
     xpk_exit(return_code)
 
 
-def deploy_coredns_manifests(args, coredns_k8s_path: str):
+def deploy_coredns_manifests(coredns_k8s_path: str):
   """Deploys CoreDNS manifests to the cluster."""
   if not os.path.isdir(coredns_k8s_path):
     xpk_print(
@@ -758,7 +754,7 @@ def deploy_coredns_manifests(args, coredns_k8s_path: str):
         f"Task: 'Deploy CoreDNS' in progress, Located at '{coredns_k8s_path}'"
     )
     return_code = run_command_with_updates(
-        command_deploy_coredns, 'Deploy CoreDNS', args
+        command_deploy_coredns, 'Deploy CoreDNS'
     )
     if return_code != 0:
       xpk_print(f'Deploy CoreDNS error {return_code}')
@@ -770,57 +766,53 @@ def deploy_coredns_manifests(args, coredns_k8s_path: str):
     xpk_exit(return_code)
 
 
-def scale_down_deployment(
-    args, deployment_name: str, namespace: str = 'kube-system'
-):
+def scale_down_deployment(deployment_name: str, namespace: str = 'kube-system'):
   """Scales down a specified Kubernetes deployment to 0 replicas."""
   command = (
       f'kubectl scale deployment {deployment_name} --replicas=0'
       f' --namespace={namespace}'
   )
   xpk_print(f"Task: 'Scaling down {deployment_name}' in progress")
   return_code = run_command_with_updates(
-      command, f'Scale down {deployment_name}', args
+      command, f'Scale down {deployment_name}'
   )
   if return_code != 0:
     xpk_print(f'Scale down {deployment_name} error {return_code}')
     xpk_exit(return_code)
   xpk_print(f'\n{deployment_name} has been scaled down.')
 
 
-def scale_up_coredns(args, replicas: int = 15, namespace: str = 'kube-system'):
+def scale_up_coredns(replicas: int = 15, namespace: str = 'kube-system'):
   """Scales up the CoreDNS deployment to a specified number of replicas."""
   command_coredns_scale = (
       f'kubectl scale deployment coredns --replicas={replicas} -n {namespace}'
   )
   xpk_print(f"Task: 'Scale CoreDNS' in progress (to {replicas} replicas)")
-  return_code = run_command_with_updates(
-      command_coredns_scale, 'Scale CoreDNS', args
-  )
+  return_code = run_command_with_updates(command_coredns_scale, 'Scale CoreDNS')
   if return_code != 0:
     xpk_print(f'Scale CoreDNS error {return_code}')
     xpk_exit(return_code)
 
 
-def check_deployment_exists(args, deployment_name: str, namespace: str) -> bool:
+def check_deployment_exists(deployment_name: str, namespace: str) -> bool:
   """Check for the existence of a specific Deployment in a given namespace."""
   # TODO: rewrite this to be more obvious, check if it is correct
   command = (
       f'kubectl get deployment {deployment_name} -n'
       f' {namespace} --ignore-not-found'
   )
   result = run_command_with_updates(
-      command, 'Waiting for kubeDNS to be checked.', args
+      command, 'Waiting for kubeDNS to be checked.'
   )
   return result != 0
 
 
 def verify_coredns_readiness(
-    args, timeout: int = 240, namespace: str = 'kube-system'
+    timeout: int = 240, namespace: str = 'kube-system'
 ):
   """Verifies CoreDNS readiness using kubectl wait commands."""
   xpk_print('Now verifying CoreDNS readiness...')
-  kube_dns_exists = check_deployment_exists(args, 'kube-dns', namespace)
+  kube_dns_exists = check_deployment_exists('kube-dns', namespace)
   if kube_dns_exists:
     # Wait for kube-dns to be fully scaled down
     command_kube_dns_wait_scaled_down = (
@@ -830,7 +822,7 @@ def verify_coredns_readiness(
     )
     xpk_print('Verifying if kube-dns has scaled down...')
     return_code_kube_dns = run_command_with_updates(
-        command_kube_dns_wait_scaled_down, 'Wait for kube-dns scale down', args
+        command_kube_dns_wait_scaled_down, 'Wait for kube-dns scale down'
     )
     if return_code_kube_dns != 0:
       xpk_print('kube-dns did not scale down successfully within the timeout.')
@@ -846,7 +838,7 @@ def verify_coredns_readiness(
   )
   xpk_print('Verifying if CoreDNS is available...')
   return_code_coredns = run_command_with_updates(
-      command_coredns_wait_available, 'Wait for coredns available', args
+      command_coredns_wait_available, 'Wait for coredns available'
   )
   if return_code_coredns != 0:
     xpk_print(
@@ -871,12 +863,9 @@ def cleanup_coredns_repo(coredns_repo_full_path: str):
     xpk_print(f'Error deleting directory {coredns_repo_full_path}: {e}')
 
 
-def update_coredns(args) -> int:
+def update_coredns() -> int:
   """Updates and deploys CoreDNS within a cluster.
 
-  Args:
-    args: user provided arguments for running the command.
-
   Returns:
     0 if successful and 1 otherwise.
   """
@@ -885,23 +874,23 @@ def update_coredns(args) -> int:
   coredns_repo_full_path = os.path.join(coredns_repo_dir, coredns_repo_dir_name)
   coredns_k8s_path = os.path.join(coredns_repo_full_path, 'kubernetes')
   # 1. Install jq
-  install_jq(args)
+  install_jq()
 
   # 2. Clone CoreDNS deployment repository
-  clone_coredns_deployment_repo(args, coredns_repo_full_path)
+  clone_coredns_deployment_repo(coredns_repo_full_path)
 
   # 3. Deploy CoreDNS to the cluster
-  deploy_coredns_manifests(args, coredns_k8s_path)
+  deploy_coredns_manifests(coredns_k8s_path)
 
   # 4. Scale down kube-dns-autoscaler
-  scale_down_deployment(args, 'kube-dns-autoscaler')
+  scale_down_deployment('kube-dns-autoscaler')
 
   # 5. Scale down kube-dns
-  scale_down_deployment(args, 'kube-dns')
+  scale_down_deployment('kube-dns')
 
   # 6. Scale up coredns and verify readiness
-  scale_up_coredns(args, replicas=15)
-  verify_coredns_readiness(args, timeout=120)
+  scale_up_coredns(replicas=15)
+  verify_coredns_readiness(timeout=120)
 
   xpk_print('The CoreDNS setup process has been completed.')
 
@@ -911,7 +900,7 @@ def update_coredns(args) -> int:
   return 0
 
 
-def coredns_deployment_exists(args, namespace: str = 'kube-system') -> bool:
+def coredns_deployment_exists(namespace: str = 'kube-system') -> bool:
   """Checks if the CoreDNS deployment exists in the given namespace.
 
   Args:
@@ -926,10 +915,10 @@ def coredns_deployment_exists(args, namespace: str = 'kube-system') -> bool:
       f' namespace: {namespace}'
   )
   return_code = run_command_with_updates(
-      command, f'Check CoreDNS deployment in {namespace}', args
+      command, f'Check CoreDNS deployment in {namespace}'
   )
   if return_code == 0:
-    verify_coredns_readiness(args)
+    verify_coredns_readiness()
     xpk_print(f"CoreDNS deployment 'coredns' found in namespace '{namespace}'.")
     return True
   else:
@@ -940,25 +929,22 @@ def coredns_deployment_exists(args, namespace: str = 'kube-system') -> bool:
     return False
 
 
-def update_coredns_if_necessary(args) -> int:
+def update_coredns_if_necessary() -> int:
   """Updates and deploys CoreDNS within the cluster if it's not already present.
 
   This function checks for the existence of the CoreDNS deployment.
   If it's not found, it proceeds to deploy and configure CoreDNS.
 
-  Args:
-    args: User-provided arguments for running the command.
-
   Returns:
     0 if successful (CoreDNS was already present or successfully deployed),
     and 1 otherwise.
   """
-  if coredns_deployment_exists(args, namespace='kube-system'):
+  if coredns_deployment_exists(namespace='kube-system'):
     xpk_print('Skipping CoreDNS deployment since it already exists.')
     return 0
   else:
     xpk_print('CoreDNS deployment not found. Proceeding with CoreDNS setup.')
-    return update_coredns(args)
+    return update_coredns()
 
 
 def create_cluster_if_necessary(
@@ -1021,7 +1007,7 @@ def run_gke_cluster_delete_command(args) -> int:
       f' --region={zone_to_region(args.zone)} --quiet'
   )
 
-  return_code = run_command_with_updates(command, 'Cluster Delete', args)
+  return_code = run_command_with_updates(command, 'Cluster Delete')
   if return_code != 0:
     xpk_print(f'Cluster delete request returned ERROR {return_code}')
     return 1
@@ -1046,7 +1032,7 @@ def run_gke_clusters_list_command(args) -> int:
       'gcloud container clusters list'
       f' --project={args.project} --region={zone_to_region(args.zone)}'
   )
-  return_code = run_command_with_updates(command, 'Cluster List', args)
+  return_code = run_command_with_updates(command, 'Cluster List')
   if return_code != 0:
     xpk_print(f'Cluster list request returned ERROR {return_code}')
     return 1
@@ -1155,7 +1141,7 @@ def run_gke_cluster_create_command(
     addons_str = ','.join(addons)
     command += f' --addons={addons_str}'
 
-  return_code = run_command_with_updates(command, 'GKE Cluster Create', args)
+  return_code = run_command_with_updates(command, 'GKE Cluster Create')
   if return_code != 0:
     xpk_print(f'GKE Cluster Create request returned ERROR {return_code}')
     return 1
@@ -1206,7 +1192,7 @@ def install_kjob(args):
     xpk_exit(err_code)
 
   xpk_print('Applying kjob CDRs')
-  err_code = apply_kjob_crds(args)
+  err_code = apply_kjob_crds()
   if err_code > 0:
     xpk_exit(err_code)
 
@@ -1217,12 +1203,12 @@ def install_kjob(args):
 
 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)
+  install_kueue_on_cluster_code = install_kueue_on_cluster()
   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)
+  wait_for_kueue_available_code = wait_for_kueue_available()
   if wait_for_kueue_available_code != 0:
     xpk_exit(wait_for_kueue_available_code)
 
@@ -1234,25 +1220,25 @@ def install_kueue(args, system: SystemCharacteristics, autoprovisioning_config):
     xpk_exit(enable_kueue_credentials_code)
 
   xpk_print('Update Kueue Controller Manager resources')
-  update_kueue_resources_code = update_kueue_resources_if_necessary(args)
+  update_kueue_resources_code = update_kueue_resources_if_necessary()
   if update_kueue_resources_code != 0:
     xpk_exit(update_kueue_resources_code)
 
 
-def prepare_gpus(args, system: SystemCharacteristics):
+def prepare_gpus(system: SystemCharacteristics):
   xpk_print('Installing NCCL Plugin for cluster')
-  install_nccl_code = install_nccl_on_cluster(args, system)
+  install_nccl_code = install_nccl_on_cluster(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)
+    install_nri_code = install_nri_on_cluster()
     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)
+    err_code = disable_mglru_on_cluster()
     if err_code > 0:
       xpk_exit(err_code)

src/xpk/commands/cluster_gcluster.py

@@ -90,7 +90,7 @@ def cluster_create(args) -> None:
 
   get_cluster_credentials(args)
 
-  err_code = apply_kjob_crds(args)
+  err_code = apply_kjob_crds()
   if err_code > 0:
     xpk_exit(err_code)
 

src/xpk/commands/common.py

@@ -42,9 +42,7 @@ def set_cluster_command(args) -> int:
       ' --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_retry(command, task, verbose=False)
   if return_code != 0:
     xpk_print(f'{task} returned ERROR {return_code}')
   return return_code