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

xpk/commands/cluster_gcluster.py

@@ -0,0 +1,185 @@
+"""
+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.blueprint.blueprint_generator import BlueprintGenerator, BlueprintGeneratorOutput, supported_device_types, a3mega_device_type, a3ultra_device_type
+from ..core.docker_manager import DockerManager
+from ..core.gcluster_manager import GclusterManager
+from ..core.core import zone_to_region, get_capacity_type
+from ..utils.console import xpk_exit, xpk_print
+from ..utils.network import all_IPs_cidr
+from ..utils.file import ensure_directory_exists
+from ..utils.objects import hash_string
+import os
+
+blueprints_path = os.path.abspath('xpkclusters/blueprints')
+gcluster_working_dir = os.path.abspath('xpkclusters/gcluster-out')
+gcloud_cfg_path = os.path.expanduser('~/.config/gcloud')
+
+
+def cluster_create(args) -> None:
+  """Function around cluster creation using Cluster toolkit.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  check_gcloud_authenticated()
+  prepare_directories()
+  gcm = prepare_gcluster_manager()
+  region = zone_to_region(args.zone)
+
+  # unique_name uses shortened hash string, so still name collision is possible
+  unique_name = get_unique_name(args.project, region, args.cluster)
+  # prefix is to prevent name collisions for blueprints and also deployments by storing them in prefix directory. Ex.: blueprints/{prefix}/cluster_name_hash
+  prefix = get_prefix_path(args.project, region)
+
+  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,
+  )
+
+  xpk_exit(0)
+
+
+def cluster_delete(args) -> None:
+  """Function around cluster delete for the clusters created by Cluster toolkit.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  check_gcloud_authenticated()
+  prepare_directories()
+  gcm = prepare_gcluster_manager()
+  region = zone_to_region(args.zone)
+
+  # unique_name uses shortened hash string, so still name collision is possible
+  unique_name = get_unique_name(args.project, region, args.cluster)
+  # prefix is to prevent name collisions for blueprints and also deployments by storing them in prefix directory. Ex.: blueprints/{prefix}/cluster_name_hash
+  prefix_path = get_prefix_path(args.project, region)
+
+  gcm.destroy_deployment(deployment_name=unique_name, prefix=prefix_path)
+
+  xpk_exit(0)
+
+
+def created_by_gcluster(args) -> bool:
+  prepare_directories()
+  region = zone_to_region(args.zone)
+  unique_name = get_unique_name(args.project, region, args.cluster)
+  prefix = get_prefix_path(args.project, region)
+  bpg = prepare_blueprint_generator()
+  return bpg.blueprint_exists(unique_name, prefix)
+
+
+def get_unique_name(project_id, region, cluster_name):
+  unique_string_hash = hash_string(
+      input_string=f'{project_id}-{region}-{cluster_name}'.lower(), length=5
+  )
+  return f'{cluster_name}-{unique_string_hash}'
+
+
+def get_prefix_path(project_id, region):
+  return f'{project_id}-{region}'.lower()
+
+
+def prepare_directories() -> None:
+  ensure_directory_exists(blueprints_path)
+  ensure_directory_exists(gcluster_working_dir)
+
+
+def check_gcloud_authenticated():
+  if not os.path.exists(gcloud_cfg_path):
+    xpk_print(
+        'Failed to find gcloud credential directory.'
+        f' {gcloud_cfg_path} {blueprints_path} {gcluster_working_dir}'
+    )
+    xpk_print(
+        'Please authenticate to gcloud ("gcloud auth application-default'
+        ' login") and then run your command.'
+    )
+    xpk_exit(1)
+
+
+def prepare_gcluster_manager() -> GclusterManager:
+  dm = DockerManager(
+      working_dir=gcluster_working_dir, gcloud_cfg_path=gcloud_cfg_path
+  )
+  dm.initialize()
+  return GclusterManager(gcluster_command_runner=dm)
+
+
+def prepare_blueprint_generator() -> BlueprintGenerator:
+  return BlueprintGenerator(storage_path=blueprints_path)
+
+
+def generate_blueprint(
+    blueprint_name, args, prefix=None
+) -> BlueprintGeneratorOutput:
+  capacity_type, return_code = get_capacity_type(args)
+  if return_code != 0:
+    xpk_print('Capacity type is invalid.')
+    xpk_exit(return_code)
+
+  bpg = prepare_blueprint_generator()
+
+  if args.device_type in supported_device_types:
+    if args.device_type == a3mega_device_type:
+      num_nodes = args.num_nodes if not args.num_nodes is None else 2
+      return bpg.generate_a3_mega_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,
+      )
+    if args.device_type == a3ultra_device_type:
+      num_nodes = args.num_nodes if not args.num_nodes is None else 2
+      return bpg.generate_a3_ultra_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/info.py

@@ -0,0 +1,245 @@
+"""
+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 ..utils.console import xpk_exit, xpk_print
+from ..core.kueue import verify_kueuectl
+from .cluster import set_cluster_command
+from ..core.commands import (
+    run_command_for_value,
+)
+from ..core.core import (
+    add_zone_and_project,
+)
+import json
+from tabulate import tabulate
+from argparse import Namespace
+
+table_fmt = 'plain'
+
+
+def info(args: Namespace) -> None:
+  """Provide info about localqueues, clusterqueues and their resources.
+
+  Args:
+    args: user provided arguments for running the command.
+  Returns:
+    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)
+
+  verify_kueuectl(args)
+  lq, cq = bool(args.localqueue), bool(args.clusterqueue)
+  if not lq and not cq:
+    lq, cq = True, True
+
+  lqs, cqs = None, None
+  if lq:
+    lqs = run_kueuectl_list_localqueue(args)
+
+  cqs = run_kueuectl_list_clusterqueue(args)
+  quotas = get_nominal_quotas(cqs)
+
+  if lq:
+    print_formatted_lqs(lqs, quotas)
+
+  if cq:
+    print_formatted_cqs(cqs, quotas)
+
+
+def get_nominal_quotas(cqs: list[dict]) -> dict[str, dict[str, str]]:
+  """Get quotas from clusterqueues.
+  This function retrieves how much of resource in each flavor is assigned to cluster queue.
+  It parses flavors of passed cluster queues.
+  Args:
+    - cqs - list of cluster queues.
+  Returns:
+    - dictionary of cluster queues resources quotas in format:
+    {cq_name:{"flavorName:resourceName":quota}}
+  """
+  try:
+    cq_list = json.loads(cqs)['items']
+  except ValueError:
+    xpk_print('Incorrect respone from list clusterqueue')
+    xpk_print(cqs)
+    xpk_exit(1)
+
+  quotas = {}
+  for cq in cq_list:
+    spec = cq['spec']
+    cq_name = cq['metadata']['name']
+    quotas[cq_name] = {}
+    for rg in spec['resourceGroups']:
+      for flavor in rg['flavors']:
+        name = flavor['name']
+        for resource in flavor['resources']:
+          key = f'{name}:{resource["name"]}'
+          quotas[cq_name][key] = resource['nominalQuota']
+  return quotas
+
+
+def print_formatted_cqs(cqs: list[dict], nominalQuotas) -> None:
+  try:
+    cq_list = json.loads(cqs)['items']
+  except ValueError:
+    xpk_print('Incorrect respone from list clusterqueue')
+    xpk_print(cqs)
+    xpk_exit(1)
+
+  cq_usages = parse_queue_lists(cq_list, nominalQuotas)
+
+  xpk_print(
+      'Cluster Queues usage \n',
+      tabulate(cq_usages, headers='keys', tablefmt=table_fmt),
+  )
+
+
+def print_formatted_lqs(lqs: list[dict], nominalQuotas) -> None:
+  try:
+    lq_list = json.loads(lqs)['items']
+  except ValueError:
+    xpk_print('Incorrect respone from list localqueue')
+    xpk_print(lqs)
+    xpk_exit(1)
+
+  lq_usages = parse_queue_lists(lq_list, nominalQuotas)
+  xpk_print(
+      'Local Queues usage \n',
+      tabulate(lq_usages, headers='keys', tablefmt=table_fmt),
+  )
+
+
+def parse_queue_lists(
+    qs: list[dict],
+    flavor_resource_quotas: dict,
+    reservation_key: str = 'flavorsReservation',
+) -> list[dict]:
+  qs_usage_list = []
+  for q in qs:
+    queue_name = q['metadata']['name']
+    q_pending_workloads = q['status']['pendingWorkloads']
+    q_admitted_workloads = q['status']['admittedWorkloads']
+    q_status = {
+        'QUEUE': queue_name,
+        'ADMITTED_WORKLOADS': q_admitted_workloads,
+        'PENDING_WORKLOADS': q_pending_workloads,
+    }
+    q_status.update(
+        get_flavors_usage(q, reservation_key, flavor_resource_quotas)
+    )
+    qs_usage_list.append(q_status)
+  return qs_usage_list
+
+
+def get_flavors_resources_reservations(
+    cq_name: str, flavors_res: list[dict]
+) -> dict[str, dict[str, str]]:
+  """Get usage of flavors resources.
+  This function parser flavorsReservation section of clusterQueue of LocalQueue.
+  Args:
+    - cq_name - name of ClusterQueue to which flavors belong.
+    - flavors_res - list of reservations made by flavors
+  Returns:
+    Dict containing usage of each resource in flavor for each flavor in cluster or local queue.
+    Dict format: {cq_name: {{flavor:resource}:reservation}}
+  """
+  reservations = {}
+  reservations[cq_name] = {}
+  for flavor_name, flavor_resources_reservation_list in flavors_res.items():
+    for resource in flavor_resources_reservation_list:
+      reservations[cq_name][f'{flavor_name}:{resource["name"]}'] = resource[
+          'total'
+      ]
+
+  return reservations
+
+
+def get_flavors_usage(
+    q_entry: dict, res_field: str, flavor_resource_quotas: dict
+) -> list[dict]:
+  """Parse q_entry to retrieve list of each resource usage in flavour.
+  Args:
+    q_entry - single entry into either LocalQueue or ClusterQueue structured as json
+    flavor_resource_quotas - nominalQuota of flavors resource usage for each clusterqueue
+  Returns:
+    list of dicts where each list entry is in format (key, entry) where:
+    - key is flavorName:resourceName
+    - entry is flavorResourceReservation/flavorResourceQuota
+  """
+  status = q_entry['status']
+  flavors_res = status[res_field]
+  queue_type = q_entry['kind']
+
+  flavors_res = {flavor['name']: flavor['resources'] for flavor in flavors_res}
+  usage_fraction = {}
+  cq_name = (
+      q_entry['metadata']['name']
+      if queue_type == 'ClusterQueue'
+      else q_entry['spec']['clusterQueue']
+  )
+
+  reservations = get_flavors_resources_reservations(cq_name, flavors_res)
+
+  for cq_name, cq_reservations in reservations.items():
+    cq_nominal_quotas = flavor_resource_quotas[cq_name]
+
+    for flavor_resource, flavor_resource_quota in cq_nominal_quotas.items():
+      flavor_resource_reservation = cq_reservations[flavor_resource]
+      usage_fraction[flavor_resource] = (
+          f'{flavor_resource_reservation}/{flavor_resource_quota}'
+      )
+  return usage_fraction
+
+
+def run_kueuectl_list_localqueue(args: Namespace) -> str:
+  """Run the kueuectl list localqueue command.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    kueuectl list localqueue formatted as json string.
+  """
+  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)
+
+  if return_code != 0:
+    xpk_print(f'Cluster info request returned ERROR {return_code}')
+    xpk_exit(return_code)
+  return val
+
+
+def run_kueuectl_list_clusterqueue(args: Namespace) -> 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)
+
+  if return_code != 0:
+    xpk_print(f'Cluster info request returned ERROR {return_code}')
+    xpk_exit(return_code)
+  return val