From 5ea54a6ed2919eb24cce1838ef99f6c1cd9e9cb1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Irving=20Mondrag=C3=B3n?= Date: Mon, 11 Nov 2024 18:43:51 +0100 Subject: [PATCH 1/5] Add kind for local testing --- README.md | 44 ++++++++ src/xpk/commands/batch.py | 9 +- src/xpk/commands/kind.py | 215 ++++++++++++++++++++++++++++++++++++++ src/xpk/parser/batch.py | 9 ++ src/xpk/parser/common.py | 19 ++++ src/xpk/parser/core.py | 9 ++ src/xpk/parser/kind.py | 94 +++++++++++++++++ 7 files changed, 395 insertions(+), 4 deletions(-) create mode 100644 src/xpk/commands/kind.py create mode 100644 src/xpk/parser/kind.py diff --git a/README.md b/README.md index 0c9f9dc8..8564e762 100644 --- a/README.md +++ b/README.md @@ -1209,6 +1209,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..00fbde42 --- /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.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 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 86b780f9..a2728e0a 100644 --- a/src/xpk/parser/core.py +++ b/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,6 +86,7 @@ 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) @@ -86,3 +94,4 @@ def default_subcommand_function( 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) 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) From 63c6a5718f2e1e10f6204c0851b8043c6f872ae4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Irving=20Mondrag=C3=B3n?= Date: Tue, 19 Nov 2024 16:10:12 +0100 Subject: [PATCH 2/5] Rename flag to use local cluster --- README.md | 4 ++-- src/xpk/parser/batch.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 8564e762..0bbf4af0 100644 --- a/README.md +++ b/README.md @@ -1246,10 +1246,10 @@ xpk interfaces seamlessly with kind to manage Kubernetes clusters locally, facil ## 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: +Local testing is available exclusively through the `batch` command of xpk with the `--kind-cluster` flag. This allows you to simulate training jobs locally: ```shell -python xpk.py batch [other-options] --local-test script +python xpk.py batch [other-options] --kind-cluster 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. diff --git a/src/xpk/parser/batch.py b/src/xpk/parser/batch.py index b66c3ac3..9369185e 100644 --- a/src/xpk/parser/batch.py +++ b/src/xpk/parser/batch.py @@ -39,7 +39,7 @@ def set_batch_parser(batch_parser): help='Cluster to which command applies.', ) batch_optional_arguments.add_argument( - '--local-test', + '--kind-cluster', type=bool, action=argparse.BooleanOptionalAction, default=False, From 157bc74f2cc23f4bfd4284c932d1702f5c77a2b9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Irving=20Mondrag=C3=B3n?= Date: Tue, 19 Nov 2024 17:38:55 +0100 Subject: [PATCH 3/5] Allow to switch local cluster --- src/xpk/commands/batch.py | 44 ++++++++++++++++++++++++++++++++++++--- 1 file changed, 41 insertions(+), 3 deletions(-) diff --git a/src/xpk/commands/batch.py b/src/xpk/commands/batch.py index d82be7e2..9269aa26 100644 --- a/src/xpk/commands/batch.py +++ b/src/xpk/commands/batch.py @@ -23,6 +23,7 @@ from ..core.app_profile import APP_PROFILE_TEMPLATE_DEFAULT_NAME from ..core.commands import ( run_command_for_value, + run_command_with_updates, ) @@ -34,11 +35,14 @@ def batch(args: Namespace) -> None: Returns: None """ - if not args.local_test: + if not args.kind_cluster: 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) + else: + set_cluster_command_code = set_local_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) @@ -57,3 +61,37 @@ def submit_job(args: Namespace) -> None: if return_code != 0: xpk_print(f'Running batch job returned ERROR {return_code}') xpk_exit(return_code) + + +def set_local_cluster_command(args) -> int: + """Run local cluster configuration command to set the kubectl config. + + Args: + args: user provided arguments for running the command. + + Returns: + 0 if successful and 1 otherwise. + """ + if not args.cluster: + command = 'kubectl config current-context' + return_code, current_context = run_command_for_value( + command, 'get current-context', args + ) + xpk_print( + 'No local cluster name specified. Using current-context' + f' `{current_context.strip()}`' + ) + return return_code + + command = ( + f'kubectl config use-context kind-{args.cluster} --namespace=default' + ) + task = f'switch to cluster {args.cluster}' + return_code = run_command_with_updates( + command, + task, + args, + ) + if return_code != 0: + xpk_print(f'{task} returned ERROR {return_code}') + return return_code From 7d19eca92a95d3f4973f4a4158674f9ab4a86a4c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Irving=20Mondrag=C3=B3n?= Date: Wed, 20 Nov 2024 15:50:29 +0100 Subject: [PATCH 4/5] Add local testing support for job command --- README.md | 2 +- src/xpk/commands/job.py | 25 ++++++++++++++++--------- src/xpk/parser/job.py | 18 ++++++++++++++++++ 3 files changed, 35 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 0bbf4af0..181b6e5f 100644 --- a/README.md +++ b/README.md @@ -1246,7 +1246,7 @@ xpk interfaces seamlessly with kind to manage Kubernetes clusters locally, facil ## Local Testing Basics -Local testing is available exclusively through the `batch` command of xpk with the `--kind-cluster` flag. This allows you to simulate training jobs locally: +Local testing is available exclusively through the `batch` and `job` commands of xpk with the `--kind-cluster` flag. This allows you to simulate training jobs locally: ```shell python xpk.py batch [other-options] --kind-cluster script diff --git a/src/xpk/commands/job.py b/src/xpk/commands/job.py index 62836816..1a68c326 100644 --- a/src/xpk/commands/job.py +++ b/src/xpk/commands/job.py @@ -21,6 +21,7 @@ run_command_with_updates, ) from .cluster import set_cluster_command +from .batch import set_local_cluster_command def job_list(args) -> None: @@ -32,15 +33,17 @@ def job_list(args) -> None: Returns: None """ - add_zone_and_project(args) - set_cluster_command_code = set_cluster_command(args) + if not args.kind_cluster: + add_zone_and_project(args) + set_cluster_command_code = set_cluster_command(args) + msg = f'Listing jobs for project {args.project} and zone {args.zone}:' + else: + set_cluster_command_code = set_local_cluster_command(args) + msg = 'Listing jobs:' + if set_cluster_command_code != 0: xpk_exit(set_cluster_command_code) - - xpk_print( - f'Listing jobs for project {args.project} and zone {args.zone}:', - flush=True, - ) + xpk_print(msg, flush=True) return_code = run_slurm_job_list_command(args) xpk_exit(return_code) @@ -67,8 +70,12 @@ def job_cancel(args) -> None: None """ xpk_print(f'Starting job cancel for job: {args.name}', flush=True) - add_zone_and_project(args) - set_cluster_command_code = set_cluster_command(args) + if not args.kind_cluster: + add_zone_and_project(args) + set_cluster_command_code = set_cluster_command(args) + else: + set_cluster_command_code = set_local_cluster_command(args) + if set_cluster_command_code != 0: xpk_exit(set_cluster_command_code) diff --git a/src/xpk/parser/job.py b/src/xpk/parser/job.py index 71e60349..bd56c6be 100644 --- a/src/xpk/parser/job.py +++ b/src/xpk/parser/job.py @@ -14,6 +14,8 @@ limitations under the License. """ +import argparse + from .common import add_shared_arguments from ..commands.job import job_list, job_cancel @@ -48,6 +50,14 @@ def set_job_parser(job_parser): required=True, ) + job_list_optional_arguments.add_argument( + '--kind-cluster', + type=bool, + action=argparse.BooleanOptionalAction, + default=False, + help='Apply command to a local test cluster.', + ) + add_shared_arguments(job_list_optional_arguments) job_list_parser.set_defaults(func=job_list) @@ -81,5 +91,13 @@ def set_job_parser(job_parser): required=True, ) + job_cancel_optional_arguments.add_argument( + '--kind-cluster', + type=bool, + action=argparse.BooleanOptionalAction, + default=False, + help='Apply command to a local test cluster.', + ) + add_shared_arguments(job_cancel_optional_arguments) job_cancel_parser.set_defaults(func=job_cancel) From 6693e4cee8f873be7d1541709492c6adab1d5480 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Irving=20Mondrag=C3=B3n?= Date: Wed, 20 Nov 2024 16:36:53 +0100 Subject: [PATCH 5/5] Set context when creating kind cluster --- src/xpk/commands/batch.py | 40 ++------------------------------------- src/xpk/commands/job.py | 2 +- src/xpk/commands/kind.py | 38 +++++++++++++++++++++++++++++++++++++ 3 files changed, 41 insertions(+), 39 deletions(-) diff --git a/src/xpk/commands/batch.py b/src/xpk/commands/batch.py index 9269aa26..8f7c5591 100644 --- a/src/xpk/commands/batch.py +++ b/src/xpk/commands/batch.py @@ -21,10 +21,8 @@ from ..core.job_template import create_job_template_instance from ..core.app_profile import create_app_profile_instance from ..core.app_profile import APP_PROFILE_TEMPLATE_DEFAULT_NAME -from ..core.commands import ( - run_command_for_value, - run_command_with_updates, -) +from ..core.commands import run_command_for_value +from .kind import set_local_cluster_command def batch(args: Namespace) -> None: @@ -61,37 +59,3 @@ def submit_job(args: Namespace) -> None: if return_code != 0: xpk_print(f'Running batch job returned ERROR {return_code}') xpk_exit(return_code) - - -def set_local_cluster_command(args) -> int: - """Run local cluster configuration command to set the kubectl config. - - Args: - args: user provided arguments for running the command. - - Returns: - 0 if successful and 1 otherwise. - """ - if not args.cluster: - command = 'kubectl config current-context' - return_code, current_context = run_command_for_value( - command, 'get current-context', args - ) - xpk_print( - 'No local cluster name specified. Using current-context' - f' `{current_context.strip()}`' - ) - return return_code - - command = ( - f'kubectl config use-context kind-{args.cluster} --namespace=default' - ) - task = f'switch to cluster {args.cluster}' - return_code = run_command_with_updates( - command, - task, - args, - ) - if return_code != 0: - xpk_print(f'{task} returned ERROR {return_code}') - return return_code diff --git a/src/xpk/commands/job.py b/src/xpk/commands/job.py index 1a68c326..8430f581 100644 --- a/src/xpk/commands/job.py +++ b/src/xpk/commands/job.py @@ -21,7 +21,7 @@ run_command_with_updates, ) from .cluster import set_cluster_command -from .batch import set_local_cluster_command +from .kind import set_local_cluster_command def job_list(args) -> None: diff --git a/src/xpk/commands/kind.py b/src/xpk/commands/kind.py index 00fbde42..07c0ab28 100644 --- a/src/xpk/commands/kind.py +++ b/src/xpk/commands/kind.py @@ -47,6 +47,10 @@ def cluster_create(args) -> None: if create_cluster_command_code != 0: xpk_exit(create_cluster_command_code) + set_cluster_command_code = set_local_cluster_command(args) + if set_cluster_command_code != 0: + xpk_exit(set_cluster_command_code) + xpk_print( 'Enabling the jobset API on our cluster, to be deprecated when Jobset is' ' globally available' @@ -213,3 +217,37 @@ def get_all_local_clusters_programmatic(args) -> tuple[list[str], int]: return [], return_code return raw_cluster_output.splitlines(), 0 + + +def set_local_cluster_command(args) -> int: + """Run local cluster configuration command to set the kubectl config. + + Args: + args: user provided arguments for running the command. + + Returns: + 0 if successful and 1 otherwise. + """ + if not args.cluster: + command = 'kubectl config current-context' + return_code, current_context = run_command_for_value( + command, 'get current-context', args + ) + xpk_print( + 'No local cluster name specified. Using current-context' + f' `{current_context.strip()}`' + ) + return return_code + + command = ( + f'kubectl config use-context kind-{args.cluster} --namespace=default' + ) + task = f'switch to cluster {args.cluster}' + return_code = run_command_with_updates( + command, + task, + args, + ) + if return_code != 0: + xpk_print(f'{task} returned ERROR {return_code}') + return return_code