xpk 0.0.1__py3-none-any.whl

xpk/commands/kind.py

@@ -0,0 +1,283 @@
+"""
+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.cluster import set_jobset_on_cluster, setup_k8s_env
+from ..core.kjob import (
+    verify_kjob_installed,
+    prepare_kjob,
+    apply_kjob_crds,
+)
+from ..core.kueue import (
+    install_kueue_on_cluster,
+    install_kueue_crs,
+    wait_for_kueue_available,
+)
+from ..core.storage import install_storage_crd
+from ..core.system_characteristics import (
+    SystemCharacteristics,
+    AcceleratorType,
+)
+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)
+
+  args.kind_cluster = True
+  err_code = prepare_kjob(args)
+  if err_code > 0:
+    xpk_exit(err_code)
+
+  k8s_client = setup_k8s_env(args)
+  install_storage_crd(k8s_client)
+
+  xpk_print('Wait for Kueue to be fully available')
+  wait_for_kueue_available_code = wait_for_kueue_available(args)
+  if wait_for_kueue_available_code != 0:
+    xpk_exit(wait_for_kueue_available_code)
+
+  args.num_slices = 1
+  args.enable_pathways = False
+  system = SystemCharacteristics(
+      'N/A',
+      1,
+      'N/A',
+      'N/A',
+      1,
+      AcceleratorType['CPU'],
+      'kind',
+  )
+
+  xpk_print('Install Kueue Custom Resources')
+  enable_kueue_credentials_code = install_kueue_crs(args, system, None)
+  if enable_kueue_credentials_code != 0:
+    xpk_exit(enable_kueue_credentials_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/kjob_common.py

@@ -0,0 +1,44 @@
+"""
+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 ..core.kjob import get_a3mega_pod_template_annotations, get_a3ultra_pod_template_annotations
+from ..core.capacity import H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
+from ..core.cluster import get_gpu_type_from_cluster
+
+
+def add_tcpxo_annotations(args, cmd: str) -> str:
+  tcpxo, interfaces, eth0 = get_a3mega_pod_template_annotations(args)
+  cmd += f" --pod-template-annotation {tcpxo} \\\n"
+  cmd += f" --pod-template-annotation {eth0} \\\n"
+  cmd += f" --pod-template-annotation {interfaces} "
+  return cmd
+
+
+def add_rdma_annotations(args, cmd) -> str:
+  eth0, interfaces = get_a3ultra_pod_template_annotations(args)
+  cmd += f" --pod-template-annotation {eth0} \\\n"
+  cmd += f" --pod-template-annotation {interfaces} \\\n"
+  return cmd
+
+
+def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
+  gpu_type = get_gpu_type_from_cluster(args)
+
+  if gpu_type == H100_MEGA_DEVICE_TYPE:
+    return add_tcpxo_annotations(args, cmd)
+  if gpu_type == H200_DEVICE_TYPE:
+    return add_rdma_annotations(args, cmd)
+  return cmd

xpk/commands/run.py

@@ -0,0 +1,128 @@
+"""
+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
+
+from ..core.cluster import create_xpk_k8s_service_account
+from ..core.commands import run_command_with_full_controls
+from ..core.gcloud_context import add_zone_and_project
+from ..core.kueue import LOCAL_QUEUE_NAME
+from ..utils.console import xpk_exit, xpk_print
+from .common import set_cluster_command
+from ..core.kjob import JobTemplateDefaults, AppProfileDefaults, prepare_kjob, Kueue_TAS_annotation, get_gcsfuse_annotation
+from .kjob_common import add_gpu_networking_annotations_to_command
+from .kind import set_local_cluster_command
+
+
+def run(args: Namespace) -> None:
+  """Run task.
+     This function runs passed script in non-blocking manner.
+  Args:
+    args: user provided arguments for running the command.
+  Returns:
+    None
+  """
+  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)
+
+  err_code = prepare_kjob(args)
+  if err_code > 0:
+    xpk_exit(err_code)
+  create_xpk_k8s_service_account()
+
+  submit_job(args)
+
+
+def submit_job(args: Namespace) -> None:
+  cmd = (
+      'kubectl kjob create slurm --profile'
+      f' {AppProfileDefaults.NAME.value} '
+      f' --localqueue {LOCAL_QUEUE_NAME} '
+      f" --pod-template-annotation '{Kueue_TAS_annotation}'"
+      f' --stream-container {JobTemplateDefaults.CONTAINER_NAME.value}'
+      f' --worker-container {JobTemplateDefaults.CONTAINER_NAME.value}'
+      ' --wait --rm  --first-node-ip'
+  )
+  cmd = add_gpu_networking_annotations_to_command(args, cmd)
+
+  gcsfuse_annotation = get_gcsfuse_annotation(args)
+  if gcsfuse_annotation is not None:
+    cmd += f' --pod-template-annotation {gcsfuse_annotation}'
+
+  if args.timeout:
+    cmd += f' --wait-timeout {args.timeout}s'
+
+  if args.ignore_unknown_flags:
+    cmd += ' --ignore-unknown-flags'
+
+  cmd += f' -- {args.script} --partition {LOCAL_QUEUE_NAME}'
+
+  if args.array is not None:
+    cmd += f' --array {args.array}'
+
+  if args.cpus_per_task is not None:
+    cmd += f' --cpus-per-task {args.cpus_per_task}'
+
+  if args.gpus_per_task is not None:
+    cmd += f' --gpus-per-task {args.gpus_per_task}'
+
+  if args.mem is not None:
+    cmd += f' --mem {args.mem}'
+
+  if args.mem_per_task is not None:
+    cmd += f' --mem-per-task {args.mem_per_task}'
+
+  if args.mem_per_cpu is not None:
+    cmd += f' --mem-per-cpu {args.mem_per_cpu}'
+
+  if args.mem_per_gpu is not None:
+    cmd += f' --mem-per-gpu {args.mem_per_gpu}'
+
+  if args.nodes is not None:
+    cmd += f' --nodes {args.nodes}'
+
+  if args.ntasks is not None:
+    cmd += f' --ntasks {args.ntasks}'
+
+  if args.output is not None:
+    cmd += f' --output {args.output}'
+
+  if args.error is not None:
+    cmd += f' --error {args.error}'
+
+  if args.input is not None:
+    cmd += f' --input {args.input}'
+
+  if args.job_name is not None:
+    cmd += f' --job-name {args.job_name}'
+
+  if args.chdir is not None:
+    cmd += f' --chdir {args.chdir}'
+
+  if args.time is not None:
+    cmd += f' --time {args.time}'
+
+  return_code = run_command_with_full_controls(cmd, 'run task', args)
+
+  if return_code != 0:
+    xpk_print(f'Running task returned ERROR {return_code}')
+    xpk_exit(return_code)

xpk/commands/shell.py

@@ -0,0 +1,140 @@
+"""
+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 ..core.cluster import get_cluster_credentials, add_zone_and_project, create_xpk_k8s_service_account
+from ..utils.console import xpk_exit, xpk_print
+from argparse import Namespace
+
+from ..core.kjob import (
+    AppProfileDefaults,
+    prepare_kjob,
+    get_pod_template_interactive_command,
+    get_gcsfuse_annotation,
+)
+
+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:
+  if not args.kind_cluster:
+    add_zone_and_project(args)
+    get_cluster_credentials(args)
+
+  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:
+  err_code = prepare_kjob(args)
+  if err_code > 0:
+    xpk_exit(err_code)
+  create_xpk_k8s_service_account()
+
+  cmd = (
+      'kubectl-kjob create interactive --profile'
+      f' {AppProfileDefaults.NAME.value} --pod-running-timeout 180s'
+  )
+
+  gcsfuse_annotation = get_gcsfuse_annotation(args)
+  if gcsfuse_annotation is not None:
+    cmd += f' --pod-template-annotation {gcsfuse_annotation}'
+
+  return run_command_with_full_controls(
+      command=cmd,
+      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' {get_pod_template_interactive_command()}'
+      ),
+      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)