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

xpk/core/workload.py

@@ -0,0 +1,133 @@
+"""
+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 import run_command_for_value
+
+
+def workload_list_awk_command(filter_key) -> str:
+  """Function returns the awk command needed from the filter specified.
+
+  Args:
+    filter_key: workload list filter to awk against
+
+  Returns:
+    awk command to use in filtering workload list.
+  """
+
+  return f" | awk -e 'NR == 1 || {filter_key} {{print $0}}'"
+
+
+def determine_workload_list_filter_by_status(args) -> str:
+  """Function to create the filtered view of workload list.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    the argument needed to filter by status of jobs in workload list.
+  """
+
+  # Argument positions related to columns created by workload list command.
+  status_arg = '$7'
+  running_vms_arg = '$5'
+  status_verbose_arg = '$9'
+  if args.filter_by_status == 'EVERYTHING':
+    return ''
+  elif args.filter_by_status == 'RUNNING':
+    # Running includes the status Admitted or Evicted, and when the number of
+    # vms running is > 0.
+    return workload_list_awk_command(
+        f'({status_arg} ~ "Admitted|Evicted" && {running_vms_arg} ~ /^[0-9]+$/'
+        f' && {running_vms_arg} > 0)'
+    )
+  elif args.filter_by_status == 'QUEUED':
+    # Queued includes the status Admitted or Evicted, and when the number of
+    # vms running is 0.
+    return workload_list_awk_command(
+        f'({status_arg} ~ "Admitted|Evicted|QuotaReserved" &&'
+        f' ({running_vms_arg} ~ "<none>" || {running_vms_arg} == 0))'
+    )
+  elif args.filter_by_status == 'FINISHED':
+    return workload_list_awk_command(f'{status_arg} == "Finished"')
+  elif args.filter_by_status == 'FAILED':
+    # Failed includes the status Finished, and when the verbose reason is failed.
+    return workload_list_awk_command(
+        f'({status_arg} == "Finished" && {status_verbose_arg} ~ "failed")'
+    )
+  elif args.filter_by_status == 'SUCCESSFUL':
+    # Failed includes the status Finished, and when the verbose reason is finished/success.
+    return workload_list_awk_command(
+        f'({status_arg} == "Finished" && {status_verbose_arg} ~ "finished")'
+    )
+  raise RuntimeError(f'Can not find filter type: {args.filter_by_status}')
+
+
+def determine_workload_list_filter_by_job(args) -> str:
+  """Function to filter view of workload list based on job name.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    the argument needed to filter job names from workload list
+  """
+  # Argument positions related to columns created by workload list command.
+  if not hasattr(args, 'filter_by_job') or args.filter_by_job is None:
+    return ''
+  else:
+    job_name_arg = '$1'
+    return workload_list_awk_command(f'{job_name_arg} ~ "{args.filter_by_job}"')
+
+
+def get_workload_list(args) -> tuple[int, str]:
+  """Function to get the list of the workloads in the cluster.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    return_code: 0 if successful and 1 otherwise.
+    return_value: workloads in the cluster matching the criteria.
+  """
+  columns = {
+      'Jobset Name': '.metadata.ownerReferences[0].name',
+      'Created Time': '.metadata.creationTimestamp',
+      'Priority': '.spec.priorityClassName',
+      'TPU VMs Needed': '.spec.podSets[0].count',
+      'TPU VMs Running/Ran': '.status.admission.podSetAssignments[-1].count',
+      'TPU VMs Done': '.status.reclaimablePods[0].count',
+      'Status': '.status.conditions[-1].type',
+      'Status Message': '.status.conditions[-1].message',
+      'Status Time': '.status.conditions[-1].lastTransitionTime',
+  }
+  s = ','.join([key + ':' + value for key, value in columns.items()])
+
+  workload_list_filter_status_cmd = determine_workload_list_filter_by_status(
+      args
+  )
+  workload_list_filter_job_cmd = determine_workload_list_filter_by_job(args)
+  command = (
+      f'kubectl get workloads -o=custom-columns="{s}" '
+      f'{workload_list_filter_status_cmd} {workload_list_filter_job_cmd}'
+  )
+
+  task = f'List Jobs with filter-by-status={args.filter_by_status}'
+  if hasattr(args, 'filter_by_job'):
+    task += f' with filter-by-job={args.filter_by_job}'
+
+  return_code, return_value = run_command_for_value(command, task, args)
+
+  return return_code, return_value

xpk/core/workload_decorators/__init__.py

@@ -0,0 +1,15 @@
+"""
+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.
+"""

xpk/core/workload_decorators/rdma_decorator.py

@@ -0,0 +1,109 @@
+"""
+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 yaml
+from ...utils.yaml import literal_string
+
+
+def decorate_jobset(jobset_manifest_str, sub_networks) -> str:
+  """
+  Decorates a JobSet manifest with the necessary components for rdma-daemon.
+
+  Args:
+    jobset_manifest_str: The JobSet manifest as a YAML string.
+
+  Returns:
+    The modified JobSet manifest as a YAML string.
+  """
+
+  manifest = yaml.safe_load(jobset_manifest_str)
+
+  for job in manifest['spec']['replicatedJobs']:
+    job_manifest = job['template']
+    job_manifest.setdefault('spec', {}).setdefault('template', {}).setdefault(
+        'metadata', {}
+    ).setdefault('annotations', {})
+    spec = (
+        job_manifest.setdefault('spec', {})
+        .setdefault('template', {})
+        .setdefault('spec', {})
+    )
+    spec.setdefault('tolerations', [])
+    spec.setdefault('volumes', [])
+
+    add_annotations(job_manifest, sub_networks)
+    add_volumes(job_manifest)
+    add_tolerations(job_manifest)
+    update_gpu_containers(job_manifest)
+
+  return yaml.dump(manifest, sort_keys=False)
+
+
+def add_annotations(job_manifest, sub_networks):
+  """Adds or updates annotations in the Pod template."""
+  annotations = job_manifest['spec']['template']['metadata']['annotations']
+  interfaces = [
+      '[',
+      '    {"interfaceName":"eth0","network":"default"},',
+      *[
+          f'    {{"interfaceName":"eth{i + 1}","network":"{sub_networks[i]}"}}{"," if i<8 else ""}'
+          for i in range(9)
+      ],
+      ']',
+  ]
+  annotations.update({
+      'networking.gke.io/default-interface': 'eth0',
+      'networking.gke.io/interfaces': literal_string('\n'.join(interfaces)),
+  })
+
+
+def add_volumes(job_manifest):
+  """Adds volumes to the Pod spec."""
+  volumes = job_manifest['spec']['template']['spec']['volumes']
+  volumes.append({
+      'name': 'library-dir-host',
+      'hostPath': {'path': '/home/kubernetes/bin/nvidia'},
+  })
+  volumes.append(
+      {'name': 'gib', 'hostPath': {'path': '/home/kubernetes/bin/gib'}}
+  )
+
+
+def add_tolerations(job_manifest):
+  """Adds tolerations to the Pod spec."""
+  tolerations = job_manifest['spec']['template']['spec']['tolerations']
+  tolerations.append({
+      'key': 'user-workload',
+      'operator': 'Equal',
+      'value': 'true',
+      'effect': 'NoSchedule',
+  })
+
+
+def update_gpu_containers(job_manifest):
+  for container in job_manifest['spec']['template']['spec']['containers']:
+    if 'nvidia.com/gpu' in container.get('resources', {}).get('limits', {}):
+      container.setdefault('env', [])
+      container['env'].append(
+          {'name': 'LD_LIBRARY_PATH', 'value': '/usr/local/nvidia/lib64'}
+      )
+      container.setdefault('volumeMounts', [])
+      container['volumeMounts'].append(
+          {'name': 'library-dir-host', 'mountPath': '/usr/local/nvidia'}
+      )
+      container['volumeMounts'].append(
+          {'name': 'gib', 'mountPath': '/usr/local/gib'}
+      )

xpk/core/workload_decorators/tcpxo_decorator.py

@@ -0,0 +1,157 @@
+"""
+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 yaml
+from ...utils.yaml import literal_string
+
+# Component version
+rxdm = 'v1.0.12'
+
+
+def decorate_jobset(jobset_manifest_str, sub_networks) -> str:
+  """
+  Decorates a JobSet manifest with the necessary components for tcpxo-daemon.
+
+  Args:
+    jobset_manifest_str: The JobSet manifest as a YAML string.
+
+  Returns:
+    The modified JobSet manifest as a YAML string.
+  """
+
+  manifest = yaml.safe_load(jobset_manifest_str)
+
+  for job in manifest['spec']['replicatedJobs']:
+    job_manifest = job['template']
+    job_manifest.setdefault('spec', {}).setdefault('template', {}).setdefault(
+        'metadata', {}
+    ).setdefault('annotations', {})
+    spec = (
+        job_manifest.setdefault('spec', {})
+        .setdefault('template', {})
+        .setdefault('spec', {})
+    )
+    spec.setdefault('tolerations', [])
+    spec.setdefault('volumes', [])
+
+    add_annotations(job_manifest, sub_networks)
+    add_volumes(job_manifest)
+    add_tolerations(job_manifest)
+    add_tcpxo_daemon_container(job_manifest)
+    update_gpu_containers(job_manifest)
+
+  return yaml.dump(manifest, sort_keys=False)
+
+
+def add_annotations(job_manifest, sub_networks):
+  """Adds or updates annotations in the Pod template."""
+  annotations = job_manifest['spec']['template']['metadata']['annotations']
+  interfaces = [
+      '[',
+      '    {"interfaceName":"eth0","network":"default"},',
+      *[
+          f'    {{"interfaceName":"eth{i + 1}","network":"{sub_networks[i]}"}}{"," if i<7 else ""}'
+          for i in range(8)
+      ],
+      ']',
+  ]
+  annotations.update({
+      'devices.gke.io/container.tcpxo-daemon': literal_string(
+          '- path: /dev/nvidia0\n'
+          '- path: /dev/nvidia1\n'
+          '- path: /dev/nvidia2\n'
+          '- path: /dev/nvidia3\n'
+          '- path: /dev/nvidia4\n'
+          '- path: /dev/nvidia5\n'
+          '- path: /dev/nvidia6\n'
+          '- path: /dev/nvidia7\n'
+          '- path: /dev/nvidiactl\n'
+          '- path: /dev/nvidia-uvm\n'
+          '- path: /dev/dmabuf_import_helper\n'
+      ),
+      'networking.gke.io/default-interface': 'eth0',
+      'networking.gke.io/interfaces': literal_string('\n'.join(interfaces)),
+  })
+
+
+def add_tolerations(job_manifest):
+  """Adds tolerations to the Pod spec."""
+  tolerations = job_manifest['spec']['template']['spec']['tolerations']
+  tolerations.append({
+      'key': 'user-workload',
+      'operator': 'Equal',
+      'value': 'true',
+      'effect': 'NoSchedule',
+  })
+
+
+def add_volumes(job_manifest):
+  """Adds volumes to the Pod spec."""
+  volumes = job_manifest['spec']['template']['spec']['volumes']
+  volumes.append({
+      'name': 'libraries',
+      'hostPath': {'path': '/home/kubernetes/bin/nvidia/lib64'},
+  })
+  volumes.append({'name': 'sys', 'hostPath': {'path': '/sys'}})
+  volumes.append({'name': 'proc-sys', 'hostPath': {'path': '/proc/sys'}})
+  volumes.append({
+      'name': 'aperture-devices',
+      'hostPath': {'path': '/dev/aperture_devices'},
+  })
+
+
+def add_tcpxo_daemon_container(job_manifest):
+  """Adds the tcpxo-daemon container to the Pod spec."""
+  tcpxo_daemon_container = {
+      'name': 'tcpxo-daemon',
+      'image': f'us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/tcpgpudmarxd-dev:{rxdm}',
+      'imagePullPolicy': 'Always',
+      'command': ['/bin/sh', '-c'],
+      'args': [
+          'set -ex\nchmod 755'
+          ' /fts/entrypoint_rxdm_container.sh\n/fts/entrypoint_rxdm_container.sh'
+          ' --num_hops=2 --num_nics=8 --uid= --alsologtostderr'
+      ],
+      'securityContext': {
+          'capabilities': {'add': ['NET_ADMIN', 'NET_BIND_SERVICE']}
+      },
+      'volumeMounts': [
+          {'name': 'libraries', 'mountPath': '/usr/local/nvidia'},
+          {'name': 'sys', 'mountPath': '/hostsysfs'},
+          {'name': 'proc-sys', 'mountPath': '/hostprocsysfs'},
+      ],
+      'env': [{'name': 'LD_LIBRARY_PATH', 'value': '/usr/local/nvidia/lib64'}],
+  }
+  job_manifest['spec']['template']['spec']['containers'].insert(
+      0, tcpxo_daemon_container
+  )
+
+
+def update_gpu_containers(job_manifest):
+  for container in job_manifest['spec']['template']['spec']['containers']:
+    if 'nvidia.com/gpu' in container.get('resources', {}).get('limits', {}):
+      container.setdefault('env', [])
+      container['env'].append(
+          {'name': 'LD_LIBRARY_PATH', 'value': '/usr/local/nvidia/lib64'}
+      )
+      container['env'].append({
+          'name': 'NCCL_FASTRAK_LLCM_DEVICE_DIRECTORY',
+          'value': '/dev/aperture_devices',
+      })
+      container.setdefault('volumeMounts', [])
+      container['volumeMounts'].append(
+          {'name': 'aperture-devices', 'mountPath': '/dev/aperture_devices'}
+      )

xpk/main.py

@@ -0,0 +1,73 @@
+"""
+Copyright 2023 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.
+"""
+
+r"""xpk (Accelerated Processing Kit).
+
+Next Steps:
+- Cluster describe is broken by Cacheimage since that counts as a workload.
+- Cluster describe: count by jobset.
+- If any instance goes down, bring down the whole job.
+- How to more gracefully handle job failures, distinguishing between software
+  and infra?
+- Look into --docker-name and --docker-image.
+  Shouldn't one string be adequate to express what we want?
+- Apply learnings from about private, region, coredns, etc:
+- Enable special preheater
+- Make Argparse logic this a function?
+  - Obvious logic that starts in main instead of here in code but args will
+    not be a universal argument.
+"""
+
+import argparse
+import sys
+
+from .parser.core import set_parser
+from .utils.console import xpk_print
+
+################### Compatibility Check ###################
+# Check that the user runs the below version or greater.
+
+major_version_supported = 3
+minor_version_supported = 10
+
+user_major_version = sys.version_info[0]
+user_minor_version = sys.version_info[1]
+if (
+    user_major_version < major_version_supported
+    or user_minor_version < minor_version_supported
+):
+  raise RuntimeError(
+      'xpk must be run with Python'
+      f' {major_version_supported}.{minor_version_supported} or greater.'
+      f' User currently is running {user_major_version}.{user_minor_version}'
+  )
+
+# Create top level parser for xpk command.
+parser = argparse.ArgumentParser(description='xpk command', prog='xpk')
+set_parser(parser=parser)
+
+xpk_print('Starting xpk', flush=True)
+main_args = parser.parse_args()
+main_args.enable_ray_cluster = False
+main_args.func(main_args)
+
+
+def main() -> None:
+  xpk_print('XPK Done.', flush=True)
+
+
+if __name__ == '__main__':
+  main()

xpk/parser/__init__.py

@@ -0,0 +1,15 @@
+"""
+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.
+"""

xpk/parser/batch.py

@@ -0,0 +1,184 @@
+"""
+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 .common import add_shared_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'
+  )
+  batch_optional_arguments.add_argument(
+      '--cluster',
+      type=str,
+      default=None,
+      help='Cluster to which command applies.',
+  )
+  batch_optional_arguments.add_argument(
+      '--kind-cluster',
+      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)
+
+  batch_optional_arguments.add_argument(
+      '--ignore-unknown-flags',
+      type=bool,
+      action=argparse.BooleanOptionalAction,
+      default=False,
+      help='Ignore all the unsupported flags in the bash script.',
+  )
+  batch_optional_arguments.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.'
+      ),
+  )
+  batch_optional_arguments.add_argument(
+      '-c',
+      '--cpus-per-task',
+      type=str,
+      default=None,
+      help='How much cpus a container inside a pod requires.',
+  )
+  batch_optional_arguments.add_argument(
+      '--gpus-per-task',
+      type=str,
+      default=None,
+      help='How much gpus a container inside a pod requires.',
+  )
+  batch_optional_arguments.add_argument(
+      '--mem',
+      type=str,
+      default=None,
+      help='How much memory a pod requires.',
+  )
+  batch_optional_arguments.add_argument(
+      '--mem-per-task',
+      type=str,
+      default=None,
+      help='How much memory a container requires.',
+  )
+  batch_optional_arguments.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.'
+      ),
+  )
+  batch_optional_arguments.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.'
+      ),
+  )
+  batch_optional_arguments.add_argument(
+      '-N',
+      '--nodes',
+      type=int,
+      default=None,
+      help='Number of pods to be used at a time.',
+  )
+  batch_optional_arguments.add_argument(
+      '-n',
+      '--ntasks',
+      type=int,
+      default=None,
+      help='Number of identical containers inside of a pod, usually 1.',
+  )
+  batch_optional_arguments.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.'
+      ),
+  )
+  batch_optional_arguments.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.'
+      ),
+  )
+  batch_optional_arguments.add_argument(
+      '--input',
+      type=str,
+      default=None,
+      help='What to pipe into the script.',
+  )
+  batch_optional_arguments.add_argument(
+      '-J',
+      '--job-name',
+      type=str,
+      default=None,
+      help='What is the job name.',
+  )
+  batch_optional_arguments.add_argument(
+      '-D',
+      '--chdir',
+      type=str,
+      default=None,
+      help='Change directory before executing the script.',
+  )
+  batch_optional_arguments.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".'
+      ),
+  )