Add kind for local testing

README.md

@@ -1201,6 +1201,50 @@ You can specify what kind of resources(clusterqueue or localqueue) you want to s
 python3 xpk.py info --cluster my-cluster --localqueue
 ```
 
+# Local testing with Kind
+
+To facilitate development and testing locally, we have integrated support for testing with `kind`. This enables you to simulate a Kubernetes environment on your local machine.
+
+## Prerequisites
+
+- Install kind on your local machine. Follow the official documentation here: [Kind Installation Guide.](https://kind.sigs.k8s.io/docs/user/quick-start#installation)
+
+## Usage
+
+xpk interfaces seamlessly with kind to manage Kubernetes clusters locally, facilitating the orchestration and management of workloads. Below are the commands for managing clusters:
+
+### Cluster Create
+*   Cluster create:
+
+    ```shell
+    python3 xpk.py kind create \
+    --cluster xpk-test
+    ```
+
+### Cluster Delete
+*   Cluster Delete:
+
+    ```shell
+    python3 xpk.py kind delete \
+    --cluster xpk-test
+    ```
+
+### Cluster List
+*   Cluster List:
+
+    ```shell
+    python3 xpk.py kind list
+    ```
+
+## Local Testing Basics
+
+Local testing is available exclusively through the `batch` command of xpk with the `--local-test` flag. This allows you to simulate training jobs locally:
+
+```shell
+python xpk.py batch [other-options] --local-test script
+```
+
+Please note that all other xpk subcommands are intended for use with cloud systems on Google Cloud Engine (GCE) and don't support local testing. This includes commands like cluster, info, inspector, etc.
 
 # Other advanced usage
 [Use a Jupyter notebook to interact with a Cloud TPU cluster](xpk-notebooks.md)

src/xpk/commands/batch.py

@@ -34,10 +34,11 @@ def batch(args: Namespace) -> None:
   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)
+  if not args.local_test:
+    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)
 
   create_job_template_instance(args)
   create_app_profile_instance(args)

src/xpk/commands/kind.py

@@ -0,0 +1,215 @@
+"""
+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.commands import (
+    run_command_for_value,
+    run_command_with_updates,
+)
+from ..core.core import (
+    set_jobset_on_cluster,
+)
+from ..core.kjob import (
+    verify_kjob_installed,
+    prepare_kjob,
+    apply_kjob_crds,
+)
+from ..core.kueue import (
+    install_kueue_on_cluster,
+)
+from ..utils.console import (xpk_exit, xpk_print)
+
+
+def cluster_create(args) -> None:
+  """Function around cluster creation.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  xpk_print(f'Starting cluster create for cluster {args.cluster}:', flush=True)
+
+  create_cluster_command_code = create_cluster_if_necessary(args)
+  if create_cluster_command_code != 0:
+    xpk_exit(create_cluster_command_code)
+
+  xpk_print(
+      'Enabling the jobset API on our cluster, to be deprecated when Jobset is'
+      ' globally available'
+  )
+  set_jobset_on_cluster_code = set_jobset_on_cluster(args)
+  if set_jobset_on_cluster_code != 0:
+    xpk_exit(set_jobset_on_cluster_code)
+
+  xpk_print('Enabling Kueue on the cluster')
+  install_kueue_on_cluster_code = install_kueue_on_cluster(args)
+  if install_kueue_on_cluster_code != 0:
+    xpk_exit(install_kueue_on_cluster_code)
+
+  xpk_print('Verifying kjob installation')
+  err_code = verify_kjob_installed(args)
+  if err_code > 0:
+    xpk_exit(err_code)
+
+  xpk_print('Applying kjob CDRs')
+  err_code = apply_kjob_crds(args)
+  if err_code > 0:
+    xpk_exit(err_code)
+
+  xpk_print('Preparing kjob')
+  err_code = prepare_kjob(args)
+  if err_code > 0:
+    xpk_exit(err_code)
+
+  xpk_print('Kind commands done! Resources are created.')
+  xpk_exit(0)
+
+
+def cluster_delete(args) -> None:
+  """Function around cluster delete.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  xpk_print(f'Starting cluster delete for cluster: {args.cluster}', flush=True)
+
+  run_kind_cluster_delete_command_code = run_kind_cluster_delete_command(args)
+  if run_kind_cluster_delete_command_code != 0:
+    xpk_exit(run_kind_cluster_delete_command_code)
+  xpk_print(f'Kind commands done! Cluster {args.cluster} deleted.')
+  xpk_exit(0)
+
+
+def cluster_list(args) -> None:
+  """Function around cluster list.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  if run_kind_clusters_list_command(args):
+    xpk_exit(1)
+  xpk_exit(0)
+
+
+def create_cluster_if_necessary(args) -> int:
+  """Creates cluster if not present in the project.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  all_clusters, return_code = get_all_local_clusters_programmatic(args)
+  if return_code > 0:
+    xpk_print('Listing all clusters failed!')
+    return 1
+  if args.cluster in all_clusters:
+    xpk_print('Skipping cluster creation since it already exists.')
+    return 0
+  else:
+    return run_kind_cluster_create_command(args)
+
+
+def run_kind_cluster_delete_command(args) -> int:
+  """Run the Delete Kind Cluster request.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = 'kind delete cluster'
+
+  if args.cluster:
+    command += f' --name={args.cluster}'
+
+  return_code = run_command_with_updates(command, 'Cluster Delete', args)
+  if return_code != 0:
+    xpk_print(f'Cluster delete request returned ERROR {return_code}')
+    return 1
+
+  return 0
+
+
+def run_kind_clusters_list_command(args) -> int:
+  """List Kind Clusters within the project and location.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = 'kind get clusters'
+  return_code = run_command_with_updates(command, 'Cluster List', args)
+  if return_code != 0:
+    xpk_print(f'Cluster list request returned ERROR {return_code}')
+    return 1
+
+  return 0
+
+
+def run_kind_cluster_create_command(args) -> int:
+  """Run the Create Kind Cluster request.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = 'kind create cluster'
+
+  if args.cluster:
+    command += f' --name={args.cluster}'
+
+  if args.k8s_version:
+    command += f' --image=kindest/node:v{args.k8s_version}'
+
+  return_code = run_command_with_updates(command, 'Kind Cluster Create', args)
+  if return_code != 0:
+    xpk_print(f'GKE Cluster Create request returned ERROR {return_code}')
+    return 1
+  return 0
+
+
+def get_all_local_clusters_programmatic(args) -> tuple[list[str], int]:
+  """Gets all the local clusters.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    List of cluster names and 0 if successful and 1 otherwise.
+  """
+  command = 'kind get clusters'
+  return_code, raw_cluster_output = run_command_for_value(
+      command, 'Find if Cluster Exists', args
+  )
+  if return_code != 0:
+    xpk_print(f'Find if Cluster Exists returned ERROR {return_code}')
+    return [], return_code
+
+  return raw_cluster_output.splitlines(), 0

src/xpk/parser/batch.py

@@ -14,6 +14,8 @@
 limitations under the License.
 """
 
+import argparse
+
 from .common import add_shared_arguments
 from ..commands.batch import batch
 
@@ -36,6 +38,13 @@ def set_batch_parser(batch_parser):
       default=None,
       help='Cluster to which command applies.',
   )
+  batch_optional_arguments.add_argument(
+      '--local-test',
+      type=bool,
+      action=argparse.BooleanOptionalAction,
+      default=False,
+      help='Apply command to a local test cluster.',
+  )
   add_shared_arguments(batch_optional_arguments)
 
   batch_parser.set_defaults(func=batch)

src/xpk/parser/common.py

@@ -50,3 +50,22 @@ def add_shared_arguments(custom_parser: argparse.ArgumentParser):
           ' branch based on the output of commands'
       ),
   )
+
+
+def add_global_arguments(custom_parser: argparse.ArgumentParser):
+  """Add global - no cloud dependent -  arguments to the parser.
+
+  Args:
+    custom_parser: parser to add global arguments to.
+  """
+  custom_parser.add_argument(
+      '--dry-run',
+      type=bool,
+      action=argparse.BooleanOptionalAction,
+      default=False,
+      help=(
+          'If given `--dry-run`, xpk will print the commands it wants to run'
+          ' but not run them. This is imperfect in cases where xpk might'
+          ' branch based on the output of commands'
+      ),
+  )

src/xpk/parser/core.py

@@ -23,6 +23,7 @@
 from .batch import set_batch_parser
 from .info import set_info_parser
 from .job import set_job_parser
+from .kind import set_kind_parser
 
 
 def set_parser(parser: argparse.ArgumentParser):
@@ -51,6 +52,10 @@ def set_parser(parser: argparse.ArgumentParser):
   job_parser = xpk_subcommands.add_parser(
       "job", help="commands around listing and cancelling jobs"
   )
+  kind_parser = xpk_subcommands.add_parser(
+      "kind",
+      help="commands around Kind cluster management",
+  )
 
   def default_subcommand_function(
       _args,
@@ -71,6 +76,8 @@ def default_subcommand_function(
     info_parser.print_help()
     job_parser.print_help()
 
+    kind_parser.print_help()
+
     return 0
 
   parser.set_defaults(func=default_subcommand_function)
@@ -79,10 +86,12 @@ def default_subcommand_function(
   batch_parser.set_defaults(func=default_subcommand_function)
   info_parser.set_defaults(func=default_subcommand_function)
   job_parser.set_defaults(func=default_subcommand_function)
+  kind_parser.set_defaults(func=default_subcommand_function)
 
   set_workload_parsers(workload_parser=workload_parser)
   set_cluster_parser(cluster_parser=cluster_parser)
   set_inspector_parser(inspector_parser=inspector_parser)
   set_batch_parser(batch_parser=batch_parser)
   set_info_parser(info_parser=info_parser)
   set_job_parser(job_parser=job_parser)
+  set_kind_parser(kind_parser=kind_parser)