Add checkpoint configuration file for MTC (#465)

Author: abhinavclemson

src/xpk/commands/cluster.py

@@ -58,6 +58,7 @@
     run_gke_node_pool_create_command,
 )
 from ..core.ray import install_ray_cluster
+from ..core.mtc import install_mtc_on_cluster
 from ..core.resources import create_cluster_configmaps
 from ..core.storage import install_storage_crd
 from ..core.system_characteristics import (
@@ -279,6 +280,12 @@ def cluster_create(args) -> None:
       xpk_print('Installation of RayCluster failed.')
       xpk_exit(return_code)
 
+  if hasattr(args, 'enable_mtc') and args.enable_mtc:
+    return_code = install_mtc_on_cluster(args, system)
+    if return_code != 0:
+      xpk_print('Installation of MTC failed.')
+      xpk_exit(return_code)
+
   xpk_print('GKE commands done! Resources are created.')
   xpk_print(
       'See your GKE Cluster here:'
@@ -815,6 +822,7 @@ def run_gke_cluster_create_command(
   addons = []
   if args.enable_gcsfuse_csi_driver:
     addons.append('GcsFuseCsiDriver')
+
   if args.enable_gcpfilestore_csi_driver:
     addons.append('GcpFilestoreCsiDriver')
 
@@ -824,6 +832,9 @@ def run_gke_cluster_create_command(
   if args.enable_pd_csi_driver:
     addons.append('GcePersistentDiskCsiDriver')
 
+  if hasattr(args, 'enable_mtc') and args.enable_mtc:
+    addons.append('HighScaleCheckpointing')
+
   if len(addons) > 0:
     addons_str = ','.join(addons)
     command += f' --addons={addons_str}'

src/xpk/core/mtc.py

@@ -0,0 +1,77 @@
+"""
+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 ..utils import templates
+from ..utils.kubectl import apply_kubectl_manifest
+from ..core.cluster import setup_k8s_env
+
+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:
+  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.
+
+  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"
+
+  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")
+  k8s_api_client = setup_k8s_env(args)
+  return_code = apply_kubectl_manifest(
+      k8s_api_client, [mtc_checkpoint_configuration_crd_data]
+  )
+
+  return return_code

src/xpk/parser/cluster.py

@@ -189,6 +189,12 @@ def set_cluster_create_parser(cluster_create_parser: ArgumentParser):
       cluster_create_tensorboard_arguments
   )
 
+  ### MTC arguments specific to "cluster create"
+  cluster_create_mtc_arguments = cluster_create_parser.add_argument_group(
+      'Optional MTC Arguments',
+      'Arguments for configuring MTC in cluster create.',
+  )
+  add_shared_cluster_create_mtc_arguments(cluster_create_mtc_arguments)
   cluster_create_parser.set_defaults(func=cluster_create)
 
 
@@ -244,6 +250,14 @@ def set_cluster_create_pathways_parser(
       cluster_create_pathways_tensorboard_arguments
   )
 
+  ### MTC arguments specific to "cluster create"
+  cluster_create_mtc_arguments = (
+      cluster_create_pathways_parser.add_argument_group(
+          'Optional MTC Arguments',
+          'Arguments for configuring MTC in cluster create.',
+      )
+  )
+  add_shared_cluster_create_mtc_arguments(cluster_create_mtc_arguments)
   cluster_create_pathways_parser.set_defaults(func=cluster_create_pathways)
 
 
@@ -313,6 +327,12 @@ def set_cluster_create_ray_parser(cluster_create_ray_parser: ArgumentParser):
       cluster_create_ray_tensorboard_arguments
   )
 
+  ### MTC arguments specific to "cluster create"
+  cluster_create_mtc_arguments = cluster_create_ray_parser.add_argument_group(
+      'Optional MTC Arguments',
+      'Arguments for configuring MTC in cluster create.',
+  )
+  add_shared_cluster_create_mtc_arguments(cluster_create_mtc_arguments)
   cluster_create_ray_parser.set_defaults(func=cluster_create_ray_cluster)
 
 
@@ -706,3 +726,43 @@ def add_shared_cluster_create_capacity_arguments(parser: ArgumentParser):
           ' See `--reservation` or `--on-demand` for other capacity types.'
       ),
   )
+
+
+def add_shared_cluster_create_mtc_arguments(parser: ArgumentParser):
+  """Add shared Multi-tier Checkpointing arguments in cluster create and Pathways cluster create.
+
+  Args:
+      List of cluster create MTC arguments parsers
+  """
+  parser.add_argument(
+      '--enable-mtc',
+      action='store_true',
+      help='Enable MTC on the cluster.',
+  )
+  parser.add_argument(
+      '--mtc-ramdisk-size',
+      type=str,
+      default=None,
+      help=(
+          '(Required if --enable-mtc is true) The size of the RAM disk to be'
+          ' used for multi-tier checkpointing. e.g. "64Mi" '
+      ),
+  )
+  parser.add_argument(
+      '--mtc-gcs-bucket',
+      type=str,
+      default=None,
+      help=(
+          '(Required if --enable-mtc is true) The GCS bucket to be used for'
+          ' multi-tier checkpointing.'
+      ),
+  )
+  parser.add_argument(
+      '--mtc-toleration-key',
+      type=str,
+      default=None,
+      help=(
+          '(Optional) The tolerance key to be used for multi-tier'
+          ' checkpointing. By default, it is set to "google.com/tpu".'
+      ),
+  )

src/xpk/templates/mtc-cpc.yaml

@@ -0,0 +1,15 @@
+apiVersion: checkpointing.gke.io/v1
+kind: CheckpointConfiguration
+metadata:
+  name: my-checkpointconfiguration
+spec:
+  cloudStorageBucketName:
+  # This field is optional
+  nodeSelector:
+    node.kubernetes.io/instance-type:
+  # This field is optional
+  tolerations:
+  - key:
+    operator: Exists
+    effect: NoSchedule
+  inMemoryVolumeSize:

src/xpk/utils/kubectl.py

@@ -20,10 +20,11 @@
 from .console import xpk_print
 
 
-def apply_kubectl_manifest(client, manifest):
+def apply_kubectl_manifest(client, manifest) -> int:
   xpk_print('Applying manifest')
   dynamic_client = DynamicClient(client)
 
+  status_code = 0
   for obj in manifest:
     api_version = obj['apiVersion']
     kind = obj['kind']
@@ -55,3 +56,5 @@ def apply_kubectl_manifest(client, manifest):
         )
       else:
         xpk_print(f'Error applying {kind}: {e}')
+        status_code = 1
+  return status_code