ray-project--ray/python/ray/train/tests/test_checkpoint.py

import os
from pathlib import Path

import pyarrow.fs
import pytest

import ray
from ray.train._checkpoint import (
    _CHECKPOINT_TEMP_DIR_PREFIX,
    _METADATA_FILE_NAME,
    Checkpoint,
    _get_del_lock_path,
    _list_existing_del_locks,
)
from ray.train._internal.storage import _exists_at_fs_path, _upload_to_fs_path
from ray.train.tests.test_new_persistence import _create_mock_custom_fs

_CHECKPOINT_CONTENT_FILE = "dummy.txt"


@pytest.fixture(params=["local", "mock", "custom_fs"])
def checkpoint(request, tmp_path):
    """Fixture that sets up a checkpoint on different filesystems."""
    checkpoint_fs_type = request.param

    checkpoint_path = tmp_path / "ckpt_dir"
    checkpoint_path.mkdir(exist_ok=True)
    (checkpoint_path / _CHECKPOINT_CONTENT_FILE).write_text("dummy")

    if checkpoint_fs_type == "local":
        yield Checkpoint.from_directory(str(checkpoint_path))
    elif checkpoint_fs_type == "mock":
        _checkpoint = Checkpoint(path="mock:///mock_bucket/ckpt_dir")
        _upload_to_fs_path(
            local_path=str(checkpoint_path),
            fs=_checkpoint.filesystem,
            fs_path=_checkpoint.path,
        )
        # The "mock://" URI doesn't persist across different instances of
        # the pyarrow.fs.MockFileSystem, so we must make sure to return
        # the checkpoint with the same filesystem instance that we uploaded
        # some mock content to.
        yield _checkpoint
    elif checkpoint_fs_type == "custom_fs":
        custom_storage_fs = _create_mock_custom_fs(tmp_path / "custom_fs")
        _upload_to_fs_path(
            local_path=str(checkpoint_path),
            fs=custom_storage_fs,
            fs_path="mock_bucket/ckpt_dir",
        )
        yield Checkpoint(path="mock_bucket/ckpt_dir", filesystem=custom_storage_fs)


def test_to_directory(checkpoint: Checkpoint):
    checkpoint_path = Path(checkpoint.to_directory())

    assert (checkpoint_path / _CHECKPOINT_CONTENT_FILE).exists()
    assert _CHECKPOINT_TEMP_DIR_PREFIX in checkpoint_path.name


def test_to_directory_with_user_specified_path(checkpoint: Checkpoint, tmp_path):
    # Test with a string
    checkpoint_path = Path(checkpoint.to_directory(str(tmp_path / "special_dir")))
    assert (checkpoint_path / _CHECKPOINT_CONTENT_FILE).exists()
    assert checkpoint_path.name == "special_dir"

    # Test with a PathLike
    checkpoint_path = Path(checkpoint.to_directory(tmp_path / "special_dir"))
    assert (checkpoint_path / _CHECKPOINT_CONTENT_FILE).exists()
    assert checkpoint_path.name == "special_dir"


def test_multiprocess_to_directory(checkpoint: Checkpoint):
    """Test the case where multiple processes are trying to checkpoint.

    Only one process should download the checkpoint, and the others should
    wait until it's finished. In the end, a single checkpoint dir should be
    shared by all processes.
    """
    if checkpoint.filesystem.type_name == "mock":
        pytest.skip("Mock filesystem cannot be pickled for use with Ray.")

    @ray.remote
    def download_checkpoint(checkpoint: Checkpoint) -> str:
        return checkpoint.to_directory()

    paths = [ray.get(download_checkpoint.remote(checkpoint)) for _ in range(5)]
    # Check that all the paths are the same (no duplicates).
    assert len(set(paths)) == 1


def test_as_directory(checkpoint: Checkpoint):
    with checkpoint.as_directory() as checkpoint_path:
        checkpoint_path = Path(checkpoint_path)
        assert (checkpoint_path / _CHECKPOINT_CONTENT_FILE).exists()

        if isinstance(checkpoint.filesystem, pyarrow.fs.LocalFileSystem):
            # We should have directly returned the local path
            assert str(checkpoint_path) == checkpoint.path
        else:
            # We should have downloaded to a temp dir.
            assert _CHECKPOINT_TEMP_DIR_PREFIX in checkpoint_path.name

    if isinstance(checkpoint.filesystem, pyarrow.fs.LocalFileSystem):
        # Checkpoint should not be deleted, if we directly gave the local path.
        assert (checkpoint_path / _CHECKPOINT_CONTENT_FILE).exists()
    else:
        # Should have been deleted, if the checkpoint downloaded to a temp dir.
        assert not checkpoint_path.exists()


def test_multiprocess_as_directory(checkpoint: Checkpoint, monkeypatch):
    """Tests that deletion lock files are created and deleted correctly,
    when multiple processes are accessing a checkpoint with `as_directory`"""
    # If this is pointing to a local path, no lockfiles will be created.
    is_local_checkpoint = isinstance(checkpoint.filesystem, pyarrow.fs.LocalFileSystem)

    monkeypatch.setattr(os, "getpid", lambda: 11111)

    with checkpoint.as_directory() as checkpoint_dir_1:
        lock_file_1 = Path(_get_del_lock_path(checkpoint_dir_1))

        # Pretend that the 2nd `as_directory` is called from a different process.
        monkeypatch.setattr(os, "getpid", lambda: 22222)

        with checkpoint.as_directory() as checkpoint_dir_2:
            lock_file_2 = Path(_get_del_lock_path(checkpoint_dir_2))

            # Should point both processes to the same canonical temp directory.
            assert checkpoint_dir_1 == checkpoint_dir_2

            if is_local_checkpoint:
                # Check that 2 different lock files were created.
                assert not lock_file_1.exists()
                assert not lock_file_2.exists()
            else:
                # Check that 2 different lock files were created.
                assert lock_file_1.exists()
                assert lock_file_2.exists()

                assert len(_list_existing_del_locks(checkpoint_dir_1)) == 2

        if not is_local_checkpoint:
            # Check that the 2nd lock file was deleted.
            assert len(_list_existing_del_locks(checkpoint_dir_1)) == 1
            assert not lock_file_2.exists()

    if not is_local_checkpoint:
        # Check that the 1st lock file was deleted.
        assert len(_list_existing_del_locks(checkpoint_dir_1)) == 0
        assert not lock_file_1.exists()

        # Check that the temp checkpoint directory was deleted.
        assert not Path(checkpoint_dir_1).exists()


def test_as_directory_lock_cleanup(checkpoint: Checkpoint):
    """Errors when accessing a checkpoint with `as_directory`
    shouldn't leave behind lock files.
    """
    with pytest.raises(RuntimeError):
        with checkpoint.as_directory() as checkpoint_dir:
            raise RuntimeError

    assert not _list_existing_del_locks(checkpoint_dir)

    is_local_checkpoint = isinstance(checkpoint.filesystem, pyarrow.fs.LocalFileSystem)
    if not is_local_checkpoint:
        assert not Path(checkpoint_dir).exists()


def test_as_directory_download_error(checkpoint: Checkpoint, monkeypatch):
    """Errors during a checkpoint download should be raised directly when accessing
    it with the `as_directory` context manager."""
    if isinstance(checkpoint.filesystem, pyarrow.fs.LocalFileSystem):
        pytest.skip(
            "Local filesystem checkpoints don't download to a temp dir, so "
            "there's no error handling to test."
        )

    error_text = "original error"

    def to_directory_error(*args, **kwargs):
        raise RuntimeError(error_text)

    monkeypatch.setattr(checkpoint, "to_directory", to_directory_error)

    with pytest.raises(RuntimeError, match=error_text):
        with checkpoint.as_directory() as _:
            pass


def test_metadata(checkpoint: Checkpoint):
    assert checkpoint.get_metadata() == {}

    # No metadata file by default.
    assert not _exists_at_fs_path(
        fs=checkpoint.filesystem,
        fs_path=str(Path(checkpoint.path) / _METADATA_FILE_NAME),
    )

    checkpoint.update_metadata({"foo": "bar"})
    assert checkpoint.get_metadata() == {"foo": "bar"}

    checkpoint.update_metadata({"foo": "baz"})
    assert checkpoint.get_metadata() == {"foo": "baz"}

    checkpoint.update_metadata({"x": 1})
    assert checkpoint.get_metadata() == {"foo": "baz", "x": 1}

    # Set metadata completely resets the metadata.
    checkpoint.set_metadata({"y": [1, 2, 3]})
    assert checkpoint.get_metadata() == {"y": [1, 2, 3]}

    # There should be a metadata file in the checkpoint directory.
    assert _exists_at_fs_path(
        fs=checkpoint.filesystem,
        fs_path=str(Path(checkpoint.path) / _METADATA_FILE_NAME),
    )

    # Non JSON serializable metadata should raise an error.
    class Test:
        pass

    with pytest.raises(TypeError):
        checkpoint.set_metadata({"non_json_serializable": Test()})


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))