fix: use fsspec for DeepSpeed checkpoint path validation to support remote URIs

Fixes #21635. _validate_checkpoint_directory was wrapping paths in
pathlib.Path, which mangles remote URI schemes (e.g. s3:// to s3:/)
and uses local-only .is_dir()/.is_file() checks that always return
False for remote paths like S3, GCS, or HDFS.

Replace with get_filesystem() + fs.isdir()/fs.isfile() from
cloud_io, which is the established pattern used by ModelCheckpoint,
TorchCheckpointIO, and CheckpointConnector.

Author: bhimrazy

src/lightning/fabric/strategies/deepspeed.py

@@ -36,6 +36,7 @@
 from lightning.fabric.strategies.ddp import DDPStrategy
 from lightning.fabric.strategies.registry import _StrategyRegistry
 from lightning.fabric.strategies.strategy import _Sharded
+from lightning.fabric.utilities.cloud_io import get_filesystem
 from lightning.fabric.utilities.distributed import log
 from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_6
 from lightning.fabric.utilities.load import _move_state_into
@@ -45,6 +46,7 @@
 
 if TYPE_CHECKING:
     from deepspeed import DeepSpeedEngine
+    from fsspec import AbstractFileSystem
     from torch.optim.lr_scheduler import _LRScheduler
 
 _DEEPSPEED_AVAILABLE = RequirementCache("deepspeed")
@@ -885,9 +887,9 @@ def _validate_device_index_selection(parallel_devices: list[torch.device]) -> No
         )
 
 
-def _is_deepspeed_checkpoint(path: Path) -> bool:
+def _is_deepspeed_checkpoint(path: str, fs: "AbstractFileSystem") -> bool:
     """Heuristic check whether the path points to a top-level DeepSpeed checkpoint directory."""
-    return path.is_dir() and (path / "checkpoint").is_dir()
+    return fs.isdir(path) and fs.isdir(f"{path.rstrip('/')}/checkpoint")
 
 
 def _validate_checkpoint_directory(path: _PATH) -> None:
@@ -903,25 +905,28 @@ def _validate_checkpoint_directory(path: _PATH) -> None:
     # ├── latest
     # └── zero_to_fp32.py
 
-    path = Path(path)
-    path_is_ds_checkpoint = _is_deepspeed_checkpoint(path)
-    default_message = f"The provided path is not a valid DeepSpeed checkpoint: {path}"
+    path_str = str(path)
+    fs = get_filesystem(path_str)
+    path_is_ds_checkpoint = _is_deepspeed_checkpoint(path_str, fs)
+    default_message = f"The provided path is not a valid DeepSpeed checkpoint: {path_str}"
 
     if not path_is_ds_checkpoint:
         # Case 1: User may have accidentally passed the subfolder "checkpoint"
-        parent_is_ds_checkpoint = _is_deepspeed_checkpoint(path.parent)
-        if parent_is_ds_checkpoint:
+        parent = path_str.rstrip("/").rsplit("/", 1)[0] if "/" in path_str else ""
+        if parent and _is_deepspeed_checkpoint(parent, fs):
             raise FileNotFoundError(
                 f"{default_message}. It looks like you passed the path to a subfolder."
-                f" Try to load using this parent directory instead: {path.parent}"
+                f" Try to load using this parent directory instead: {parent}"
             )
         # Case 2: User may have accidentally passed the path to a file inside the "checkpoint" subfolder
-        parent_parent_is_ds_checkpoint = path.is_file() and _is_deepspeed_checkpoint(path.parent.parent)
-        if parent_parent_is_ds_checkpoint:
-            raise FileNotFoundError(
-                f"{default_message}. It looks like you passed the path to a file inside a DeepSpeed checkpoint folder."
-                f" Try to load using this parent directory instead: {path.parent.parent}"
-            )
+        if parent and fs.isfile(path_str):
+            grandparent = parent.rstrip("/").rsplit("/", 1)[0] if "/" in parent else ""
+            if grandparent and _is_deepspeed_checkpoint(grandparent, fs):
+                raise FileNotFoundError(
+                    f"{default_message}. It looks like you passed the path to a file inside a DeepSpeed"
+                    f" checkpoint folder."
+                    f" Try to load using this parent directory instead: {grandparent}"
+                )
         raise FileNotFoundError(default_message)
 
 

tests/tests_fabric/strategies/test_deepspeed.py

@@ -264,6 +264,40 @@ def test_deepspeed_load_checkpoint_validate_path(tmp_path):
         strategy.load_checkpoint(path=checkpoint_path, state={"model": Mock()})
 
 
+def test_validate_checkpoint_directory_remote_uri():
+    """Test that _validate_checkpoint_directory works with remote filesystem URIs (e.g., S3, HDFS)."""
+    from lightning.fabric.strategies.deepspeed import _validate_checkpoint_directory
+
+    mock_fs = Mock()
+    mock_fs.isdir = Mock(side_effect=lambda p: p in ("s3://bucket/ckpt", "s3://bucket/ckpt/checkpoint"))
+    mock_fs.isfile = Mock(return_value=False)
+
+    with mock.patch("lightning.fabric.strategies.deepspeed.get_filesystem", return_value=mock_fs):
+        # Should not raise when the remote path is a valid DeepSpeed checkpoint
+        _validate_checkpoint_directory("s3://bucket/ckpt")
+
+    # Verify the URI was NOT mangled (s3:// must stay as s3://, not s3:/)
+    mock_fs.isdir.assert_any_call("s3://bucket/ckpt")
+    mock_fs.isdir.assert_any_call("s3://bucket/ckpt/checkpoint")
+
+
+def test_validate_checkpoint_directory_remote_uri_subfolder_suggestion():
+    """Test that the subfolder suggestion works with remote URIs."""
+    from lightning.fabric.strategies.deepspeed import _validate_checkpoint_directory
+
+    mock_fs = Mock()
+    mock_fs.isdir = Mock(
+        side_effect=lambda p: p in ("s3://bucket/ckpt", "s3://bucket/ckpt/checkpoint"),
+    )
+    mock_fs.isfile = Mock(return_value=False)
+
+    with (
+        mock.patch("lightning.fabric.strategies.deepspeed.get_filesystem", return_value=mock_fs),
+        pytest.raises(FileNotFoundError, match="Try to load using this parent directory instead: s3://bucket/ckpt"),
+    ):
+        _validate_checkpoint_directory("s3://bucket/ckpt/checkpoint")
+
+
 @RunIf(deepspeed=True)
 def test_deepspeed_load_checkpoint_no_state(tmp_path):
     """Test that DeepSpeed can't load the full state without access to a model instance from the user."""