MSC Checkpointing Changes (#789)

* Working changes to be cleaned up.

* Rename msc_config.yaml

* Fixed pytorch test issue by removing MSC Cache

* Updated project dependencies

* Find MSC config using absolute path.

* Re-added cuda test parameter.

* Add test to read from public S3 bucket using MSC.

* Revert save_checkpoint_freq value.

* Remove temporary printing

* Remove unnecessary dependency

* Switched to use consistent mechanism for detecting msc URIs

* Moved fsspec.filesystem logic into filesystem.py

* Change to cache for non-file protocols when reading non-modulus models.

* Moved code to generate checkpoint directory.directory

* Added get_checkpoint_dir import

* Address review feedback.

* Changes from code review.

* Addressed file test issue from review.

* Fix to file existence check.

* Fix merge conflicts due to project name change.

* Updated CHANGELOG.

* Added Multi-Storage Client to allow checkpointing to/from Object Storage

Signed-off-by: Chris Hawes <[email protected]>

* Addressed issues identified by pre-commit.

* Update filesystem.py

* Update __init__.py

* Update Dockerfile

---------

Signed-off-by: Chris Hawes <[email protected]>
Co-authored-by: Nicholas Geneva <[email protected]>

Author: chris-hawes

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - Added ReGen score-based data assimilation example
+- Added Multi-Storage Client to allow checkpointing to/from Object Storage
 
 ### Changed
 

examples/generative/corrdiff/train.py

@@ -25,7 +25,11 @@
 from physicsnemo.launch.logging import PythonLogger, RankZeroLoggingWrapper
 from physicsnemo.metrics.diffusion import RegressionLoss, ResLoss, RegressionLossCE
 from physicsnemo.launch.logging import PythonLogger, RankZeroLoggingWrapper
-from physicsnemo.launch.utils import load_checkpoint, save_checkpoint
+from physicsnemo.launch.utils import (
+    load_checkpoint,
+    save_checkpoint,
+    get_checkpoint_dir,
+)
 from datasets.dataset import init_train_valid_datasets_from_config
 from helpers.train_helpers import (
     set_patch_shape,
@@ -66,8 +70,8 @@ def main(cfg: DictConfig) -> None:
     enable_amp = fp_optimizations.startswith("amp")
     amp_dtype = torch.float16 if (fp_optimizations == "amp-fp16") else torch.bfloat16
     logger.info(f"Saving the outputs in {os.getcwd()}")
-    checkpoint_dir = os.path.join(
-        cfg.training.io.get("checkpoint_dir", "."), f"checkpoints_{cfg.model.name}"
+    checkpoint_dir = get_checkpoint_dir(
+        str(cfg.training.io.get("checkpoint_dir", ".")), cfg.model.name
     )
     if cfg.training.hp.batch_size_per_gpu == "auto":
         cfg.training.hp.batch_size_per_gpu = (

physicsnemo/launch/utils/__init__.py

@@ -14,4 +14,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .checkpoint import load_checkpoint, save_checkpoint
+from .checkpoint import get_checkpoint_dir, load_checkpoint, save_checkpoint

physicsnemo/launch/utils/checkpoint.py

@@ -14,11 +14,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import glob
+import os
 import re
 from pathlib import Path
 from typing import Any, Dict, List, NewType, Optional, Union
 
+import fsspec
+import fsspec.utils
 import torch
 from torch.cuda.amp import GradScaler
 from torch.optim.lr_scheduler import _LRScheduler
@@ -27,6 +29,7 @@
 from physicsnemo.distributed import DistributedManager
 from physicsnemo.launch.logging import PythonLogger
 from physicsnemo.utils.capture import _StaticCapture
+from physicsnemo.utils.filesystem import LOCAL_CACHE, _download_cached
 
 optimizer = NewType("optimizer", torch.optim)
 scheduler = NewType("scheduler", _LRScheduler)
@@ -86,10 +89,13 @@ def _get_checkpoint_filename(
         else 0
     )
 
-    # Input file name
-    checkpoint_filename = str(
-        Path(path).resolve() / f"{base_name}.{model_parallel_rank}"
-    )
+    # Determine input file name. Get absolute file path if Posix path.
+    # pathlib does not support custom schemes (eg: msc://...) so only perform resolve() for Posix.
+    protocol = fsspec.utils.get_protocol(path)
+    fs = fsspec.filesystem(protocol)
+    if protocol == "file":
+        path = str(Path(path).resolve())
+    checkpoint_filename = f"{path}/{base_name}.{model_parallel_rank}"
 
     # File extension for PhysicsNeMo models or PyTorch models
     file_extension = ".mdlus" if model_type == "mdlus" else ".pt"
@@ -101,20 +107,21 @@ def _get_checkpoint_filename(
     # Otherwise try loading the latest epoch or rolling checkpoint
     else:
         file_names = [
-            Path(fname).name
-            for fname in glob.glob(
-                checkpoint_filename + "*" + file_extension, recursive=False
-            )
+            fname for fname in fs.glob(checkpoint_filename + "*" + file_extension)
         ]
 
         if len(file_names) > 0:
             # If checkpoint from a null index save exists load that
             # This is the most likely line to error since it will fail with
             # invalid checkpoint names
+
+            # Remove protocol prefix if present to allow generic matching
+            _, path_without_protocol = fsspec.core.split_protocol(path)
             file_idx = [
                 int(
                     re.sub(
-                        f"^{base_name}.{model_parallel_rank}.|" + file_extension,
+                        f"^{path_without_protocol}/{base_name}.{model_parallel_rank}.|"
+                        + file_extension,
                         "",
                         fname,
                     )
@@ -212,8 +219,11 @@ def save_checkpoint(
     metadata : Optional[Dict[str, Any]], optional
         Additional metadata to save, by default None
     """
-    # Create checkpoint directory if it does not exist
-    if not Path(path).is_dir():
+    protocol = fsspec.utils.get_protocol(path)
+    fs = fsspec.filesystem(protocol)
+    # Create checkpoint directory if it does not exist.
+    # Only applicable to Posix filesystems ("file" protocol), not object stores.
+    if protocol == "file" and not Path(path).is_dir():
         checkpoint_logging.warning(
             f"Output directory {path} does not exist, will " "attempt to create"
         )
@@ -239,7 +249,8 @@ def save_checkpoint(
             if isinstance(model, physicsnemo.models.Module):
                 model.save(file_name)
             else:
-                torch.save(model.state_dict(), file_name)
+                with fs.open(file_name, "wb") as fp:
+                    torch.save(model.state_dict(), fp)
             checkpoint_logging.success(f"Saved model state dictionary: {file_name}")
 
     # == Saving training checkpoint ==
@@ -270,10 +281,11 @@ def save_checkpoint(
 
     # Save checkpoint to memory
     if bool(checkpoint_dict):
-        torch.save(
-            checkpoint_dict,
-            output_filename,
-        )
+        with fs.open(output_filename, "wb") as fp:
+            torch.save(
+                checkpoint_dict,
+                fp,
+            )
         checkpoint_logging.success(f"Saved training checkpoint: {output_filename}")
 
 
@@ -318,8 +330,14 @@ def load_checkpoint(
     int
         Loaded epoch
     """
+    fs = fsspec.filesystem(fsspec.utils.get_protocol(path))
     # Check if checkpoint directory exists
-    if not Path(path).is_dir():
+    if fs.exists(path):
+        if fs.isfile(path):
+            raise FileNotFoundError(
+                f"Provided checkpoint directory {path} is a file, not directory"
+            )
+    else:
         checkpoint_logging.warning(
             f"Provided checkpoint directory {path} does not exist, skipping load"
         )
@@ -340,7 +358,7 @@ def load_checkpoint(
             file_name = _get_checkpoint_filename(
                 path, name, index=epoch, model_type=model_type
             )
-            if not Path(file_name).exists():
+            if not fs.exists(file_name):
                 checkpoint_logging.error(
                     f"Could not find valid model file {file_name}, skipping load"
                 )
@@ -349,21 +367,22 @@ def load_checkpoint(
             if isinstance(model, physicsnemo.models.Module):
                 model.load(file_name)
             else:
-                model.load_state_dict(torch.load(file_name, map_location=device))
-
+                file_to_load = _cache_if_needed(file_name)
+                model.load_state_dict(torch.load(file_to_load, map_location=device))
             checkpoint_logging.success(
                 f"Loaded model state dictionary {file_name} to device {device}"
             )
 
     # == Loading training checkpoint ==
     checkpoint_filename = _get_checkpoint_filename(path, index=epoch, model_type="pt")
-    if not Path(checkpoint_filename).is_file():
+    if not fs.exists(checkpoint_filename):
         checkpoint_logging.warning(
             "Could not find valid checkpoint file, skipping load"
         )
         return 0
 
-    checkpoint_dict = torch.load(checkpoint_filename, map_location=device)
+    file_to_load = _cache_if_needed(checkpoint_filename)
+    checkpoint_dict = torch.load(file_to_load, map_location=device)
     checkpoint_logging.success(
         f"Loaded checkpoint file {checkpoint_filename} to device {device}"
     )
@@ -397,3 +416,41 @@ def load_checkpoint(
         metadata_dict[key] = value
 
     return epoch
+
+
+def get_checkpoint_dir(base_dir: str, model_name: str) -> str:
+    """Get a checkpoint directory based on a given base directory and model name
+
+    Parameters
+    ----------
+    base_dir : str
+        Path to the base directory where checkpoints are stored
+    model_name: str, optional
+        Name of the model which is generating the checkpoint
+
+    Returns
+    -------
+    str
+        Checkpoint directory
+    """
+    top_level_dir = f"checkpoints_{model_name}"
+    protocol = fsspec.utils.get_protocol(base_dir)
+    if protocol == "msc":
+        if not base_dir.endswith("/"):
+            base_dir += "/"
+        return base_dir + top_level_dir
+    else:
+        return os.path.join(base_dir, top_level_dir)
+
+
+# Read via cache and return the cached path for non-file protocols, otherwise just return the path
+def _cache_if_needed(path: str) -> str:
+    protocol = fsspec.utils.get_protocol(path)
+    if protocol == "file":
+        return path
+    else:
+        return _download_cached(
+            path,
+            recursive=False,
+            local_cache_path=os.path.join(LOCAL_CACHE, f"checkpoint_pid_{os.getpid()}"),
+        )

physicsnemo/utils/filesystem.py

@@ -24,6 +24,7 @@
 
 import fsspec
 import fsspec.implementations.cached
+import fsspec.utils
 import requests
 import s3fs
 from tqdm import tqdm
@@ -46,7 +47,7 @@ def _get_fs(path):
     if path.startswith("s3://"):
         return s3fs.S3FileSystem(client_kwargs=dict(endpoint_url="https://pbss.s8k.io"))
     else:
-        return fsspec.filesystem("file")
+        return fsspec.filesystem(fsspec.utils.get_protocol(path))
 
 
 def _download_ngc_model_file(path: str, out_path: str, timeout: int = 300) -> str:
@@ -175,8 +176,8 @@ def _download_cached(
     # TODO watch for race condition here
     if not os.path.exists(cache_path):
         logger.debug("Downloading %s to cache: %s", path, cache_path)
-        if path.startswith("s3://"):
-            fs = _get_fs(path)
+        if url.scheme in ("s3", "msc"):
+            fs = fsspec.filesystem(fsspec.utils.get_protocol(path))
             fs.get(path, cache_path, recursive=recursive)
         elif path.startswith("ngc://models/"):
             path = _download_ngc_model_file(path, cache_path)

pyproject.toml

@@ -57,6 +57,7 @@ dev = [
     "interrogate==1.5.0",
     "coverage==6.5.0",
     "ruff==0.0.290",
+    "moto[s3]>=5.0.28",
 ]
 
 makani = [
@@ -76,7 +77,7 @@ fignet = [
 ]
 
 storage = [
-    "multi-storage-client>=0.14.0",
+    "multi-storage-client[boto3]>=0.14.0",
 ]
 
 shardtensor = [

test/utils/msc_config_checkpoint.yaml

@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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
+#
+#     http://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.
+
+
+# This is an example MSC configuration file for testing checkpoint logic.
+profiles:
+  checkpoint-test:
+    storage_provider:
+      type: s3
+      options:
+        region_name: us-east-1
+        base_path: checkpoint-test-bucket
+    credentials_provider:
+      type: S3Credentials
+      options:
+        access_key: "access-key-id"
+        secret_key: "secret-access-key"

test/utils/msc_config_public_read.yaml

@@ -0,0 +1,31 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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
+#
+#     http://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.
+
+
+# This is an example MSC configuration file for accessing the CMIP6 archive on AWS:
+# https://registry.opendata.aws/cmip6/
+profiles:
+  cmip6-pds:
+    storage_provider:
+      type: s3
+      options:
+        region_name: us-west-2
+        base_path: cmip6-pds
+        signature_version: UNSIGNED
+cache:
+  location: /tmp/.cache
+  size_mb: 5000
+