xpk 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

xpk/parser/common.py

@@ -0,0 +1,259 @@
+"""
+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.
+"""
+
+import argparse
+
+
+def add_shared_arguments(
+    custom_parser: argparse.ArgumentParser, required=False
+) -> None:
+  """Add shared arguments to the parser.
+
+  Args:
+    custom_parser: parser to add shared arguments to.
+  """
+  custom_parser.add_argument(
+      '--project',
+      type=str,
+      default=None,
+      help='GCE project name, defaults to "gcloud config project."',
+      required=required,
+  )
+  custom_parser.add_argument(
+      '--zone',
+      type=str,
+      default=None,
+      help=(
+          'GCE zone, e.g. us-central2-b, defaults to "gcloud config '
+          'compute/zone." Only one of --zone or --region is allowed in a '
+          'command.'
+      ),
+      required=required,
+  )
+  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'
+      ),
+      required=required,
+  )
+
+
+def add_cluster_arguments(
+    custom_parser: argparse.ArgumentParser, required=False
+) -> None:
+  """Add cluster argument to the parser.
+
+  Args:
+    custom_parser: parser to add shared arguments to.
+  """
+  custom_parser.add_argument(
+      '--cluster',
+      type=str,
+      default=None,
+      help='The name of the cluster.',
+      required=required,
+  )
+
+
+def add_kind_cluster_arguments(custom_parser: argparse.ArgumentParser) -> None:
+  """Add kind cluster arguments to the parser.
+
+  Args:
+    custom_parser: parser to add shared arguments to.
+  """
+  custom_parser.add_argument(
+      '--kind-cluster',
+      type=bool,
+      action=argparse.BooleanOptionalAction,
+      default=False,
+      help='Apply command to a local test cluster.',
+  )
+
+
+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'
+      ),
+  )
+
+
+def add_slurm_arguments(custom_parser: argparse.ArgumentParser):
+  """Add Slurm job arguments to the parser.
+
+  Args:
+    custom_parser: parser to add global arguments to.
+  """
+  custom_parser.add_argument(
+      '--ignore-unknown-flags',
+      type=bool,
+      action=argparse.BooleanOptionalAction,
+      default=False,
+      help='Ignore all the unsupported flags in the bash script.',
+  )
+  custom_parser.add_argument(
+      '-a',
+      '--array',
+      type=str,
+      default=None,
+      help=(
+          'Submit a job array, multiple jobs to be executed with identical'
+          ' parameters. The indexes specification identifies what array index'
+          ' values should be used. For example, "--array=0-15" or'
+          ' "--array=0,6,16-32". Multiple values may be specified using a comma'
+          ' separated list and/or a range of values with a "-" separator. For'
+          ' example "--array=0-15%%4" will limit the number of simultaneously'
+          ' running tasks from this job array to 4. The minimum index value is'
+          ' 0. The maximum index value is 2147483647.'
+      ),
+  )
+  custom_parser.add_argument(
+      '-c',
+      '--cpus-per-task',
+      type=str,
+      default=None,
+      help='How much cpus a container inside a pod requires.',
+  )
+  custom_parser.add_argument(
+      '--gpus-per-task',
+      type=str,
+      default=None,
+      help='How much gpus a container inside a pod requires.',
+  )
+  custom_parser.add_argument(
+      '--mem',
+      type=str,
+      default=None,
+      help='How much memory a pod requires.',
+  )
+  custom_parser.add_argument(
+      '--mem-per-task',
+      type=str,
+      default=None,
+      help='How much memory a container requires.',
+  )
+  custom_parser.add_argument(
+      '--mem-per-cpu',
+      type=str,
+      default=None,
+      help=(
+          'How much memory a container requires, it multiplies the number '
+          'of requested cpus per task by mem-per-cpu.'
+      ),
+  )
+  custom_parser.add_argument(
+      '--mem-per-gpu',
+      type=str,
+      default=None,
+      help=(
+          'How much memory a container requires, it multiplies the number '
+          'of requested gpus per task by mem-per-gpu.'
+      ),
+  )
+  custom_parser.add_argument(
+      '-N',
+      '--nodes',
+      type=int,
+      default=None,
+      help='Number of pods to be used at a time.',
+  )
+  custom_parser.add_argument(
+      '-n',
+      '--ntasks',
+      type=int,
+      default=None,
+      help='Number of identical containers inside of a pod, usually 1.',
+  )
+  custom_parser.add_argument(
+      '-o',
+      '--output',
+      type=str,
+      default=None,
+      help=(
+          'Where to redirect the standard output stream of a task. If not'
+          ' passed it proceeds to stdout, and is available via kubectl logs.'
+      ),
+  )
+  custom_parser.add_argument(
+      '-e',
+      '--error',
+      type=str,
+      default=None,
+      help=(
+          'Where to redirect std error stream of a task. If not passed it'
+          ' proceeds to stdout, and is available via kubectl logs.'
+      ),
+  )
+  custom_parser.add_argument(
+      '--input',
+      type=str,
+      default=None,
+      help='What to pipe into the script.',
+  )
+  custom_parser.add_argument(
+      '-J',
+      '--job-name',
+      type=str,
+      default=None,
+      help='What is the job name.',
+  )
+  custom_parser.add_argument(
+      '-D',
+      '--chdir',
+      type=str,
+      default=None,
+      help='Change directory before executing the script.',
+  )
+  custom_parser.add_argument(
+      '-t',
+      '--time',
+      type=str,
+      default=None,
+      help=(
+          'Set a limit on the total run time of the job. '
+          'A time limit of zero requests that no time limit be imposed. '
+          'Acceptable time formats include "minutes", "minutes:seconds", '
+          '"hours:minutes:seconds", "days-hours", "days-hours:minutes" '
+          'and "days-hours:minutes:seconds".'
+      ),
+  )
+  custom_parser.add_argument(
+      '--priority',
+      type=str,
+      default='medium',
+      choices=['very-low', 'low', 'medium', 'high', 'very-high'],
+      help=(
+          'A priority, one of `very-low`, `low`, `medium`, `high` or'
+          ' `very-high`. Defaults to `medium`.'
+      ),
+  )

xpk/parser/config.py

@@ -0,0 +1,49 @@
+"""
+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.config import get_config, set_config
+from ..core.config import DEFAULT_KEYS
+from .common import add_shared_arguments
+
+
+def set_config_parsers(config_parser):
+  add_shared_arguments(config_parser)
+
+  config_subcommands = config_parser.add_subparsers(
+      title='config subcommands', dest='xpk_config_subcommands'
+  )
+  config_set_parser = config_subcommands.add_parser(
+      'set', help='set config key'
+  )
+  config_get_parser = config_subcommands.add_parser(
+      'get', help='get config key'
+  )
+  config_set_parser.add_argument(
+      'set_config_args',
+      help=f"""Pair of (key, value) to be set in config. Allowed keys are: {DEFAULT_KEYS}.
+      Command usage: `xpk config set key value`""",
+      type=str,
+      nargs=2,
+  )
+  config_get_parser.add_argument(
+      'get_config_key',
+      help=f"""Get key value from config. Allowed keys are: {DEFAULT_KEYS} .
+      Command usage: `xpk config get key`""",
+      type=str,
+      nargs=1,
+  )
+  config_set_parser.set_defaults(func=set_config)
+  config_get_parser.set_defaults(func=get_config)

xpk/parser/core.py

@@ -0,0 +1,135 @@
+"""
+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.
+"""
+
+import argparse
+
+from .config import set_config_parsers
+
+from ..utils.console import xpk_print
+from .cluster import set_cluster_parser
+from .inspector import set_inspector_parser
+from .storage import set_storage_parser
+from .workload import set_workload_parsers
+from .batch import set_batch_parser
+from .job import set_job_parser
+from .info import set_info_parser
+from .kind import set_kind_parser
+from .shell import set_shell_parser
+from .version import set_version_parser
+from .run import set_run_parser
+
+
+def set_parser(parser: argparse.ArgumentParser):
+  xpk_subcommands = parser.add_subparsers(
+      title="xpk subcommands", dest="xpk_subcommands", help="Top level commands"
+  )
+  workload_parser = xpk_subcommands.add_parser(
+      "workload", help="Commands around workload management"
+  )
+  storage_parser = xpk_subcommands.add_parser(
+      "storage", help="Commands around storage management"
+  )
+  cluster_parser = xpk_subcommands.add_parser(
+      "cluster",
+      help="Commands around creating, deleting, and viewing clusters.",
+  )
+  inspector_parser = xpk_subcommands.add_parser(
+      "inspector",
+      help="Commands around investigating workload, and Kueue failures.",
+  )
+  info_parser = xpk_subcommands.add_parser(
+      "info",
+      help="Commands around listing kueue clusterqueues and localqueues.",
+  )
+  batch_parser = xpk_subcommands.add_parser(
+      "batch",
+      help="commands around running batch job",
+  )
+  job_parser = xpk_subcommands.add_parser(
+      "job", help="commands around listing, cancelling and investigating jobs"
+  )
+  kind_parser = xpk_subcommands.add_parser(
+      "kind",
+      help="commands around Kind cluster management",
+  )
+  shell_parser = xpk_subcommands.add_parser(
+      "shell", help="Commands around configuring and using interactive shell."
+  )
+  version_parser = xpk_subcommands.add_parser(
+      "version", help="Command to get xpk version"
+  )
+
+  config_parser = xpk_subcommands.add_parser(
+      "config", help="Commands to set and retrieve values from xpk config."
+  )
+
+  run_parser = xpk_subcommands.add_parser(
+      "run",
+      help="Command to run parallel jobs",
+  )
+
+  def default_subcommand_function(
+      _args,
+  ) -> int:  # args is unused, so pylint: disable=invalid-name
+    """Default subcommand function.
+
+    Args:
+      _args: user provided arguments for running the command.
+
+    Returns:
+      0 if successful and 1 otherwise.
+    """
+    xpk_print("Welcome to XPK! See below for overall commands:", flush=True)
+    parser.print_help()
+    cluster_parser.print_help()
+    workload_parser.print_help()
+    batch_parser.print_help()
+    info_parser.print_help()
+    job_parser.print_help()
+    shell_parser.print_help()
+    version_parser.print_help()
+    kind_parser.print_help()
+    config_parser.print_help()
+    run_parser.print_help()
+
+    storage_parser.print_help()
+    return 0
+
+  parser.set_defaults(func=default_subcommand_function)
+  workload_parser.set_defaults(func=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)
+  job_parser.set_defaults(func=default_subcommand_function)
+  kind_parser.set_defaults(func=default_subcommand_function)
+  shell_parser.set_defaults(func=default_subcommand_function)
+  storage_parser.set_defaults(func=default_subcommand_function)
+  version_parser.set_defaults(func=default_subcommand_function)
+  config_parser.set_defaults(func=default_subcommand_function)
+  run_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_job_parser(job_parser=job_parser)
+  set_kind_parser(kind_parser=kind_parser)
+  set_shell_parser(shell_parser=shell_parser)
+  set_storage_parser(storage_parser=storage_parser)
+  set_version_parser(version_parser=version_parser)
+  set_config_parsers(config_parser=config_parser)
+  set_run_parser(run_parser=run_parser)

xpk/parser/info.py

@@ -0,0 +1,64 @@
+"""
+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.info import info
+from .common import add_shared_arguments
+from .validators import name_type
+import argparse
+
+
+def set_info_parser(info_parser: argparse.ArgumentParser) -> None:
+  info_required_arguments = info_parser.add_argument_group(
+      'Required Arguments', 'Arguments required for info.'
+  )
+  info_optional_arguments = info_parser.add_argument_group(
+      'Optional Arguments', 'Arguments optional for info.'
+  )
+
+  info_required_arguments.add_argument(
+      '--cluster',
+      type=name_type,
+      default=None,
+      help='Cluster to which command applies.',
+      required=True,
+  )
+
+  info_optional_arguments.add_argument(
+      '--namespace',
+      type=str,
+      default='',
+      help='Namespace to which resources and queues belong',
+  )
+
+  queues_flitering_group = (
+      info_optional_arguments.add_mutually_exclusive_group()
+  )
+
+  queues_flitering_group.add_argument(
+      '--clusterqueue',
+      action='store_true',
+      default=None,
+      help='Show only clusterqueues resources and usage',
+  )
+
+  queues_flitering_group.add_argument(
+      '--localqueue',
+      action='store_true',
+      default=None,
+      help='Show only localqueues resources and usage',
+  )
+  add_shared_arguments(info_optional_arguments)
+  info_parser.set_defaults(func=info)

xpk/parser/inspector.py

@@ -0,0 +1,65 @@
+"""
+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.inspector import inspector
+from .validators import name_type
+from .common import add_shared_arguments
+
+
+def set_inspector_parser(inspector_parser):
+  inspector_parser.add_subparsers(
+      title='inspector subcommands',
+      dest='xpk_inspector_subcommands',
+      help='Investigate workload, and Kueue failures.',
+  )
+
+  inspector_parser_required_arguments = inspector_parser.add_argument_group(
+      'inspector Built-in Arguments', 'Arguments required for `inspector`.'
+  )
+  inspector_parser_optional_arguments = inspector_parser.add_argument_group(
+      'Optional Arguments', 'Arguments optional for `inspector`.'
+  )
+
+  ### "inspector" Required arguments
+
+  inspector_parser_required_arguments.add_argument(
+      '--cluster',
+      type=name_type,
+      default=None,
+      help='The name of the cluster to investigate.',
+      required=True,
+  )
+
+  ### "inspector" Optional Arguments
+  add_shared_arguments(inspector_parser_optional_arguments)
+
+  inspector_parser_optional_arguments.add_argument(
+      '--workload',
+      type=name_type,
+      default=None,
+      help='The name of the workload to investigate.',
+  )
+
+  inspector_parser_optional_arguments.add_argument(
+      '--print-to-terminal',
+      action='store_true',
+      help=(
+          'Prints inspector output to terminal. A user can always look at the'
+          ' returned file.'
+      ),
+  )
+
+  inspector_parser.set_defaults(func=inspector)