Add tests

Author: leoleoasd

src/pytorch_lightning/trainer/connectors/checkpoint_connector.py

@@ -19,9 +19,9 @@
 from typing import Any, Dict, Optional
 
 import torch
+from fsspec.core import url_to_fs
 from torch import Tensor
 from torchmetrics import Metric
-from fsspec.core import uri_to_fs
 
 import pytorch_lightning as pl
 from lightning_lite.plugins.environments.slurm import SLURMEnvironment
@@ -575,7 +575,7 @@ def __max_ckpt_version_in_folder(dir_path: _PATH, name_key: str = "ckpt_") -> Op
         """
 
         # check directory existence
-        fs, uri = uri_to_fs(dir_path)
+        fs, uri = url_to_fs(dir_path)
         if not fs.exists(dir_path):
             return None
 

tests/tests_pytorch/trainer/connectors/test_checkpoint_connector.py

@@ -14,8 +14,10 @@
 import os
 from unittest import mock
 
+import fsspec.registry
 import pytest
 import torch
+from fsspec.implementations.arrow import ArrowFSWrapper
 
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint
@@ -102,6 +104,38 @@ def test_hpc_max_ckpt_version(tmpdir):
     )
 
 
+def test_max_ckpt_version_for_fsspec(tmpdir):
+    """Test that the CheckpointConnector is able to find the hpc checkpoint file with the highest version."""
+
+    class MockFileSystem(ArrowFSWrapper):
+        """A wrapper on top of the pyarrow.fs.HadoopFileSystem to connect it's interface with fsspec."""
+
+        protocol = "mock"
+
+        def __init__(self, **kwargs):
+            from pyarrow.fs import FileSystem
+
+            fs = FileSystem.from_uri("mock://")
+            super().__init__(fs=fs, **kwargs)
+
+    fsspec.registry.register_implementation("mock", MockFileSystem)
+
+    model = BoringModel()
+    trainer = Trainer(default_root_dir="mock://" + tmpdir, max_steps=1)
+    trainer.fit(model)
+    trainer.save_checkpoint(tmpdir / "hpc_ckpt.ckpt")
+    trainer.save_checkpoint(tmpdir / "hpc_ckpt_0.ckpt")
+    trainer.save_checkpoint(tmpdir / "hpc_ckpt_3.ckpt")
+    trainer.save_checkpoint(tmpdir / "hpc_ckpt_33.ckpt")
+
+    assert trainer._checkpoint_connector._hpc_resume_path == str(tmpdir / "hpc_ckpt_33.ckpt")
+    assert trainer._checkpoint_connector._CheckpointConnector__max_ckpt_version_in_folder(tmpdir) == 33
+    assert (
+        trainer._checkpoint_connector._CheckpointConnector__max_ckpt_version_in_folder(tmpdir / "not" / "existing")
+        is None
+    )
+
+
 def test_loops_restore(tmpdir):
     """Test that required loop state_dict is loaded correctly by checkpoint connector."""
     model = BoringModel()