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

xpk/commands/cluster_gcluster.py

@@ -16,6 +16,12 @@ limitations under the License.
 
 import os
 
+from ..utils.execution_context import is_dry_run
+from ..core.kueue_manager import KueueConfig, KueueManager
+from ..core.nap import enable_autoprovisioning_on_cluster
+from ..core.scheduling import get_total_chips_requested_from_args
+from ..core.system_characteristics import get_system_characteristics
+
 from ..core.blueprint.blueprint_generator import (
     BlueprintGenerator,
     BlueprintGeneratorOutput,
@@ -75,22 +81,29 @@ def cluster_create(args) -> None:
   bp = generate_blueprint(blueprint_name=unique_name, args=args, prefix=prefix)
 
   # staging: sending the blueprint file(s) to gcluster's working directory
-  bp_staged_path = gcm.stage_files(
-      blueprint_file=bp.blueprint_file,
-      blueprint_dependencies=bp.blueprint_dependencies,
-      prefix=prefix,
-  )
-  gcm.deploy(
-      blueprint_path=bp_staged_path,
-      deployment_name=unique_name,
-      prefix=prefix,
-  )
-  if args.cluster_state_gcs_bucket is not None:
-    gcm.upload_state()
+  if is_dry_run():
+    xpk_print(f'Blueprint file: {bp.blueprint_file}')
+  else:
+    bp_staged_path = gcm.stage_files(
+        blueprint_file=bp.blueprint_file,
+        blueprint_dependencies=bp.blueprint_dependencies,
+        prefix=prefix,
+    )
+    gcm.deploy(
+        blueprint_path=bp_staged_path,
+        deployment_name=unique_name,
+        prefix=prefix,
+    )
+    if args.cluster_state_gcs_bucket is not None:
+      gcm.upload_state()
 
   get_cluster_credentials(args)
 
-  err_code = apply_kjob_crds(args)
+  err_code = __install_kueue(args)
+  if err_code > 0:
+    xpk_exit(err_code)
+
+  err_code = apply_kjob_crds()
   if err_code > 0:
     xpk_exit(err_code)
 
@@ -101,6 +114,57 @@ def cluster_create(args) -> None:
   xpk_exit(0)
 
 
+def __install_kueue(args) -> int:
+  system, return_code = get_system_characteristics(args)
+
+  if return_code > 0 or system is None:
+    xpk_print('Fetching system characteristics failed!')
+    return return_code
+
+  # 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:
+      return return_code
+
+  autoprovisioning_enabled = False
+  if autoprovisioning_config:
+    # Determine total resources available based on autoprovisioning max chips.
+    autoprovisioning_enabled = True
+    total_chips = autoprovisioning_config.maximum_chips
+  else:
+    # Determine total chips based on user specified topology.
+    total_chips = get_total_chips_requested_from_args(args, system)
+  kueue_manager = KueueManager()
+
+  tolerations = [{
+      'key': 'components.gke.io/gke-managed-components',
+      'operator': 'Equal',
+      'value': 'true',
+      'effect': 'NoSchedule',
+  }]
+
+  kueue_manager.install_or_upgrade(
+      KueueConfig(
+          system,
+          total_chips=total_chips,
+          autoprovisioning_enabled=autoprovisioning_enabled,
+          num_slices=args.num_slices,
+          memory_limit=args.memory_limit,
+          cpu_limit=args.cpu_limit,
+          is_pathways_cluster=args.enable_pathways,
+          flex=args.flex,
+      ),
+      tolerations=tolerations,
+  )
+  return 0
+
+
 def cluster_delete(args) -> None:
   """Function around cluster delete for the clusters created by Cluster toolkit.
 
@@ -213,7 +277,6 @@ def validate_state_gcs_bucket(args):
   err_code, _ = run_command_for_value(
       bucket_validate_cmd,
       'Validate remote state bucket existence.',
-      global_args=args,
   )
   if err_code != 0:
     xpk_exit(err_code)

xpk/commands/cluster_gcluster_test.py

@@ -0,0 +1,177 @@
+"""
+Copyright 2024 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 unittest.mock import MagicMock, patch
+
+import pytest
+
+from xpk.commands.cluster_gcluster import cluster_create
+from xpk.core.kueue_manager import KueueConfig
+from xpk.core.system_characteristics import AcceleratorType, SystemCharacteristics
+
+
+@pytest.fixture
+def mock_args():
+  """Provides a mock for args."""
+  args = MagicMock()
+  args.enable_autoprovisioning = False
+  args.num_slices = 1
+  args.memory_limit = "200G"
+  args.cpu_limit = "50"
+  args.enable_pathways = False
+  args.flex = False
+  args.project = "test-project"
+  args.cluster = "test-cluster"
+  args.zone = "us-central1-c"
+  args.cluster_state_gcs_bucket = None
+  return args
+
+
+@pytest.fixture
+def mock_cluster_create_deps(request):
+  """Mocks dependencies for cluster_create."""
+  with (
+      patch("xpk.commands.cluster_gcluster.xpk_exit") as mock_exit,
+      patch("xpk.commands.cluster_gcluster.prepare_kjob") as mock_prep_kjob,
+      patch("xpk.commands.cluster_gcluster.apply_kjob_crds") as mock_apply_kjob,
+      patch(
+          "xpk.commands.cluster_gcluster.get_cluster_credentials"
+      ) as mock_get_creds,
+      patch("xpk.commands.cluster_gcluster.generate_blueprint") as mock_gen_bp,
+      patch(
+          "xpk.commands.cluster_gcluster.prepare_gcluster_manager"
+      ) as mock_prep_gcm,
+      patch(
+          "xpk.commands.cluster_gcluster.prepare_directories"
+      ) as mock_prep_dirs,
+      patch(
+          "xpk.commands.cluster_gcluster.check_gcloud_authenticated"
+      ) as mock_check_auth,
+      patch(
+          "xpk.commands.cluster_gcluster.get_system_characteristics"
+      ) as mock_get_sys_char,
+      patch("xpk.commands.cluster_gcluster.KueueManager") as mock_kueue_manager,
+  ):
+    yield {
+        "xpk_exit": mock_exit,
+        "prepare_kjob": mock_prep_kjob,
+        "apply_kjob_crds": mock_apply_kjob,
+        "get_cluster_credentials": mock_get_creds,
+        "generate_blueprint": mock_gen_bp,
+        "prepare_gcluster_manager": mock_prep_gcm,
+        "prepare_directories": mock_prep_dirs,
+        "check_gcloud_authenticated": mock_check_auth,
+        "get_system_characteristics": mock_get_sys_char,
+        "KueueManager": mock_kueue_manager,
+    }
+
+
+@patch("xpk.commands.cluster_gcluster.get_total_chips_requested_from_args")
+def test_install_kueue_standard(
+    mock_get_total_chips, mock_args, mock_cluster_create_deps
+):
+  """Tests __install_kueue for a standard installation."""
+  mock_cluster_create_deps["prepare_kjob"].return_value = 0
+  mock_cluster_create_deps["apply_kjob_crds"].return_value = 0
+
+  mock_system = SystemCharacteristics(
+      topology="N/A",
+      vms_per_slice=1,
+      gke_accelerator="nvidia-h100-mega-80gb",
+      gce_machine_type="a3-megagpu-8g",
+      chips_per_vm=8,
+      accelerator_type=AcceleratorType["GPU"],
+      device_type="h100-mega-80gb-8",
+      supports_sub_slicing=False,
+  )
+  mock_cluster_create_deps["get_system_characteristics"].return_value = (
+      mock_system,
+      0,
+  )
+  mock_get_total_chips.return_value = 16
+
+  cluster_create(mock_args)
+
+  mock_cluster_create_deps["xpk_exit"].assert_called_with(0)
+  mock_kueue_manager = mock_cluster_create_deps["KueueManager"]
+  mock_kueue_manager.return_value.install_or_upgrade.assert_called_once()
+  call_args, call_kwargs = (
+      mock_kueue_manager.return_value.install_or_upgrade.call_args
+  )
+  kueue_config: KueueConfig = call_args[0]
+
+  assert kueue_config.system == mock_system
+  assert kueue_config.total_chips == 16
+  assert not kueue_config.autoprovisioning_enabled
+  assert "tolerations" in call_kwargs
+  tolerations = call_kwargs["tolerations"]
+  assert any(
+      t.get("key") == "components.gke.io/gke-managed-components"
+      and t.get("effect") == "NoSchedule"
+      for t in tolerations
+  )
+
+
+@patch("xpk.commands.cluster_gcluster.enable_autoprovisioning_on_cluster")
+def test_install_kueue_with_autoprovisioning(
+    mock_enable_autoprovisioning, mock_args, mock_cluster_create_deps
+):
+  """Tests __install_kueue with autoprovisioning enabled."""
+  mock_cluster_create_deps["prepare_kjob"].return_value = 0
+  mock_cluster_create_deps["apply_kjob_crds"].return_value = 0
+
+  mock_args.enable_autoprovisioning = True
+  mock_system = SystemCharacteristics(
+      topology="N/A",
+      vms_per_slice=1,
+      gke_accelerator="nvidia-h100-mega-80gb",
+      gce_machine_type="a3-megagpu-8g",
+      chips_per_vm=8,
+      accelerator_type=AcceleratorType["GPU"],
+      device_type="h100-mega-80gb-8",
+      supports_sub_slicing=False,
+  )
+  mock_cluster_create_deps["get_system_characteristics"].return_value = (
+      mock_system,
+      0,
+  )
+
+  mock_autoprovisioning_config = MagicMock()
+  mock_autoprovisioning_config.maximum_chips = 128
+  mock_enable_autoprovisioning.return_value = (mock_autoprovisioning_config, 0)
+
+  cluster_create(mock_args)
+
+  mock_cluster_create_deps["xpk_exit"].assert_called_with(0)
+  mock_enable_autoprovisioning.assert_called_once_with(mock_args, mock_system)
+  mock_kueue_manager = mock_cluster_create_deps["KueueManager"]
+  mock_kueue_manager.return_value.install_or_upgrade.assert_called_once()
+
+  call_args, call_kwargs = (
+      mock_kueue_manager.return_value.install_or_upgrade.call_args
+  )
+  kueue_config: KueueConfig = call_args[0]
+
+  assert kueue_config.system == mock_system
+  assert kueue_config.total_chips == 128
+  assert kueue_config.autoprovisioning_enabled
+  assert "tolerations" in call_kwargs
+  tolerations = call_kwargs["tolerations"]
+  assert any(
+      t.get("key") == "components.gke.io/gke-managed-components"
+      and t.get("effect") == "NoSchedule"
+      for t in tolerations
+  )

xpk/commands/common.py

@@ -16,7 +16,7 @@ limitations under the License.
 
 from ..core.commands import run_command_with_updates_retry
 from ..core.capacity import H100_MEGA_DEVICE_TYPE, CapacityType
-from ..core.gcloud_context import zone_to_region
+from ..core.gcloud_context import get_cluster_location
 from ..utils.console import xpk_print, xpk_exit
 from ..utils.execution_context import is_dry_run
 from ..core.system_characteristics import (
@@ -35,16 +35,12 @@ 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'
+      f' {args.cluster} --location={get_cluster_location(args.project, args.cluster, args.zone)} --dns-endpoint'
+      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_retry(command, task, verbose=False)
   if return_code != 0:
     xpk_print(f'{task} returned ERROR {return_code}')
   return return_code
@@ -53,16 +49,8 @@ def set_cluster_command(args) -> int:
 def is_TAS_possible(
     system_characteristics: SystemCharacteristics | None,
     capacity_type: CapacityType | None,
-    flex: bool,
 ) -> 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.
-  """
+  """Check cluster's machine_type and capacity type to determine if Kueue TAS is possible"""
 
   if is_dry_run():
     return True
@@ -75,13 +63,7 @@ def is_TAS_possible(
     xpk_print('capacity_type data was not found in configmaps.')
     xpk_exit(1)
 
-  if not flex:
-    return False
-
-  if (
-      system_characteristics.device_type == H100_MEGA_DEVICE_TYPE
-      and capacity_type != CapacityType.RESERVATION
-  ):
-    return False
-
-  return True
+  return (
+      system_characteristics.device_type != H100_MEGA_DEVICE_TYPE
+      or capacity_type == CapacityType.RESERVATION
+  )

xpk/commands/info.py

@@ -22,8 +22,8 @@ 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 ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
 
 table_fmt = 'plain'
 
@@ -36,10 +36,15 @@ def info(args: Namespace) -> None:
   Returns:
     None
   """
+  if should_validate_dependencies(args):
+    validate_dependencies_list([
+        SystemDependency.KUBECTL,
+        SystemDependency.GCLOUD,
+        SystemDependency.KUEUECTL,
+    ])
   add_zone_and_project(args)
   get_cluster_credentials(args)
 
-  verify_kueuectl(args)
   lq, cq = bool(args.localqueue), bool(args.clusterqueue)
   if not lq and not cq:
     lq, cq = True, True
@@ -48,7 +53,7 @@ def info(args: Namespace) -> None:
   if lq:
     lqs = run_kueuectl_list_localqueue(args)
 
-  cqs = run_kueuectl_list_clusterqueue(args)
+  cqs = run_kueuectl_list_clusterqueue()
   quotas = get_nominal_quotas(cqs)
 
   if lq and lqs is not None:
@@ -214,7 +219,7 @@ def run_kueuectl_list_localqueue(args: Namespace) -> str:
   command = 'kubectl kueue list localqueue -o json'
   if args.namespace != '':
     command += f' --namespace {args.namespace}'
-  return_code, val = run_command_for_value(command, 'list localqueue', args)
+  return_code, val = run_command_for_value(command, 'list localqueue')
 
   if return_code != 0:
     xpk_print(f'Cluster info request returned ERROR {return_code}')
@@ -222,18 +227,15 @@ def run_kueuectl_list_localqueue(args: Namespace) -> str:
   return val
 
 
-def run_kueuectl_list_clusterqueue(args: Namespace) -> str:
+def run_kueuectl_list_clusterqueue() -> str:
   """Run the kueuectl list clusterqueue command.
 
-  Args:
-    args: user provided arguments for running the command.
-
   Returns:
     kueuectl list clusterqueue formatted as json string
   """
   command = 'kubectl kueue list clusterqueue -o json'
 
-  return_code, val = run_command_for_value(command, 'list clusterqueue', args)
+  return_code, val = run_command_for_value(command, 'list clusterqueue')
 
   if return_code != 0:
     xpk_print(f'Cluster info request returned ERROR {return_code}')

xpk/commands/inspector.py

@@ -16,11 +16,12 @@ limitations under the License.
 
 from ..core.cluster import get_cluster_credentials
 from ..core.commands import run_command_for_value
-from ..core.gcloud_context import add_zone_and_project, zone_to_region
-from ..core.kueue import CLUSTER_QUEUE_NAME, LOCAL_QUEUE_NAME
+from ..core.gcloud_context import add_zone_and_project, get_cluster_location
+from ..core.kueue_manager import CLUSTER_QUEUE_NAME, LOCAL_QUEUE_NAME
 from ..core.resources import CLUSTER_METADATA_CONFIGMAP, CLUSTER_RESOURCES_CONFIGMAP
 from ..utils.console import xpk_exit, xpk_print
 from ..utils.file import append_tmp_file, write_tmp_file
+from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
 from .workload import get_workload_list
 
 
@@ -41,7 +42,7 @@ def inspector_run_command_helper(
   prefix = f'Command: {command}\nCommand Description: {command_description}\n'
   postfix = '========================================================'
   return_code, command_output = run_command_for_value(
-      command, f'{command_description}', args
+      command, f'{command_description}'
   )
 
   if return_code != 0:
@@ -116,7 +117,10 @@ def inspector(args) -> None:
   # Future Improvements for inspector:
   # 2. List what is next in Queue.
   # 3. Split inspector into different subcommands to parse info easier.
-
+  if should_validate_dependencies(args):
+    validate_dependencies_list(
+        [SystemDependency.KUBECTL, SystemDependency.GCLOUD]
+    )
   final_return_code = 0
   xpk_print(args)
 
@@ -138,8 +142,9 @@ def inspector(args) -> None:
       (
           (
               'gcloud beta container clusters list --project'
-              f' {args.project} --region {zone_to_region(args.zone)} | grep -e'
-              f' NAME -e {args.cluster}'
+              f' {args.project} --location'
+              f' {get_cluster_location(args.project, args.cluster, args.zone)} |'
+              f' grep -e NAME -e {args.cluster}'
           ),
           'GKE: Cluster Details',
       ),
@@ -160,7 +165,7 @@ def inspector(args) -> None:
       (
           (
               f'gcloud beta container node-pools list --cluster {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)}'
           ),
           'GKE: Node pool Details',
       ),
@@ -309,19 +314,25 @@ def inspector(args) -> None:
     workload_links = [(
         f'Cloud Console for the workload {args.workload}',
         # pylint: disable=line-too-long
-        f'https://console.cloud.google.com/kubernetes/service/{zone_to_region(args.zone)}/{args.cluster}/default/{args.workload}/details?project={args.project}',
+        (
+            f'https://console.cloud.google.com/kubernetes/service/{get_cluster_location(args.project, args.cluster, args.zone)}/{args.cluster}/default/{args.workload}/details?project={args.project}'
+        ),
     )]
 
   links = [
       (
           'Cloud Console for the GKE Cluster',
           # 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}',
+          (
+              f'https://console.cloud.google.com/kubernetes/clusters/details/{get_cluster_location(args.project, args.cluster, args.zone)}/{args.cluster}/details?project={args.project}'
+          ),
       ),
       (
           'Cloud Console for all workloads in GKE Cluster',
           # pylint: disable=line-too-long
-          f'https://console.cloud.google.com/kubernetes/workload/overview?project={args.project}&pageState=((gke%2F{zone_to_region(args.zone)}%2F{args.cluster}))',
+          (
+              f'https://console.cloud.google.com/kubernetes/workload/overview?project={args.project}&pageState=((gke%2F{get_cluster_location(args.project, args.cluster, args.zone)}%2F{args.cluster}))'
+          ),
       ),
       (
           'Cloud Console for IAM Permissions',

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