xpk 0.17.2__py3-none-any.whl → 1.0.0__py3-none-any.whl

xpk/commands/kind.py

@@ -1,286 +0,0 @@
-"""
-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.kueue_manager import (KueueConfig, KueueManager)
-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.scheduling import get_total_chips_requested_from_args
-from ..core.storage import install_storage_crd
-from ..core.system_characteristics import (
-    SystemCharacteristics,
-    AcceleratorType,
-    DockerPlatform,
-)
-from ..utils.console import (xpk_exit, xpk_print)
-from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
-
-
-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.
-  """
-  if should_validate_dependencies(args):
-    validate_dependencies_list([
-        SystemDependency.KUBECTL,
-        SystemDependency.KJOB,
-        SystemDependency.GCLOUD,
-    ])
-  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('Verifying kjob installation')
-  err_code = verify_kjob_installed()
-  if err_code > 0:
-    xpk_exit(err_code)
-
-  xpk_print('Applying kjob CDRs')
-  err_code = apply_kjob_crds()
-  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)
-
-  args.num_slices = 1
-  args.enable_pathways = False
-  system = SystemCharacteristics(
-      'N/A',
-      1,
-      'N/A',
-      'N/A',
-      1,
-      AcceleratorType.CPU,
-      'kind',
-      supports_sub_slicing=False,
-      supports_super_slicing=False,
-      docker_platform=DockerPlatform.ARM,
-  )
-
-  kueue_manager = KueueManager(project='', zone='')
-  kueue_manager.install_or_upgrade(
-      KueueConfig(
-          system,
-          total_chips=get_total_chips_requested_from_args(args, system),
-          autoprovisioning_enabled=False,
-          num_slices=args.num_slices,
-          memory_limit='',
-          cpu_limit=0,
-          is_pathways_cluster=False,
-          flex=False,
-          configure_sub_slicing=False,
-          configure_super_slicing=False,
-      ),
-  )
-
-  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.
-  """
-  if should_validate_dependencies(args):
-    validate_dependencies_list([SystemDependency.GCLOUD])
-  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.
-
-  Returns:
-    0 if successful and 1 otherwise.
-  """
-  if should_validate_dependencies(args):
-    validate_dependencies_list([SystemDependency.GCLOUD])
-  if run_kind_clusters_list_command():
-    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()
-  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')
-  if return_code != 0:
-    xpk_print(f'Cluster delete request returned ERROR {return_code}')
-    return 1
-
-  return 0
-
-
-def run_kind_clusters_list_command() -> int:
-  """List Kind Clusters within the project and location.
-
-  Returns:
-    0 if successful and 1 otherwise.
-  """
-  command = 'kind get clusters'
-  return_code = run_command_with_updates(command, 'Cluster List')
-  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')
-  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() -> tuple[list[str], int]:
-  """Gets all the local clusters.
-
-  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'
-  )
-  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'
-    )
-    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,
-  )
-  if return_code != 0:
-    xpk_print(f'{task} returned ERROR {return_code}')
-  return return_code

xpk/commands/kjob_common.py

@@ -1,60 +0,0 @@
-"""
-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.capacity import (
-    B200_DEVICE_TYPE,
-    H100_MEGA_DEVICE_TYPE,
-    H200_DEVICE_TYPE,
-)
-from ..core.cluster import get_gpu_type_from_cluster
-from ..core.kjob import (
-    get_a3mega_pod_template_annotations,
-    get_a3ultra_pod_template_annotations,
-    get_a4_pod_template_annotations,
-    Kueue_TAS_annotation,
-)
-from .common import is_TAS_possible
-from ..core.resources import get_cluster_capacity_type, get_cluster_system_characteristics
-
-
-def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
-  gpu_type = get_gpu_type_from_cluster(args)
-
-  annotations: tuple
-  if gpu_type == H100_MEGA_DEVICE_TYPE:
-    annotations = get_a3mega_pod_template_annotations()
-  elif gpu_type == H200_DEVICE_TYPE:
-    annotations = get_a3ultra_pod_template_annotations()
-  elif gpu_type == B200_DEVICE_TYPE:
-    annotations = get_a4_pod_template_annotations()
-  else:
-    annotations = tuple()
-
-  flags = [
-      f" --pod-template-annotation {annotation} " for annotation in annotations
-  ]
-  cmd += "\\\n".join(flags)
-
-  return cmd
-
-
-def add_TAS_annotations_to_command(args, cmd: str) -> str:
-  system_characteristics = get_cluster_system_characteristics(args)
-  capacity_type = get_cluster_capacity_type(args)
-  if is_TAS_possible(system_characteristics, capacity_type):
-    cmd += f" --pod-template-annotation {Kueue_TAS_annotation}"
-
-  return cmd

xpk/commands/run.py

@@ -1,140 +0,0 @@
-"""
-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 (
-    setup_k8s_service_accounts,
-    get_cluster_credentials,
-)
-from ..core.commands import run_command_with_full_controls
-from ..core.gcloud_context import add_zone_and_project
-from ..core.kjob import (
-    AppProfileDefaults,
-    JobTemplateDefaults,
-    get_storage_annotations,
-    prepare_kjob,
-)
-from ..core.kueue_manager import LOCAL_QUEUE_NAME
-from ..utils.console import xpk_exit, xpk_print
-from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
-from .kind import set_local_cluster_command
-from .kjob_common import add_gpu_networking_annotations_to_command, add_TAS_annotations_to_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 should_validate_dependencies(args):
-    validate_dependencies_list([
-        SystemDependency.KUBECTL,
-        SystemDependency.KJOB,
-        SystemDependency.GCLOUD,
-    ])
-  if not args.kind_cluster:
-    add_zone_and_project(args)
-    get_cluster_credentials(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)
-  setup_k8s_service_accounts()
-
-  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' --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)
-  cmd = add_TAS_annotations_to_command(args, cmd)
-
-  for annotation in get_storage_annotations(args):
-    cmd += f' --pod-template-annotation {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')
-
-  if return_code != 0:
-    xpk_print(f'Running task returned ERROR {return_code}')
-    xpk_exit(return_code)

xpk/commands/shell.py

@@ -1,142 +0,0 @@
-"""
-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, setup_k8s_service_accounts
-from ..utils.console import xpk_exit, xpk_print
-from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
-from argparse import Namespace
-
-from ..core.kjob import (
-    AppProfileDefaults,
-    prepare_kjob,
-    get_pod_template_interactive_command,
-    get_storage_annotations,
-)
-
-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.
-  """
-  if should_validate_dependencies(args):
-    validate_dependencies_list([
-        SystemDependency.KUBECTL,
-        SystemDependency.KJOB,
-        SystemDependency.GCLOUD,
-    ])
-  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)
-
-  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.',
-  )
-  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)
-  setup_k8s_service_accounts()
-
-  cmd = (
-      'kubectl-kjob create interactive --profile'
-      f' {AppProfileDefaults.NAME.value} --pod-running-timeout 180s'
-  )
-
-  for annotation in get_storage_annotations(args):
-    cmd += f' --pod-template-annotation {annotation}'
-
-  return run_command_with_full_controls(
-      command=cmd,
-      task='Creating new interactive shell and entering it',
-      instructions=exit_instructions,
-  )
-
-
-def connect_to_existing_interactive_shell(pod_name: str) -> 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',
-      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.
-  """
-  if should_validate_dependencies(args):
-    validate_dependencies_list(
-        [SystemDependency.KUBECTL, SystemDependency.GCLOUD]
-    )
-  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.',
-  )
-  if return_code != 0:
-    xpk_exit(return_code)
-
-  xpk_print('The shell was deleted successfully.')
-  xpk_exit(0)

xpk/parser/batch.py

@@ -1,43 +0,0 @@
-"""
-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 .common import (
-    add_shared_arguments,
-    add_slurm_arguments,
-    add_cluster_arguments,
-    add_kind_cluster_arguments,
-)
-from ..commands.batch import batch
-
-
-def set_batch_parser(batch_parser):
-  batch_required_arguments = batch_parser.add_argument_group(
-      'batch Built-in Arguments', 'Arguments required for `batch`.'
-  )
-  batch_optional_arguments = batch_parser.add_argument_group(
-      'Optional Arguments', 'Arguments optional for `batch`.'
-  )
-
-  ### "batch" Required arguments
-  batch_required_arguments.add_argument(
-      'script', help='script with batch task to run'
-  )
-
-  add_cluster_arguments(batch_optional_arguments)
-  add_kind_cluster_arguments(batch_optional_arguments)
-  add_shared_arguments(batch_optional_arguments)
-  add_slurm_arguments(batch_optional_arguments)
-  batch_parser.set_defaults(func=batch)