xpk 0.5.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

xpk/core/kueue.py

@@ -0,0 +1,352 @@
+"""
+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 argparse import Namespace
+from packaging.version import Version
+import packaging
+from ..utils.file import write_tmp_file
+from ..utils.console import xpk_print, xpk_exit
+from .commands import run_command_with_updates, run_command_with_updates_retry, run_command_for_value
+from .core import (
+    AutoprovisioningConfig,
+    create_accelerator_label,
+    create_machine_label,
+    get_total_chips_requested_from_args,
+)
+from .pathways import add_pw_resource_flavors, add_pw_resources_to_kueue
+from .system_characteristics import (
+    AcceleratorTypeToAcceleratorCharacteristics,
+    SystemCharacteristics,
+)
+
+KUEUE_VERSION = 'v0.10.0'
+CLUSTER_QUEUE_NAME = 'cluster-queue'
+LOCAL_QUEUE_NAME = 'multislice-queue'
+WAIT_FOR_KUEUE_TIMEOUT = '5m'
+
+packaging.version.VERSION_PATTERN = r'^v\d+\.\d+\.\d+$'
+
+cluster_set_crd_yaml = """apiVersion: kueue.x-k8s.io/v1beta1
+kind: ResourceFlavor
+metadata:
+  name: {cluster_hardware_name}
+spec:
+  nodeLabels:
+    {accelerator_label}
+    {machine_label}
+---
+{pw_resource_flavors}
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: ClusterQueue
+metadata:
+  name: {cluster_queue_name}
+spec:
+  preemption:
+      reclaimWithinCohort: Never # Don't preempt other queues in the cohort.
+      withinClusterQueue: LowerPriority
+  namespaceSelector: {{}} # match all.
+  resourceGroups:
+  {covered_resources_config}
+  {pw_resources_kueue}
+---
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: LocalQueue
+metadata:
+  namespace: default
+  name: {local_queue_name}
+spec:
+  clusterQueue: {cluster_queue_name}
+---
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: very-low
+value: 100
+globalDefault: false
+description: "Very Low"
+---
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: low
+value: 250
+globalDefault: false
+description: "Low"
+---
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: medium
+value: 500
+globalDefault: false
+description: "Medium"
+---
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: high
+value: 750
+globalDefault: false
+description: "High"
+---
+apiVersion: scheduling.k8s.io/v1
+kind: PriorityClass
+metadata:
+  name: very-high
+value: 1000
+globalDefault: false
+description: "Very High"
+"""
+
+cluster_preheat_yml = """
+apiVersion: apps/v1
+kind: DaemonSet
+metadata:
+  name: {cachekey}
+  labels:
+    k8s-app: {cachekey}
+spec:
+  selector:
+    matchLabels:
+      k8s-app: {cachekey}
+  updateStrategy:
+    type: RollingUpdate
+  template:
+    metadata:
+      labels:
+        name: {cachekey}
+        k8s-app: {cachekey}
+    spec:
+      affinity:
+        nodeAffinity:
+          requiredDuringSchedulingIgnoredDuringExecution:
+            nodeSelectorTerms:
+            - matchExpressions:
+              - key: {nodeSelectorKey}
+                operator: Exists
+      tolerations:
+      - operator: "Exists"
+      containers:
+      - image: {image_name}
+        name: {cachekey}
+        command: [ "sleep", "inf" ]
+"""
+
+
+def verify_kueuectl(args: Namespace) -> None:
+  """Verify if kueuectl is installed.
+  Args:
+    args: user provided arguments.
+  Returns:
+    None
+  """
+  xpk_print('Veryfing kueuectl installation')
+
+  command = 'kubectl kueue version'
+  task = 'Verify kueuectl installation on cluster'
+  verify_kueuectl_installed_code, _ = run_command_for_value(command, task, args)
+
+  if verify_kueuectl_installed_code == 0:
+    xpk_print('kueuectl found')
+
+  if verify_kueuectl_installed_code != 0:
+    xpk_print(
+        'kueuectl not found. Please follow'
+        ' https://kueue.sigs.k8s.io/docs/reference/kubectl-kueue/installation/'
+        ' to install kueuectl.'
+    )
+    xpk_exit(verify_kueuectl_installed_code)
+
+
+def delete_multikueueconfigs_definitions(args) -> int:
+  command = 'kubectl delete crd multikueueconfigs.kueue.x-k8s.io'
+  task = 'Delete multikueueconfigs crds'
+  return_code = run_command_with_updates_retry(command, task, args)
+  if return_code != 0:
+    xpk_print(f'{task} returned ERROR {return_code}')
+  return return_code
+
+
+def delete_multikueueclusters_definitions(args) -> int:
+  command = 'kubectl delete crd multikueueclusters.kueue.x-k8s.io'
+  task = 'Delete multikueueclusters crds'
+  return_code = run_command_with_updates_retry(command, task, args)
+  if return_code != 0:
+    xpk_print(f'{task} returned ERROR {return_code}')
+  return return_code
+
+
+def get_kueue_version(args) -> (int, str):
+  command = 'kubectl kueue version'
+  task = 'Get kueue version on server'
+  return_code, val = run_command_for_value(command, task, args)
+  if return_code != 0:
+    return return_code, ''
+  lines = val.splitlines()
+  if len(lines) == 1:
+    return 1, ''
+  server_version_line = lines[1]
+  manager_image_version = server_version_line.split(':')[-1]
+  return return_code, manager_image_version
+
+
+def install_kueue_on_cluster(args) -> int:
+  """Install Kueue on the cluster.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+
+  err_code, kueue_version_installed = get_kueue_version(args)
+  if err_code == 0:
+    if Version(kueue_version_installed) < Version('v0.9.0') and Version(
+        KUEUE_VERSION
+    ) >= Version('v0.9.0'):
+      xpk_print('Upgrading kueue on cluster from version < 0.9.0.')
+      upgrade_code = delete_multikueueclusters_definitions(args)
+      if upgrade_code != 0:
+        return upgrade_code
+      upgrade_code = delete_multikueueconfigs_definitions(args)
+      if upgrade_code != 0:
+        return upgrade_code
+
+  command = (
+      'kubectl apply --server-side --force-conflicts -f'
+      f' https://github.com/kubernetes-sigs/kueue/releases/download/{KUEUE_VERSION}/manifests.yaml'
+  )
+  task = 'Set Kueue On Cluster'
+  return_code = run_command_with_updates_retry(command, task, args)
+  if return_code != 0:
+    xpk_print(f'{task} returned ERROR {return_code}')
+  return return_code
+
+
+def wait_for_kueue_available(args: Namespace) -> int:
+  """Wait for Kueue to be fully available.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = (
+      'kubectl wait deploy/kueue-controller-manager -nkueue-system'
+      f' --for=condition=available --timeout={WAIT_FOR_KUEUE_TIMEOUT}'
+  )
+  task = 'Wait for Kueue to be available'
+  return_code = run_command_with_updates(command, task, args)
+  if return_code != 0:
+    xpk_print(f'{task} returned ERROR {return_code}')
+  return return_code
+
+
+def install_kueue_crs(
+    args,
+    system: SystemCharacteristics,
+    autoprovisioning_config: AutoprovisioningConfig | None,
+) -> int:
+  """Install Kueue Custom Resources.
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system level arguments.
+    autoprovisioning_config: Autoprovisioning config to configure kueue with if
+        autoprovisioning is enabled.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  device_type = system.device_type
+  cluster_hardware_name = f'{args.num_slices}x{device_type}'
+  resource_type = AcceleratorTypeToAcceleratorCharacteristics[
+      system.accelerator_type
+  ].resource_type
+
+  autoprovisioning_enabled = False
+  if autoprovisioning_config:
+    # Determine total resources available based on autoprovisioning max chips.
+    autoprovisioning_enabled = True
+    total_chips = autoprovisioning_config.maximum_chips
+    cluster_hardware_name = f'{system.gke_accelerator}'
+  else:
+    # Determine total chips based on user specified topology.
+    total_chips = get_total_chips_requested_from_args(args, system)
+
+  covered_resources_config = get_kueue_covered_resources_config(
+      cluster_hardware_name=cluster_hardware_name,
+      resource_type=resource_type,
+      total_chips=total_chips,
+  )
+  yml_string = cluster_set_crd_yaml.format(
+      system=system,
+      cluster_hardware_name=cluster_hardware_name,
+      accelerator_label=create_accelerator_label(
+          system.accelerator_type, system
+      ),
+      machine_label=create_machine_label(
+          system.accelerator_type, system, autoprovisioning_enabled
+      ),
+      covered_resources_config=covered_resources_config,
+      resource_type=AcceleratorTypeToAcceleratorCharacteristics[
+          system.accelerator_type
+      ].resource_type,
+      pw_resource_flavors=add_pw_resource_flavors(args),
+      pw_resources_kueue=add_pw_resources_to_kueue(args),
+      cluster_queue_name=CLUSTER_QUEUE_NAME,
+      local_queue_name=LOCAL_QUEUE_NAME,
+  )
+
+  tmp = write_tmp_file(yml_string)
+  command = f'kubectl apply -f {str(tmp.file.name)}'
+
+  task = 'Applying Kueue Custom Resources'
+  return_code = run_command_with_updates_retry(command, task, args)
+  if return_code != 0:
+    xpk_print(f'{task} returned ERROR {return_code}')
+  return return_code
+
+
+def get_kueue_covered_resources_config(
+    cluster_hardware_name, resource_type, total_chips
+) -> str:
+  """Gets Kueue covered resources configuration.
+
+  Args:
+    cluster_hardware_name: cluster hardware name.
+    resource_type: resource type of tpu or gpu.
+    total_chips: total number of chips for the specific resource type.
+
+  Returns:
+    A string of Kueue covered resources configuration.
+  """
+  config_format = """
+  - coveredResources: ["{resource_type}"]
+    flavors:
+    - name: {cluster_hardware_name}
+      resources:
+      - name: "{resource_type}"
+        nominalQuota: {total_chips}
+  """
+  config_string = config_format.format(
+      cluster_hardware_name=cluster_hardware_name,
+      resource_type=resource_type,
+      total_chips=total_chips,
+  )
+  return config_string

xpk/core/nap.py

@@ -0,0 +1,349 @@
+"""
+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 ..core.core import (
+    AUTOPROVISIONING_CONFIG_VALUE,
+    CAPACITY_TYPE_CONFIG_KEY,
+    CLUSTER_METADATA_CONFIGMAP,
+    CLUSTER_RESOURCES_CONFIGMAP,
+    RESERVATION_CONFIG_KEY,
+    AutoprovisioningConfig,
+    CapacityType,
+    get_all_nodepools_programmatic,
+    get_capacity_node_selectors_from_capacity_type,
+    get_capacity_type,
+    get_cluster_configmap,
+    get_total_chips_requested_from_args,
+    verify_reservation_exists,
+    zone_to_region,
+)
+from ..utils.objects import get_value_from_map
+from ..utils.file import write_tmp_file
+from ..utils.console import xpk_print
+from .commands import run_command_with_updates, run_commands
+from .system_characteristics import AcceleratorType, SystemCharacteristics
+
+autoprovisioning_config_file = """
+management:
+  autoRepair: true
+  autoUpgrade: true
+autoprovisioningLocations:
+  {zones}
+{resource_limits}
+"""
+
+autoprovisioning_resource_limits = """
+resourceLimits:
+- resourceType: 'cpu'
+  {cpu_limits}
+- resourceType: 'memory'
+  {memory_limits}
+{custom_resource_type}
+"""
+
+autoprovisioning_custom_resource_type = """
+- resourceType: {resource_type}
+  minimum: {minimum}
+  maximum: {maximum}
+"""
+
+
+def enable_autoprovisioning_on_cluster(
+    args, system: SystemCharacteristics | None
+) -> tuple[AutoprovisioningConfig | None, int]:
+  """Enable autoprovisioning on the cluster.
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system characteristics.
+
+  Returns:
+    Autoprovisioning Config or None.
+    0 if successful and 1 otherwise.
+  """
+  if not system:
+    return None, 1
+
+  # TODO(@vbarr): Disable NAP if they call xpk cluster create again without --enable-autoprovisioning.
+  # TODO(@vbarr): Support Pathways.
+  # TODO(@vbarr): Support timeout period for idle np before they are deleted.
+  # TODO(@vbarr): Support for hot idle configuration (timeout period is infinity).
+  return_code = 0
+  if system.accelerator_type == AcceleratorType['CPU']:
+    xpk_print("Error: XPK NAP doesn't support Accelerators of Types: CPUs.")
+    return None, 1
+
+  autoprovisioning_config, return_code = create_autoprovisioning_config(
+      args, system
+  )
+  if return_code != 0 or not autoprovisioning_config:
+    xpk_print('Unable to create autoprovisioning config.')
+    return autoprovisioning_config, return_code
+
+  command = (
+      'gcloud container clusters update'
+      f' {args.cluster} --project={args.project}'
+      f' --region={zone_to_region(args.zone)} --enable-autoprovisioning'
+      ' --autoprovisioning-config-file'
+      f' {autoprovisioning_config.config_filename}'
+  )
+  task = 'Update cluster with autoprovisioning enabled'
+  return_code = run_command_with_updates(command, task, args)
+  if return_code != 0:
+    xpk_print(f'{task} request returned ERROR {return_code}')
+    return autoprovisioning_config, return_code
+
+  # Update created accelerator node pools to support autoprovisioning.
+  existing_node_pool_names, return_code = get_all_nodepools_programmatic(args)
+  if return_code != 0:
+    xpk_print('Listing all node pools failed!')
+    return autoprovisioning_config, return_code
+
+  desired_node_pool_names = [
+      f'{args.cluster}-np-{slice_num}' for slice_num in range(args.num_slices)
+  ]
+
+  commands = []
+  task_names = []
+  for node_pool_name in desired_node_pool_names:
+    if node_pool_name not in existing_node_pool_names:
+      # Ignore node pools that are not created yet, and not of the accelerator type.
+      continue
+    commands.append(
+        f'gcloud container node-pools update {node_pool_name}'
+        f' --cluster {args.cluster}'
+        f' --project={args.project}'
+        f' --region={zone_to_region(args.zone)}'
+        ' --enable-autoprovisioning'
+        ' --enable-autoscaling'
+    )
+    task_name = (
+        f'Update node pool {node_pool_name} with autoprovisioning support.'
+    )
+    task_names.append(task_name)
+
+  for i, command in enumerate(commands):
+    xpk_print(f'To complete {task_names[i]} we are executing {command}')
+  max_return_code = run_commands(
+      commands,
+      'Update node pools with autoprovisioning support',
+      task_names,
+      dry_run=args.dry_run,
+  )
+  if max_return_code != 0:
+    xpk_print(
+        'Update node pools with autoprovisioning support returned ERROR:'
+        f' {max_return_code}'
+    )
+    return None, max_return_code
+  return autoprovisioning_config, return_code
+
+
+def create_autoprovisioning_config(
+    args, system: SystemCharacteristics
+) -> tuple[AutoprovisioningConfig | None, int]:
+  """Create autoprovisioning config based on template file and user args
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system characteristics.
+
+  Returns:
+    tuple[AutoprovisioningConfig, int]
+    AutoprovisioningConfig: config used to enable autoprovisioning
+    int: return code
+  """
+
+  # CPU Limits and Memory Limits are for user jobs only. The default node pool
+  # is not controlled by NAP.
+  cpu_limits = """
+  minimum: 1
+  maximum: 10000
+  """
+  memory_limits = """
+  minimum: 1
+  maximum: 10000
+  """
+
+  # By default, the maximum chips is set to be the current number of resources used
+  # in the cluster. The minimum is set to zero.
+  minimum = 0
+  maximum = get_total_chips_requested_from_args(args, system)
+  xpk_print(f'Default Chips quota is minimum: {minimum}, maximum: {maximum}.')
+
+  # Check for user overrides.
+  if args.autoprovisioning_min_chips:
+    minimum = args.autoprovisioning_min_chips
+    xpk_print(
+        f'User provided min chip quota of {minimum}. Overriding defaults.'
+    )
+  if args.autoprovisioning_max_chips:
+    maximum = args.autoprovisioning_max_chips
+    xpk_print(
+        f'User provided max chip quota of {maximum}. Overriding defaults.'
+    )
+
+  # Check for edge cases in min and max chip values.
+  if minimum < 0:
+    xpk_print(
+        f'Error: Minimum chips is set to {minimum}, and must be zero or'
+        ' greater.'
+    )
+    return None, 1
+  if maximum <= minimum or maximum < 0:
+    xpk_print(
+        f'Error: Maximum chips is set to {maximum}, and must be greater than'
+        f' zero and greater or equal to minimum: {minimum}.Use'
+        ' --autoprovisioning-max-chips=$MAX_CHIPS to adjust.'
+    )
+    return None, 1
+  xpk_print(
+      f'Chips quota is minimum: {minimum}, maximum: {maximum}. XPK will'
+      f' autoprovision {maximum - minimum} chips based on incoming workload'
+      f' requests, keeping at least {minimum} available at all times, and'
+      f' maximum of {maximum}. If the difference ({maximum - minimum} chips) is'
+      ' small, rescaling will not work well.'
+  )
+
+  custom_resource_string = autoprovisioning_custom_resource_type.format(
+      resource_type=system.gke_accelerator,
+      minimum=minimum,
+      maximum=maximum,
+  )
+
+  resource_limits = autoprovisioning_resource_limits.format(
+      cpu_limits=cpu_limits,
+      memory_limits=memory_limits,
+      custom_resource_type=custom_resource_string,
+  )
+
+  yml_string = autoprovisioning_config_file.format(
+      resource_limits=resource_limits,
+      zones=f'- {args.zone}',
+  )
+  autoprovisioning_config = AutoprovisioningConfig(
+      config_filename=write_tmp_file(yml_string).name,
+      minimum_chips=minimum,
+      maximum_chips=maximum,
+  )
+  return autoprovisioning_config, 0
+
+
+def is_autoprovisioning_enabled(
+    args, system: SystemCharacteristics
+) -> tuple[bool, int]:
+  """Determine if autoprovisioning is enabled.
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system characteristics.
+
+  Returns:
+    bool is true if autoprovisioning is enabled, false otherwise.
+    int of 0 if successful and 1 otherwise.
+  """
+  resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
+  cluster_config_map = get_cluster_configmap(args, resources_configmap_name)
+
+  if cluster_config_map is None:
+    xpk_print(
+        f'Unable to find config map: {resources_configmap_name}.'
+        ' Autoprovisioning is not enabled.'
+    )
+    return False, 0
+
+  return_code, autoprovisioning_value = get_value_from_map(
+      system.gke_accelerator, cluster_config_map
+  )
+  if return_code != 0:
+    xpk_print(
+        'gke_accelerator type not found in config map:'
+        f' {resources_configmap_name}. Autoprovisioning is not enabled.'
+    )
+    return False, 0
+
+  if autoprovisioning_value == AUTOPROVISIONING_CONFIG_VALUE:
+    xpk_print('Autoprovisioning is Enabled.')
+    return True, 0
+  else:
+    xpk_print(
+        'Error: Autoprovisioning not enabled but should be so exiting xpk.'
+        f' Value should be {AUTOPROVISIONING_CONFIG_VALUE} but instead found'
+        f' value of {autoprovisioning_value}'
+    )
+    return False, 1
+
+
+def get_autoprovisioning_node_selector_args(args) -> tuple[str, int]:
+  """Determine the capacity type when autoprovisioning is enabled.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    Tuple with string of autoprovisioning node selector args and
+    int of 0 if successful and 1 otherwise.
+  """
+  return_code = 0
+  node_selector_args = ''
+  # If the user doesn't specify args, then use the cluster settings.
+  capacity_type, return_code = get_capacity_type(args)
+  capacity_type_str = capacity_type.name
+  if return_code != 0:
+    xpk_print('Unable to get capacity type.')
+    return node_selector_args, return_code
+
+  if capacity_type_str == CapacityType.UNKNOWN.name:
+    # Use default settings from cluster creation.
+    metadata_configmap_name = f'{args.cluster}-{CLUSTER_METADATA_CONFIGMAP}'
+    cluster_config_map = get_cluster_configmap(args, metadata_configmap_name)
+
+    # Error out if the metadata config map doesn't exist, and is attempting to use
+    # autoprovisioning.
+    if cluster_config_map is None:
+      xpk_print(
+          'Unable to find config map. Please specify a capacity type'
+          ' --on-demand, --spot, --reservation=$RESERVATION_ID) to continue'
+          ' to use autoprovisioning (--enable-autoprovisioning).'
+      )
+      return node_selector_args, 1
+
+    return_code, capacity_type_str = get_value_from_map(
+        CAPACITY_TYPE_CONFIG_KEY, cluster_config_map
+    )
+    if return_code != 0:
+      return node_selector_args, return_code
+
+    if capacity_type_str == CapacityType.RESERVATION.name:
+      return_code, args.reservation = get_value_from_map(
+          RESERVATION_CONFIG_KEY, cluster_config_map
+      )
+      if return_code != 0:
+        return node_selector_args, return_code
+      return_code = verify_reservation_exists(args)
+      if return_code > 0:
+        xpk_print('Unable to verify reservation name saved in config map.')
+        return node_selector_args, return_code
+
+  # Check if reservation id is valid. Shared function with cluster creation.
+  node_selector_args, return_code = (
+      get_capacity_node_selectors_from_capacity_type(args, capacity_type_str)
+  )
+  if return_code != 0:
+    xpk_print('Unable to get node selectors from capacity type.')
+    return node_selector_args, return_code
+
+  return node_selector_args, return_code