Repository: Lightning-AI/pytorch-lightning
File: tests/tests_pytorch/trainer/connectors/test_checkpoint_connector.py (Python)

Code before:
    assert trainer._checkpoint_connector._CheckpointConnector__max_ckpt_version_in_folder(tmp_path) == 33
    assert (
        trainer._checkpoint_connector._CheckpointConnector__max_ckpt_version_in_folder(tmp_path / "not" / "existing")
        is None
    )


def test_ckpt_for_fsspec(tmpdir):
    """Test that the _CheckpointConnector is able to write to fsspec file systems."""
    model = BoringModel()
    # hardcoding dir since `tmp_path` can be windows path
    trainer = Trainer(
        default_root_dir="memory://test_ckpt_for_fsspec", limit_train_batches=1, limit_val_batches=1, max_epochs=1
    )
    trainer.fit(model)
    trainer.save_checkpoint("memory://test_ckpt_for_fsspec/hpc_ckpt.ckpt")
    trainer.save_checkpoint("memory://test_ckpt_for_fsspec/hpc_ckpt_0.ckpt")
    trainer.save_checkpoint("memory://test_ckpt_for_fsspec/hpc_ckpt_3.ckpt")
    trainer.save_checkpoint("memory://test_ckpt_for_fsspec/hpc_ckpt_33.ckpt")

    # Simulate the behavior of fsspec when writing to a local file system but other device.
    with (
        mock.patch("os.rename", side_effect=OSError(errno.EXDEV, "Invalid cross-device link")),
        mock.patch("os.chmod", side_effect=PermissionError("Operation not permitted")),
    ):
        trainer.save_checkpoint(tmpdir + "/test_ckpt_for_fsspec/hpc_ckpt_18.ckpt")

    assert trainer._checkpoint_connector._hpc_resume_path == "memory://test_ckpt_for_fsspec/hpc_ckpt_33.ckpt"
    assert (
        trainer._checkpoint_connector._CheckpointConnector__max_ckpt_version_in_folder("memory://test_ckpt_for_fsspec")
        == 33
    )
    assert (
        trainer._checkpoint_connector._CheckpointConnector__max_ckpt_version_in_folder("memory://not_existing") is None
    )


def test_loops_restore(tmp_path):
    """Test that required loop state_dict is loaded correctly by checkpoint connector."""
    model = BoringModel()
    checkpoint_callback = ModelCheckpoint(dirpath=tmp_path, save_last=True)
    trainer_args = {
        "default_root_dir": tmp_path,
        "max_epochs": 1,
        "limit_train_batches": 1,
        "limit_val_batches": 1,
        "logger": False,
        "callbacks": [checkpoint_callback],
        "num_sanity_val_steps": 0,
    }
    trainer = Trainer(**trainer_args)

Review comment (suggestion):
I think we shall add a new test covering this case, which will be conditioned by `fsspec[http] >=2025.5.0`, but there is no reason to be such aggressive and deny other past versions

Code after:

    assert trainer._checkpoint_connector._hpc_resume_path == str(tmp_path / "hpc_ckpt_33.ckpt")
    assert trainer._checkpoint_connector._CheckpointConnector__max_ckpt_version_in_folder(tmp_path) == 33
    assert (
        trainer._checkpoint_connector._CheckpointConnector__max_ckpt_version_in_folder(tmp_path / "not" / "existing")
        is None
    )


def test_local_cross_device_checkpoint(tmpdir):
    """Test that the _CheckpointConnector can write local cross-device files or raises an error if fsspec<2025.5.0."""
    model = BoringModel()
    # hardcoding dir since `tmp_path` can be windows path
    trainer = Trainer(
        default_root_dir="memory://test_ckpt_for_fsspec", limit_train_batches=1, limit_val_batches=1, max_epochs=1
    )
    trainer.fit(model)
    # Simulate the behavior of fsspec when writing to a local file system but other device.
    with (
        mock.patch("os.rename", side_effect=OSError(errno.EXDEV, "Invalid cross-device link")),
        mock.patch("os.chmod", side_effect=PermissionError("Operation not permitted")),
    ):
        if fsspec.__version__ < "2025.5.0":
            with pytest.raises(
                RuntimeError,
                match=re.escape(
                    'Upgrade fsspec to enable cross-device local checkpoints: pip install "fsspec[http]>=2025.5.0"'
                ),
            ):
                trainer.save_checkpoint(tmpdir + "/test_ckpt_for_fsspec/hpc_ckpt.ckpt")
        else:
            trainer.save_checkpoint(tmpdir + "/test_ckpt_for_fsspec/hpc_ckpt.ckpt")


def test_ckpt_for_fsspec():
    """Test that the _CheckpointConnector is able to write to fsspec file systems."""
    model = BoringModel()
    # hardcoding dir since `tmp_path` can be windows path
    trainer = Trainer(
        default_root_dir="memory://test_ckpt_for_fsspec", limit_train_batches=1, limit_val_batches=1, max_epochs=1
    )
    trainer.fit(model)
    trainer.save_checkpoint("memory://test_ckpt_for_fsspec/hpc_ckpt.ckpt")
    trainer.save_checkpoint("memory://test_ckpt_for_fsspec/hpc_ckpt_0.ckpt")
    trainer.save_checkpoint("memory://test_ckpt_for_fsspec/hpc_ckpt_3.ckpt")
    trainer.save_checkpoint("memory://test_ckpt_for_fsspec/hpc_ckpt_33.ckpt")

    assert trainer._checkpoint_connector._hpc_resume_path == "memory://test_ckpt_for_fsspec/hpc_ckpt_33.ckpt"
    assert (
        trainer._checkpoint_connector._CheckpointConnector__max_ckpt_version_in_folder("memory://test_ckpt_for_fsspec")
        == 33