xpk 0.0.1__py3-none-any.whl

xpk/commands/cluster_gcluster.py

@@ -0,0 +1,269 @@
+"""
+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.
+"""
+
+import os
+
+from ..core.remote_state.remote_state_client import RemoteStateClient
+from ..core.remote_state.fuse_remote_state import FuseStateClient
+from ..core.blueprint.blueprint_generator import (
+    BlueprintGenerator,
+    BlueprintGeneratorOutput,
+    a3mega_device_type,
+    a3ultra_device_type,
+    supported_device_types,
+)
+from ..core.commands import run_command_for_value
+from ..core.capacity import get_capacity_type
+from ..core.docker_manager import DockerManager
+from ..core.gcloud_context import zone_to_region
+from ..core.gcluster_manager import GclusterManager
+from ..utils.console import xpk_exit, xpk_print
+from ..utils.file import ensure_directory_exists
+from ..utils.network import all_IPs_cidr
+from ..utils.objects import hash_string
+from ..core.cluster import get_cluster_credentials
+from ..core.kjob import apply_kjob_crds, prepare_kjob
+
+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()
+  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)
+  remote_state_client = None
+  if args.cluster_state_gcs_bucket is not None:
+    remote_state_client = FuseStateClient(
+        bucket=args.cluster_state_gcs_bucket,
+        state_directory=os.path.join(blueprints_path, prefix, unique_name),
+        prefix=prefix,
+        cluster=args.cluster,
+        deployment_name=unique_name,
+    )
+  gcm = prepare_gcluster_manager(remote_state_client)
+
+  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()
+
+  get_cluster_credentials(args)
+
+  err_code = apply_kjob_crds(args)
+  if err_code > 0:
+    xpk_exit(err_code)
+
+  err_code = prepare_kjob(args)
+  if err_code > 0:
+    xpk_exit(err_code)
+
+  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()
+  region = zone_to_region(args.zone)
+  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)
+  remote_state_client = None
+  if args.cluster_state_gcs_bucket is not None:
+    remote_state_client = FuseStateClient(
+        bucket=args.cluster_state_gcs_bucket,
+        state_directory=os.path.join(blueprints_path, prefix, unique_name),
+        prefix=prefix,
+        cluster=args.cluster,
+        deployment_name=unique_name,
+    )
+  gcm = prepare_gcluster_manager(remote_state_client)
+
+  # 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)
+  if args.cluster_state_gcs_bucket is not None:
+    gcm.download_state()
+
+    bp = BlueprintGeneratorOutput(
+        blueprint_file=os.path.join(blueprints_path, prefix, unique_name)
+        + '.yaml',
+        blueprint_dependencies=os.path.join(
+            blueprints_path, prefix, unique_name
+        ),
+    )
+
+    gcm.stage_files(
+        blueprint_file=bp.blueprint_file,
+        blueprint_dependencies=bp.blueprint_dependencies,
+        prefix=prefix,
+    )
+  gcm.destroy_deployment(deployment_name=unique_name, prefix=prefix)
+
+  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(
+    remote_state_client: RemoteStateClient | None,
+) -> GclusterManager:
+  dm = DockerManager(
+      working_dir=gcluster_working_dir, gcloud_cfg_path=gcloud_cfg_path
+  )
+  dm.initialize()
+  return GclusterManager(
+      gcluster_command_runner=dm, remote_state_client=remote_state_client
+  )
+
+
+def prepare_blueprint_generator() -> BlueprintGenerator:
+  return BlueprintGenerator(storage_path=blueprints_path)
+
+
+def validate_state_gcs_bucket(args):
+  bucket_validate_cmd = (
+      f'gcloud storage buckets describe gs://{args.cluster_state_gcs_bucket}'
+  )
+  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)
+
+
+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.cluster_state_gcs_bucket is not None:
+    validate_state_gcs_bucket(args)
+
+  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,
+          gcs_bucket=args.cluster_state_gcs_bucket,
+      )
+    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,
+          enable_filestore_csi_driver=args.enable_gcpfilestore_csi_driver,
+          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,
+          gcs_bucket=args.cluster_state_gcs_bucket,
+      )
+  return None

xpk/commands/common.py

@@ -0,0 +1,44 @@
+"""
+Copyright 2025 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.commands import run_command_with_updates_retry
+from ..core.gcloud_context import zone_to_region
+from ..utils.console import xpk_print
+
+
+def set_cluster_command(args) -> int:
+  """Run cluster configuration command to set the kubectl config.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = (
+      'gcloud container clusters get-credentials'
+      f' {args.cluster} --region={zone_to_region(args.zone)}'
+      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
+  )
+  if return_code != 0:
+    xpk_print(f'{task} returned ERROR {return_code}')
+  return return_code

xpk/commands/config.py

@@ -0,0 +1,29 @@
+"""
+Copyright 2025 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.config import XpkConfig
+from ..utils.console import xpk_print
+
+xpk_cfg = XpkConfig()
+
+
+def set_config(args):
+  xpk_cfg.set(args.set_config_args[0], args.set_config_args[1])
+
+
+def get_config(args):
+  value = xpk_cfg.get(args.get_config_key[0])
+  xpk_print(value)

xpk/commands/info.py

@@ -0,0 +1,243 @@
+"""
+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.
+"""
+
+import json
+from argparse import Namespace
+
+from tabulate import tabulate
+
+from ..core.commands import run_command_for_value
+from ..core.gcloud_context import add_zone_and_project
+from ..core.kueue import verify_kueuectl
+from ..utils.console import xpk_exit, xpk_print
+from .common import set_cluster_command
+
+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