Multi-tier checkpointing + orbax replicator

Dockerfile

@@ -90,7 +90,7 @@ ENTRYPOINT ["/opt/apache/beam/boot"]
 
 FROM base AS tpu
 
-ARG EXTRAS=
+ARG EXTRAS=orbax
 
 ENV UV_FIND_LINKS=https://storage.googleapis.com/jax-releases/libtpu_releases.html
 # Ensure we install the TPU version, even if building locally.

axlearn/cloud/gcp/jobset_utils.py

@@ -243,6 +243,7 @@ class Config(BaseReplicatedJob.Config):
             env_vars: Optional env vars to set.
             gcsfuse_mount: Optional configs for the GCS FUSE sidecar and volume mount.
                 See `GCSFuseMount` for details.
+            mtc_mount: Optionally mount multi-tier checkpointing volume.
             host_mounts: List of volumes from host to mount into the container.
                 See `HostMount` for details.
             service_account: Optional service account to execute the job as.
@@ -256,6 +257,7 @@ class Config(BaseReplicatedJob.Config):
         accelerator: AcceleratorConfig = AcceleratorConfig()
         env_vars: dict[str, str] = {}
         gcsfuse_mount: Optional[GCSFuseMount] = None
+        mtc_mount: Optional[bool] = None
         host_mounts: Optional[Sequence[HostMount]] = None
         service_account: Optional[str] = None
         # This config is made Optional for backwards compatibility. Default is False.
@@ -377,6 +379,7 @@ class Config(SingleReplicatedJob.Config):
                 to attach to the node pool. This is needed to support multiple NIC.
                 Refer to GKE TPU provisioner for more context:
                 https://github.com/GoogleCloudPlatform/ai-on-gke/blob/5f256eed7075a5cb8e73cd72328aea46237b8ce6/tpu-provisioner/internal/cloud/common.go#L29-L31
+            mtc_mount: Whether or not to mount volume for multi-tier checkpointing
         """
 
         reservation: Optional[str] = None
@@ -386,6 +389,7 @@ class Config(SingleReplicatedJob.Config):
         enable_tpu_smart_repair: bool = False
         priority_class: Optional[str] = None
         additional_node_networks: Optional[str] = None
+        mtc_mount: Optional[bool] = None
 
     @classmethod
     def define_flags(cls, fv: flags.FlagValues):
@@ -410,6 +414,12 @@ def define_flags(cls, fv: flags.FlagValues):
             "The GKE PriorityClass for the job.",
             **common_kwargs,
         )
+        flags.DEFINE_boolean(
+            "mtc_mount",
+            None,
+            "Whether to mount checkpoint volume",
+            **common_kwargs,
+        )
 
     @classmethod
     def from_flags(cls, fv: flags.FlagValues, **kwargs) -> Config:
@@ -433,6 +443,7 @@ def from_flags(cls, fv: flags.FlagValues, **kwargs) -> Config:
         cfg.additional_node_networks = gcp_settings(
             "additional_node_networks", required=False, fv=fv
         )
+        cfg.mtc_mount = fv.mtc_mount
         return cfg
 
     def __init__(self, cfg: Config, *, bundler: Bundler):
@@ -476,6 +487,11 @@ def _build_container(self) -> Nested[Any]:
                 volume_mounts, spec=VolumeMount(name="shared-memory", mount_path="/dev/shm")
             )
 
+        if cfg.mtc_mount:
+            self._maybe_add_volume_mount(
+                volume_mounts, spec=VolumeMount(name="checkpoint", mount_path="/checkpoint")
+            )
+
         if cfg.host_mounts:
             for mount in cfg.host_mounts:
                 self._maybe_add_volume_mount(volume_mounts, spec=mount)
@@ -577,6 +593,17 @@ def _build_pod(self) -> Nested[Any]:
         annotations, labels, selector, volumes, tolerations = {}, {}, {}, [], []
 
         volumes.append(dict(name="shared-output", emptyDir={}))
+
+        if cfg.mtc_mount:
+            volumes.append(
+                dict(
+                    name="checkpoint",
+                    csi=dict(
+                        driver="multitier-checkpoint.csi.storage.gke.io",
+                    ),
+                )
+            )
+
         if cfg.gcsfuse_mount:
             # Increases the shared memory volumes when enabled gcsfuse. This is useful when grain
             # prefetch is enabled.