Fix FSDP remote checkpointing tests by using local:// instead of memory://

The `memory://` fsspec filesystem is process-local and cannot be used across separately-spawned ranks in distributed tests. Switch to `local://` which simulates a remote filesystem but is backed by a shared disk, and broadcast the path from rank 0.

Additionally:
- Skip FSDP GPU tests on Windows.
- Conditionally import `_fsspec_filesystem` in reader/writer selection tests to support older PyTorch versions.

TAG=agy
CONV=90d4b033-0f6c-4895-bf1c-4184200442cd

Author: zhixiangli

tests/tests_fabric/strategies/test_fsdp.py

@@ -499,17 +499,21 @@ def test_is_sharded_checkpoint_remote_memory():
 
 def test_distributed_checkpoint_reader_writer_selection(tmp_path):
     from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter
-    from torch.distributed.checkpoint._fsspec_filesystem import FsspecReader, FsspecWriter
 
     from lightning.fabric.strategies.fsdp import (
         _get_distributed_checkpoint_reader,
         _get_distributed_checkpoint_writer,
     )
 
+    # the local file-system writer/reader are always available
     assert isinstance(_get_distributed_checkpoint_writer(str(tmp_path)), FileSystemWriter)
     assert isinstance(_get_distributed_checkpoint_reader(str(tmp_path)), FileSystemReader)
-    assert isinstance(_get_distributed_checkpoint_writer("memory:///w/ckpt"), FsspecWriter)
-    assert isinstance(_get_distributed_checkpoint_reader("memory:///w/ckpt"), FsspecReader)
+
+    # the fsspec writer/reader live in a private torch module; skip the remote assertions if the
+    # installed torch does not ship it (keeps the oldest-supported-torch CI green)
+    fsspec_fs = pytest.importorskip("torch.distributed.checkpoint._fsspec_filesystem")
+    assert isinstance(_get_distributed_checkpoint_writer("memory:///w/ckpt"), fsspec_fs.FsspecWriter)
+    assert isinstance(_get_distributed_checkpoint_reader("memory:///w/ckpt"), fsspec_fs.FsspecReader)
 
 
 def test_distributed_checkpoint_save_load_remote_memory():
@@ -581,10 +585,10 @@ def test_load_raw_module_state_from_path_remote(monkeypatch):
     assert "via_lazy" not in called
 
 
-@RunIf(min_cuda_gpus=2, standalone=True)
+@RunIf(min_cuda_gpus=2, skip_windows=True, standalone=True)
 @pytest.mark.parametrize("state_dict_type", ["full", "sharded"])
-def test_fsdp_remote_memory_roundtrip(state_dict_type):
-    """Save to and load from an fsspec ``memory://`` URL for both state-dict types."""
+def test_fsdp_remote_memory_roundtrip(state_dict_type, tmp_path):
+    """Save to and load from an fsspec remote-style URL for both state-dict types."""
     from lightning.fabric import Fabric
 
     fabric = Fabric(accelerator="cuda", devices=2, strategy=FSDPStrategy(state_dict_type=state_dict_type))
@@ -595,7 +599,11 @@ def test_fsdp_remote_memory_roundtrip(state_dict_type):
     optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
     optimizer = fabric.setup_optimizers(optimizer)
 
-    url = "memory:///ckpt-" + state_dict_type
+    # `local://` is treated as a remote fsspec URL (it exercises the FsspecWriter/Reader path) but is
+    # backed by real shared disk visible to every rank. `memory://` cannot be used here because its
+    # store is process-local, so the separately-spawned ranks would not see each other's shards.
+    # Broadcast rank 0's path so all ranks save/load from the same location.
+    url = fabric.broadcast("local://" + str(tmp_path / ("ckpt-" + state_dict_type)))
     fabric.save(url, {"model": model, "optimizer": optimizer, "step": 1})
 
     loaded = {"model": model, "optimizer": optimizer, "step": 0}

tests/tests_fabric/utilities/test_cloud_io.py

@@ -182,7 +182,7 @@ def test_checkpoint_join_does_not_corrupt_remote_urls():
 def test_is_checkpoint_dir_local(tmp_path):
     d = tmp_path / "adir"
     d.mkdir()
-    f = tmp_path / "afile"
+    f = tmp_path / "a_file"
     f.write_text("x")
     assert _is_checkpoint_dir(d) is True
     assert _is_checkpoint_dir(f) is False

tests/tests_pytorch/strategies/test_fsdp.py

@@ -1004,10 +1004,10 @@ def test_pl_save_checkpoint_does_not_corrupt_remote_path(monkeypatch):
     assert captured["path"] == "gs://bucket/run/ckpt"
 
 
-@RunIf(min_cuda_gpus=2, standalone=True)
+@RunIf(min_cuda_gpus=2, skip_windows=True, standalone=True)
 @pytest.mark.parametrize("state_dict_type", ["full", "sharded"])
-def test_pl_fsdp_remote_memory_roundtrip(state_dict_type):
-    """Save to and load from an fsspec ``memory://`` URL via the Trainer for both state-dict types."""
+def test_pl_fsdp_remote_memory_roundtrip(state_dict_type, tmp_path):
+    """Save to and load from an fsspec remote-style URL via the Trainer for both state-dict types."""
     from lightning.pytorch import Trainer
     from lightning.pytorch.demos.boring_classes import BoringModel
 
@@ -1021,7 +1021,10 @@ def test_pl_fsdp_remote_memory_roundtrip(state_dict_type):
         logger=False,
     )
     trainer.fit(model)
-    url = "memory:///pl-ckpt-" + state_dict_type
+    # `local://` exercises the remote fsspec path but is backed by shared disk visible to all ranks;
+    # `memory://` is process-local and unusable across the separately-spawned ranks. Broadcast rank 0's
+    # path so every rank uses the same location.
+    url = trainer.strategy.broadcast("local://" + str(tmp_path / ("pl-ckpt-" + state_dict_type)))
     trainer.save_checkpoint(url)
     # loading back must not raise and must restore the path uncorrupted
     trainer.strategy.load_checkpoint(url)