xpk 0.4.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

xpk/commands/kind.py

@@ -0,0 +1,253 @@
+"""
+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)
+
+  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'
+  )
+  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
+
+
+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

xpk/commands/shell.py

@@ -0,0 +1,120 @@
+"""
+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_with_full_controls, run_command_for_value, run_command_with_updates
+from ..utils.console import xpk_exit, xpk_print
+from argparse import Namespace
+
+from ..core.kjob import AppProfileDefaults, PodTemplateDefaults
+
+
+exit_instructions = 'To exit the shell input "exit".'
+
+
+def shell(args: Namespace):
+  """Enter interactive shell.
+  Args:
+    args: user provided arguments for running the command.
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  exisitng_shell_pod_name = get_existing_shell_pod_name(args)
+
+  if exisitng_shell_pod_name is None:
+    return_code = connect_to_new_interactive_shell(args)
+  else:
+    return_code = connect_to_existing_interactive_shell(
+        exisitng_shell_pod_name, args
+    )
+
+  if return_code != 0:
+    xpk_print(f'The command failed with code {return_code}.')
+    xpk_exit(return_code)
+
+  xpk_exit(0)
+
+
+def get_existing_shell_pod_name(args: Namespace) -> str | None:
+  return_code, shell_name = run_command_for_value(
+      command=(
+          'kubectl get pods --no-headers --field-selector status.phase=Running'
+          ' -o custom-columns=":metadata.name"'
+      ),
+      task='Get existing interactive shell pod name.',
+      global_args=args,
+  )
+  if return_code != 0:
+    xpk_print(
+        f'Encounter an error with a code {return_code} when checking for'
+        ' existing running shell.'
+    )
+    xpk_exit(return_code)
+
+  pod_names = shell_name.strip().split('\n')
+  kjob_pod_names = [
+      name for name in pod_names if AppProfileDefaults.NAME.value in name
+  ]
+  shell_pod_names = [name for name in kjob_pod_names if 'interactive' in name]
+
+  return shell_pod_names[0] if shell_pod_names else None
+
+
+def connect_to_new_interactive_shell(args: Namespace) -> int:
+  return run_command_with_full_controls(
+      command=(
+          'kubectl-kjob create interactive --profile'
+          f' {AppProfileDefaults.NAME.value} --pod-running-timeout 30s'
+      ),
+      task='Creating new interactive shell and entering it',
+      global_args=args,
+      instructions=exit_instructions,
+  )
+
+
+def connect_to_existing_interactive_shell(
+    pod_name: str, args: Namespace
+) -> int:
+  return run_command_with_full_controls(
+      command=(
+          f'kubectl exec --stdin --tty {pod_name} --'
+          f' {PodTemplateDefaults.INTERACTIVE_COMMAND.value}'
+      ),
+      task='Entering existing interactive shell',
+      global_args=args,
+      instructions=exit_instructions,
+  )
+
+
+def shell_stop(args: Namespace):
+  """Stop the running interactive shell by deleting the pod.
+  Args:
+    args: user provided arguments for running the command.
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  exisitng_shell_pod_name = get_existing_shell_pod_name(args)
+
+  if exisitng_shell_pod_name is None:
+    xpk_print('There is no shell running to stop')
+    xpk_exit(0)
+
+  return_code = run_command_with_updates(
+      command=f'kubectl delete pod {exisitng_shell_pod_name}',
+      task='Deleting the existing shell.',
+      global_args=args,
+  )
+  if return_code != 0:
+    xpk_exit(return_code)
+
+  xpk_print('The shell was deleted successfully.')
+  xpk_exit(0)

xpk/commands/version.py

@@ -0,0 +1,39 @@
+"""
+Copyright 2025 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 argparse import Namespace
+import os
+
+from ..core.commands import run_command_for_value
+
+XPK_VERSION = 'v0.6.0'
+
+from ..utils.console import xpk_exit, xpk_print
+
+
+def version(args: Namespace) -> None:
+  """Get version of xpk."""
+  xpk_print('xpk_version:', XPK_VERSION)
+  if os.path.exists(os.path.join(os.getcwd(), '.git')):
+    code, xpk_version = run_command_for_value(
+        'git rev-parse HEAD',
+        task='Get latest hash',
+        global_args=args,
+        quiet=True,
+    )
+    if code != 0:
+      xpk_exit(code)
+    xpk_print('git commit:', xpk_version.strip('\n'))