xpk 0.8.0__py3-none-any.whl → 0.10.0__py3-none-any.whl

xpk/core/kueue.py

@@ -16,11 +16,13 @@ limitations under the License.
 
 from argparse import Namespace
 
+import math
 import packaging
 from packaging.version import Version
 
 from ..utils.console import xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
+from .capacity import B200_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
 from .commands import (
     run_command_for_value,
     run_command_with_updates,
@@ -38,13 +40,28 @@ from .system_characteristics import (
     SystemCharacteristics,
 )
 
-KUEUE_VERSION = 'v0.10.0'
+KUEUE_VERSION = 'v0.12.2'
 CLUSTER_QUEUE_NAME = 'cluster-queue'
 LOCAL_QUEUE_NAME = 'multislice-queue'
 WAIT_FOR_KUEUE_TIMEOUT = '5m'
+MEMORY_SIZE_PER_VM = 1.2
+MIN_MEMORY_LIMIT_SIZE = 4096
 
 packaging.version.VERSION_PATTERN = r'^v\d+\.\d+\.\d+$'
 
+topology_yaml = """apiVersion: kueue.x-k8s.io/v1alpha1
+kind: Topology
+metadata:
+  name: "gke-default"
+spec:
+  levels:
+  - nodeLabel: "cloud.google.com/gce-topology-block"
+  - nodeLabel: "cloud.google.com/gce-topology-subblock"
+  - nodeLabel: "cloud.google.com/gce-topology-host"
+  - nodeLabel: "kubernetes.io/hostname"
+---
+"""
+
 cluster_set_crd_yaml = """apiVersion: kueue.x-k8s.io/v1beta1
 kind: ResourceFlavor
 metadata:
@@ -53,6 +70,27 @@ spec:
   nodeLabels:
     {accelerator_label}
     {machine_label}
+  {topology_label}
+---
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: AdmissionCheck
+metadata:
+  name: dws-prov
+spec:
+  controllerName: kueue.x-k8s.io/provisioning-request
+  parameters:
+    apiGroup: kueue.x-k8s.io
+    kind: ProvisioningRequestConfig
+    name: dws-config
+---
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: ProvisioningRequestConfig
+metadata:
+  name: dws-config
+spec:
+  provisioningClassName: queued-provisioning.gke.io
+  managedResources:
+  - {managed_resource}
 ---
 {pw_resource_flavors}
 apiVersion: kueue.x-k8s.io/v1beta1
@@ -67,6 +105,7 @@ spec:
   resourceGroups:
   {covered_resources_config}
   {pw_resources_kueue}
+  {admission_checks}
 ---
 apiVersion: kueue.x-k8s.io/v1beta1
 kind: LocalQueue
@@ -151,6 +190,99 @@ spec:
         command: [ "sleep", "inf" ]
 """
 
+kueue_controller_manager_yml = """
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  labels:
+    app.kubernetes.io/component: controller
+    app.kubernetes.io/name: kueue
+    control-plane: controller-manager
+  name: kueue-controller-manager
+  namespace: kueue-system
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      control-plane: controller-manager
+  template:
+    metadata:
+      annotations:
+        kubectl.kubernetes.io/default-container: manager
+      labels:
+        app.kubernetes.io/component: controller
+        app.kubernetes.io/name: kueue
+        control-plane: controller-manager
+    spec:
+      containers:
+      - args:
+        - --config=/controller_manager_config.yaml
+        - --zap-log-level=2
+        command:
+        - /manager
+        image: registry.k8s.io/kueue/kueue:v0.10.0
+        imagePullPolicy: Always
+        livenessProbe:
+          httpGet:
+            path: /healthz
+            port: 8081
+          initialDelaySeconds: 15
+          periodSeconds: 20
+        name: manager
+        ports:
+        - containerPort: 8082
+          name: visibility
+          protocol: TCP
+        - containerPort: 9443
+          name: webhook-server
+          protocol: TCP
+        readinessProbe:
+          httpGet:
+            path: /readyz
+            port: 8081
+          initialDelaySeconds: 5
+          periodSeconds: 10
+        resources:
+          limits:
+            cpu: 500m
+            memory: {memory_limit_size}
+          requests:
+            cpu: 500m
+            memory: 512Mi
+        securityContext:
+          allowPrivilegeEscalation: false
+        volumeMounts:
+        - mountPath: /tmp/k8s-webhook-server/serving-certs
+          name: cert
+          readOnly: true
+        - mountPath: /controller_manager_config.yaml
+          name: manager-config
+          subPath: controller_manager_config.yaml
+      - args:
+        - --secure-listen-address=0.0.0.0:8443
+        - --upstream=http://127.0.0.1:8080/
+        - --logtostderr=true
+        - --v=10
+        image: registry.k8s.io/kubebuilder/kube-rbac-proxy:v0.16.0
+        name: kube-rbac-proxy
+        ports:
+        - containerPort: 8443
+          name: https
+          protocol: TCP
+      securityContext:
+        runAsNonRoot: true
+      serviceAccountName: kueue-controller-manager
+      terminationGracePeriodSeconds: 10
+      volumes:
+      - name: cert
+        secret:
+          defaultMode: 420
+          secretName: kueue-webhook-server-cert
+      - configMap:
+          name: kueue-manager-config
+        name: manager-config
+"""
+
 
 def verify_kueuectl(args: Namespace) -> None:
   """Verify if kueuectl is installed.
@@ -267,6 +399,7 @@ def install_kueue_crs(
     args,
     system: SystemCharacteristics,
     autoprovisioning_config: AutoprovisioningConfig | None,
+    flex_with_tpu=False,
 ) -> int:
   """Install Kueue Custom Resources.
 
@@ -294,12 +427,29 @@ def install_kueue_crs(
   else:
     # Determine total chips based on user specified topology.
     total_chips = get_total_chips_requested_from_args(args, system)
+  if args.flex and flex_with_tpu is False:
+    admission_checks = """
+  admissionChecks:
+  - dws-prov
+"""
+  else:
+    admission_checks = ''
 
   covered_resources_config = get_kueue_covered_resources_config(
       cluster_hardware_name=cluster_hardware_name,
       resource_type=resource_type,
       total_chips=total_chips,
   )
+  topology_label = ''
+  if system.device_type in [
+      H100_MEGA_DEVICE_TYPE,
+      H200_DEVICE_TYPE,
+      B200_DEVICE_TYPE,
+  ]:
+    topology_label = 'topologyName: "gke-default"'
+  res_type = AcceleratorTypeToAcceleratorCharacteristics[
+      system.accelerator_type
+  ].resource_type
   yml_string = cluster_set_crd_yaml.format(
       system=system,
       cluster_hardware_name=cluster_hardware_name,
@@ -309,15 +459,22 @@ def install_kueue_crs(
       machine_label=create_machine_label(
           system.accelerator_type, system, autoprovisioning_enabled
       ),
+      topology_label=topology_label,
       covered_resources_config=covered_resources_config,
-      resource_type=AcceleratorTypeToAcceleratorCharacteristics[
-          system.accelerator_type
-      ].resource_type,
+      resource_type=res_type,
       pw_resource_flavors=add_pw_resource_flavors(args),
       pw_resources_kueue=add_pw_resources_to_kueue(args),
+      admission_checks=admission_checks,
+      managed_resource=res_type,
       cluster_queue_name=CLUSTER_QUEUE_NAME,
       local_queue_name=LOCAL_QUEUE_NAME,
   )
+  if system.device_type in [
+      H100_MEGA_DEVICE_TYPE,
+      H200_DEVICE_TYPE,
+      B200_DEVICE_TYPE,
+  ]:
+    yml_string = topology_yaml + yml_string
 
   tmp = write_tmp_file(yml_string)
   command = f'kubectl apply -f {str(tmp.file.name)}'
@@ -356,3 +513,36 @@ def get_kueue_covered_resources_config(
       total_chips=total_chips,
   )
   return config_string
+
+
+def update_kueue_resources_if_necessary(args):
+  """Update the kueue manifest to increase the resources for the kueue controller manager.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  # Get total number of nodes
+  cmd_total_node_num = 'kubectl get node --no-headers | wc -l'
+  return_code, out = run_command_for_value(
+      cmd_total_node_num, 'Count total nodes', args
+  )
+  if return_code != 0:
+    xpk_exit(1)
+  # 1.2MiB per VM or 4GiB (whichever is greater).
+  new_memory_limit = (
+      f'{max(math.ceil(int(out) * MEMORY_SIZE_PER_VM), MIN_MEMORY_LIMIT_SIZE)}Mi'
+  )
+  yml_string = kueue_controller_manager_yml.format(
+      memory_limit_size=new_memory_limit,
+  )
+  tmp = write_tmp_file(yml_string)
+  command = f'kubectl apply -f {str(tmp.file.name)}'
+
+  task = 'Updating Kueue Controller Manager 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

xpk/core/mtc.py

@@ -0,0 +1,195 @@
+"""
+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 requests
+import yaml
+
+from ..core.cluster import JOBSET_VERSION
+from ..core.cluster import setup_k8s_env
+from ..utils import templates
+from ..utils.console import xpk_exit
+from ..utils.console import xpk_print
+from ..utils.kubectl import apply_kubectl_manifest
+
+
+MTC_CPC_PATH = "/../templates/mtc-cpc.yaml"
+
+
+def create_mtc_cpc(
+    mtc_gcs_bucket: str,
+    mtc_machine_type: str,
+    mtc_toleration_key: str,
+    mtc_ramdisk_size: str,
+) -> dict:
+  """Create MTC Checkpoint Configuration.
+
+  Args:
+    mtc_gcs_bucket: GCS bucket for MTC
+    mtc_machine_type: Machine type for MTC
+    mtc_toleration_key: Toleration key for MTC
+    mtc_ramdisk_size: Ramdisk size for MTC
+
+  Returns:
+    MTC Checkpoint Configuration
+  """
+  data = templates.load(MTC_CPC_PATH)
+
+  data["spec"]["cloudStorageBucketName"] = mtc_gcs_bucket
+  data["spec"]["nodeSelector"][
+      "node.kubernetes.io/instance-type"
+  ] = mtc_machine_type
+  data["spec"]["tolerations"][0]["key"] = mtc_toleration_key
+  data["spec"]["inMemoryVolumeSize"] = mtc_ramdisk_size
+
+  return data
+
+
+def install_mtc_on_cluster(args, system) -> int:
+  """Install MTC on the cluster.
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system related information.
+
+  Returns:
+    return code of the command.
+  """
+  if args.mtc_gcs_bucket is None:
+    xpk_print("MTC GCS bucket is required.")
+    xpk_exit(1)
+  if args.mtc_gcs_bucket.startswith("gs://"):
+    args.mtc_gcs_bucket = args.mtc_gcs_bucket.replace("gs://", "")
+
+  if args.mtc_ramdisk_size is None:
+    xpk_print("MTC ramdisk size is required.")
+    xpk_exit(1)
+
+  if args.mtc_toleration_key is None:
+    args.mtc_toleration_key = "google.com/tpu"
+
+  k8s_api_client = setup_k8s_env(args)
+  jobset_manifest = update_jobset_manifest()
+  if jobset_manifest is None:
+    xpk_print(
+        "Updated jobset manifest is empty, not updating the jobset controller."
+    )
+
+  xpk_print("Applying Jobset with MTC Configuration")
+  return_code = apply_kubectl_manifest(k8s_api_client, [jobset_manifest])
+  if return_code != 0:
+    return return_code
+
+  mtc_checkpoint_configuration_crd_data = create_mtc_cpc(
+      args.mtc_gcs_bucket,
+      system.gce_machine_type,
+      args.mtc_toleration_key,
+      args.mtc_ramdisk_size,
+  )
+  xpk_print("Applying MTC Checkpoint Configuration")
+  return_code = apply_kubectl_manifest(
+      k8s_api_client, [mtc_checkpoint_configuration_crd_data]
+  )
+
+  return return_code
+
+
+def update_jobset_manifest():
+  """Update the jobset manifest to increase the resources for the jobset controller manager.
+
+  Returns:
+    The updated jobset manifest.
+  """
+  manifest_url = f"https://github.com/kubernetes-sigs/jobset/releases/download/{JOBSET_VERSION}/manifests.yaml"
+  manifest_content = None
+  # Fetch the manifest content
+  try:
+    response = requests.get(manifest_url, timeout=10)
+    response.raise_for_status()  # Raise an exception for HTTP errors
+    manifest_content = response.text
+  except requests.exceptions.Timeout as e:
+    xpk_print(f"Error: Request to {manifest_url} after 10 seconds: {e}")
+    xpk_exit(1)
+  except requests.exceptions.RequestException as e:
+    xpk_print(f"Error fetching manifest from {manifest_url}: {e}")
+    xpk_exit(1)
+
+  if manifest_content is None:
+    xpk_print("Manifest content not found.")
+    xpk_exit(1)
+
+  # Load all YAML documents from the manifest
+  yaml_data_list = list(yaml.safe_load_all(manifest_content))
+  # Iterate through the yaml_data to find the Deployment for
+  # jobset-controller-manager
+  update_manifest = False
+  for yaml_data in yaml_data_list:
+    if (
+        yaml_data
+        and yaml_data.get("apiVersion") == "apps/v1"
+        and yaml_data.get("kind") == "Deployment"
+        and yaml_data.get("metadata", {}).get("name")
+        == "jobset-controller-manager"
+    ):
+      # Found the Deployment, now modify the resources
+      containers = yaml_data["spec"]["template"]["spec"]["containers"]
+      for container in containers:
+        if container["name"] == "manager":
+          # Update resource limits and requests
+          current_cpu_request = (
+              container["resources"].get("requests", {}).get("cpu", "0m")
+          )
+          current_memory_request = (
+              container["resources"].get("requests", {}).get("memory", "0Mi")
+          )
+          current_memory_limit = (
+              container["resources"].get("limits", {}).get("memory", "0Mi")
+          )
+
+          # Define new values for comparison
+          new_cpu_request = "1000m"
+          new_memory_request = "1Gi"
+          new_memory_limit = "2Gi"
+
+          if parse_resource_value(current_cpu_request) < parse_resource_value(
+              new_cpu_request
+          ):
+            container["resources"]["requests"]["cpu"] = new_cpu_request
+            update_manifest = True
+          if parse_resource_value(
+              current_memory_request
+          ) < parse_resource_value(new_memory_request):
+            container["resources"]["requests"]["memory"] = new_memory_request
+            update_manifest = True
+          if parse_resource_value(current_memory_limit) < parse_resource_value(
+              new_memory_limit
+          ):
+            container["resources"]["limits"]["memory"] = new_memory_limit
+            update_manifest = True
+          break
+      if update_manifest:
+        xpk_print("Jobset controller updation required.")
+        return yaml_data
+  xpk_print("Jobset controller no updation required.")
+
+
+def parse_resource_value(value) -> int:
+  if value.endswith("m"):
+    return int(value[:-1])
+  if value.endswith("Mi"):
+    return int(value[:-2])
+  if value.endswith("Gi"):
+    return int(value[:-2]) * 1024
+  return int(value)

xpk/core/network.py

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from ..utils.console import xpk_print
+from ..utils.console import xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
 from .commands import run_command_for_value, run_command_with_updates
 from .gcloud_context import zone_to_region
@@ -235,6 +235,28 @@ def create_cluster_network_config(args) -> int:
   return 0
 
 
+def get_cluster_subnetworks(args) -> list[str]:
+  """Gets the list of cluster networks.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    list[str]: list of cluster networks
+  """
+  command = 'kubectl get GKENetworkParamSet'
+  return_code, stdout = run_command_for_value(
+      command, 'Get Cluster Networks', args
+  )
+  if return_code != 0:
+    xpk_print('GKE Cluster Get NetworkParamSet failed')
+    xpk_exit(return_code)
+
+  networks = [line.split()[0] for line in stdout.splitlines()][1:]
+
+  return networks
+
+
 def set_up_cluster_network_for_a3(args) -> int:
   """Set up GKE Cluster networks, subnets and firewall rules for A3.
   Note: there are 4 NICs for GPU-GPU bw and 1 NIC for host in an A3 node.

xpk/core/nodepool.py

@@ -77,8 +77,12 @@ def run_gke_node_pool_create_command(
     if return_code > 0:
       xpk_print('Listing all reservations failed!')
     return_code = 1
+  if system.accelerator_type == AcceleratorType['TPU']:
+    max_nodes = system.vms_per_slice
+  else:
+    max_nodes = 1000
   capacity_args, return_code = get_capacity_arguments_from_capacity_type(
-      args, capacity_type
+      args, capacity_type, max_nodes
   )
   if return_code > 0:
     xpk_print('Parsing capacity arguments failed!')
@@ -275,7 +279,10 @@ def run_gke_node_pool_create_command(
     )
     if system.accelerator_type == AcceleratorType['TPU']:
       command += f' --node-version={gke_node_pool_version}'
-      command += f' --num-nodes={system.vms_per_slice}'
+      if capacity_type == CapacityType.FLEX_START:
+        command += ' --num-nodes=0'
+      else:
+        command += f' --num-nodes={system.vms_per_slice}'
       command += ' --placement-type=COMPACT  --max-pods-per-node 15'
       command += (
           f' --scopes=storage-full,gke-default,{CLOUD_PLATFORM_AUTH_SCOPE_URL}'
@@ -284,7 +291,10 @@ def run_gke_node_pool_create_command(
       command += f' {args.custom_tpu_nodepool_arguments}'
     elif system.accelerator_type == AcceleratorType['GPU']:
       subnet_prefix = f'{args.cluster}-{zone_to_region(args.zone)}'
-      command += f' --num-nodes={args.num_nodes}'
+      if capacity_type == CapacityType.FLEX_START:
+        command += ' --num-nodes=0'
+      else:
+        command += f' --num-nodes={args.num_nodes}'
       command += (
           ' --accelerator'
           f' type={system.gke_accelerator},count={str(system.chips_per_vm)},gpu-driver-version=latest'
@@ -298,7 +308,10 @@ def run_gke_node_pool_create_command(
           )
         command += ' --max-pods-per-node=32'
     elif system.accelerator_type == AcceleratorType['CPU']:
-      command += f' --num-nodes={system.vms_per_slice}'
+      if capacity_type == CapacityType.FLEX_START:
+        command += ' --num-nodes=0'
+      else:
+        command += f' --num-nodes={system.vms_per_slice}'
       command += (
           f' --scopes=storage-full,gke-default,{CLOUD_PLATFORM_AUTH_SCOPE_URL}'
       )

xpk/core/pathways.py

@@ -19,8 +19,7 @@ from ..core.docker_container import get_user_workload_container
 from ..core.gcloud_context import zone_to_region
 from ..core.nodepool import get_all_nodepools_programmatic
 from ..utils.console import xpk_exit, xpk_print
-from .config import AcceleratorType
-from .system_characteristics import SystemCharacteristics
+from .system_characteristics import AcceleratorType, SystemCharacteristics
 
 
 def add_pw_resource_flavors(args):
@@ -211,7 +210,7 @@ def append_custom_pathways_worker(args) -> str:
   """
   yaml = """"""
   if args.server_image or args.custom_pathways_worker_args:
-    yaml = """- componentType: pathways_worker"""
+    yaml = """- componentType: worker"""
   indentation = (
       ' ' * 8
   )  # Currently 8, based on the YAML, may need to update in the future.

xpk/core/resources.py

@@ -236,3 +236,24 @@ def get_cluster_system_characteristics(args) -> SystemCharacteristics | None:
       return system
 
   return None
+
+
+def get_cluster_capacity_type(args) -> CapacityType | None:
+  """Get systemCharcteristics based on the cluster resources configMap
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    returns system characteristics
+  """
+  metadata_configmap_name = f'{args.cluster}-{CLUSTER_METADATA_CONFIGMAP}'
+  cluster_config_map = get_cluster_configmap(args, metadata_configmap_name)
+
+  if cluster_config_map is None:
+    return None
+
+  capacityValue = cluster_config_map.get('capacity_type')
+  if capacityValue is not None:
+    return CapacityType[capacityValue.upper()]
+
+  return None