ray-project--ray/python/ray/air/tests/test_experiment_restore.py

import json
import os
import shutil
import signal
import subprocess
import sys
import time
from pathlib import Path

import numpy as np
import pandas as pd
import pytest

from ray.tune.analysis import ExperimentAnalysis
from ray.tune.result_grid import ResultGrid

_RUN_SCRIPT_FILENAME = "_test_experiment_restore_run.py"


def _kill_process_if_needed(
    process: subprocess.Popen, timeout_s: float = 10, poll_interval_s: float = 1.0
):
    """Kills a process if it hasn't finished in `timeout_s` seconds.
    Polls every `poll_interval_s` seconds to check if the process is still running."""
    kill_timeout = time.monotonic() + timeout_s
    while process.poll() is None and time.monotonic() < kill_timeout:
        time.sleep(poll_interval_s)
    if process.poll() is None:
        process.terminate()


def _print_message(message):
    sep = "=" * 50
    print(f"\n{sep}\n{message}\n{sep}\n")


@pytest.mark.parametrize("runner_type", ["tuner", "trainer"])
def test_experiment_restore(tmp_path, runner_type):
    """
    This is an integration stress test for experiment restoration.


    Test setup:

    - For Tuner.restore:
        - 8 trials, with a max of 2 running concurrently (--> 4 rounds of trials)
        - Each iteration takes 0.5 seconds
        - Each trial runs for 8 iterations --> 4 seconds
        - Each round of 2 trials should take 4 seconds
        - Without any interrupts/restoration:
            - Minimum runtime: 4 rounds * 4 seconds / round = 16 seconds
        - The test will stop the script with a SIGINT at a random time between
        6-10 iterations each restore.

    - For Trainer.restore:
        - 1 trial with 4 workers
        - Each iteration takes 0.5 seconds
        - Runs for 32 iterations --> Minimum runtime = 16 seconds
        - The test will stop the script with a SIGINT at a random time between
        6-10 iterations after each restore.

    Requirements:
    - Req 1: Training progress persisted
        - The experiment should progress monotonically.
        (The training iteration shouldn't go backward at any point)
        - Trials shouldn't start from scratch.
    - Req 2: Searcher state saved/restored correctly
    - Req 3: Callback state saved/restored correctly
    """

    np.random.seed(2023)

    script_path = Path(__file__).parent / _RUN_SCRIPT_FILENAME

    # Args to pass into the script as environment variables
    exp_name = f"{runner_type}_restore_integration_test"
    callback_dump_file = tmp_path / f"{runner_type}-callback_dump_file.json"
    storage_path = tmp_path / "ray_results"
    if storage_path.exists():
        shutil.rmtree(storage_path)

    csv_file = str(tmp_path / "dummy_data.csv")
    dummy_df = pd.DataFrame({"x": np.arange(128), "y": 2 * np.arange(128)})
    dummy_df.to_csv(csv_file)

    run_started_marker = tmp_path / "run_started_marker"

    time_per_iter_s = 0.5
    max_concurrent = 2

    if runner_type == "tuner":
        iters_per_trial = 8
        num_trials = 8
    elif runner_type == "trainer":
        iters_per_trial = 32
        num_trials = 1

    total_iters = iters_per_trial * num_trials

    env = os.environ.copy()
    env.update(
        {
            "RUNNER_TYPE": runner_type,
            "STORAGE_PATH": str(storage_path),
            "EXP_NAME": exp_name,
            "CALLBACK_DUMP_FILE": str(callback_dump_file),
            "RUN_STARTED_MARKER": str(run_started_marker),
            "TIME_PER_ITER_S": str(time_per_iter_s),
            "ITERATIONS_PER_TRIAL": str(iters_per_trial),
            "NUM_TRIALS": str(num_trials),
            "MAX_CONCURRENT_TRIALS": str(max_concurrent),
            "CSV_DATA_FILE": csv_file,
        }
    )

    # Variables used in the loop
    return_code = None
    total_runtime = 0
    run_iter = 0
    progress = 0
    progress_history = []

    poll_interval_s = 0.1
    test_start_time = time.monotonic()

    while True:
        run_started_marker.write_text("", encoding="utf-8")

        run = subprocess.Popen([sys.executable, script_path], env=env)
        run_iter += 1

        _print_message(f"Started run #{run_iter} w/ PID = {run.pid}")

        # Start the timer after the first trial has entered its training loop.
        while run.poll() is None and run_started_marker.exists():
            time.sleep(poll_interval_s)

        # If the run already finished, then exit immediately.
        if run.poll() is not None:
            return_code = run.poll()
            break

        timeout_s = np.random.uniform(6 * time_per_iter_s, 10 * time_per_iter_s)

        _print_message(
            "Training has started...\n"
            f"Interrupting after {timeout_s:.2f} seconds\n"
            f"Currently at {total_runtime:.2f} seconds"
        )

        # Sleep for a random amount of time, then stop the run.
        start_time = time.monotonic()
        time.sleep(timeout_s)
        total_runtime += time.monotonic() - start_time

        return_code = run.poll()
        if return_code is None:
            # Send "SIGINT" to stop the run
            _print_message(f"Sending SIGUSR1 to run #{run_iter} w/ PID = {run.pid}")
            run.send_signal(signal.SIGUSR1)

            # Make sure the process is stopped forcefully after a timeout.
            _kill_process_if_needed(run)
        else:
            _print_message("Run has already terminated!")
            break

        # Check up on the results.
        results = ResultGrid(ExperimentAnalysis(str(storage_path / exp_name)))
        iters = [result.metrics.get("training_iteration", 0) for result in results]
        progress = sum(iters) / total_iters
        progress_history.append(progress)
        _print_message(
            f"Number of trials = {len(results)}\n"
            f"% completion = {progress} ({sum(iters)} iters / {total_iters})\n"
            f"Currently at {total_runtime:.2f} seconds"
        )

    _print_message(
        f"Total number of restorations = {run_iter}\n"
        f"Total runtime = {total_runtime:.2f}\n"
        f"Return code = {return_code}"
    )
    test_end_time = time.monotonic()

    assert progress == 1.0

    # The script shouldn't have errored. (It should have finished by this point.)
    assert return_code == 0, (
        f"The script errored with return code: {return_code}.\n"
        f"Check the `{_RUN_SCRIPT_FILENAME}` script for any issues. "
    )

    # Req 1: training progress persisted
    # Check that progress increases monotonically (we never go backwards/start from 0)
    assert np.all(np.diff(progress_history) >= 0), (
        "Expected progress to increase monotonically. Instead, got:\n"
        "{progress_history}"
    )

    # Req 2: searcher state
    results = ResultGrid(ExperimentAnalysis(str(storage_path / exp_name)))
    # Check that all trials have unique ids assigned by the searcher (if applicable)
    ids = [result.config.get("id", -1) for result in results]
    ids = [id for id in ids if id >= 0]
    if ids:
        assert sorted(ids) == list(range(1, num_trials + 1)), (
            "Expected the searcher to assign increasing id for each trial, but got:"
            f"{ids}"
        )

    # Req 3: callback state
    with open(callback_dump_file, "r") as f:
        callback_state = json.load(f)

    trial_iters = callback_state["trial_iters"]
    for iters in trial_iters.values():
        # Check that the callback has data for each trial, for all iters
        # NOTE: There may be some duplicate data, due to the fact that
        # the callback will be updated on every `on_trial_result` hook,
        # but the trial may crash before the corresponding checkpoint gets processed.
        assert sorted(set(iters)) == list(
            range(1, iters_per_trial + 1)
        ), f"Expected data from all iterations, but got: {iters}"

    _print_message(f"Success! Test took {test_end_time - test_start_time:.2f} seconds.")


if __name__ == "__main__":
    import pytest

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