xpk 0.7.2__py3-none-any.whl → 0.8.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.
@@ -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.
 
@@ -80,12 +80,12 @@ def get_interfaces_entry(sub_networks: list[str]) -> tuple[str, str]:
   return 'networking.gke.io/interfaces', literal_string('\n'.join(interfaces))
 
 
-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']
   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/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.
 
@@ -105,7 +105,7 @@ 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']
   tcpxo_deamon_key, tcpxo_deamon_paths = get_tcpxo_deamon_entry()

xpk/parser/cluster.py

@@ -107,6 +107,7 @@ def set_cluster_parser(cluster_parser):
           ' enable cluster to accept Pathways workloads.'
       ),
   )
+
   ### Autoprovisioning arguments specific to "cluster create"
   cluster_create_autoprovisioning_arguments = (
       cluster_create_parser.add_argument_group(
@@ -462,7 +463,7 @@ def add_shared_cluster_create_optional_arguments(args_parsers):
     custom_parser.add_argument(
         '--pathways-gce-machine-type',
         type=str,
-        default='n1-standard-32',
+        default='n2-standard-64',
         help='The CPU type for Pathways CPU nodepools',
     )
     custom_parser.add_argument(
@@ -580,14 +581,22 @@ def add_shared_cluster_create_optional_arguments(args_parsers):
             ' Identity is enabled by default.'
         ),
     )
-
     custom_parser.add_argument(
         '--enable-gcpfilestore-csi-driver',
         action='store_true',
-        help=(
-            'Enable GCPFilestore driver on the cluster. This enables Workload'
-            ' Identity Federation.'
-        ),
+        help='Enable GCPFilestore driver on the cluster.',
+    )
+
+    custom_parser.add_argument(
+        '--enable-parallelstore-csi-driver',
+        action='store_true',
+        help='Enable Parallelstore CSI driver on the cluster.',
+    )
+
+    custom_parser.add_argument(
+        '--enable-pd-csi-driver',
+        action='store_true',
+        help='Enable PersistentDisk CSI driver on the cluster.',
     )
 
 

xpk/parser/storage.py

@@ -73,7 +73,7 @@ def add_storage_attach_parser(
           'The type of storage. Currently supported types: ["gcsfuse",'
           ' "gcpfilestore"]'
       ),
-      choices=['gcsfuse', 'gcpfilestore'],
+      choices=['gcsfuse', 'gcpfilestore', 'parallelstore', 'pd'],
       required=True,
   )
   add_cluster_arguments(req_args, required=True)
@@ -146,13 +146,19 @@ def add_storage_attach_parser(
 
   opt_args = storage_attach_parser.add_argument_group(
       'Optional Arguments',
-      'Optional arguments for storage create.',
+      'Optional arguments for storage attach.',
   )
   opt_args.add_argument(
       '--manifest',
       type=str,
       help='Path to manifest file containing volume definitions',
   )
+  opt_args.add_argument(
+      '--mount-options',
+      type=str,
+      help='Comma-separated list of mountOptions for PersistentVolume',
+      default='implicit-dirs',
+  )
   add_kind_cluster_arguments(opt_args)
 
 
@@ -184,7 +190,6 @@ def add_storage_create_parser(
       ),
       required=True,
   )
-
   req_args.add_argument(
       '--type',
       type=str,
@@ -248,6 +253,12 @@ def add_storage_create_parser(
       type=str,
       help='Path to manifest file containing volume definitions',
   )
+  opt_args.add_argument(
+      '--mount-options',
+      type=str,
+      help='Comma-separated list of mountOptions for PersistentVolume',
+      default='',
+  )
 
   add_kind_cluster_arguments(opt_args)
 

xpk/parser/workload.py

@@ -134,6 +134,24 @@ def set_workload_parsers(workload_parser):
           ' to use `gke.io/topology-aware-auto`.'
       ),
   )
+  workload_create_parser_optional_arguments.add_argument(
+      '--ramdisk-directory',
+      type=str,
+      default='',
+      help=(
+          'The directory of the locally mounted RAM disk. This is only to'
+          ' be used with the CSI driver provided by GKE.'
+      ),
+  )
+  workload_create_parser_optional_arguments.add_argument(
+      '--mtc-enabled',
+      action='store_true',
+      help=(
+          'The workload can use multi-tier checkpointing controllers when the'
+          ' --ramdisk-directory argument is used with this additional'
+          ' argument.'
+      ),
+  )
   workload_create_parser_optional_arguments.add_argument(
       '--debug-dump-gcs',
       type=str,
@@ -161,6 +179,19 @@ def set_workload_parsers(workload_parser):
           ' create Pathways workloads.'
       ),
   )
+  workload_create_parser_optional_arguments.add_argument(
+      '--restart-on-exit-codes',
+      type=str,
+      default=None,
+      help=(
+          'Adding this argument specifies additional user-defined exit codes'
+          ' that allow restarting the workload when --max-restarts is set to'
+          ' a value greater than 0. By default, workloads restart on exit'
+          ' codes 42 and 127-255. Any exit codes provided through this flag'
+          ' will be included alongside the default codes for restarting'
+          ' conditions.'
+      ),
+  )
 
   # Autoprovisioning workload arguments
   workload_create_autoprovisioning_arguments.add_argument(
@@ -244,9 +275,7 @@ def set_workload_parsers(workload_parser):
   workload_create_pathways_parser_optional_arguments.add_argument(
       '--proxy-server-image',
       type=str,
-      default=(
-          'us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest'
-      ),
+      default='',
       help=(
           'Please provide the proxy server image for Pathways. This arg can'
           ' only be used in `xpk workload create-pathways`.'
@@ -255,7 +284,7 @@ def set_workload_parsers(workload_parser):
   workload_create_pathways_parser_optional_arguments.add_argument(
       '--server-image',
       type=str,
-      default='us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest',
+      default='',
       help=(
           'Please provide the server image for Pathways. This arg can only be'
           ' used in `xpk workload create-pathways`.'
@@ -293,7 +322,7 @@ def set_workload_parsers(workload_parser):
   workload_create_pathways_parser_optional_arguments.add_argument(
       '--custom-pathways-server-args',
       type=str,
-      default=None,
+      default='',
       help=(
           'Provide custom Pathways server args as follows -'
           " --custom-pathways-server-args='--arg_1=xxx --arg2=yyy'"
@@ -304,7 +333,7 @@ def set_workload_parsers(workload_parser):
   workload_create_pathways_parser_optional_arguments.add_argument(
       '--custom-pathways-proxy-server-args',
       type=str,
-      default=None,
+      default='',
       help=(
           'Provide custom Pathways proxy server args as follows -'
           " --custom-pathways-proxy-server-args='--arg_1=xxx --arg2=yyy'"
@@ -315,7 +344,7 @@ def set_workload_parsers(workload_parser):
   workload_create_pathways_parser_optional_arguments.add_argument(
       '--custom-pathways-worker-args',
       type=str,
-      default=None,
+      default='',
       help=(
           'Provide custom Pathways worker args as follows -'
           " --custom-pathways-worker-args='--arg_1=xxx --arg2=yyy'"
@@ -323,6 +352,27 @@ def set_workload_parsers(workload_parser):
       required=False,
   )
 
+  workload_create_pathways_parser_optional_arguments.add_argument(
+      '--elastic-slices',
+      type=int,
+      default=0,
+      help=(
+          'Enable elastic slices in Pathways and specify'
+          ' the number of slices the workload could lose.'
+      ),
+      required=False,
+  )
+  workload_create_pathways_parser_optional_arguments.add_argument(
+      '--max-slice-restarts',
+      type=int,
+      default=1,
+      help=(
+          'Specify the maximum times the workers in a slice can be'
+          ' restarted. Used with --elastic-slices for Pathways workloads.'
+      ),
+      required=False,
+  )
+
   add_shared_workload_create_required_arguments([
       workload_create_parser_required_arguments,
       workload_create_pathways_parser_required_arguments,
@@ -583,9 +633,9 @@ def add_shared_workload_create_optional_arguments(args_parsers):
         ),
     )
     custom_parser.add_argument(
-        '--remote-python-sidecar-image',
+        '--colocated-python-sidecar-image',
         type=str,
-        default=None,
+        default='',
         help='Remote Python sidecar server image.',
     )
     custom_parser.add_argument(
@@ -596,28 +646,6 @@ def add_shared_workload_create_optional_arguments(args_parsers):
             ' the workload.'
         ),
     )
-    custom_parser.add_argument(
-        '--restart-on-exit-codes',
-        type=str,
-        default=None,
-        help=(
-            'Adding this argument specifies additional user-defined exit codes'
-            ' that allow restarting the workload when --max-restarts is set to'
-            ' a value greater than 0. By default, workloads restart on exit'
-            ' codes 42 and 127-255. Any exit codes provided through this flag'
-            ' will be included alongside the default codes for restarting'
-            ' conditions.'
-        ),
-    )
-    custom_parser.add_argument(
-        '--ramdisk-directory',
-        type=str,
-        default='',
-        help=(
-            'The directory of the locally mounted RAM disk. This is only to'
-            ' be used with the CSI driver provided by GKE.'
-        ),
-    )
 
 
 def add_shared_workload_create_env_arguments(args_parsers):

{xpk-0.7.2.dist-info → xpk-0.8.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xpk
-Version: 0.7.2
+Version: 0.8.0
 Summary: xpk helps Cloud developers to orchestrate training jobs on accelerators on GKE.
 Author-email: XPK team <xpk-code-reviewers@google.com>
 License: Apache-2.0
@@ -88,9 +88,11 @@ and the following GPU types:
 and the following CPU types:
 * n2-standard-32
 
-xpk also supports Google Cloud Storage solutions:
+xpk also supports [Google Cloud Storage solutions](#storage):
 * [Cloud Storage FUSE](#fuse)
 * [Filestore](#filestore)
+* [Parallelstore](#parallelstore)
+* [Block storage (Persistent Disk, Hyperdisk)](#block-storage-persistent-disk-hyperdisk)
 
 # Permissions needed on Cloud Console:
 
@@ -253,6 +255,7 @@ all zones.
     --num-slices=4 --on-demand \
     --tpu-type=v5litepod-16
     ```
+    Note that Pathways clusters need a CPU nodepool of n2-standard-64 or higher.
 
 *   Cluster Create for Ray:
     A cluster with KubeRay enabled and a RayCluster can be created using `cluster create-ray`.
@@ -475,7 +478,11 @@ Currently, the below flags/arguments are supported for A3-Mega and A3-Ultra mach
 
 
 ## Storage
-Currently XPK supports two types of storages: Cloud Storage FUSE and Google Cloud Filestore.
+Currently XPK supports the below types of storages:
+- [Cloud Storage FUSE](#fuse)
+- [Google Cloud Filestore](#filestore)
+- [Google Cloud Parallelstore](#parallelstore)
+- [Google Cloud Block storages (Persistent Disk, Hyperdisk)](#block-storage-persistent-disk-hyperdisk)
 
 ### FUSE
 A FUSE adapter lets you mount and access Cloud Storage buckets as local file systems, so applications can read and write objects in your bucket using standard file system semantics.
@@ -499,11 +506,12 @@ Parameters:
 - `--readonly` - if set to true, workload can only read from storage.
 - `--size` - size of the storage in Gb.
 - `--bucket` - name of the storage bucket. If not set then the name of the storage is used as a bucket name.
+- `--mount-options` - comma-separated list of additional mount options for PersistentVolume ([reference](https://cloud.google.com/kubernetes-engine/docs/how-to/cloud-storage-fuse-csi-driver-perf#mount-options)).
 - `--manifest` - path to the manifest file containing PersistentVolume and PresistentVolumeClaim definitions. If set, then values from manifest override the following parameters: `--size` and `--bucket`.
 
 ### Filestore
 
-A Filestore adapter lets you mount and access [Filestore instances](https://cloud.google.com/filestore/) as local file systems, so applications can read and write objects in your volumes using standard file system semantics.
+A Filestore adapter lets you mount and access [Filestore instances](https://cloud.google.com/filestore/) as local file systems, so applications can read and write files in your volumes using standard file system semantics.
 
 To create and attach a GCP Filestore instance to your cluster use `xpk storage create` command with `--type=gcpfilestore`:
 
@@ -537,6 +545,54 @@ Commands `xpk storage create` and `xpk storage attach` with `--type=gcpfilestore
 - `--instance` - the name of the Filestore instance. If not set then the name parameter is used as an instance name. Useful when connecting multiple volumes from the same Filestore instance.
 - `--manifest` - path to the manifest file containing PersistentVolume, PresistentVolumeClaim and StorageClass definitions. If set, then values from manifest override the following parameters: `--access-mode`, `--size` and `--volume`.
 
+### Parallelstore
+
+A Parallelstore adapter lets you mount and access [Parallelstore instances](https://cloud.google.com/parallelstore/) as local file systems, so applications can read and write files in your volumes using standard file system semantics.
+
+To use the GCS Parallelstore with XPK you need to create a [Parallelstore Instance](https://console.cloud.google.com/parallelstore/).
+
+Once it's ready you can use `xpk storage attach` with `--type=parallelstore` command to attach a Parallelstore instance to your cluster. Currently, attaching a Parallelstore is supported only by providing a manifest file.
+
+```shell
+python3 xpk.py storage attach test-parallelstore-storage --type=parallelstore \
+  --project=$PROJECT --cluster=$CLUSTER --zone=$ZONE \
+  --mount-point='/test-mount-point' --readonly=false \
+  --auto-mount=true \
+  --manifest='./examples/storage/parallelstore-manifest-attach.yaml'
+```
+
+Parameters:
+
+- `--type` - type of the storage `parallelstore`
+- `--auto-mount` - if set to true all workloads will have this storage mounted by default.
+- `--mount-point` - the path on which this storage should be mounted for a workload.
+- `--readonly` - if set to true, workload can only read from storage.
+- `--manifest` - path to the manifest file containing PersistentVolume and PresistentVolumeClaim definitions.
+
+### Block storage (Persistent Disk, Hyperdisk)
+
+A PersistentDisk adapter lets you mount and access Google Cloud Block storage solutions ([Persistent Disk](https://cloud.google.com/kubernetes-engine/docs/concepts/storage-overview#pd), [Hyperdisk](https://cloud.google.com/kubernetes-engine/docs/concepts/storage-overview#hyperdisk)) as local file systems, so applications can read and write files in your volumes using standard file system semantics.
+
+To use the GCE PersistentDisk with XPK you need to create a [disk in GCE](https://cloud.google.com/compute/docs/disks). Please consider that the disk type you are creating is [compatible with the VMs](https://cloud.google.com/compute/docs/machine-resource#machine_type_comparison) in the default and accelerator nodepools.
+
+Once it's ready you can use `xpk storage attach` with `--type=pd` command to attach a PersistentDisk instance to your cluster. Currently, attaching a PersistentDisk is supported only by providing a manifest file.
+
+```shell
+python3 xpk.py storage attach test-pd-storage --type=pd \
+  --project=$PROJECT --cluster=$CLUSTER --zone=$ZONE \
+  --mount-point='/test-mount-point' --readonly=false \
+  --auto-mount=true \
+  --manifest='./examples/storage/pd-manifest-attach.yaml'
+```
+
+Parameters:
+
+- `--type` - type of the storage `pd`
+- `--auto-mount` - if set to true all workloads will have this storage mounted by default.
+- `--mount-point` - the path on which this storage should be mounted for a workload.
+- `--readonly` - if set to true, workload can only read from storage.
+- `--manifest` - path to the manifest file containing PersistentVolume and PresistentVolumeClaim definitions.
+
 ### List attached storages
 
 ```shell