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

xpk/core/vertex.py

@@ -0,0 +1,105 @@
+"""
+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 ..utils.console import xpk_print
+from .resources import CLUSTER_METADATA_CONFIGMAP, get_cluster_configmap
+
+DEFAULT_VERTEX_TENSORBOARD_NAME = 'tb-instance'
+
+
+def create_vertex_tensorboard(args) -> dict:
+  """Creates a Tensorboard instance in Vertex AI.
+
+  Args:
+    args: user provided arguments.
+
+  Returns:
+    dict containing Tensorboard instance name, id and location.
+  """
+  from cloud_accelerator_diagnostics import (  # pylint: disable=import-outside-toplevel
+      tensorboard,
+  )
+
+  tensorboard_config = {}
+  tensorboard_name = args.tensorboard_name
+  if tensorboard_name is None:
+    tensorboard_name = f'{args.cluster}-{DEFAULT_VERTEX_TENSORBOARD_NAME}'
+  instance_id = tensorboard.create_instance(  # pylint: disable=used-before-assignment
+      project=args.project,
+      location=args.tensorboard_region,
+      tensorboard_name=tensorboard_name,
+  )
+  if instance_id:
+    xpk_print(
+        f'Tensorboard instance {tensorboard_name} is successfully created.'
+    )
+    tensorboard_config['tensorboard_region'] = args.tensorboard_region
+    tensorboard_config['tensorboard_name'] = tensorboard_name
+    tensorboard_config['tensorboard_id'] = instance_id
+  return tensorboard_config
+
+
+def create_vertex_experiment(args) -> dict | None:
+  """Creates an Experiment in Vertex AI.
+
+  Args:
+    args: user provided arguments.
+
+  Returns:
+    map containing Vertex Tensorboard configurations.
+  """
+  from cloud_accelerator_diagnostics import (  # pylint: disable=import-outside-toplevel
+      tensorboard,
+  )
+
+  metadata_configmap_name = f'{args.cluster}-{CLUSTER_METADATA_CONFIGMAP}'
+  cluster_config_map = get_cluster_configmap(args, metadata_configmap_name)
+
+  if cluster_config_map is None or 'tensorboard_name' not in cluster_config_map:
+    xpk_print(
+        'No Vertex Tensorboard instance has been created in cluster create. Run'
+        ' `xpk cluster create --create-vertex-tensorboard` before running `xpk'
+        ' workload create --use-vertex-tensorboard` to create a Vertex'
+        ' Tensorboard instance. Alternatively, use `xpk cluster create-pathways'
+        ' --create-vertex-tensorboard` before running `xpk workload'
+        ' create-pathways --use-vertex-tensorboard`.'
+    )
+    return None
+
+  tensorboard_config = {}
+  tensorboard_config['tensorboard_project'] = args.project
+  tensorboard_config['tensorboard_region'] = cluster_config_map[
+      'tensorboard_region'
+  ]
+  tensorboard_config['tensorboard_name'] = cluster_config_map[
+      'tensorboard_name'
+  ]
+  experiment_name = args.experiment_name
+  if experiment_name is None:
+    experiment_name = f'{args.cluster}-{args.workload}'
+  tensorboard_config['experiment_name'] = experiment_name
+
+  _, tensorboard_url = tensorboard.create_experiment(
+      project=args.project,
+      location=tensorboard_config['tensorboard_region'],
+      experiment_name=experiment_name,
+      tensorboard_name=tensorboard_config['tensorboard_name'],
+  )
+  if tensorboard_url is None:
+    return None
+
+  xpk_print(f'You can view Vertex Tensorboard at: {tensorboard_url}')
+  return tensorboard_config

xpk/core/workload.py

@@ -0,0 +1,341 @@
+"""
+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 ..utils.console import xpk_exit, xpk_print
+from .capacity import H100_DEVICE_TYPE, H100_MEGA_DEVICE_TYPE
+from .commands import run_command_for_value
+from .gcloud_context import zone_to_region
+from .system_characteristics import SystemCharacteristics
+
+
+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
+
+
+def check_if_workload_exists(args) -> bool:
+  """Check if workload exists.
+
+  Args:
+     args: user provided arguments for running the command.
+
+  Returns:
+    returns true if workload exist, otherwise returns false.
+  """
+  columns = {
+      'Jobset': '.metadata.ownerReferences[0].name',
+  }
+
+  s = ','.join([key + ':' + value for key, value in columns.items()])
+
+  command = f"kubectl get workloads -o=custom-columns='{s}'"
+  return_code, return_msg = run_command_for_value(
+      command, 'Check if Workload Already Exists', args
+  )
+
+  if return_code != 0:
+    xpk_print(f'List Job request returned ERROR {return_code}')
+    xpk_exit(return_code)
+
+  lines = return_msg.split('\n')
+  new_workload_name = args.workload
+  for line in lines:
+    if line == new_workload_name:
+      return True
+  return False
+
+
+def wait_for_job_completion(args) -> int:
+  """Function to wait for job completion.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    return_code: 0 if successful, 124 if timeout, 125 if unsuccessful job, 1 otherwise
+  """
+  # Check that the workload exists
+  args.workload = args.wait_for_job_completion
+  workload_exists = check_if_workload_exists(args)
+  if not workload_exists:
+    xpk_print(f'Workload named {args.workload} does not exist.')
+    return 1
+
+  # Get the full workload name
+  get_workload_name_cmd = f'kubectl get workloads | grep jobset-{args.workload}'
+  return_code, return_value = run_command_for_value(
+      get_workload_name_cmd, 'Get full workload name', args
+  )
+  if return_code != 0:
+    xpk_print(f'Get full workload name request returned ERROR {return_code}')
+    return return_code
+  full_workload_name = return_value.split(' ')[0]
+
+  # Call kubectl wait on the workload using the full workload name
+  timeout_val = args.timeout if args.timeout is not None else -1
+  timeout_msg = (
+      f'{timeout_val}s' if timeout_val != -1 else 'max timeout (1 week)'
+  )
+  wait_cmd = (
+      "kubectl  wait --for jsonpath='.status.conditions[-1].type'=Finished"
+      f' workload {full_workload_name} --timeout={timeout_val}s'
+  )
+  return_code, return_value = run_command_for_value(
+      wait_cmd,
+      f'Wait for workload to finish with timeout of {timeout_msg}',
+      args,
+      print_timer=True,
+  )
+  if return_code != 0:
+    if 'timed out' in return_value:
+      xpk_print(
+          f'Timed out waiting for your workload after {timeout_msg}, see your'
+          ' workload here:'
+          # pylint: disable=line-too-long
+          f' https://console.cloud.google.com/kubernetes/service/{zone_to_region(args.zone)}/{args.cluster}/default/{args.workload}/details?project={args.project}'
+      )
+      return 124
+    else:
+      xpk_print(f'{return_value}')
+      xpk_print(f'Wait for workload returned ERROR {return_code}')
+      return return_code
+  xpk_print(
+      'Finished waiting for your workload, see your workload here:'
+      # pylint: disable=line-too-long
+      f' https://console.cloud.google.com/kubernetes/service/{zone_to_region(args.zone)}/{args.cluster}/default/{args.workload}/details?project={args.project}'
+  )
+  status_cmd = (
+      f'kubectl get jobset {args.workload} -o'
+      " jsonpath='{.status.conditions[-1].type}'"
+  )
+  return_code, return_value = run_command_for_value(
+      status_cmd, 'Get jobset status', args
+  )
+  if return_code != 0:
+    xpk_print(f'Get workload status request returned ERROR {return_code}')
+    return return_code
+  xpk_print(f'Your workload finished with status: {return_value}')
+  if return_value != 'Completed':
+    xpk_print('Your workload did not complete successfully')
+    return 125
+  return 0
+
+
+def get_gpu_volume(system: SystemCharacteristics) -> str:
+  """Get gpu volume based on user provided arguments.
+
+  Args:
+    system: system characteristics.
+
+  Returns:
+    str: yaml containing gpu volume
+  """
+  gpu_volume = ''
+  if system.device_type == H100_DEVICE_TYPE:
+    gpu_volume = """- name: nvidia-install-dir-host
+                hostPath:
+                  path: /home/kubernetes/bin/nvidia/lib64
+              - name: tcpd-socket
+                hostPath:
+                  path: /run/tcpx
+              - name: shared-memory
+                emptyDir:
+                  medium: "Memory"
+                  sizeLimit: 200Gi
+              - name: workload-terminated-volume
+                emptyDir:
+              - name: tcpx-nccl-plugin-volume
+                emptyDir:"""
+  elif system.device_type == H100_MEGA_DEVICE_TYPE:
+    gpu_volume = """- name: nvidia-install-dir-host
+                hostPath:
+                  path: /home/kubernetes/bin/nvidia/lib64
+              - name: shared-memory
+                emptyDir:
+                  medium: "Memory"
+                  sizeLimit: 1Gi
+              - name: workload-terminated-volume
+                emptyDir:"""
+  return gpu_volume
+
+
+def get_gpu_rxdm_image(system: SystemCharacteristics) -> str:
+  """Get config of rxdm based on user provided arguments.
+
+  Args:
+    system: system characteristics.
+
+  Returns:
+    str: yaml containing the rxdm name and image
+  """
+  gpu_rxdm_image = ''
+  if system.device_type == H100_DEVICE_TYPE:
+    gpu_rxdm_image = """- name: tcpd-daemon
+                image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/tcpgpudmarxd-dev:v2.0.9"""
+  elif system.device_type == H100_MEGA_DEVICE_TYPE:
+    gpu_rxdm_image = """- name: fastrak-daemon
+                image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/tcpgpudmarxd-dev:v1.0.9"""
+  return gpu_rxdm_image
+
+
+def get_gpu_rxdm_cmd(system: SystemCharacteristics) -> str:
+  """Get rxdm command based on user provided arguments.
+
+  Args:
+    system: system characteristics.
+
+  Returns:
+    str: command of running rxdm container
+  """
+  gpu_rxdm_cmd = ''
+  if system.device_type == H100_DEVICE_TYPE:
+    gpu_rxdm_cmd = (
+        '/tcpgpudmarxd/build/app/tcpgpudmarxd --gpu_nic_preset a3vm'
+        ' --gpu_shmem_type fd --setup_param "--verbose 128 2 0"'
+    )
+  elif system.device_type == H100_MEGA_DEVICE_TYPE:
+    gpu_rxdm_cmd = (
+        'set -ex; chmod 755 /fts/entrypoint_rxdm_container.sh;'
+        ' /fts/entrypoint_rxdm_container.sh --num_hops=2 --num_nics=8 --uid='
+        ' --alsologtostderr'
+    )
+  return gpu_rxdm_cmd
+
+
+def get_gpu_tcp_volume(system: SystemCharacteristics) -> str:
+  """Get gpu tcp volume based on user provided arguments.
+
+  Args:
+    system: system characteristics.
+
+  Returns:
+    str: yaml containing gpu tcp volume
+  """
+  gpu_tcp_volume = ''
+  if system.device_type == H100_DEVICE_TYPE:
+    gpu_tcp_volume = """- name: tcpd-socket
+                  mountPath: /tmp"""
+  return gpu_tcp_volume

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,129 @@
+"""
+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_kjob_template(job_manifest) -> str:
+  spec = (
+      job_manifest.setdefault('spec', {})
+      .setdefault('template', {})
+      .setdefault('spec', {})
+  )
+  spec.setdefault('tolerations', [])
+  spec.setdefault('volumes', [])
+
+  add_volumes(job_manifest)
+  add_tolerations(job_manifest)
+  update_gpu_containers(job_manifest)
+  return job_manifest
+
+
+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 get_interfaces_entry(sub_networks: list[str]) -> tuple[str, str]:
+  interfaces = [
+      '[',
+      '    {"interfaceName":"eth0","network":"default"},',
+      *[
+          f'    {{"interfaceName":"eth{i + 1}","network":"{sub_networks[i]}"}}{"," if i<8 else ""}'
+          for i in range(9)
+      ],
+      ']',
+  ]
+  return 'networking.gke.io/interfaces', literal_string('\n'.join(interfaces))
+
+
+def add_annotations(job_manifest, sub_networks):
+  """Adds or updates annotations in the Pod template."""
+  annotations = job_manifest['spec']['template']['metadata']['annotations']
+  interfaces_key, interfaces_value = get_interfaces_entry(sub_networks)
+  annotations.update({
+      'networking.gke.io/default-interface': "'eth0'",
+      interfaces_key: interfaces_value,
+  })
+
+
+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/storage_decorator.py

@@ -0,0 +1,52 @@
+"""
+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 ...core.storage import GCS_FUSE_TYPE, get_storage_volumes_yaml_dict, GCS_FUSE_ANNOTATION
+
+
+def decorate_jobset(jobset_manifest_str, storages) -> str:
+  """
+  Decorates a JobSet manifest with the necessary storages.
+
+  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)
+  storage_volumes = get_storage_volumes_yaml_dict(storages)
+  for job in manifest['spec']['replicatedJobs']:
+    job_manifest = job['template']
+    add_annotations(job_manifest, storages)
+    add_volumes(job_manifest, storage_volumes)
+  return yaml.dump(manifest, sort_keys=False)
+
+
+def add_annotations(job_manifest, storages):
+  """Adds or updates storage annotations in the Pod template."""
+  annotations = job_manifest['spec']['template']['metadata']['annotations']
+  gcs_present = [storage.type == GCS_FUSE_TYPE for storage in storages]
+  if gcs_present:
+    annotations.update(GCS_FUSE_ANNOTATION)
+
+
+def add_volumes(job_manifest, storage_volumes):
+  volumes = job_manifest['spec']['template']['spec']['volumes']
+  volumes.extend(storage_volumes)