xpk 0.9.0__tar.gz → 0.10.0__tar.gz

{xpk-0.9.0/src/xpk.egg-info → xpk-0.10.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xpk
-Version: 0.9.0
+Version: 0.10.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
@@ -259,6 +259,13 @@ all zones.
     --num-slices=4 --spot
     ```
 
+* Cluster Create (DWS flex queued capacity):
+    ```shell
+        python3 xpk.py cluster create \
+        --cluster xpk-test --tpu-type=v5litepod-16 \
+        --num-slices=4 --flex
+    ```
+
 * Cluster Create for Pathways:
     Pathways compatible cluster can be created using `cluster create-pathways`.
     ```shell
@@ -495,6 +502,7 @@ Currently, the below flags/arguments are supported for A3 Mega, A3 Ultra and A4
   * `--reservation`
   * `--spot`
   * `--on-demand` (A3 Mega only)
+  * `--flex`
 
 ## Running XPK on existing clusters
 
@@ -518,9 +526,10 @@ Currently XPK supports the below types of storages:
 - [Google Cloud Filestore](#filestore)
 - [Google Cloud Parallelstore](#parallelstore)
 - [Google Cloud Block storages (Persistent Disk, Hyperdisk)](#block-storage-persistent-disk-hyperdisk)
+- [Google Cloud Managed Lustre](#managed-lustre)
 
 ### 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.
+A FUSE adapter lets you mount and access Cloud Storage buckets as local file systems, so workloads can read and write objects in your bucket using standard file system semantics.
 
 To use the GCS FUSE with XPK you need to create a [Storage Bucket](https://console.cloud.google.com/storage/).
 
@@ -547,7 +556,7 @@ Parameters:
 
 ### 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 files 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 workloads 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`:
 
@@ -583,7 +592,7 @@ Commands `xpk storage create` and `xpk storage attach` with `--type=gcpfilestore
 
 ### 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.
+A Parallelstore adapter lets you mount and access [Parallelstore instances](https://cloud.google.com/parallelstore/) as local file systems, so workloads 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/).
 
@@ -607,7 +616,7 @@ Parameters:
 
 ### 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.
+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 workloads 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.
 
@@ -629,6 +638,30 @@ Parameters:
 - `--readonly` - if set to true, workload can only read from storage.
 - `--manifest` - path to the manifest file containing PersistentVolume and PresistentVolumeClaim definitions.
 
+### Managed Lustre
+
+A Managed Lustre adaptor lets you mount and access [Google Cloud Managed Lustre instances](https://cloud.google.com/kubernetes-engine/docs/concepts/managed-lustre) as local file systems, so workloads can read and write files in your volumes using standard file system semantics.
+
+To use the GCP Managed Lustre with XPK you need to create [an instance](https://cloud.google.com/managed-lustre/docs/create-instance). Please make sure you enable GKE support when creating the instance (gcloud ex. `--gke-support-enabled`).
+
+Once it's ready you can use `xpk storage attach` with `--type=lustre` command to attach a Managed Lustre instance to your cluster. Currently, attaching a Managed Lustre instance is supported only by providing a manifest file.
+
+```shell
+python3 xpk.py storage attach test-lustre-storage --type=lustre \
+  --project=$PROJECT --cluster=$CLUSTER --zone=$ZONE \
+  --mount-point='/test-mount-point' --readonly=false \
+  --auto-mount=true \
+  --manifest='./examples/storage/lustre-manifest-attach.yaml'
+```
+
+Parameters:
+
+- `--type` - type of the storage `lustre`
+- `--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
@@ -670,8 +703,14 @@ python3 xpk.py storage delete test-fs-instance \
     python3 xpk.py workload create \
     --workload xpk-test-workload --command "echo goodbye" \
     --cluster xpk-test \
-    --tpu-type=v5litepod-16 --projet=$PROJECT
+    --tpu-type=v5litepod-16 --project=$PROJECT
     ```
+*   Workload create(DWS flex with queued provisioning):
+        ```shell
+    python3 xpk.py workload create \
+    --workload xpk-test-workload --command "echo goodbye" \
+    --cluster xpk-test  --flex \
+    --tpu-type=v5litepod-16 --project=$PROJECT
 
 *   Workload Create for Pathways:
     Pathways workload can be submitted using `workload create-pathways` on a Pathways enabled cluster (created with `cluster create-pathways`)

{xpk-0.9.0 → xpk-0.10.0}/README.md

@@ -227,6 +227,13 @@ all zones.
     --num-slices=4 --spot
     ```
 
+* Cluster Create (DWS flex queued capacity):
+    ```shell
+        python3 xpk.py cluster create \
+        --cluster xpk-test --tpu-type=v5litepod-16 \
+        --num-slices=4 --flex
+    ```
+
 * Cluster Create for Pathways:
     Pathways compatible cluster can be created using `cluster create-pathways`.
     ```shell
@@ -463,6 +470,7 @@ Currently, the below flags/arguments are supported for A3 Mega, A3 Ultra and A4
   * `--reservation`
   * `--spot`
   * `--on-demand` (A3 Mega only)
+  * `--flex`
 
 ## Running XPK on existing clusters
 
@@ -486,9 +494,10 @@ Currently XPK supports the below types of storages:
 - [Google Cloud Filestore](#filestore)
 - [Google Cloud Parallelstore](#parallelstore)
 - [Google Cloud Block storages (Persistent Disk, Hyperdisk)](#block-storage-persistent-disk-hyperdisk)
+- [Google Cloud Managed Lustre](#managed-lustre)
 
 ### 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.
+A FUSE adapter lets you mount and access Cloud Storage buckets as local file systems, so workloads can read and write objects in your bucket using standard file system semantics.
 
 To use the GCS FUSE with XPK you need to create a [Storage Bucket](https://console.cloud.google.com/storage/).
 
@@ -515,7 +524,7 @@ Parameters:
 
 ### 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 files 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 workloads 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`:
 
@@ -551,7 +560,7 @@ Commands `xpk storage create` and `xpk storage attach` with `--type=gcpfilestore
 
 ### 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.
+A Parallelstore adapter lets you mount and access [Parallelstore instances](https://cloud.google.com/parallelstore/) as local file systems, so workloads 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/).
 
@@ -575,7 +584,7 @@ Parameters:
 
 ### 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.
+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 workloads 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.
 
@@ -597,6 +606,30 @@ Parameters:
 - `--readonly` - if set to true, workload can only read from storage.
 - `--manifest` - path to the manifest file containing PersistentVolume and PresistentVolumeClaim definitions.
 
+### Managed Lustre
+
+A Managed Lustre adaptor lets you mount and access [Google Cloud Managed Lustre instances](https://cloud.google.com/kubernetes-engine/docs/concepts/managed-lustre) as local file systems, so workloads can read and write files in your volumes using standard file system semantics.
+
+To use the GCP Managed Lustre with XPK you need to create [an instance](https://cloud.google.com/managed-lustre/docs/create-instance). Please make sure you enable GKE support when creating the instance (gcloud ex. `--gke-support-enabled`).
+
+Once it's ready you can use `xpk storage attach` with `--type=lustre` command to attach a Managed Lustre instance to your cluster. Currently, attaching a Managed Lustre instance is supported only by providing a manifest file.
+
+```shell
+python3 xpk.py storage attach test-lustre-storage --type=lustre \
+  --project=$PROJECT --cluster=$CLUSTER --zone=$ZONE \
+  --mount-point='/test-mount-point' --readonly=false \
+  --auto-mount=true \
+  --manifest='./examples/storage/lustre-manifest-attach.yaml'
+```
+
+Parameters:
+
+- `--type` - type of the storage `lustre`
+- `--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
@@ -638,8 +671,14 @@ python3 xpk.py storage delete test-fs-instance \
     python3 xpk.py workload create \
     --workload xpk-test-workload --command "echo goodbye" \
     --cluster xpk-test \
-    --tpu-type=v5litepod-16 --projet=$PROJECT
+    --tpu-type=v5litepod-16 --project=$PROJECT
     ```
+*   Workload create(DWS flex with queued provisioning):
+        ```shell
+    python3 xpk.py workload create \
+    --workload xpk-test-workload --command "echo goodbye" \
+    --cluster xpk-test  --flex \
+    --tpu-type=v5litepod-16 --project=$PROJECT
 
 *   Workload Create for Pathways:
     Pathways workload can be submitted using `workload create-pathways` on a Pathways enabled cluster (created with `cluster create-pathways`)

{xpk-0.9.0 → xpk-0.10.0}/src/xpk/commands/batch.py

@@ -18,7 +18,7 @@ import re
 from argparse import Namespace
 
 from ..core.cluster import (
-    create_xpk_k8s_service_account,
+    setup_k8s_service_accounts,
     get_cluster_credentials,
 )
 from ..core.commands import run_command_for_value
@@ -54,14 +54,14 @@ def batch(args: Namespace) -> None:
   err_code = prepare_kjob(args)
   if err_code > 0:
     xpk_exit(err_code)
-  create_xpk_k8s_service_account()
+  setup_k8s_service_accounts()
 
   submit_job(args)
 
 
 def submit_job(args: Namespace) -> None:
 
-  create_xpk_k8s_service_account()
+  setup_k8s_service_accounts()
 
   cmd = (
       'kubectl kjob create slurm'

{xpk-0.9.0 → xpk-0.10.0}/src/xpk/commands/cluster.py

@@ -31,6 +31,7 @@ from ..core.cluster import (
     update_cluster_with_gcsfuse_driver_if_necessary,
     update_cluster_with_parallelstore_driver_if_necessary,
     update_cluster_with_pd_driver_if_necessary,
+    update_cluster_with_lustre_driver_if_necessary,
     update_cluster_with_workload_identity_if_necessary,
 )
 from ..core.cluster_private import authorize_private_cluster_access_if_necessary
@@ -42,12 +43,14 @@ from ..core.gcloud_context import (
     get_gke_server_config,
     zone_to_region,
 )
+from ..core.jobset import update_jobset_resources_if_necessary
 from ..core.kjob import apply_kjob_crds, prepare_kjob, verify_kjob_installed
 from ..core.kueue import (
     cluster_preheat_yml,
     install_kueue_crs,
     install_kueue_on_cluster,
     wait_for_kueue_available,
+    update_kueue_resources_if_necessary,
 )
 from ..core.nap import enable_autoprovisioning_on_cluster
 from ..core.network import (
@@ -170,7 +173,6 @@ def cluster_adapt(args) -> None:
   install_kueue(args, system, autoprovisioning_config)
 
   install_kjob(args)
-
   if system.accelerator_type == AcceleratorType['GPU']:
     prepare_gpus(args, system)
 
@@ -308,6 +310,9 @@ def cluster_create(args) -> None:
   set_jobset_on_cluster_code = set_jobset_on_cluster(args)
   if set_jobset_on_cluster_code != 0:
     xpk_exit(set_jobset_on_cluster_code)
+  update_jobset_resources_code = update_jobset_resources_if_necessary(args)
+  if update_jobset_resources_code != 0:
+    xpk_exit(update_jobset_resources_code)
 
   set_pathways_job_on_cluster_code = set_pathways_job_on_cluster(args)
   if set_pathways_job_on_cluster_code != 0:
@@ -879,6 +884,10 @@ def run_gke_cluster_create_command(
   if args.enable_pd_csi_driver:
     addons.append('GcePersistentDiskCsiDriver')
 
+  if args.enable_lustre_csi_driver:
+    addons.append('LustreCsiDriver')
+    command += ' --enable-legacy-lustre-port'
+
   if hasattr(args, 'enable_mtc') and args.enable_mtc:
     addons.append('HighScaleCheckpointing')
 
@@ -922,6 +931,13 @@ def install_storage_csis(args):
     if update_cluster_command_code != 0:
       xpk_exit(update_cluster_command_code)
 
+  if args.enable_lustre_csi_driver:
+    update_cluster_command_code = (
+        update_cluster_with_lustre_driver_if_necessary(args)
+    )
+    if update_cluster_command_code != 0:
+      xpk_exit(update_cluster_command_code)
+
 
 def install_kjob(args):
   xpk_print('Verifying kjob installation')
@@ -957,6 +973,11 @@ def install_kueue(args, system: SystemCharacteristics, autoprovisioning_config):
   if enable_kueue_credentials_code != 0:
     xpk_exit(enable_kueue_credentials_code)
 
+  xpk_print('Update Kueue Controller Manager resources')
+  update_kueue_resources_code = update_kueue_resources_if_necessary(args)
+  if update_kueue_resources_code != 0:
+    xpk_exit(update_kueue_resources_code)
+
 
 def prepare_gpus(args, system: SystemCharacteristics):
   xpk_print('Installing NCCL Plugin for cluster')

{xpk-0.9.0 → xpk-0.10.0}/src/xpk/commands/cluster_gcluster.py

@@ -37,6 +37,7 @@ from ..utils.console import xpk_exit, xpk_print
 from ..utils.file import ensure_directory_exists
 from ..utils.network import all_IPs_cidr
 from ..utils.objects import hash_string
+from ..core.capacity import get_reservation_maintenance_interval, get_reservation_placement_policy
 
 blueprints_path = os.path.abspath('xpkclusters/blueprints')
 gcluster_working_dir = os.path.abspath('xpkclusters/gcluster-out')
@@ -234,6 +235,30 @@ def generate_blueprint(
   if args.device_type in supported_device_types:
     if args.device_type == a3mega_device_type:
       num_nodes = args.num_nodes if not args.num_nodes is None else 2
+
+      maintenance_interval = (
+          get_reservation_maintenance_interval(
+              args.reservation, args.zone, args.project
+          )
+          if args.reservation is not None
+          else 'PERIODIC'
+      )
+      placement_policy_name = (
+          get_reservation_placement_policy(
+              args.reservation, args.zone, args.project
+          )
+          if args.reservation is not None
+          else None
+      )
+      placement_policy = (
+          {
+              'type': 'COMPACT',
+              'name': placement_policy_name.split('/')[-1],
+          }
+          if placement_policy_name is not None
+          and len(placement_policy_name) > 0
+          else None
+      )
       return bpg.generate_a3_mega_blueprint(
           blueprint_name=blueprint_name,
           prefix=prefix,
@@ -243,6 +268,8 @@ def generate_blueprint(
           zone=args.zone,
           auth_cidr=all_IPs_cidr,
           num_nodes=num_nodes,
+          reservation_maintenance_interval=maintenance_interval,
+          reservation_placement_policy=placement_policy,
           reservation=args.reservation if args.reservation else None,
           capacity_type=capacity_type,
           system_node_pool_machine_type=args.default_pool_cpu_machine_type,

{xpk-0.9.0 → xpk-0.10.0}/src/xpk/commands/common.py

@@ -15,10 +15,12 @@ limitations under the License.
 """
 
 from ..core.commands import run_command_with_updates_retry
-from ..core.resources import get_cluster_capacity_type, get_cluster_system_characteristics
 from ..core.capacity import H100_MEGA_DEVICE_TYPE, CapacityType
 from ..core.gcloud_context import zone_to_region
 from ..utils.console import xpk_print, xpk_exit
+from ..core.system_characteristics import (
+    SystemCharacteristics,
+)
 
 
 def set_cluster_command(args) -> int:
@@ -47,7 +49,11 @@ def set_cluster_command(args) -> int:
   return return_code
 
 
-def is_TAS_possible(args) -> bool:
+def is_TAS_possible(
+    system_characteristics: SystemCharacteristics,
+    capacity_type: CapacityType,
+    flex: bool,
+) -> bool:
   """Check cluster's machine_type and capacity type to determine if Kueue TAS is possible
 
   Args:
@@ -56,8 +62,6 @@ def is_TAS_possible(args) -> bool:
   Returns:
     True if possible and False otherwise.
   """
-  system_characteristics = get_cluster_system_characteristics(args)
-  capacity_type = get_cluster_capacity_type(args)
 
   if system_characteristics is None:
     xpk_print('system_characteristics data was not found in configmaps.')
@@ -67,9 +71,12 @@ def is_TAS_possible(args) -> bool:
     xpk_print('capacity_type data was not found in configmaps.')
     xpk_exit(1)
 
+  if flex:
+    return False
+
   if (
       system_characteristics.device_type == H100_MEGA_DEVICE_TYPE
-      and capacity_type == CapacityType.SPOT
+      and capacity_type != CapacityType.RESERVATION
   ):
     return False
 

{xpk-0.9.0 → xpk-0.10.0}/src/xpk/commands/kjob_common.py

@@ -27,6 +27,7 @@ from ..core.kjob import (
     Kueue_TAS_annotation,
 )
 from .common import is_TAS_possible
+from ..core.resources import get_cluster_capacity_type, get_cluster_system_characteristics
 
 
 def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
@@ -50,7 +51,9 @@ def add_gpu_networking_annotations_to_command(args, cmd: str) -> str:
 
 
 def add_TAS_annotations_to_command(args, cmd: str) -> str:
-  if is_TAS_possible(args):
+  system_characteristics = get_cluster_system_characteristics(args)
+  capacity_type = get_cluster_capacity_type(args)
+  if is_TAS_possible(system_characteristics, capacity_type, flex=False):
     cmd += f" --pod-template-annotation {Kueue_TAS_annotation}"
 
   return cmd

{xpk-0.9.0 → xpk-0.10.0}/src/xpk/commands/run.py

@@ -17,7 +17,7 @@ limitations under the License.
 from argparse import Namespace
 
 from ..core.cluster import (
-    create_xpk_k8s_service_account,
+    setup_k8s_service_accounts,
     get_cluster_credentials,
 )
 from ..core.commands import run_command_with_full_controls
@@ -53,7 +53,7 @@ def run(args: Namespace) -> None:
   err_code = prepare_kjob(args)
   if err_code > 0:
     xpk_exit(err_code)
-  create_xpk_k8s_service_account()
+  setup_k8s_service_accounts()
 
   submit_job(args)
 

{xpk-0.9.0 → xpk-0.10.0}/src/xpk/commands/shell.py

@@ -12,7 +12,7 @@ limitations under the License.
 """
 
 from ..core.commands import run_command_with_full_controls, run_command_for_value, run_command_with_updates
-from ..core.cluster import get_cluster_credentials, add_zone_and_project, create_xpk_k8s_service_account
+from ..core.cluster import get_cluster_credentials, add_zone_and_project, setup_k8s_service_accounts
 from ..utils.console import xpk_exit, xpk_print
 from argparse import Namespace
 
@@ -82,7 +82,7 @@ def connect_to_new_interactive_shell(args: Namespace) -> int:
   err_code = prepare_kjob(args)
   if err_code > 0:
     xpk_exit(err_code)
-  create_xpk_k8s_service_account()
+  setup_k8s_service_accounts()
 
   cmd = (
       'kubectl-kjob create interactive --profile'

{xpk-0.9.0 → xpk-0.10.0}/src/xpk/commands/storage.py

@@ -29,6 +29,7 @@ from ..core.cluster import (
     setup_k8s_env,
     update_cluster_with_parallelstore_driver_if_necessary,
     update_cluster_with_pd_driver_if_necessary,
+    update_cluster_with_lustre_driver_if_necessary,
     update_cluster_with_gcpfilestore_driver_if_necessary,
     update_cluster_with_gcsfuse_driver_if_necessary,
     update_cluster_with_workload_identity_if_necessary,
@@ -45,6 +46,7 @@ from ..core.storage import (
     GCS_FUSE_TYPE,
     GCE_PD_TYPE,
     PARALLELSTORE_TYPE,
+    LUSTRE_TYPE,
     STORAGE_CRD_PLURAL,
     XPK_API_GROUP_NAME,
     XPK_API_GROUP_VERSION,
@@ -183,11 +185,11 @@ def storage_attach(args: Namespace) -> None:
           args.prefetch_metadata,
       )
 
-  elif args.type in [PARALLELSTORE_TYPE, GCE_PD_TYPE]:
+  elif args.type in [PARALLELSTORE_TYPE, GCE_PD_TYPE, LUSTRE_TYPE]:
     if args.manifest is None:
       xpk_print(
-          "Parallelstore and PersistentDisk are currently supported only with"
-          " --manifest"
+          "Parallelstore, PersistentDisk, and Lustre are currently supported"
+          " only with --manifest"
       )
       xpk_exit(1)
 
@@ -234,6 +236,11 @@ def enable_csi_drivers_if_necessary(args: Namespace) -> None:
     if return_code > 0:
       xpk_exit(return_code)
 
+  if args.type == LUSTRE_TYPE:
+    return_code = update_cluster_with_lustre_driver_if_necessary(args)
+    if return_code > 0:
+      xpk_exit(return_code)
+
 
 def storage_list(args: Namespace) -> None:
   k8s_api_client = setup_k8s_env(args)