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

xpk/core/remote_state/remote_state_client.py

@@ -0,0 +1,38 @@
+"""
+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 abc import ABC, abstractmethod
+
+
+class RemoteStateClient(ABC):
+  """This is a base class that defines methods a class for managing remote cluster state.
+  Args:
+      ABC (_type_): _description_
+  """
+
+  @abstractmethod
+  def upload_state(self) -> None:
+    """Upload state to remote storage"""
+    return None
+
+  @abstractmethod
+  def download_state(self) -> None:
+    """Download state from remote storage"""
+    return None
+
+  @abstractmethod
+  def check_remote_state_exists(self) -> bool:
+    return False

xpk/core/resources.py

@@ -0,0 +1,238 @@
+"""
+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 dataclasses import dataclass
+
+from ..utils.console import xpk_print
+from ..utils.file import write_tmp_file
+from .capacity import (
+    AUTOPROVISIONING_CONFIG_MAXIMUM_KEY,
+    AUTOPROVISIONING_CONFIG_MINIMUM_KEY,
+    AUTOPROVISIONING_CONFIG_VALUE,
+    CAPACITY_TYPE_CONFIG_KEY,
+    RESERVATION_CONFIG_KEY,
+    CapacityType,
+    get_capacity_type,
+)
+from .commands import run_command_for_value, run_commands
+from .config import XPK_CURRENT_VERSION
+from .system_characteristics import AcceleratorType, get_system_characteristics_by_device_type, SystemCharacteristics
+
+CLUSTER_RESOURCES_CONFIGMAP = 'resources-configmap'
+CLUSTER_METADATA_CONFIGMAP = 'metadata-configmap'
+
+CLUSTER_CONFIGMAP_YAML = """kind: ConfigMap
+apiVersion: v1
+metadata:
+  name: {name}
+data:
+  {data}
+"""
+
+
+@dataclass
+class AutoprovisioningConfig:
+  config_filename: str
+  minimum_chips: int
+  maximum_chips: int
+
+
+def get_cluster_configmap(args, configmap_name) -> dict[str, str] | None:
+  """Run the Get GKE Cluster ConfigMap request.
+
+  Args:
+    args: user provided arguments for running the command.
+    configmap_name: name of the configmap.
+
+  Returns:
+    key:value pairs stored in cluster ConfigMap.
+  """
+  command = (
+      'kubectl get configmap'
+      f' {configmap_name} -o=custom-columns="ConfigData:data" --no-headers=true'
+  )
+
+  return_code, return_value = run_command_for_value(
+      command, 'GKE Cluster Get ConfigMap', args
+  )
+  if return_code != 0:
+    xpk_print(f'GKE Cluster Get ConfigMap request returned ERROR {return_code}')
+    return None
+
+  config_map = {}
+  return_value = return_value.strip()
+
+  if return_value:
+    # Format of ConfigMap: map[key1:value1 key2:value2]
+    return_value = return_value[return_value.index('map') :]
+    configs = return_value[4:-1].split(' ')
+
+    for config in configs:
+      key, value = config.strip().split(':')
+      config_map[key] = value
+  return config_map
+
+
+def create_cluster_configmaps(
+    args,
+    system,
+    tensorboard_config: dict,
+    autoprovisioning_config: AutoprovisioningConfig | None,
+) -> int:
+  """Run the Create GKE Cluster ConfigMap request.
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system characteristics.
+    tensorboard_config: map that contains Vertex Tensorboard name, id and location
+    autoprovisioning_config: Config used in autoprovisioning.
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  configmap_yml = {}
+
+  # ConfigMap to store resources available in the cluster.
+  device_type = system.device_type
+  if system.accelerator_type == AcceleratorType['GPU']:
+    resources_data = f'{device_type}: "{int(args.num_nodes)}"'
+  elif (
+      not args.enable_pathways
+      and args.enable_autoprovisioning
+      and autoprovisioning_config
+  ):
+    # Currently autoprovisioning is not supported with Pathways.
+    # Auto provisioning will have variable topologies for a gke accelerator type.
+    resources_data = (
+        f'{system.gke_accelerator}: {AUTOPROVISIONING_CONFIG_VALUE}'
+    )
+    resources_data += (
+        f'\n  {AUTOPROVISIONING_CONFIG_MINIMUM_KEY}:'
+        f' "{autoprovisioning_config.minimum_chips}"'
+    )
+    resources_data += (
+        f'\n  {AUTOPROVISIONING_CONFIG_MAXIMUM_KEY}:'
+        f' "{autoprovisioning_config.maximum_chips}"'
+    )
+  else:
+    resources_data = (
+        f'{device_type}: "{int(args.num_slices) * system.vms_per_slice}"'
+    )
+  resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
+  resources_yml = CLUSTER_CONFIGMAP_YAML.format(
+      args=args, name=resources_configmap_name, data=resources_data
+  )
+  configmap_yml[resources_configmap_name] = resources_yml
+
+  # ConfigMap to store cluster metadata.
+  # XPK Version.
+  metadata = f'xpk_version: {XPK_CURRENT_VERSION}'
+  # Vertex Tensorboard information
+  for key, value in tensorboard_config.items():
+    metadata += f'\n  {key}: "{value}"'
+  # Capacity Type.
+  capacity_type, return_code = get_capacity_type(args)
+  if return_code != 0:
+    xpk_print('Unable to determine capacity type.')
+    return return_code
+  metadata += f'\n  {CAPACITY_TYPE_CONFIG_KEY}: {capacity_type.name}'
+  # Reservation ID if applicable.
+  if capacity_type == CapacityType.RESERVATION:
+    metadata += f'\n  {RESERVATION_CONFIG_KEY}: {args.reservation}'
+  metadata_configmap_name = f'{args.cluster}-{CLUSTER_METADATA_CONFIGMAP}'
+  metadata_yml = CLUSTER_CONFIGMAP_YAML.format(
+      args=args, name=metadata_configmap_name, data=metadata
+  )
+  configmap_yml[metadata_configmap_name] = metadata_yml
+  return create_or_update_cluster_configmap(configmap_yml)
+
+
+def create_or_update_cluster_configmap(configmap_yml: dict) -> int:
+  """
+  Args:
+    configmap_yml: dict containing ConfigMap name and yml string.
+
+  Returns:
+    0 if successful, 1 otherwise.
+  """
+  commands = []
+  task_names = []
+  for configmap_name, yml_string in configmap_yml.items():
+    tmp = write_tmp_file(yml_string)
+    command = f'kubectl apply -f {str(tmp.file.name)}'
+    commands.append(command)
+    task_name = f'ConfigMap CreateOrUpdate-{configmap_name}'
+    task_names.append(task_name)
+
+  return_code = run_commands(
+      commands, 'GKE Cluster CreateOrUpdate ConfigMap(s)', task_names
+  )
+  if return_code != 0:
+    xpk_print(
+        'GKE Cluster Create/Update ConfigMap(s) request returned ERROR'
+        f' {return_code}'
+    )
+    return 1
+  return 0
+
+
+def check_cluster_resources(args, system) -> tuple[bool, bool]:
+  """Check if cluster has resources of a specified device_type/gke_accelerator.
+  This check will be skipped if <args.cluster>-<_CLUSTER_RESOURCES_CONFIGMAP> ConfigMap doesn't exist for the cluster.
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system characteristics.
+
+  Returns:
+    Tuple of bool, bool
+    True if resources in the cluster should be checked, False otherwise.
+    True if device_type/gke_accelerator exists in the cluster, False otherwise.
+  """
+  resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
+  resources_config_map = get_cluster_configmap(args, resources_configmap_name)
+  if resources_config_map is None:
+    xpk_print(
+        f'No ConfigMap exist for cluster with the name {resources_config_map}.'
+        ' Cluster resources check will be skipped.'
+    )
+    return False, False
+  if system.device_type in resources_config_map:
+    return True, True
+  elif system.gke_accelerator in resources_config_map:
+    return True, True
+  return True, False
+
+
+def get_cluster_system_characteristics(args) -> SystemCharacteristics | None:
+  """Get systemCharcteristics based on the cluster resources configMap
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    returns system characteristics
+  """
+  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:
+    return None
+
+  for key in cluster_config_map:
+    system, result_code = get_system_characteristics_by_device_type(key)
+    if result_code == 0:
+      return system
+
+  return None

xpk/core/scheduling.py

@@ -0,0 +1,253 @@
+"""
+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 ..utils.console import xpk_print
+from .capacity import AUTOPROVISIONING_CONFIG_MAXIMUM_KEY, AUTOPROVISIONING_CONFIG_VALUE
+from .resources import CLUSTER_RESOURCES_CONFIGMAP, get_cluster_configmap
+from .system_characteristics import (
+    AcceleratorType,
+    AcceleratorTypeToAcceleratorCharacteristics,
+    SystemCharacteristics,
+)
+
+
+def check_if_workload_can_schedule(args, system: SystemCharacteristics) -> bool:
+  """Check if workload can schedule based on the cluster resources (tpu_type and maximum VM in cluster).
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system characteristics
+
+  Returns:
+    returns true if workload can schedule, otherwise returns false.
+  """
+  resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
+  cluster_config_map = get_cluster_configmap(args, resources_configmap_name)
+
+  # Prevents workload creation failure for existing clusters with no ConfigMap
+  if cluster_config_map is None:
+    xpk_print(
+        'No ConfigMap exist for cluster with the name'
+        f' {resources_configmap_name}.'
+    )
+    return True
+
+  # Check for gke accelerator type:
+  missing_gke_accelerator_type = False
+  if not cluster_config_map.get(system.gke_accelerator):
+    xpk_print(
+        f'Gke Accelerator Type Check: {args.workload} is requesting'
+        f' {system.gke_accelerator} but cluster only contains'
+        f' {cluster_config_map.keys()}. '
+    )
+    missing_gke_accelerator_type = True
+  elif (
+      cluster_config_map[system.gke_accelerator]
+      == AUTOPROVISIONING_CONFIG_VALUE
+  ):
+    # Run total chip check when in autoprovisioning mode.
+    max_chips_in_cluster = int(
+        cluster_config_map[AUTOPROVISIONING_CONFIG_MAXIMUM_KEY]
+    )
+    num_chips_in_workload = get_total_chips_requested_from_args(args, system)
+
+    if num_chips_in_workload > max_chips_in_cluster:
+      xpk_print(
+          f'{args.workload} is requesting {num_chips_in_workload} chips but'
+          f' the cluster {args.cluster} supports up to {max_chips_in_cluster}.'
+          '  Resize the cluster to support more chips with'
+          ' `xpk cluster create --autoprovisioning-max-chips=X ...`'
+      )
+      return False
+    return True
+
+  # Check for device type
+  missing_device_type = False
+  device_type = system.device_type
+  if device_type not in cluster_config_map:
+    xpk_print(
+        f'Device Type Check: {args.workload} is requesting {device_type} but '
+        f'cluster only contains {cluster_config_map.keys()}. '
+    )
+    missing_device_type = True
+
+  if missing_device_type and missing_gke_accelerator_type:
+    xpk_print(
+        'Both Device Type and GKE Accelerator Type checks failed.'
+        f' XPK will not create the workload {args.workload}.'
+    )
+    return False
+  else:
+    # Check if the size of the workload will fit in the cluster.
+    max_vm_in_cluster = int(cluster_config_map[device_type])
+    if system.accelerator_type == AcceleratorType['GPU']:
+      vm_required_by_workload = args.num_nodes
+    else:
+      vm_required_by_workload = args.num_slices * system.vms_per_slice
+    if vm_required_by_workload > max_vm_in_cluster:
+      xpk_print(
+          f'{args.workload} is requesting {args.num_slices} slice/slices of'
+          f' {device_type}, which is {vm_required_by_workload} VMs, but the'
+          f' cluster only contains {max_vm_in_cluster} VMs of {device_type}.'
+          ' XPK will not create this workload.'
+      )
+      return False
+
+  return True
+
+
+def get_total_chips_requested_from_args(
+    args, system: SystemCharacteristics
+) -> int:
+  """Return the total chips requested based on user args.
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system characteristics.
+
+  Returns:
+    num of chips for the current request.
+  """
+  if system.accelerator_type == AcceleratorType['GPU']:
+    num_chips = system.vms_per_slice * system.chips_per_vm * args.num_nodes
+  else:
+    num_chips = system.vms_per_slice * system.chips_per_vm * args.num_slices
+
+  return int(num_chips)
+
+
+def get_cpu_affinity(accelerator_type) -> str:
+  """Generate affinity rules for CPU nodepools, so that workload pods are
+  not scheduled on the default pool machines.
+  Args:
+    accelerator_type: TPU / GPU / CPU
+
+  Returns:
+    str: yaml containing affinity constraints
+  """
+  yaml = """affinity:
+                nodeAffinity:
+                  requiredDuringSchedulingIgnoredDuringExecution:
+                    nodeSelectorTerms:
+                    - matchExpressions:
+                      - key: cloud.google.com/gke-nodepool
+                        operator: NotIn
+                        values:
+                        - default-pool
+"""
+  if accelerator_type == AcceleratorType['CPU']:
+    return yaml
+  return ''
+
+
+def get_gpu_scheduler(
+    args, system: SystemCharacteristics, autoprovisioning_args: str
+) -> tuple[str, int]:
+  """Get gpu scheduler configuration.
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system characteristics.
+    autoprovisioning_args: a string of arguments for Autoprovisioning.
+
+  Returns:
+    str: yaml containing gpu scheduler configuration
+    int of 0 if successful and 1 otherwise.
+  """
+  gpu_scheduler = ''
+  return_code = 0
+
+  if args.scheduler == 'gke.io/topology-aware-auto':
+    gpu_scheduler = f"""schedulingGates:
+              - name: "{args.scheduler}-{args.workload}"
+              """
+  elif args.scheduler == 'default-scheduler':
+    gpu_scheduler_yaml = """schedulerName: {scheduler_name}
+              affinity:
+                nodeAffinity:
+                  requiredDuringSchedulingIgnoredDuringExecution:
+                    nodeSelectorTerms:
+                    - matchExpressions:
+                      - key: cloud.google.com/gke-accelerator
+                        operator: Exists
+                      - key: cloud.google.com/gke-nodepool
+                        operator: In
+                        values: [{node_pool_name}]
+              nodeSelector:
+                {accelerator_label}
+                {machine_label}
+                {autoprovisioning_args}
+              """
+    gpu_scheduler = gpu_scheduler_yaml.format(
+        scheduler_name=args.scheduler,
+        accelerator_label=create_accelerator_label(
+            system.accelerator_type, system
+        ),
+        machine_label=create_machine_label(system.accelerator_type, system),
+        node_pool_name=f'{args.cluster}-np-0',
+        autoprovisioning_args=autoprovisioning_args,
+    )
+  else:
+    return_code = 1
+    xpk_print(
+        '--scheduler needs to be set as either `default-scheduler`'
+        ' or `gke.io/topology-aware-auto` in order to schedule the'
+        ' workloads on GPUs.'
+    )
+
+  return gpu_scheduler, return_code
+
+
+def create_accelerator_label(accelerator_type, system) -> str:
+  """Generates accelerator label.
+
+  Args:
+    accelerator_type: type of accelerator.
+    system: system characteristics.
+
+  Returns:
+    The accelerator label.
+  """
+  if accelerator_type == AcceleratorType['CPU']:
+    return ''
+  return (
+      f'{AcceleratorTypeToAcceleratorCharacteristics[accelerator_type].accelerator_label}:'
+      f' {system.gke_accelerator}'
+  )
+
+
+def create_machine_label(
+    accelerator_type, system, autoprovisioning_enabled: bool = False
+) -> str:
+  """Generates machine label.
+
+  Args:
+    accelerator_type: type of accelerator.
+    system: system characteristics.
+    autoprovisioning_enabled: describes autoprovisioning enablement.
+
+  Returns:
+    The machine label.
+  """
+  if (
+      accelerator_type == AcceleratorType['TPU']
+      and not autoprovisioning_enabled
+  ):
+    return (
+        f'{AcceleratorTypeToAcceleratorCharacteristics[accelerator_type].machine_label}:'
+        f' {system.topology}'
+    )
+  return ''