import multiprocessing
import os
import shutil
import signal
import time
from unittest.mock import create_autospec

import pytest

import ray
import ray.cloudpickle as ray_pickle
from ray._common.test_utils import simulate_s3_bucket
from ray.air._internal.uri_utils import URI
from ray.tests.client_test_utils import create_remote_signal_actor
from ray.train import Checkpoint, CheckpointConfig, RunConfig, ScalingConfig
from ray.train.tests.util import create_dict_checkpoint, load_dict_checkpoint
from ray.train.v2.api.context import LocalTrainContext
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
from ray.train.v2.api.exceptions import WorkerGroupError
from ray.train.v2.api.report_config import (
    CheckpointConsistencyMode,
    CheckpointUploadMode,
)
from ray.train.v2.api.reported_checkpoint import ReportedCheckpointStatus
from ray.train.v2.api.validation_config import (
    ValidationConfig,
    ValidationTaskConfig,
)


@pytest.fixture(scope="module", autouse=True)
def ray_start_4_cpus():
    ray.init(num_cpus=4)
    yield
    ray.shutdown()


def test_report_mixed_checkpoint_upload_modes(tmp_path):
    """Run all 10 possible pairs (e.g. (SYNC, ASYNC)) of checkpoint upload modes between 2 workers."""

    def get_checkpoint_iteration(checkpoint):
        if not checkpoint:
            return -1
        return int(checkpoint.path.split("_")[-1])

    def train_fn():
        # When reporting with async checkpointing, write the checkpoint to
        # tmp_path, which stays alive for the duration of the test, instead of
        # tempfile.TemporaryDirectory(), which might get deleted before the
        # async checkpoint upload completes.

        # Run all 10 possible pairs of checkpoint upload modes
        rank = ray.train.get_context().get_world_rank()
        if rank == 0:
            ASYNC_ITERATIONS = [0, 1, 2, 3]
            SYNC_ITERATIONS = [4, 5, 6]
            NO_UPLOAD_ITERATIONS = [7, 8]
            NO_CHECKPOINT_ITERATIONS = [9]
        else:
            ASYNC_ITERATIONS = [0]
            SYNC_ITERATIONS = [1, 4]
            NO_UPLOAD_ITERATIONS = [2, 5, 7]
            NO_CHECKPOINT_ITERATIONS = [3, 6, 8, 9]

        prev_latest_checkpoint_iteration = -1
        for i in range(10):
            # Set variables
            if i in ASYNC_ITERATIONS:
                checkpoint_upload_mode = CheckpointUploadMode.ASYNC
            elif i in SYNC_ITERATIONS:
                checkpoint_upload_mode = CheckpointUploadMode.SYNC
            else:
                checkpoint_upload_mode = CheckpointUploadMode.NO_UPLOAD
            metrics = {"metric": f"iteration_{i}_shard_{rank}"}

            # Create and report checkpoint
            if i in NO_CHECKPOINT_ITERATIONS:
                ray.train.report(
                    metrics=metrics,
                    checkpoint=None,
                    validation=False,
                )
                assert prev_latest_checkpoint_iteration <= get_checkpoint_iteration(
                    ray.train.get_checkpoint()
                )
            else:
                # Create remote or local checkpoint_dir
                checkpoint_dir_name = f"checkpoint_iteration_{i}"
                if i in NO_UPLOAD_ITERATIONS:
                    checkpoint_dir = (
                        ray.train.get_context()
                        .get_storage()
                        .build_checkpoint_path_from_name(checkpoint_dir_name)
                    )
                else:
                    checkpoint_dir = os.path.join(
                        tmp_path, checkpoint_dir_name, f"_{rank}"
                    )

                # Create and report that remote or local checkpoint
                os.makedirs(checkpoint_dir, exist_ok=True)
                with open(os.path.join(checkpoint_dir, f"shard_{rank}"), "wb") as f:
                    ray_pickle.dump(f"iteration_{i}_shard_{rank}", f)
                checkpoint = Checkpoint(checkpoint_dir)
                ray.train.report(
                    metrics=metrics,
                    checkpoint=checkpoint,
                    checkpoint_upload_mode=checkpoint_upload_mode,
                    checkpoint_dir_name=checkpoint_dir_name,
                )

                # Check the status of latest_checkpoint
                latest_checkpoint = ray.train.get_checkpoint()
                if i in NO_UPLOAD_ITERATIONS:
                    assert latest_checkpoint == checkpoint
                elif i in SYNC_ITERATIONS:
                    assert checkpoint_dir_name in latest_checkpoint.path
                else:
                    assert prev_latest_checkpoint_iteration <= get_checkpoint_iteration(
                        latest_checkpoint
                    )

                prev_latest_checkpoint_iteration = get_checkpoint_iteration(
                    latest_checkpoint
                )

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=2),
        run_config=RunConfig(storage_path=str(tmp_path)),
    )
    result = trainer.fit()
    # Note that the (checkpoint=None, checkpoint=None) pair does not produce any checkpoint
    assert len(result.best_checkpoints) == 9
    for i, (checkpoint, metrics) in enumerate(result.best_checkpoints):
        assert checkpoint.path.endswith(f"checkpoint_iteration_{i}")
        assert metrics["metric"] == f"iteration_{i}_shard_0"


@pytest.mark.parametrize(
    "delete_local_checkpoint_after_upload,checkpoint_upload_mode",
    [
        (True, CheckpointUploadMode.ASYNC),
        (False, CheckpointUploadMode.ASYNC),
        (True, CheckpointUploadMode.SYNC),
        (False, CheckpointUploadMode.SYNC),
        (True, CheckpointUploadMode.NO_UPLOAD),
        (False, CheckpointUploadMode.NO_UPLOAD),
    ],
)
def test_report_delete_local_checkpoint_after_upload(
    tmp_path,
    delete_local_checkpoint_after_upload,
    checkpoint_upload_mode,
):
    """Check that the local checkpoint is deleted after upload."""

    def train_fn():
        rank = ray.train.get_context().get_world_rank()
        if rank == 0:
            if checkpoint_upload_mode == CheckpointUploadMode.NO_UPLOAD:
                checkpoint_dir = (
                    ray.train.get_context()
                    .get_storage()
                    .build_checkpoint_path_from_name("my_checkpoint_dir")
                )
            else:
                checkpoint_dir = os.path.join(
                    tmp_path,
                    "my_checkpoint_dir",
                )
            os.makedirs(checkpoint_dir, exist_ok=True)
            with open(os.path.join(checkpoint_dir, "shard_0"), "wb") as f:
                ray_pickle.dump("some_checkpoint_contents", f)
            checkpoint = Checkpoint(checkpoint_dir)
            ray.train.report(
                {},
                checkpoint,
                checkpoint_upload_mode=checkpoint_upload_mode,
                delete_local_checkpoint_after_upload=delete_local_checkpoint_after_upload,
            )
        else:
            ray.train.report(
                {},
                None,
            )

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=2),
        run_config=RunConfig(storage_path=str(tmp_path)),
    )
    trainer.fit()
    if (
        delete_local_checkpoint_after_upload
        or checkpoint_upload_mode == CheckpointUploadMode.NO_UPLOAD
    ):
        assert not os.path.exists(os.path.join(tmp_path, "my_checkpoint_dir"))
    else:
        assert os.path.exists(os.path.join(tmp_path, "my_checkpoint_dir"))


def test_report_checkpoint_upload_error(monkeypatch, tmp_path):
    """Check that the trainer shuts down when an error occurs during checkpoint upload."""

    def train_fn():

        if ray.train.get_context().get_world_rank() == 0:

            # Mock persist_current_checkpoint to raise an error
            mock_persist_current_checkpoint = create_autospec(
                ray.train.get_context().get_storage().persist_current_checkpoint
            )
            mock_persist_current_checkpoint.side_effect = ValueError("error")
            monkeypatch.setattr(
                ray.train.get_context().get_storage(),
                "persist_current_checkpoint",
                mock_persist_current_checkpoint,
            )

            # Report minimal valid checkpoint
            local_checkpoint_dir = os.path.join(tmp_path, "local_checkpoint_dir")
            os.makedirs(local_checkpoint_dir, exist_ok=True)
            ray.train.report(
                {},
                Checkpoint.from_directory(local_checkpoint_dir),
                checkpoint_upload_mode=CheckpointUploadMode.ASYNC,
            )
        else:
            ray.train.report(
                {}, None, checkpoint_upload_mode=CheckpointUploadMode.ASYNC
            )

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=2),
        run_config=RunConfig(storage_path=str(tmp_path)),
    )
    with pytest.raises(WorkerGroupError, match="error") as exc_info:
        trainer.fit()
    assert isinstance(exc_info.value.worker_failures[0], ValueError)


@pytest.mark.parametrize(
    "kwarg",
    (
        # both of these cases can cause the local checkpoint to be deleted after upload
        dict(delete_local_checkpoint_after_upload=True),
        dict(checkpoint_upload_mode=CheckpointUploadMode.ASYNC),
    ),
    ids=["delete_local_checkpoint_after_upload=True", "checkpoint_upload_mode=ASYNC"],
)
def test_report_checkpoint_delete_storage_path(kwarg, tmp_path):
    """Test that the trainer raises an error if the Checkpoint path is contains the storage_path."""
    # Test in `tmp_path` in case the test fails which means that the tmp_path.parent might be deleted
    base_dir = tmp_path / "test_base"
    storage_dir = base_dir / "storage"
    os.makedirs(storage_dir, exist_ok=True)

    def train_fn_equal_storage_path():
        ray.train.report(
            {},
            Checkpoint(str(storage_dir)),
            **kwarg,
        )

    def train_fn_within_storage_path():
        ray.train.report(
            {},
            Checkpoint(str(base_dir)),
            **kwarg,
        )

    for train_fn in [train_fn_equal_storage_path, train_fn_within_storage_path]:
        trainer = DataParallelTrainer(
            train_fn, run_config=RunConfig(storage_path=str(storage_dir))
        )
        with pytest.raises(WorkerGroupError, match="error") as exc_info:
            trainer.fit()
        assert isinstance(exc_info.value.worker_failures[0], ValueError)
        assert (
            exc_info.value.worker_failures[0]
            .args[0]
            .startswith("Ray Train's experiment directory")
        )
        # full error message:
        #   Ray Train's experiment directory (<file path>) is contained within the checkpoint path (<file path>)
        #   and `ray.train.report(delete_local_checkpoint_after_upload=True)`.
        #   As a result, this would delete the experiment directory.
        #   Please write the checkpoint to a subdirectory of the experiment directory
        #   or use `delete_local_checkpoint_after_upload=False`.


@pytest.mark.parametrize(
    "kwarg",
    (
        # both of these cases can cause the local checkpoint to be deleted after upload
        dict(delete_local_checkpoint_after_upload=True),
        dict(checkpoint_upload_mode=CheckpointUploadMode.ASYNC),
    ),
    ids=["delete_local_checkpoint_after_upload=True", "checkpoint_upload_mode=ASYNC"],
)
def test_report_checkpoint_delete_s3_storage_path(kwarg):
    """Test that the trainer raises an error if a s3 checkpoint path is contains a s3 storage_path."""
    port, region = 5002, "us-west-2"
    with simulate_s3_bucket(port=port, region=region) as s3_uri:
        import boto3

        s3 = boto3.client(
            "s3", region_name=region, endpoint_url=f"http://localhost:{port}"
        )
        # Bucket name will be autogenerated/unique per test
        bucket_name = URI(s3_uri).name
        s3.create_bucket(
            Bucket=bucket_name,
            CreateBucketConfiguration={"LocationConstraint": region},
        )
        # Use URI(s3_uri) / "storage" to correctly insert the path before query params.
        s3_storage_path = str(URI(s3_uri) / "storage")

        def train_fn_equal_storage_path():
            ray.train.report(
                {},
                Checkpoint(s3_storage_path),
                **kwarg,
            )

        def train_fn_within_storage_path():
            # s3_uri is the bucket root, which is a parent of s3_storage_path.
            ray.train.report(
                {},
                Checkpoint(s3_uri),
                **kwarg,
            )

        for train_fn in [train_fn_equal_storage_path, train_fn_within_storage_path]:
            trainer = DataParallelTrainer(
                train_fn, run_config=RunConfig(storage_path=s3_storage_path)
            )
            with pytest.raises(WorkerGroupError, match="error") as exc_info:
                trainer.fit()
            assert isinstance(exc_info.value.worker_failures[0], ValueError)
            assert (
                exc_info.value.worker_failures[0]
                .args[0]
                .startswith("Ray Train's experiment directory")
            )
            # full error message:
            #   Ray Train's experiment directory (<file path>) is contained within the checkpoint path (<file path>)
            #   and `ray.train.report(delete_local_checkpoint_after_upload=True)`.
            #   As a result, this would delete the experiment directory.
            #   Please write the checkpoint to a subdirectory of the experiment directory
            #   or use `delete_local_checkpoint_after_upload=False`.


def test_report_validation_without_validation_fn():
    def train_fn():
        with create_dict_checkpoint({}) as checkpoint:
            ray.train.report(metrics={}, checkpoint=checkpoint, validation=True)

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=1),
    )
    with pytest.raises(
        WorkerGroupError,
        match="`validation_config` was not set on the trainer, but a validation was requested.",
    ) as exc_info:
        trainer.fit()
    assert isinstance(exc_info.value.worker_failures[0], ValueError)


def test_report_validation_without_checkpoint():
    def train_fn():
        ray.train.report(metrics={}, validation=True)

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=1),
    )
    with pytest.raises(
        WorkerGroupError, match="Validation requires a checkpoint to be provided."
    ) as exc_info:
        trainer.fit()
    assert isinstance(exc_info.value.worker_failures[0], ValueError)


def test_report_validation_fn_keeps_correct_checkpoints(tmp_path):
    def validation_fn(checkpoint, new_score=None):
        if new_score:
            return {"score": new_score}
        else:
            return {}

    def train_fn():
        rank = ray.train.get_context().get_world_rank()
        checkpoint_dir = os.path.join(
            tmp_path,
            "my_checkpoint_dir",
        )
        os.makedirs(checkpoint_dir, exist_ok=True)
        with open(os.path.join(checkpoint_dir, f"shard_{rank}"), "wb") as f:
            ray_pickle.dump("some_checkpoint_contents", f)
        ray.train.report(
            metrics={"score": 1},
            checkpoint=Checkpoint(checkpoint_dir),
            checkpoint_upload_mode=CheckpointUploadMode.ASYNC,
            delete_local_checkpoint_after_upload=False,
            validation=ValidationTaskConfig(fn_kwargs={}),
        )
        with create_dict_checkpoint({}) as cp2:
            ray.train.report(
                metrics={"score": 3},
                checkpoint=cp2,
                checkpoint_upload_mode=CheckpointUploadMode.SYNC,
                validation=True,
            )
        with create_dict_checkpoint({}) as cp3:
            ray.train.report(
                metrics={"score": 2},
                checkpoint=cp3,
                checkpoint_upload_mode=CheckpointUploadMode.SYNC,
                validation=ValidationTaskConfig(fn_kwargs={"new_score": 5}),
            )

    trainer = DataParallelTrainer(
        train_fn,
        validation_config=ValidationConfig(fn=validation_fn),
        scaling_config=ScalingConfig(num_workers=2),
        run_config=RunConfig(
            storage_path=str(tmp_path),
            checkpoint_config=CheckpointConfig(
                num_to_keep=2, checkpoint_score_attribute="score"
            ),
        ),
    )
    result = trainer.fit()
    assert result.error is None
    assert result.checkpoint == result.best_checkpoints[1][0]
    assert len(result.best_checkpoints) == 2
    assert result.best_checkpoints[0][1] == {"score": 3}
    assert result.best_checkpoints[1][1] == {"score": 5}


@pytest.mark.parametrize("num_validation_workers", [0, 1])
def test_report_validation_fn_with_trainer_train_fn_report(num_validation_workers):
    """Test implementing the validation_fn with train_fn that reports metrics."""

    def eval_only_train_fn(config_dict):
        if isinstance(ray.train.get_context(), LocalTrainContext):
            checkpoint = config_dict["checkpoint"]
        else:
            checkpoint = ray.train.Checkpoint(
                ray.train.get_context()
                .get_storage()
                .build_checkpoint_path_from_name("placeholder")
            )

        ray.train.report(
            metrics={"validation": ray.train.get_context().get_world_rank()},
            checkpoint=checkpoint,
            checkpoint_upload_mode=CheckpointUploadMode.NO_UPLOAD,
        )

    def validation_fn(checkpoint: ray.train.Checkpoint):
        validation_trainer = DataParallelTrainer(
            eval_only_train_fn,
            train_loop_config={"checkpoint": checkpoint},
            scaling_config=ScalingConfig(num_workers=num_validation_workers),
        )
        validation_results = validation_trainer.fit()
        return validation_results.metrics

    def train_fn(config: dict):
        with create_dict_checkpoint({}) as cp:
            ray.train.report(
                metrics={"training": ray.train.get_context().get_world_rank()},
                checkpoint=cp,
                validation=True,
            )

    trainer = DataParallelTrainer(
        train_fn,
        validation_config=ValidationConfig(fn=validation_fn),
    )
    results = trainer.fit()
    assert results.error is None
    assert results.metrics == {"training": 0, "validation": 0}


@pytest.mark.parametrize("num_validation_workers", [0, 1])
def test_report_validation_fn_with_trainer_train_fn_return(num_validation_workers):
    """Test implementing the validation_fn with train_fn returns metrics."""

    def eval_only_train_fn(config_dict):
        return {"validation": ray.train.get_context().get_world_rank()}

    def validation_fn(checkpoint: ray.train.Checkpoint):
        validation_trainer = DataParallelTrainer(
            eval_only_train_fn,
            scaling_config=ScalingConfig(num_workers=num_validation_workers),
        )
        validation_results = validation_trainer.fit()
        return validation_results.return_value

    def train_fn(config: dict):
        with create_dict_checkpoint({}) as cp:
            ray.train.report(
                metrics={"training": ray.train.get_context().get_world_rank()},
                checkpoint=cp,
                validation=True,
            )

    trainer = DataParallelTrainer(
        train_fn,
        validation_config=ValidationConfig(fn=validation_fn),
    )
    results = trainer.fit()
    assert results.error is None
    assert results.metrics == {"training": 0, "validation": 0}
    assert results.return_value is None


def test_report_validation_fn_overrides_default_kwargs(tmp_path):
    def validation_fn(checkpoint, validation_score, other_key):
        return {"validation_score": validation_score, "other_key": other_key}

    def train_fn():
        with create_dict_checkpoint({}) as cp:
            ray.train.report(
                metrics={},
                checkpoint=cp,
                validation=ValidationTaskConfig(fn_kwargs={"validation_score": 2}),
            )

    trainer = DataParallelTrainer(
        train_fn,
        validation_config=ValidationConfig(
            fn=validation_fn,
            task_config=ValidationTaskConfig(
                fn_kwargs={"validation_score": 1, "other_key": "other_value"}
            ),
        ),
        run_config=RunConfig(storage_path=str(tmp_path)),
        scaling_config=ScalingConfig(num_workers=1),
    )
    result = trainer.fit()
    assert result.best_checkpoints[0][1] == {
        "validation_score": 2,
        "other_key": "other_value",
    }


def test_report_validation_fn_error(tmp_path):
    def validation_fn(checkpoint, rank=None, iteration=None):
        if rank == 0 and iteration == 0:
            raise ValueError("validation failed")
        return {"validation_score": iteration}

    def train_fn():
        rank = ray.train.get_context().get_world_rank()
        with create_dict_checkpoint({}) as cp1:
            ray.train.report(
                metrics={"training_score": 0},
                checkpoint=cp1,
                validation=ValidationTaskConfig(
                    fn_kwargs={"rank": rank, "iteration": 0}
                ),
            )
        with create_dict_checkpoint({}) as cp2:
            ray.train.report(
                metrics={"training_score": 1},
                checkpoint=cp2,
                validation=ValidationTaskConfig(
                    fn_kwargs={"rank": rank, "iteration": 1}
                ),
            )

        reported_checkpoints = ray.train.get_all_reported_checkpoints()
        assert len(reported_checkpoints) == 2
        assert (
            reported_checkpoints[0].status == ReportedCheckpointStatus.VALIDATION_FAILED
        )
        assert reported_checkpoints[0].metrics == {"training_score": 0}
        assert reported_checkpoints[1].status == ReportedCheckpointStatus.VALIDATED
        assert reported_checkpoints[1].metrics == {
            "training_score": 1,
            "validation_score": 1,
        }

    trainer = DataParallelTrainer(
        train_fn,
        validation_config=ValidationConfig(fn=validation_fn),
        scaling_config=ScalingConfig(num_workers=2),
        run_config=RunConfig(storage_path=str(tmp_path)),
    )
    result = trainer.fit()
    assert result.error is None
    assert len(result.best_checkpoints) == 2
    assert result.best_checkpoints[0][1] == {"training_score": 0}
    assert result.best_checkpoints[1][1] == {"training_score": 1, "validation_score": 1}


def test_report_validation_fn_timeout(tmp_path):
    def validation_fn(checkpoint):
        while True:
            time.sleep(1)

    def train_fn():
        with create_dict_checkpoint({}) as cp:
            ray.train.report(
                metrics={"training_score": 0}, checkpoint=cp, validation=True
            )

        reported_checkpoints = ray.train.get_all_reported_checkpoints()
        assert len(reported_checkpoints) == 1
        assert (
            reported_checkpoints[0].status
            == ReportedCheckpointStatus.VALIDATION_TIMEOUT
        )
        assert reported_checkpoints[0].metrics == {"training_score": 0}

    trainer = DataParallelTrainer(
        train_fn,
        validation_config=ValidationConfig(
            fn=validation_fn, task_config=ValidationTaskConfig(timeout_s=2)
        ),
        scaling_config=ScalingConfig(num_workers=1),
        run_config=RunConfig(storage_path=str(tmp_path)),
    )
    result = trainer.fit()
    assert result.error is None
    assert len(result.best_checkpoints) == 1
    assert result.best_checkpoints[0][1] == {"training_score": 0}


def test_report_validation_fn_success_after_retry():
    @ray.remote
    class Counter:
        def __init__(self):
            self.value = 0

        def increment(self):
            self.value += 1
            return self.value

    counter = Counter.remote()

    def validation_fn(checkpoint):
        if ray.get(counter.increment.remote()) < 2:
            raise ValueError("validation failed")
        return {"score": 100}

    def train_fn():
        with create_dict_checkpoint({}) as cp:
            ray.train.report(
                metrics={},
                checkpoint=cp,
                validation=True,
            )

        reported_checkpoints = ray.train.get_all_reported_checkpoints()
        assert len(reported_checkpoints) == 1
        assert reported_checkpoints[0].status == ReportedCheckpointStatus.VALIDATED

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=1),
        validation_config=ValidationConfig(
            fn=validation_fn,
            ray_remote_kwargs={"max_retries": 1, "retry_exceptions": [ValueError]},
        ),
    )
    result = trainer.fit()
    assert result.best_checkpoints[0][1] == {"score": 100}


def _run_first_trainer_for_resumption(storage_path, validation_task_config):
    """Subprocess target: run a trainer with a stalling validation, then get SIGINT'd."""
    # Lives outside the test because multiprocessing cannot pickle nested functions.
    ray.init(address="auto")

    def validation_fn_stall(checkpoint, score):
        signal_actor = ray.get_actor(
            "validation_resumption_signal", namespace="test_validation_resumption"
        )
        ray.get(signal_actor.send.remote())
        while True:
            time.sleep(1)

    def train_fn():
        with create_dict_checkpoint({}) as cp:
            ray.train.report(
                metrics={},
                checkpoint=cp,
                validation=validation_task_config,
            )

    trainer = DataParallelTrainer(
        train_fn,
        validation_config=ValidationConfig(
            fn=validation_fn_stall,
            task_config=ValidationTaskConfig(fn_kwargs={"score": 1}),
        ),
        scaling_config=ScalingConfig(num_workers=1),
        run_config=RunConfig(
            name="validation_fn_resumption", storage_path=storage_path
        ),
    )
    trainer.fit()


@pytest.mark.parametrize(
    "validation_task_config, expected_score",
    [
        (True, 1),
        (ValidationTaskConfig(fn_kwargs={"score": 2}), 2),
    ],
)
def test_report_validation_fn_resumption(
    tmp_path, validation_task_config, expected_score
):
    """A train_func call a validation_fn that stalls and the trainer is cancelled.
    Does the resumed trainer restart the validation?"""
    signal_actor = (
        create_remote_signal_actor(ray)
        .options(
            name="validation_resumption_signal",
            namespace="test_validation_resumption",
        )
        .remote()
    )

    multiprocessing.set_start_method("spawn", force=True)
    process = multiprocessing.Process(
        target=_run_first_trainer_for_resumption,
        args=(str(tmp_path), validation_task_config),
    )
    process.start()

    # Wait for validation to start, then SIGINT the trainer process.
    ray.get(signal_actor.wait.remote())
    os.kill(process.pid, signal.SIGINT)
    process.join()

    def validation_fn_finish(checkpoint, score):
        return {"score": score}

    def train_fn_second():
        rc = ray.train.get_all_reported_checkpoints(
            consistency_mode=CheckpointConsistencyMode.VALIDATED
        )
        assert len(rc) == 1
        assert rc[0].status == ReportedCheckpointStatus.VALIDATED
        assert rc[0].metrics == {"score": expected_score}

    # Run second trainer that should finish interrupted validations.
    trainer = DataParallelTrainer(
        train_fn_second,
        validation_config=ValidationConfig(
            fn=validation_fn_finish,
            task_config=ValidationTaskConfig(fn_kwargs={"score": 1}),
        ),
        scaling_config=ScalingConfig(num_workers=1),
        run_config=RunConfig(
            name="validation_fn_resumption", storage_path=str(tmp_path)
        ),
    )
    result = trainer.fit()
    assert result.metrics == {"score": expected_score}


@pytest.mark.parametrize(
    "validation_task_config, expected_score",
    [
        (True, 1),
        (ValidationTaskConfig(fn_kwargs={"score": 2}), 2),
    ],
)
def test_report_validation_fn_resumption_on_train_fn_error(
    tmp_path, validation_task_config, expected_score
):
    """Train run where train_fn fails after reporting a checkpoint with pending validation.
    The validation only returns after train_fn signals failure. before_controller_shutdown
    drains the validation, persisting the validated metrics. The second run sees them."""
    signal_actor = create_remote_signal_actor(ray).remote()

    def validation_fn(checkpoint, score):
        # Block until train_fn has signaled and sleep to ensure that the train_func has closed.
        ray.get(signal_actor.wait.remote())
        time.sleep(1)
        return {"score": score}

    def train_fn_first():
        with create_dict_checkpoint({}) as cp:
            ray.train.report(
                metrics={},
                checkpoint=cp,
                validation=validation_task_config,
            )
        try:
            raise RuntimeError("train_fn failed intentionally")
        finally:
            signal_actor.send.remote()

    def train_fn_second():
        rc = ray.train.get_all_reported_checkpoints(
            consistency_mode=CheckpointConsistencyMode.VALIDATED
        )
        assert len(rc) == 1
        assert rc[0].status == ReportedCheckpointStatus.VALIDATED
        assert rc[0].metrics == {"score": expected_score}

    run_config = RunConfig(
        name="validation_fn_resumption_on_train_fn_error",
        storage_path=str(tmp_path),
    )
    validation_config = ValidationConfig(
        fn=validation_fn,
        task_config=ValidationTaskConfig(fn_kwargs={"score": 1}),
    )

    with pytest.raises(WorkerGroupError):
        DataParallelTrainer(
            train_fn_first,
            validation_config=validation_config,
            run_config=run_config,
        ).fit()

    result = DataParallelTrainer(
        train_fn_second,
        validation_config=validation_config,
        run_config=run_config,
    ).fit()
    assert result.metrics == {"score": expected_score}


def test_report_validation_fn_resumption_checkpoint_status(tmp_path):
    def validation_fn(checkpoint, name):
        if name == "timeout":
            while True:
                time.sleep(1)
        elif name == "error":
            raise ValueError("validation error")
        else:
            return {"validation": name}

    def train_fn_first():
        with create_dict_checkpoint({}) as cp:
            ray.train.report(
                metrics={"score": 0},
                checkpoint=cp,
                validation=ValidationTaskConfig(fn_kwargs={"name": "success"}),
            )

        with create_dict_checkpoint({}) as cp:
            ray.train.report(
                metrics={"score": 1},
                checkpoint=cp,
                validation=ValidationTaskConfig(
                    fn_kwargs={"name": "timeout"}, timeout_s=1
                ),
            )

        with create_dict_checkpoint({}) as cp:
            ray.train.report(
                metrics={"score": 2},
                checkpoint=cp,
                validation=ValidationTaskConfig(fn_kwargs={"name": "error"}),
            )

        with create_dict_checkpoint({}) as cp:
            ray.train.report(
                metrics={"score": 3},
                checkpoint=cp,
                validation=ValidationTaskConfig(fn_kwargs={"name": "success"}),
            )

        reported_checkpoints = ray.train.get_all_reported_checkpoints()
        assert len(reported_checkpoints) == 4
        assert reported_checkpoints[0].status == ReportedCheckpointStatus.VALIDATED
        assert (
            reported_checkpoints[1].status
            == ReportedCheckpointStatus.VALIDATION_TIMEOUT
        )
        assert (
            reported_checkpoints[2].status == ReportedCheckpointStatus.VALIDATION_FAILED
        )
        assert reported_checkpoints[3].status == ReportedCheckpointStatus.VALIDATED
        assert reported_checkpoints[3].metrics == {"score": 3, "validation": "success"}

        raise RuntimeError("train_fn failed intentionally")

    def train_fn_second():
        reported_checkpoints = ray.train.get_all_reported_checkpoints()
        assert len(reported_checkpoints) == 4
        assert reported_checkpoints[0].status == ReportedCheckpointStatus.VALIDATED
        assert (
            reported_checkpoints[1].status
            == ReportedCheckpointStatus.VALIDATION_TIMEOUT
        )
        assert (
            reported_checkpoints[2].status == ReportedCheckpointStatus.VALIDATION_FAILED
        )
        assert reported_checkpoints[3].status == ReportedCheckpointStatus.VALIDATED

    with pytest.raises(WorkerGroupError):
        DataParallelTrainer(
            train_fn_first,
            run_config=RunConfig(
                "test-trainer-resumption-with-checkpoint-status",
                storage_path=str(tmp_path),
            ),
            validation_config=ValidationConfig(fn=validation_fn),
        ).fit()

    result = DataParallelTrainer(
        train_fn_second,
        run_config=RunConfig(
            "test-trainer-resumption-with-checkpoint-status", storage_path=str(tmp_path)
        ),
    ).fit()
    assert len(result.best_checkpoints) == 4


def test_multiple_workers_return_value_only_worker_zero():
    """Check that the `return_value` is of worker 0."""

    def train_fn():
        return (
            ray.train.get_context().get_world_size(),
            ray.train.get_context().get_world_rank(),
        )

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=3),
    )
    result = trainer.fit()
    assert result.return_value == (3, 0)


def test_report_checkpoint_upload_fn(tmp_path):
    def checkpoint_upload_fn(checkpoint, checkpoint_dir_name):
        full_checkpoint_path = (
            ray.train.get_context()
            .get_storage()
            .build_checkpoint_path_from_name(checkpoint_dir_name)
        )
        shutil.copytree(checkpoint.path, full_checkpoint_path)
        return Checkpoint.from_directory(full_checkpoint_path)

    def train_fn():
        if ray.train.get_context().get_world_rank() == 0:
            with create_dict_checkpoint(
                {"checkpoint_key": "checkpoint_value"}
            ) as checkpoint:
                ray.train.report(
                    metrics={},
                    checkpoint=checkpoint,
                    checkpoint_dir_name="my_checkpoint_dir_name",
                    checkpoint_upload_fn=checkpoint_upload_fn,
                )
        else:
            ray.train.report(metrics={}, checkpoint=None)

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=2),
        run_config=RunConfig(storage_path=str(tmp_path)),
    )
    result = trainer.fit()
    assert load_dict_checkpoint(result.checkpoint) == {
        "checkpoint_key": "checkpoint_value"
    }


def test_checkpoint_upload_fn_returns_checkpoint(tmp_path):
    def train_fn():
        with create_dict_checkpoint({}) as checkpoint:
            ray.train.report(
                metrics={},
                checkpoint=checkpoint,
                checkpoint_upload_fn=lambda x, y: None,
            )

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=1),
        run_config=RunConfig(storage_path=str(tmp_path)),
    )
    with pytest.raises(
        WorkerGroupError,
        match="checkpoint_upload_fn must return a `ray.train.Checkpoint`",
    ):
        trainer.fit()


def test_report_get_all_reported_checkpoints(tmp_path):
    """Check that get_all_reported_checkpoints returns checkpoints depending on # report calls."""

    def train_fn():
        if ray.train.get_context().get_world_rank() == 0:
            ray.train.report(metrics={}, checkpoint=None)
            with create_dict_checkpoint({}) as checkpoint:
                ray.train.report(metrics={}, checkpoint=checkpoint)

            reported_checkpoints = ray.train.get_all_reported_checkpoints()
            assert len(reported_checkpoints) == 1
            assert reported_checkpoints[0].status == ReportedCheckpointStatus.COMMITTED

            with create_dict_checkpoint({}) as checkpoint:
                ray.train.report(metrics={}, checkpoint=checkpoint)
        else:
            ray.train.report(metrics={}, checkpoint=None)
            ray.train.report(metrics={}, checkpoint=None)
            ray.train.report(metrics={}, checkpoint=None)

            reported_checkpoints = ray.train.get_all_reported_checkpoints()
            assert len(reported_checkpoints) == 2
            assert all(
                rc.status == ReportedCheckpointStatus.COMMITTED
                for rc in reported_checkpoints
            )

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=2),
        run_config=RunConfig(storage_path=str(tmp_path)),
    )
    trainer.fit()


def test_get_all_reported_checkpoints_all_consistency_modes(tmp_path):
    signal_actor = create_remote_signal_actor(ray).remote()

    def validation_fn(checkpoint, validation_score):
        ray.get(signal_actor.wait.remote())
        return {
            "validation_score": validation_score,
        }

    def train_fn(config):
        signal_actor = config["signal_actor"]

        if ray.train.get_context().get_world_rank() == 0:
            # Assert that we get committed checkpoints
            with create_dict_checkpoint({}) as cp1:
                # The validation check will hang until signal_actor.send.remote() is called
                ray.train.report(
                    metrics={"training_score": 1},
                    checkpoint=cp1,
                    validation=True,
                )

            # Check with ConsistencyMode.COMMITTED
            reported_checkpoints = ray.train.get_all_reported_checkpoints(
                consistency_mode=CheckpointConsistencyMode.COMMITTED
            )
            assert len(reported_checkpoints) == 1
            assert (
                reported_checkpoints[0].status
                == ReportedCheckpointStatus.PENDING_VALIDATION
            )
            assert reported_checkpoints[0].metrics == {"training_score": 1}

            # Check with ConsistencyMode.VALIDATED with timeout as the validation is hanging currently
            reported_checkpoints = ray.train.get_all_reported_checkpoints(
                consistency_mode=CheckpointConsistencyMode.VALIDATED,
                timeout_s=2,
            )
            assert len(reported_checkpoints) == 1
            assert (
                reported_checkpoints[0].status
                == ReportedCheckpointStatus.PENDING_VALIDATION
            )
            assert reported_checkpoints[0].metrics == {"training_score": 1}

            # Assert that we get validated checkpoints
            signal_actor.send.remote()
            reported_checkpoints = ray.train.get_all_reported_checkpoints(
                consistency_mode=CheckpointConsistencyMode.VALIDATED
            )
            assert len(reported_checkpoints) == 1
            assert reported_checkpoints[0].status == ReportedCheckpointStatus.VALIDATED
            assert reported_checkpoints[0].metrics == {
                "training_score": 1,
                "validation_score": 100,
            }
        else:
            ray.train.report(metrics={}, checkpoint=None)

    trainer = DataParallelTrainer(
        train_fn,
        validation_config=ValidationConfig(
            fn=validation_fn,
            task_config=ValidationTaskConfig(fn_kwargs={"validation_score": 100}),
        ),
        scaling_config=ScalingConfig(num_workers=2),
        train_loop_config={"signal_actor": signal_actor},
        run_config=RunConfig(storage_path=str(tmp_path)),
    )
    trainer.fit()


def test_hanging_checkpoint_upload_fn(tmp_path):
    """Test hanging async checkpoint upload fn with `get_all_reported_checkpoints` timeout."""
    signal_actor = create_remote_signal_actor(ray).remote()

    def checkpoint_fn(checkpoint, checkpoint_name):
        ray.get(signal_actor.wait.remote())
        full_checkpoint_path = (
            ray.train.get_context()
            .get_storage()
            .build_checkpoint_path_from_name(checkpoint_name)
        )
        shutil.copytree(checkpoint.path, full_checkpoint_path)
        return Checkpoint.from_directory(full_checkpoint_path)

    def train_fn(config):
        signal_actor = config["signal_actor"]

        with create_dict_checkpoint({}) as checkpoint:
            ray.train.report(
                metrics={"training_score": 1},
                checkpoint=checkpoint,
                checkpoint_upload_mode=CheckpointUploadMode.ASYNC,
                checkpoint_upload_fn=checkpoint_fn,
                checkpoint_dir_name="my_checkpoint",
            )

            # This will hang without a timeout.
            reported_checkpoints = ray.train.get_all_reported_checkpoints(
                consistency_mode=CheckpointConsistencyMode.COMMITTED, timeout_s=1
            )
            # As the checkpoint hasn't been committed yet, then the length is zero.
            assert len(reported_checkpoints) == 0

            signal_actor.send.remote()

            # Now the checkpoint has been completed, a timeout isn't necessary
            reported_checkpoints = ray.train.get_all_reported_checkpoints(
                consistency_mode=CheckpointConsistencyMode.COMMITTED
            )
            assert len(reported_checkpoints) == 1
            assert reported_checkpoints[0].status == ReportedCheckpointStatus.COMMITTED
            assert reported_checkpoints[0].metrics == {"training_score": 1}
            assert "my_checkpoint" in reported_checkpoints[0].checkpoint.path

    trainer = DataParallelTrainer(
        train_fn,
        train_loop_config={"signal_actor": signal_actor},
        run_config=RunConfig(
            storage_path=str(tmp_path),
            checkpoint_config=CheckpointConfig(num_to_keep=1),
        ),
    )
    trainer.fit()


def test_get_all_reported_checkpoints_empty_reports():
    def train_fn():
        ray.train.report(metrics={}, checkpoint=None)
        assert len(ray.train.get_all_reported_checkpoints()) == 0

    trainer = DataParallelTrainer(
        train_fn,
        scaling_config=ScalingConfig(num_workers=2),
    )
    trainer.fit()


if __name__ == "__main__":
    import sys

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