diff --git a/README.md b/README.md index a8d4acb8..3e9056cc 100644 --- a/README.md +++ b/README.md @@ -1172,6 +1172,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) diff --git a/src/xpk/commands/batch.py b/src/xpk/commands/batch.py index 7001ec2c..d82be7e2 100644 --- a/src/xpk/commands/batch.py +++ b/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) diff --git a/src/xpk/commands/kind.py b/src/xpk/commands/kind.py new file mode 100644 index 00000000..a5b3bd6b --- /dev/null +++ b/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 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 diff --git a/src/xpk/parser/batch.py b/src/xpk/parser/batch.py index 1a0f6275..b66c3ac3 100644 --- a/src/xpk/parser/batch.py +++ b/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) diff --git a/src/xpk/parser/common.py b/src/xpk/parser/common.py index 390f7bd2..f3bfe50e 100644 --- a/src/xpk/parser/common.py +++ b/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' + ), + ) diff --git a/src/xpk/parser/core.py b/src/xpk/parser/core.py index 14b93f8f..7df53cc0 100644 --- a/src/xpk/parser/core.py +++ b/src/xpk/parser/core.py @@ -22,6 +22,7 @@ from .workload import set_workload_parsers from .batch import set_batch_parser from .info import set_info_parser +from .kind import set_kind_parser def set_parser(parser: argparse.ArgumentParser): @@ -47,6 +48,10 @@ def set_parser(parser: argparse.ArgumentParser): "batch", help="Run batch job.", ) + kind_parser = xpk_subcommands.add_parser( + "kind", + help="commands around Kind cluster management", + ) def default_subcommand_function( _args, @@ -65,6 +70,8 @@ def default_subcommand_function( workload_parser.print_help() batch_parser.print_help() info_parser.print_help() + kind_parser.print_help() + return 0 parser.set_defaults(func=default_subcommand_function) @@ -72,9 +79,11 @@ def default_subcommand_function( cluster_parser.set_defaults(func=default_subcommand_function) batch_parser.set_defaults(func=default_subcommand_function) info_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_kind_parser(kind_parser=kind_parser) diff --git a/src/xpk/parser/kind.py b/src/xpk/parser/kind.py new file mode 100644 index 00000000..1554fb0c --- /dev/null +++ b/src/xpk/parser/kind.py @@ -0,0 +1,94 @@ +""" +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 ..commands.kind import ( + cluster_create, + cluster_delete, + cluster_list, +) +from .common import add_global_arguments + + +def set_kind_parser(kind_parser): + cluster_subcommands = kind_parser.add_subparsers( + title='kind subcommands', + dest='xpk_kind_subcommands', + help=( + 'These are commands related to kind management. Look at help for' + ' specific subcommands for more details.' + ), + ) + + ### "cluster create" command parser ### + cluster_create_parser = cluster_subcommands.add_parser( + 'create', help='Create local clusters.' + ) + + ### Optional Arguments + cluster_create_parser.add_argument( + '--cluster', + type=str, + default='kind', + help=( + 'The name of the cluster. Will be used as the prefix for internal' + ' objects in the cluster.' + ), + required=False, + ) + + cluster_create_parser.add_argument( + '--k8s-version', + type=str, + default='', + help='The Kubernetes version of the cluster.', + required=False, + ) + + add_global_arguments(cluster_create_parser) + cluster_create_parser.set_defaults(func=cluster_create) + + ### "cluster delete" command parser ### + cluster_delete_parser = cluster_subcommands.add_parser( + 'delete', + help='Delete cloud clusters.', + ) + + cluster_delete_required_arguments = cluster_delete_parser.add_argument_group( + 'Required Arguments', + 'Arguments required for cluster delete.', + ) + + ### Required arguments + cluster_delete_required_arguments.add_argument( + '--cluster', + type=str, + default=None, + help='The name of the cluster to be deleted.', + required=True, + ) + + ### Optional Arguments + add_global_arguments(cluster_delete_parser) + cluster_delete_parser.set_defaults(func=cluster_delete) + + # "cluster list" command parser. + cluster_list_parser = cluster_subcommands.add_parser( + 'list', help='List cloud clusters.' + ) + + ### Optional Arguments + add_global_arguments(cluster_list_parser) + cluster_list_parser.set_defaults(func=cluster_list)