Fix Google Cloud Storage checkpointing (#18088)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Adrian Wälchli <[email protected]>

Author: 3 people

src/lightning/fabric/loggers/csv_logs.py

@@ -21,7 +21,7 @@
 from torch import Tensor
 
 from lightning.fabric.loggers.logger import Logger, rank_zero_experiment
-from lightning.fabric.utilities.cloud_io import get_filesystem
+from lightning.fabric.utilities.cloud_io import _is_dir, get_filesystem
 from lightning.fabric.utilities.logger import _add_prefix
 from lightning.fabric.utilities.rank_zero import rank_zero_only, rank_zero_warn
 from lightning.fabric.utilities.types import _PATH
@@ -155,15 +155,15 @@ def finalize(self, status: str) -> None:
     def _get_next_version(self) -> int:
         versions_root = os.path.join(self._root_dir, self.name)
 
-        if not self._fs.isdir(versions_root):
+        if not _is_dir(self._fs, versions_root, strict=True):
             log.warning("Missing logger folder: %s", versions_root)
             return 0
 
         existing_versions = []
         for d in self._fs.listdir(versions_root):
             full_path = d["name"]
             name = os.path.basename(full_path)
-            if self._fs.isdir(full_path) and name.startswith("version_"):
+            if _is_dir(self._fs, full_path) and name.startswith("version_"):
                 existing_versions.append(int(name.split("_")[1]))
 
         if len(existing_versions) == 0:

src/lightning/fabric/loggers/tensorboard.py

@@ -22,7 +22,7 @@
 from torch.nn import Module
 
 from lightning.fabric.loggers.logger import Logger, rank_zero_experiment
-from lightning.fabric.utilities.cloud_io import get_filesystem
+from lightning.fabric.utilities.cloud_io import _is_dir, get_filesystem
 from lightning.fabric.utilities.logger import _add_prefix, _convert_params, _flatten_dict
 from lightning.fabric.utilities.logger import _sanitize_params as _utils_sanitize_params
 from lightning.fabric.utilities.rank_zero import rank_zero_only, rank_zero_warn
@@ -294,7 +294,7 @@ def _get_next_version(self) -> int:
         for listing in listdir_info:
             d = listing["name"]
             bn = os.path.basename(d)
-            if self._fs.isdir(d) and bn.startswith("version_"):
+            if _is_dir(self._fs, d) and bn.startswith("version_"):
                 dir_ver = bn.split("_")[1].replace("/", "")
                 existing_versions.append(int(dir_ver))
         if len(existing_versions) == 0:

src/lightning/fabric/utilities/cloud_io.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Utilities related to data saving/loading."""
-
 import io
 from pathlib import Path
 from typing import Any, Dict, IO, Union
@@ -21,6 +20,7 @@
 import torch
 from fsspec.core import url_to_fs
 from fsspec.implementations.local import AbstractFileSystem
+from lightning_utilities.core.imports import module_available
 
 from lightning.fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
 
@@ -70,3 +70,54 @@ def _atomic_save(checkpoint: Dict[str, Any], filepath: Union[str, Path]) -> None
     torch.save(checkpoint, bytesbuffer)
     with fsspec.open(filepath, "wb") as f:
         f.write(bytesbuffer.getvalue())
+
+
+def _is_object_storage(fs: AbstractFileSystem) -> bool:
+    if module_available("adlfs"):
+        from adlfs import AzureBlobFileSystem
+
+        if isinstance(fs, AzureBlobFileSystem):
+            return True
+
+    if module_available("gcsfs"):
+        from adlfs import GCSFileSystem
+
+        if isinstance(fs, GCSFileSystem):
+            return True
+
+    if module_available("s3fs"):
+        from s3fs import S3FileSystem
+
+        if isinstance(fs, S3FileSystem):
+            return True
+
+    return False
+
+
+def _is_dir(fs: AbstractFileSystem, path: Union[str, Path], strict: bool = False) -> bool:
+    """Check if a path is directory-like.
+
+    This function determines if a given path is considered directory-like, taking into account the behavior
+    specific to object storage platforms. For other filesystems, it behaves similarly to the standard `fs.isdir`
+    method.
+
+    Args:
+        fs: The filesystem to check the path against.
+        path: The path or URL to be checked.
+        strict: A flag specific to Object Storage platforms. If set to ``False``, any non-existing path is considered
+            as a valid directory-like path. In such cases, the directory (and any non-existing parent directories)
+            will be created on the fly. Defaults to False.
+    """
+    # Object storage fsspec's are inconsistent with other file systems because they do not have real directories,
+    # see for instance https://gcsfs.readthedocs.io/en/latest/api.html?highlight=makedirs#gcsfs.core.GCSFileSystem.mkdir
+    # In particular, `fs.makedirs` is a no-op so we use `strict=False` to consider any path as valid, except if the
+    # path already exists but is a file
+    if _is_object_storage(fs):
+        if strict:
+            return fs.isdir(path)
+
+        # Check if the path is not already taken by a file. If not, it is considered a valid directory-like path
+        # because the directory (and all non-existing parent directories) will be created on the fly.
+        return not fs.isfile(path)
+
+    return fs.isdir(path)

src/lightning/pytorch/CHANGELOG.md

@@ -174,6 +174,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - `LightningCLI` not saving correctly `seed_everything` when `run=True` and `seed_everything=True` ([#18056](https://github.com/Lightning-AI/lightning/pull/18056))
 
+- Fixed a `Missing folder` exception when using a Google Storage URL as a `default_root_dir` ([#18088](https://github.com/Lightning-AI/lightning/pull/18088))
 
 - Fixed an issue that prevented the use of custom logger classes without an `experiment` property defined ([#18093](https://github.com/Lightning-AI/lightning/pull/18093))
 

src/lightning/pytorch/callbacks/model_checkpoint.py

@@ -32,7 +32,7 @@
 from torch import Tensor
 
 import lightning.pytorch as pl
-from lightning.fabric.utilities.cloud_io import get_filesystem
+from lightning.fabric.utilities.cloud_io import _is_dir, get_filesystem
 from lightning.fabric.utilities.types import _PATH
 from lightning.pytorch.callbacks import Checkpoint
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
@@ -620,7 +620,7 @@ def _find_last_checkpoints(self, trainer: "pl.Trainer") -> Set[str]:
         return set()
 
     def __warn_if_dir_not_empty(self, dirpath: _PATH) -> None:
-        if self.save_top_k != 0 and self._fs.isdir(dirpath) and len(self._fs.ls(dirpath)) > 0:
+        if self.save_top_k != 0 and _is_dir(self._fs, dirpath, strict=True) and len(self._fs.ls(dirpath)) > 0:
             rank_zero_warn(f"Checkpoint directory {dirpath} exists and is not empty.")
 
     def _get_metric_interpolated_filepath_name(

src/lightning/pytorch/core/saving.py

@@ -30,6 +30,7 @@
 from lightning_utilities.core.apply_func import apply_to_collection
 
 import lightning.pytorch as pl
+from lightning.fabric.utilities.cloud_io import _is_dir
 from lightning.fabric.utilities.cloud_io import _load as pl_load
 from lightning.fabric.utilities.cloud_io import get_filesystem
 from lightning.fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
@@ -252,7 +253,7 @@ def load_hparams_from_tags_csv(tags_csv: _PATH) -> Dict[str, Any]:
 
 def save_hparams_to_tags_csv(tags_csv: _PATH, hparams: Union[dict, Namespace]) -> None:
     fs = get_filesystem(tags_csv)
-    if not fs.isdir(os.path.dirname(tags_csv)):
+    if not _is_dir(fs, os.path.dirname(tags_csv)):
         raise RuntimeError(f"Missing folder: {os.path.dirname(tags_csv)}.")
 
     if isinstance(hparams, Namespace):
@@ -306,7 +307,7 @@ def save_hparams_to_yaml(config_yaml: _PATH, hparams: Union[dict, Namespace], us
 
     """
     fs = get_filesystem(config_yaml)
-    if not fs.isdir(os.path.dirname(config_yaml)):
+    if not _is_dir(fs, os.path.dirname(config_yaml)):
         raise RuntimeError(f"Missing folder: {os.path.dirname(config_yaml)}.")
 
     # convert Namespace or AD to dict

src/lightning/pytorch/loggers/tensorboard.py

@@ -26,6 +26,7 @@
 import lightning.pytorch as pl
 from lightning.fabric.loggers.tensorboard import _TENSORBOARD_AVAILABLE, _TENSORBOARDX_AVAILABLE
 from lightning.fabric.loggers.tensorboard import TensorBoardLogger as FabricTensorBoardLogger
+from lightning.fabric.utilities.cloud_io import _is_dir
 from lightning.fabric.utilities.logger import _convert_params
 from lightning.fabric.utilities.types import _PATH
 from lightning.pytorch.callbacks import ModelCheckpoint
@@ -217,7 +218,7 @@ def save(self) -> None:
         hparams_file = os.path.join(dir_path, self.NAME_HPARAMS_FILE)
 
         # save the metatags file if it doesn't exist and the log directory exists
-        if self._fs.isdir(dir_path) and not self._fs.isfile(hparams_file):
+        if _is_dir(self._fs, dir_path) and not self._fs.isfile(hparams_file):
             save_hparams_to_yaml(hparams_file, self.hparams)
 
     @rank_zero_only
@@ -248,7 +249,7 @@ def _get_next_version(self) -> int:
         for listing in listdir_info:
             d = listing["name"]
             bn = os.path.basename(d)
-            if self._fs.isdir(d) and bn.startswith("version_"):
+            if _is_dir(self._fs, d) and bn.startswith("version_"):
                 dir_ver = bn.split("_")[1].replace("/", "")
                 existing_versions.append(int(dir_ver))
         if len(existing_versions) == 0:

src/lightning/pytorch/trainer/connectors/checkpoint_connector.py

@@ -23,7 +23,7 @@
 
 import lightning.pytorch as pl
 from lightning.fabric.plugins.environments.slurm import SLURMEnvironment
-from lightning.fabric.utilities.cloud_io import get_filesystem
+from lightning.fabric.utilities.cloud_io import _is_dir, get_filesystem
 from lightning.fabric.utilities.types import _PATH
 from lightning.pytorch.callbacks import ModelCheckpoint
 from lightning.pytorch.plugins.precision import MixedPrecisionPlugin
@@ -55,7 +55,7 @@ def _hpc_resume_path(self) -> Optional[str]:
         dir_path_hpc = self.trainer.default_root_dir
         dir_path_hpc = str(dir_path_hpc)
         fs, path = url_to_fs(dir_path_hpc)
-        if not fs.isdir(path):
+        if not _is_dir(fs, path):
             return None
         max_version = self.__max_ckpt_version_in_folder(dir_path_hpc, "hpc_ckpt_")
         if max_version is not None:

tests/tests_fabric/utilities/test_cloud_io.py

@@ -15,8 +15,9 @@
 
 import fsspec
 from fsspec.implementations.local import LocalFileSystem
+from fsspec.spec import AbstractFileSystem
 
-from lightning.fabric.utilities.cloud_io import get_filesystem
+from lightning.fabric.utilities.cloud_io import _is_dir, get_filesystem
 
 
 def test_get_filesystem_custom_filesystem():
@@ -32,3 +33,61 @@ class DummyFileSystem(LocalFileSystem):
 
 def test_get_filesystem_local_filesystem():
     assert isinstance(get_filesystem("tmpdir/tmp_file"), LocalFileSystem)
+
+
+def test_is_dir_with_local_filesystem(tmp_path):
+    fs = LocalFileSystem()
+    tmp_existing_directory = tmp_path
+    tmp_non_existing_directory = tmp_path / "non_existing"
+
+    assert _is_dir(fs, tmp_existing_directory)
+    assert not _is_dir(fs, tmp_non_existing_directory)
+
+
+def test_is_dir_with_object_storage_filesystem():
+    class MockAzureBlobFileSystem(AbstractFileSystem):
+        def isdir(self, path):
+            return path.startswith("azure://") and not path.endswith(".txt")
+
+        def isfile(self, path):
+            return path.startswith("azure://") and path.endswith(".txt")
+
+    class MockGCSFileSystem(AbstractFileSystem):
+        def isdir(self, path):
+            return path.startswith("gcs://") and not path.endswith(".txt")
+
+        def isfile(self, path):
+            return path.startswith("gcs://") and path.endswith(".txt")
+
+    class MockS3FileSystem(AbstractFileSystem):
+        def isdir(self, path):
+            return path.startswith("s3://") and not path.endswith(".txt")
+
+        def isfile(self, path):
+            return path.startswith("s3://") and path.endswith(".txt")
+
+    fsspec.register_implementation("azure", MockAzureBlobFileSystem, clobber=True)
+    fsspec.register_implementation("gcs", MockGCSFileSystem, clobber=True)
+    fsspec.register_implementation("s3", MockS3FileSystem, clobber=True)
+
+    azure_directory = "azure://container/directory/"
+    azure_file = "azure://container/file.txt"
+    gcs_directory = "gcs://bucket/directory/"
+    gcs_file = "gcs://bucket/file.txt"
+    s3_directory = "s3://bucket/directory/"
+    s3_file = "s3://bucket/file.txt"
+
+    assert _is_dir(get_filesystem(azure_directory), azure_directory)
+    assert _is_dir(get_filesystem(azure_directory), azure_directory, strict=True)
+    assert not _is_dir(get_filesystem(azure_directory), azure_file)
+    assert not _is_dir(get_filesystem(azure_directory), azure_file, strict=True)
+
+    assert _is_dir(get_filesystem(gcs_directory), gcs_directory)
+    assert _is_dir(get_filesystem(gcs_directory), gcs_directory, strict=True)
+    assert not _is_dir(get_filesystem(gcs_directory), gcs_file)
+    assert not _is_dir(get_filesystem(gcs_directory), gcs_file, strict=True)
+
+    assert _is_dir(get_filesystem(s3_directory), s3_directory)
+    assert _is_dir(get_filesystem(s3_directory), s3_directory, strict=True)
+    assert not _is_dir(get_filesystem(s3_directory), s3_file)
+    assert not _is_dir(get_filesystem(s3_directory), s3_file, strict=True)