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

xpk/commands/job.py

@@ -25,6 +25,7 @@ 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 ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
 from .kind import set_local_cluster_command
 
 
@@ -59,12 +60,16 @@ def job_info(args):
   Returns:
     None
   """
+  if should_validate_dependencies(args):
+    validate_dependencies_list([
+        SystemDependency.KUBECTL,
+        SystemDependency.KJOB,
+        SystemDependency.GCLOUD,
+    ])
   job_name = args.name
 
   desc_command = f'kubectl-kjob describe slurm {job_name}'
-  desc_code, desc_text = run_command_for_value(
-      desc_command, 'Getting job data', args
-  )
+  desc_code, desc_text = run_command_for_value(desc_command, 'Getting job data')
   if desc_code != 0:
     xpk_print(f'Data info request returned ERROR {desc_code}')
     xpk_exit(desc_code)
@@ -76,7 +81,6 @@ def job_info(args):
   job_code, job_text = run_command_for_value(
       job_command,
       'Getting job info',
-      args,
       dry_run_return_val=JOBS_DRY_RUN_YAML,
   )
   if job_code != 0:
@@ -87,7 +91,6 @@ def job_info(args):
   pods_code, pods_text = run_command_for_value(
       pods_command,
       'Getting pods list',
-      args,
       dry_run_return_val=PODS_DRY_RUN_RESULT,
   )
   if pods_code != 0:
@@ -171,6 +174,12 @@ def job_list(args) -> None:
   Returns:
     None
   """
+  if should_validate_dependencies(args):
+    validate_dependencies_list([
+        SystemDependency.KUBECTL,
+        SystemDependency.KJOB,
+        SystemDependency.GCLOUD,
+    ])
   if not args.kind_cluster:
     add_zone_and_project(args)
     get_cluster_credentials(args)
@@ -183,14 +192,14 @@ def job_list(args) -> None:
 
   xpk_print(msg, flush=True)
 
-  return_code = run_slurm_job_list_command(args)
+  return_code = run_slurm_job_list_command()
   xpk_exit(return_code)
 
 
-def run_slurm_job_list_command(args) -> int:
+def run_slurm_job_list_command() -> int:
   cmd = f'kubectl-kjob list slurm  --profile {AppProfileDefaults.NAME.value}'
 
-  return_code = run_command_with_updates(cmd, 'list jobs', args)
+  return_code = run_command_with_updates(cmd, 'list jobs')
   if return_code != 0:
     xpk_print(f'Listing jobs returned ERROR {return_code}')
   return return_code
@@ -205,6 +214,13 @@ def job_cancel(args) -> None:
   Returns:
     None
   """
+  if should_validate_dependencies(args):
+    validate_dependencies_list([
+        SystemDependency.KUBECTL,
+        SystemDependency.KJOB,
+        SystemDependency.GCLOUD,
+    ])
+
   xpk_print(f'Starting job cancel for job: {args.name}', flush=True)
   if not args.kind_cluster:
     add_zone_and_project(args)
@@ -222,7 +238,7 @@ def run_slurm_job_delete_command(args) -> int:
   list_of_jobs = ' '.join(args.name)
   cmd = f'kubectl-kjob delete slurm {list_of_jobs}'
 
-  return_code = run_command_with_updates(cmd, 'delete job', args)
+  return_code = run_command_with_updates(cmd, 'delete job')
   if return_code != 0:
     xpk_print(f'Delete job request returned ERROR {return_code}')
   return return_code

xpk/commands/kind.py

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from ..core.kueue_manager import (KueueConfig, KueueManager)
 from ..core.commands import (
     run_command_for_value,
     run_command_with_updates,
@@ -24,17 +25,14 @@ from ..core.kjob import (
     prepare_kjob,
     apply_kjob_crds,
 )
-from ..core.kueue import (
-    install_kueue_on_cluster,
-    install_kueue_crs,
-    wait_for_kueue_available,
-)
+from ..core.scheduling import get_total_chips_requested_from_args
 from ..core.storage import install_storage_crd
 from ..core.system_characteristics import (
     SystemCharacteristics,
     AcceleratorType,
 )
 from ..utils.console import (xpk_exit, xpk_print)
+from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
 
 
 def cluster_create(args) -> None:
@@ -46,6 +44,12 @@ def cluster_create(args) -> None:
   Returns:
     0 if successful and 1 otherwise.
   """
+  if should_validate_dependencies(args):
+    validate_dependencies_list([
+        SystemDependency.KUBECTL,
+        SystemDependency.KJOB,
+        SystemDependency.GCLOUD,
+    ])
   xpk_print(f'Starting cluster create for cluster {args.cluster}:', flush=True)
 
   create_cluster_command_code = create_cluster_if_necessary(args)
@@ -64,18 +68,13 @@ 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)
-
   xpk_print('Verifying kjob installation')
-  err_code = verify_kjob_installed(args)
+  err_code = verify_kjob_installed()
   if err_code > 0:
     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)
 
@@ -87,11 +86,6 @@ def cluster_create(args) -> None:
   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)
-
   args.num_slices = 1
   args.enable_pathways = False
   system = SystemCharacteristics(
@@ -102,12 +96,23 @@ def cluster_create(args) -> None:
       1,
       AcceleratorType['CPU'],
       'kind',
+      supports_sub_slicing=False,
   )
 
-  xpk_print('Install Kueue Custom Resources')
-  enable_kueue_credentials_code = install_kueue_crs(args, system, None)
-  if enable_kueue_credentials_code != 0:
-    xpk_exit(enable_kueue_credentials_code)
+  kueue_manager = KueueManager()
+  kueue_manager.install_or_upgrade(
+      KueueConfig(
+          system,
+          total_chips=get_total_chips_requested_from_args(args, system),
+          autoprovisioning_enabled=False,
+          num_slices=args.num_slices,
+          memory_limit='',
+          cpu_limit=0,
+          is_pathways_cluster=False,
+          flex=False,
+          configure_sub_slicing=False,
+      ),
+  )
 
   xpk_print('Kind commands done! Resources are created.')
   xpk_exit(0)
@@ -122,6 +127,8 @@ def cluster_delete(args) -> None:
   Returns:
     0 if successful and 1 otherwise.
   """
+  if should_validate_dependencies(args):
+    validate_dependencies_list([SystemDependency.GCLOUD])
   xpk_print(f'Starting cluster delete for cluster: {args.cluster}', flush=True)
 
   run_kind_cluster_delete_command_code = run_kind_cluster_delete_command(args)
@@ -134,13 +141,12 @@ def cluster_delete(args) -> None:
 def cluster_list(args) -> None:
   """Function around cluster list.
 
-  Args:
-    args: user provided arguments for running the command.
-
   Returns:
     0 if successful and 1 otherwise.
   """
-  if run_kind_clusters_list_command(args):
+  if should_validate_dependencies(args):
+    validate_dependencies_list([SystemDependency.GCLOUD])
+  if run_kind_clusters_list_command():
     xpk_exit(1)
   xpk_exit(0)
 
@@ -154,7 +160,7 @@ def create_cluster_if_necessary(args) -> int:
   Returns:
     0 if successful and 1 otherwise.
   """
-  all_clusters, return_code = get_all_local_clusters_programmatic(args)
+  all_clusters, return_code = get_all_local_clusters_programmatic()
   if return_code > 0:
     xpk_print('Listing all clusters failed!')
     return 1
@@ -179,7 +185,7 @@ def run_kind_cluster_delete_command(args) -> int:
   if args.cluster:
     command += f' --name={args.cluster}'
 
-  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
@@ -187,17 +193,14 @@ def run_kind_cluster_delete_command(args) -> int:
   return 0
 
 
-def run_kind_clusters_list_command(args) -> int:
+def run_kind_clusters_list_command() -> int:
   """List Kind Clusters within the project and location.
 
-  Args:
-    args: user provided arguments for running the command.
-
   Returns:
     0 if successful and 1 otherwise.
   """
   command = 'kind get clusters'
-  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
@@ -222,25 +225,22 @@ def run_kind_cluster_create_command(args) -> int:
   if args.k8s_version:
     command += f' --image=kindest/node:v{args.k8s_version}'
 
-  return_code = run_command_with_updates(command, 'Kind Cluster Create', args)
+  return_code = run_command_with_updates(command, 'Kind Cluster Create')
   if return_code != 0:
     xpk_print(f'GKE Cluster Create request returned ERROR {return_code}')
     return 1
   return 0
 
 
-def get_all_local_clusters_programmatic(args) -> tuple[list[str], int]:
+def get_all_local_clusters_programmatic() -> tuple[list[str], int]:
   """Gets all the local clusters.
 
-  Args:
-    args: user provided arguments for running the command.
-
   Returns:
     List of cluster names and 0 if successful and 1 otherwise.
   """
   command = 'kind get clusters'
   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}')
@@ -261,7 +261,7 @@ def set_local_cluster_command(args) -> int:
   if not args.cluster:
     command = 'kubectl config current-context'
     return_code, current_context = run_command_for_value(
-        command, 'get current-context', args
+        command, 'get current-context'
     )
     xpk_print(
         'No local cluster name specified. Using current-context'
@@ -276,7 +276,6 @@ def set_local_cluster_command(args) -> int:
   return_code = run_command_with_updates(
       command,
       task,
-      args,
   )
   if return_code != 0:
     xpk_print(f'{task} returned ERROR {return_code}')

xpk/commands/kjob_common.py

@@ -35,11 +35,11 @@ def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
 
   annotations: tuple
   if gpu_type == H100_MEGA_DEVICE_TYPE:
-    annotations = get_a3mega_pod_template_annotations(args)
+    annotations = get_a3mega_pod_template_annotations()
   elif gpu_type == H200_DEVICE_TYPE:
-    annotations = get_a3ultra_pod_template_annotations(args)
+    annotations = get_a3ultra_pod_template_annotations()
   elif gpu_type == B200_DEVICE_TYPE:
-    annotations = get_a4_pod_template_annotations(args)
+    annotations = get_a4_pod_template_annotations()
   else:
     annotations = tuple()
 
@@ -54,7 +54,7 @@ def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
 def add_TAS_annotations_to_command(args, cmd: str) -> str:
   system_characteristics = get_cluster_system_characteristics(args)
   capacity_type = get_cluster_capacity_type(args)
-  if is_TAS_possible(system_characteristics, capacity_type, flex=False):
+  if is_TAS_possible(system_characteristics, capacity_type):
     cmd += f" --pod-template-annotation {Kueue_TAS_annotation}"
 
   return cmd

xpk/commands/run.py

@@ -28,8 +28,9 @@ from ..core.kjob import (
     get_storage_annotations,
     prepare_kjob,
 )
-from ..core.kueue import LOCAL_QUEUE_NAME
+from ..core.kueue_manager import LOCAL_QUEUE_NAME
 from ..utils.console import xpk_exit, xpk_print
+from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
 from .kind import set_local_cluster_command
 from .kjob_common import add_gpu_networking_annotations_to_command, add_TAS_annotations_to_command
 
@@ -42,6 +43,12 @@ def run(args: Namespace) -> None:
   Returns:
     None
   """
+  if should_validate_dependencies(args):
+    validate_dependencies_list([
+        SystemDependency.KUBECTL,
+        SystemDependency.KJOB,
+        SystemDependency.GCLOUD,
+    ])
   if not args.kind_cluster:
     add_zone_and_project(args)
     get_cluster_credentials(args)
@@ -126,7 +133,7 @@ def submit_job(args: Namespace) -> None:
   if args.time is not None:
     cmd += f' --time {args.time}'
 
-  return_code = run_command_with_full_controls(cmd, 'run task', args)
+  return_code = run_command_with_full_controls(cmd, 'run task')
 
   if return_code != 0:
     xpk_print(f'Running task returned ERROR {return_code}')

xpk/commands/shell.py

@@ -14,6 +14,7 @@ limitations under the License.
 from ..core.commands import run_command_with_full_controls, run_command_for_value, run_command_with_updates
 from ..core.cluster import get_cluster_credentials, add_zone_and_project, setup_k8s_service_accounts
 from ..utils.console import xpk_exit, xpk_print
+from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
 from argparse import Namespace
 
 from ..core.kjob import (
@@ -33,14 +34,18 @@ def shell(args: Namespace):
   Returns:
     0 if successful and 1 otherwise.
   """
+  if should_validate_dependencies(args):
+    validate_dependencies_list([
+        SystemDependency.KUBECTL,
+        SystemDependency.KJOB,
+        SystemDependency.GCLOUD,
+    ])
   exisitng_shell_pod_name = get_existing_shell_pod_name(args)
 
   if exisitng_shell_pod_name is None:
     return_code = connect_to_new_interactive_shell(args)
   else:
-    return_code = connect_to_existing_interactive_shell(
-        exisitng_shell_pod_name, args
-    )
+    return_code = connect_to_existing_interactive_shell(exisitng_shell_pod_name)
 
   if return_code != 0:
     xpk_print(f'The command failed with code {return_code}.')
@@ -60,7 +65,6 @@ def get_existing_shell_pod_name(args: Namespace) -> str | None:
           ' -o custom-columns=":metadata.name"'
       ),
       task='Get existing interactive shell pod name.',
-      global_args=args,
   )
   if return_code != 0:
     xpk_print(
@@ -95,21 +99,17 @@ def connect_to_new_interactive_shell(args: Namespace) -> int:
   return run_command_with_full_controls(
       command=cmd,
       task='Creating new interactive shell and entering it',
-      global_args=args,
       instructions=exit_instructions,
   )
 
 
-def connect_to_existing_interactive_shell(
-    pod_name: str, args: Namespace
-) -> int:
+def connect_to_existing_interactive_shell(pod_name: str) -> int:
   return run_command_with_full_controls(
       command=(
           f'kubectl exec --stdin --tty {pod_name} --'
           f' {get_pod_template_interactive_command()}'
       ),
       task='Entering existing interactive shell',
-      global_args=args,
       instructions=exit_instructions,
   )
 
@@ -121,6 +121,10 @@ def shell_stop(args: Namespace):
   Returns:
     0 if successful and 1 otherwise.
   """
+  if should_validate_dependencies(args):
+    validate_dependencies_list(
+        [SystemDependency.KUBECTL, SystemDependency.GCLOUD]
+    )
   exisitng_shell_pod_name = get_existing_shell_pod_name(args)
 
   if exisitng_shell_pod_name is None:
@@ -130,7 +134,6 @@ def shell_stop(args: Namespace):
   return_code = run_command_with_updates(
       command=f'kubectl delete pod {exisitng_shell_pod_name}',
       task='Deleting the existing shell.',
-      global_args=args,
   )
   if return_code != 0:
     xpk_exit(return_code)

xpk/commands/storage.py

@@ -59,9 +59,14 @@ from ..core.storage import (
 from ..utils.console import get_user_input, xpk_exit, xpk_print
 from ..utils.kubectl import apply_kubectl_manifest
 from ..utils.execution_context import is_dry_run
+from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
 
 
 def storage_create(args: Namespace) -> None:
+  if should_validate_dependencies(args):
+    validate_dependencies_list(
+        [SystemDependency.KUBECTL, SystemDependency.GCLOUD]
+    )
   add_zone_and_project(args)
   if args.type == GCP_FILESTORE_TYPE:
     if args.instance is None:
@@ -107,6 +112,10 @@ def storage_create(args: Namespace) -> None:
 
 
 def storage_delete(args: Namespace) -> None:
+  if should_validate_dependencies(args):
+    validate_dependencies_list(
+        [SystemDependency.KUBECTL, SystemDependency.GCLOUD]
+    )
   add_zone_and_project(args)
   k8s_api_client = setup_k8s_env(args)
   storages = list_storages(k8s_api_client)
@@ -141,6 +150,10 @@ def storage_delete(args: Namespace) -> None:
 
 
 def storage_attach(args: Namespace) -> None:
+  if should_validate_dependencies(args):
+    validate_dependencies_list(
+        [SystemDependency.KUBECTL, SystemDependency.GCLOUD]
+    )
   add_zone_and_project(args)
   manifest: list[dict] = [{}]
   if args.type == GCP_FILESTORE_TYPE:
@@ -244,6 +257,10 @@ def enable_csi_drivers_if_necessary(args: Namespace) -> None:
 
 
 def storage_list(args: Namespace) -> None:
+  if should_validate_dependencies(args):
+    validate_dependencies_list(
+        [SystemDependency.KUBECTL, SystemDependency.GCLOUD]
+    )
   storages = []
   if not is_dry_run():
     k8s_api_client = setup_k8s_env(args)
@@ -252,6 +269,10 @@ def storage_list(args: Namespace) -> None:
 
 
 def storage_detach(args: Namespace) -> None:
+  if should_validate_dependencies(args):
+    validate_dependencies_list(
+        [SystemDependency.KUBECTL, SystemDependency.GCLOUD]
+    )
   k8s_api_client = setup_k8s_env(args)
   storage = get_storage(k8s_api_client, args.name)
   delete_storage_resources(k8s_api_client, storage)

xpk/commands/version.py

@@ -18,10 +18,6 @@ from ..core.config import __version__
 from ..utils.console import xpk_print
 
 
-def get_xpk_version() -> str:
-  return __version__
-
-
 def version(args) -> None:  # pylint: disable=unused-argument
   """Get version of xpk."""
   xpk_print('xpk_version:', __version__)