Add the option to use Multi-tier checkpointing in workloads

src/xpk/core/core.py

@@ -2109,10 +2109,14 @@ def get_volumes(args, system: SystemCharacteristics) -> str:
                 name: dshm-2"""
 
   if args.ramdisk_directory != '':
-    volumes += """
+    driver = 'phase1-checkpoint.csi.storage.gke.io'
+    if args.mtc_enabled:
+      driver = 'multitier-checkpoint.csi.storage.gke.io'
+
+    volumes += f"""
               - name: cache
                 csi:
-                  driver: phase1-checkpoint.csi.storage.gke.io"""
+                  driver: {driver}"""
 
   if (
       system.accelerator_type == AcceleratorType['TPU']

src/xpk/parser/workload.py

@@ -598,6 +598,15 @@ def add_shared_workload_create_optional_arguments(args_parsers):
             ' be used with the CSI driver provided by GKE.'
         ),
     )
+    custom_parser.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.'
+        ),
+    )
 
 
 def add_shared_workload_create_env_arguments(args_parsers):