xpk 0.7.2__py3-none-any.whl → 0.9.0__py3-none-any.whl

xpk/core/workload.py

@@ -14,12 +14,19 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import yaml
+
+from ..utils import templates
 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 .storage import Storage, get_storage_volume_mounts_for_gpu
 from .system_characteristics import SystemCharacteristics
 
+RXDM_CONTAINER_A3HIGH_PATH = '/../templates/rxdm_container_a3high.yaml'
+RXDM_CONTAINER_A3MEGA_PATH = '/../templates/rxdm_container_a3mega.yaml'
+
 
 def workload_list_awk_command(filter_key) -> str:
   """Function returns the awk command needed from the filter specified.
@@ -124,7 +131,7 @@ def get_workload_list(args) -> tuple[int, str]:
   )
   workload_list_filter_job_cmd = determine_workload_list_filter_by_job(args)
   command = (
-      f'kubectl get workloads -o=custom-columns="{s}" '
+      f'kubectl get workloads --ignore-not-found -o=custom-columns="{s}" '
       f'{workload_list_filter_status_cmd} {workload_list_filter_job_cmd}'
   )
 
@@ -244,98 +251,36 @@ def wait_for_job_completion(args) -> int:
   return 0
 
 
-def get_gpu_volume(system: SystemCharacteristics) -> str:
-  """Get gpu volume based on user provided arguments.
+def add_gpu_rxdm_container(
+    jobset_manifest_str: str,
+    system: SystemCharacteristics,
+    all_storages: list[Storage],
+) -> str:
+  """Add gpu rxdm container to jobset manifest based on user provided arguments.
 
   Args:
+    jobset_manifest_str: the JobSet manifest as a YAML string.
     system: system characteristics.
+    all_storages: list of all storages.
 
   Returns:
-    str: yaml containing gpu volume
+    str: the modified JobSet manifest as a YAML string.
   """
-  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:"""
+    gpu_rxdm_container = templates.load(RXDM_CONTAINER_A3HIGH_PATH)
   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
+    gpu_rxdm_container = templates.load(RXDM_CONTAINER_A3MEGA_PATH)
+  else:
+    return jobset_manifest_str
 
+  storage_volume_mounts = get_storage_volume_mounts_for_gpu(all_storages)
+  gpu_rxdm_container['volumeMounts'].extend(storage_volume_mounts)
 
-def get_gpu_rxdm_cmd(system: SystemCharacteristics) -> str:
-  """Get rxdm command based on user provided arguments.
+  manifest = yaml.safe_load(jobset_manifest_str)
 
-  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"'
+  for job in manifest['spec']['replicatedJobs']:
+    job['template']['spec']['template']['spec']['containers'].append(
+        gpu_rxdm_container
     )
-  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
+  return yaml.dump(manifest, sort_keys=False)

xpk/core/workload_decorators/rdma_decorator.py

@@ -33,7 +33,7 @@ def decorate_kjob_template(job_manifest) -> str:
   return job_manifest
 
 
-def decorate_jobset(jobset_manifest_str, sub_networks) -> str:
+def decorate_jobset(jobset_manifest_str: str, sub_networks: list[str]) -> str:
   """
   Decorates a JobSet manifest with the necessary components for rdma-daemon.
 
@@ -68,24 +68,20 @@ def decorate_jobset(jobset_manifest_str, sub_networks) -> str:
 
 
 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):
+  entries = ',\n'.join([
+      f'    {{"interfaceName":"eth{i}","network":"{network}"}}'
+      for i, network in enumerate(sub_networks)
+  ])
+  interfaces = f'[\n{entries}\n]'
+  return 'networking.gke.io/interfaces', literal_string(interfaces)
+
+
+def add_annotations(job_manifest: dict, sub_networks: list[str]):
   """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/default-interface': 'eth0',
       interfaces_key: interfaces_value,
   })
 

xpk/core/workload_decorators/storage_decorator.py

@@ -16,7 +16,7 @@ limitations under the License.
 
 import yaml
 
-from ...core.storage import GCS_FUSE_TYPE, get_storage_volumes_yaml_dict, GCS_FUSE_ANNOTATION
+from ...core.storage import GCS_FUSE_TYPE, PARALLELSTORE_TYPE, get_storage_volumes_yaml_dict, GCS_FUSE_ANNOTATIONS, PARALLELSTORE_ANNOTATIONS
 
 
 def decorate_jobset(jobset_manifest_str, storages) -> str:
@@ -42,9 +42,14 @@ def decorate_jobset(jobset_manifest_str, storages) -> str:
 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]
+  gcs_present = any(storage.type == GCS_FUSE_TYPE for storage in storages)
   if gcs_present:
-    annotations.update(GCS_FUSE_ANNOTATION)
+    annotations.update(GCS_FUSE_ANNOTATIONS)
+  parallelstore_present = any(
+      storage.type == PARALLELSTORE_TYPE for storage in storages
+  )
+  if parallelstore_present:
+    annotations.update(PARALLELSTORE_ANNOTATIONS)
 
 
 def add_volumes(job_manifest, storage_volumes):

xpk/core/workload_decorators/tcpx_decorator.py

@@ -0,0 +1,179 @@
+"""
+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
+tcpx = 'v2.0.11'
+
+
+def decorate_kjob_template(job_manifest: dict) -> dict:
+  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) -> dict:
+  add_annotations(job_manifest)
+  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: str) -> 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 = decorate_job(job_manifest)
+  return yaml.dump(manifest, sort_keys=False)
+
+
+def get_interfaces_annotation() -> dict:
+  interfaces = [
+      '[',
+      '    {"interfaceName":"eth0","network":"default"},',
+      '    {"interfaceName":"eth1","network":"vpc1"},',
+      '    {"interfaceName":"eth2","network":"vpc2"},',
+      '    {"interfaceName":"eth3","network":"vpc3"},',
+      '    {"interfaceName":"eth4","network":"vpc4"}',
+      ']',
+  ]
+  return {'networking.gke.io/interfaces': literal_string('\n'.join(interfaces))}
+
+
+def get_tcpx_deamon_annotation() -> dict:
+  return {
+      'devices.gke.io/container.tcpx-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'
+      )
+  }
+
+
+def add_annotations(job_manifest: dict):
+  """Adds or updates annotations in the Pod template."""
+  annotations: dict = (
+      job_manifest.setdefault('spec', {})
+      .setdefault('template', {})
+      .setdefault('metadata', {})
+      .setdefault('annotations', {})
+  )
+  annotations.update(get_tcpx_deamon_annotation())
+  annotations.update({'networking.gke.io/default-interface': 'eth0'})
+  annotations.update(get_interfaces_annotation())
+
+
+def add_tolerations(job_manifest: dict):
+  """Adds tolerations to the Pod spec."""
+  tolerations: list = (
+      job_manifest.setdefault('spec', {})
+      .setdefault('template', {})
+      .setdefault('spec', {})
+      .setdefault('tolerations', [])
+  )
+  tolerations.append({
+      'key': 'user-workload',
+      'operator': 'Equal',
+      'value': 'true',
+      'effect': 'NoSchedule',
+  })
+
+
+def add_volumes(job_manifest: dict):
+  """Adds volumes to the Pod spec."""
+  volumes: list = (
+      job_manifest.setdefault('spec', {})
+      .setdefault('template', {})
+      .setdefault('spec', {})
+      .setdefault('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'}})
+
+
+def add_tcpxo_daemon_container(job_manifest):
+  """Adds the tcpxo-daemon container to the Pod spec."""
+  tcpxo_daemon_container = {
+      'name': 'tcpx-daemon',
+      'image': f'us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/tcpgpudmarxd-dev:{tcpx}',
+      'imagePullPolicy': 'Always',
+      'restartPolicy': 'Always',
+      'command': [
+          '/tcpgpudmarxd/build/app/tcpgpudmarxd',
+          '--gpu_nic_preset',
+          'a3vm',
+          '--gpu_shmem_type',
+          'fd',
+          '--uds_path',
+          '/run/tcpx',
+          '--setup_param',
+          '"--verbose 128 2 0 "',
+      ],
+      'securityContext': {'capabilities': {'add': ['NET_ADMIN']}},
+      'volumeMounts': [
+          {'name': 'libraries', 'mountPath': '/usr/local/nvidia/lib64'},
+          {'name': 'tcpx-socket', 'mountPath': '/run/tcpx'},
+          {'name': 'sys', 'mountPath': '/hostsysfs'},
+          {'name': 'proc-sys', 'mountPath': '/hostprocsysfs'},
+      ],
+      'env': [{'name': 'LD_LIBRARY_PATH', 'value': '/usr/local/nvidia/lib64'}],
+  }
+  spec = job_manifest['spec']['template']['spec']
+  spec.setdefault('initContainers', [])
+  spec['initContainers'].append(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', {}):
+      env: list = container.setdefault('env', [])
+      env.append(
+          {'name': 'LD_LIBRARY_PATH', 'value': '/usr/local/nvidia/lib64'}
+      )
+      volumeMounts: list = container.setdefault('volumeMounts', [])
+      volumeMounts.append({'name': 'tcpx-socket', 'mountPath': '/tmp'})
+      volumeMounts.append(
+          {'name': 'libraries', 'mountPath': '/usr/local/nvidia/lib64'}
+      )

xpk/core/workload_decorators/tcpxo_decorator.py

@@ -57,7 +57,7 @@ def decorate_job(job_manifest: dict, sub_networks: list[str]) -> dict:
   return job_manifest
 
 
-def decorate_jobset(jobset_manifest_str, sub_networks) -> str:
+def decorate_jobset(jobset_manifest_str: str, sub_networks: list[str]) -> str:
   """
   Decorates a JobSet manifest with the necessary components for tcpxo-daemon.
 
@@ -77,16 +77,12 @@ def decorate_jobset(jobset_manifest_str, sub_networks) -> str:
 
 
 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<7 else ""}'
-          for i in range(8)
-      ],
-      ']',
-  ]
-  return 'networking.gke.io/interfaces', literal_string('\n'.join(interfaces))
+  entries = ',\n'.join([
+      f'    {{"interfaceName":"eth{i}","network":"{network}"}}'
+      for i, network in enumerate(sub_networks)
+  ])
+  interfaces = f'[\n{entries}\n]'
+  return 'networking.gke.io/interfaces', literal_string(interfaces)
 
 
 def get_tcpxo_deamon_entry() -> tuple[str, str]:
@@ -105,9 +101,13 @@ def get_tcpxo_deamon_entry() -> tuple[str, str]:
   )
 
 
-def add_annotations(job_manifest, sub_networks):
+def add_annotations(job_manifest: dict, sub_networks: list[str]):
   """Adds or updates annotations in the Pod template."""
-  annotations = job_manifest['spec']['template']['metadata']['annotations']
+  metadata = job_manifest['spec']['template']['metadata']
+  annotations = metadata.get('annotations')
+  if annotations is None:
+    annotations = {}
+    metadata['annotations'] = annotations
   tcpxo_deamon_key, tcpxo_deamon_paths = get_tcpxo_deamon_entry()
   interfaces_key, interfaces_value = get_interfaces_entry(sub_networks)
   annotations.update({
@@ -149,6 +149,7 @@ def add_tcpxo_daemon_container(job_manifest):
       'name': 'tcpxo-daemon',
       'image': f'us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/tcpgpudmarxd-dev:{rxdm}',
       'imagePullPolicy': 'Always',
+      'restartPolicy': 'Always',
       'command': ['/bin/sh', '-c'],
       'args': [
           'set -ex\nchmod 755'
@@ -165,9 +166,9 @@ def add_tcpxo_daemon_container(job_manifest):
       ],
       'env': [{'name': 'LD_LIBRARY_PATH', 'value': '/usr/local/nvidia/lib64'}],
   }
-  job_manifest['spec']['template']['spec']['containers'].append(
-      tcpxo_daemon_container
-  )
+  spec = job_manifest['spec']['template']['spec']
+  spec.setdefault('initContainers', [])
+  spec['initContainers'].append(tcpxo_daemon_container)
 
 
 def update_gpu_containers(job_manifest):