xpk 0.6.0__py3-none-any.whl → 0.7.1__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

@@ -14,7 +14,11 @@ 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:
@@ -129,5 +133,209 @@ def get_workload_list(args) -> tuple[int, str]:
     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/rdma_decorator.py

@@ -18,6 +18,21 @@ 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.
@@ -52,9 +67,7 @@ def decorate_jobset(jobset_manifest_str, sub_networks) -> str:
   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']
+def get_interfaces_entry(sub_networks: list[str]) -> tuple[str, str]:
   interfaces = [
       '[',
       '    {"interfaceName":"eth0","network":"default"},',
@@ -64,9 +77,16 @@ def add_annotations(job_manifest, sub_networks):
       ],
       ']',
   ]
+  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',
-      'networking.gke.io/interfaces': literal_string('\n'.join(interfaces)),
+      'networking.gke.io/default-interface': "'eth0'",
+      interfaces_key: interfaces_value,
   })
 
 

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)

xpk/core/workload_decorators/tcpxo_decorator.py

@@ -21,6 +21,42 @@ from ...utils.yaml import literal_string
 rxdm = 'v1.0.12'
 
 
+def decorate_kjob_template(job_manifest: dict) -> dict:
+  spec = (
+      job_manifest.setdefault('spec', {})
+      .setdefault('template', {})
+      .setdefault('spec', {})
+  )
+  spec.setdefault('tolerations', [])
+  spec.setdefault('volumes', [])
+
+  add_volumes(job_manifest)
+  add_tolerations(job_manifest)
+  add_tcpxo_daemon_container(job_manifest)
+  update_gpu_containers(job_manifest)
+  return job_manifest
+
+
+def decorate_job(job_manifest: dict, sub_networks: list[str]) -> dict:
+  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 job_manifest
+
+
 def decorate_jobset(jobset_manifest_str, sub_networks) -> str:
   """
   Decorates a JobSet manifest with the necessary components for tcpxo-daemon.
@@ -36,29 +72,11 @@ def decorate_jobset(jobset_manifest_str, sub_networks) -> 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)
-
+    job_manifest = decorate_job(job_manifest, sub_networks)
   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']
+def get_interfaces_entry(sub_networks: list[str]) -> tuple[str, str]:
   interfaces = [
       '[',
       '    {"interfaceName":"eth0","network":"default"},',
@@ -68,22 +86,34 @@ def add_annotations(job_manifest, sub_networks):
       ],
       ']',
   ]
+  return 'networking.gke.io/interfaces', literal_string('\n'.join(interfaces))
+
+
+def get_tcpxo_deamon_entry() -> tuple[str, str]:
+  return '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'
+  )
+
+
+def add_annotations(job_manifest, sub_networks):
+  """Adds or updates annotations in the Pod template."""
+  annotations = job_manifest['spec']['template']['metadata']['annotations']
+  tcpxo_deamon_key, tcpxo_deamon_paths = get_tcpxo_deamon_entry()
+  interfaces_key, interfaces_value = get_interfaces_entry(sub_networks)
   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'
-      ),
+      tcpxo_deamon_key: tcpxo_deamon_paths,
       'networking.gke.io/default-interface': 'eth0',
-      'networking.gke.io/interfaces': literal_string('\n'.join(interfaces)),
+      interfaces_key: interfaces_value,
   })
 
 
@@ -103,7 +133,7 @@ def add_volumes(job_manifest):
   volumes = job_manifest['spec']['template']['spec']['volumes']
   volumes.append({
       'name': 'libraries',
-      'hostPath': {'path': '/home/kubernetes/bin/nvidia/lib64'},
+      'hostPath': {'path': '/home/kubernetes/bin/nvidia'},
   })
   volumes.append({'name': 'sys', 'hostPath': {'path': '/sys'}})
   volumes.append({'name': 'proc-sys', 'hostPath': {'path': '/proc/sys'}})
@@ -135,8 +165,8 @@ def add_tcpxo_daemon_container(job_manifest):
       ],
       'env': [{'name': 'LD_LIBRARY_PATH', 'value': '/usr/local/nvidia/lib64'}],
   }
-  job_manifest['spec']['template']['spec']['containers'].insert(
-      0, tcpxo_daemon_container
+  job_manifest['spec']['template']['spec']['containers'].append(
+      tcpxo_daemon_container
   )
 
 
@@ -155,3 +185,6 @@ def update_gpu_containers(job_manifest):
       container['volumeMounts'].append(
           {'name': 'aperture-devices', 'mountPath': '/dev/aperture_devices'}
       )
+      container['volumeMounts'].append(
+          {'name': 'libraries', 'mountPath': '/usr/local/nvidia'}
+      )

xpk/main.py

@@ -36,10 +36,11 @@ import sys
 
 from .parser.core import set_parser
 from .utils.console import xpk_print
-
+from .utils.validation import validate_dependencies
 ################### Compatibility Check ###################
 # Check that the user runs the below version or greater.
 
+
 major_version_supported = 3
 minor_version_supported = 10
 
@@ -60,6 +61,7 @@ parser = argparse.ArgumentParser(description='xpk command', prog='xpk')
 set_parser(parser=parser)
 
 xpk_print('Starting xpk', flush=True)
+validate_dependencies()
 main_args = parser.parse_args()
 main_args.enable_ray_cluster = False
 main_args.func(main_args)