noaa-goes-visualization/tests/test_gpu_watchdog.py

"""The watchdog's verdict, which decides whether a run is still worth waiting for.

These cases are the ones that have actually cost GPU hours on this project: a
finished run mistaken for a live one, a live run mistaken for a finished one,
and an unreachable host mistaken for either.
"""

import gpu_train


def runs(**entries):
    """{name: (step, psnr, age_minutes)} -> the shape _watchdog builds."""
    return {name: {"step": step, "psnr": psnr, "age": minutes * 60}
            for name, (step, psnr, minutes) in entries.items()}


def test_container_up_and_history_moving_is_training():
    state, lines = gpu_train.watchdog_verdict(
        ["suvi-run-123", "llama-build"], runs(v6=(1500, 46.2, 4)), now=0)
    assert state == "training"
    assert "step 1500" in lines[0] and "46.2" in lines[0]


def test_container_up_but_history_frozen_is_stalled():
    """The expensive case: the box looks busy and is doing nothing."""
    state, lines = gpu_train.watchdog_verdict(
        ["suvi-run-123"], runs(v6=(1500, 46.2, 90)), now=0, busy_age=90 * 60)
    assert state == "stalled"
    assert "STALLED" in lines[0] and "90 min" in lines[0]


def test_a_freshly_launched_run_is_not_a_stall():
    """A new arm has written nothing, so the newest history on the box belongs
    to the *previous* run and looks hours old.  Calling that a stall fires a
    false alarm on every single launch."""
    state, lines = gpu_train.watchdog_verdict(
        ["suvi-run-999"], runs(e1_attention=(5000, 47.46, 113)), now=0,
        busy_age=3 * 60)
    assert state == "training"
    assert "starting up" in lines[0]


def test_a_long_lived_container_with_frozen_history_still_stalls():
    """The startup exemption must not swallow a genuine wedge later on."""
    state, lines = gpu_train.watchdog_verdict(
        ["suvi-run-999"], runs(v6=(1500, 46.2, 200)), now=0, busy_age=400 * 60)
    assert state == "stalled"


def test_no_container_is_idle_however_recent_the_history():
    """An arm that just finished leaves a fresh history and a dead container;
    that is idle, not training, or the next arm never launches."""
    state, lines = gpu_train.watchdog_verdict(
        ["llama-build"], runs(e1_attention=(5000, 47.4, 1)), now=0)
    assert state == "idle"
    assert "idle" in lines[0] and "step 5000" in lines[0]


def test_unreachable_host_is_unknown_not_idle():
    """A dropped connection previously read as 'finished' and launched a second
    run on top of a live one.  Unknown must never be actionable."""
    state, lines = gpu_train.watchdog_verdict(None, runs(v6=(10, 40.0, 1)), now=0)
    assert state == "unknown"
    assert "unreachable" in lines[0]


def test_the_most_recently_written_run_is_the_one_reported():
    """Fifteen stale run directories sit beside the live one."""
    state, lines = gpu_train.watchdog_verdict(
        ["suvi-run-9"],
        runs(v5=(9000, 44.0, 6000), v6=(4000, 46.9, 2000), e1_relative=(750, 45.1, 3)),
        now=0)
    assert state == "training" and "e1_relative" in lines[0]


def test_container_up_with_no_history_yet_is_training():
    """Startup: the image is pulling and torch is importing.  Calling that idle
    would relaunch the arm that is already coming up."""
    state, _ = gpu_train.watchdog_verdict(["suvi-run-1"], {}, now=0)
    assert state == "training"


def test_empty_box_with_no_history_is_idle():
    state, _ = gpu_train.watchdog_verdict([], {}, now=0)
    assert state == "idle"


def test_stall_threshold_is_respected_at_its_boundary():
    just_under = gpu_train.watchdog_verdict(
        ["suvi-run-1"], runs(v6=(100, 45.0, 29)), now=0, stall_minutes=30)
    just_over = gpu_train.watchdog_verdict(
        ["suvi-run-1"], runs(v6=(100, 45.0, 31)), now=0, stall_minutes=30)
    assert just_under[0] == "training"
    assert just_over[0] == "stalled"


def test_a_run_without_a_psnr_still_reports():
    """The untrained row and crashed-before-eval rows have no val_psnr."""
    state, lines = gpu_train.watchdog_verdict(
        ["suvi-run-1"], {"v6": {"step": 0, "psnr": None, "age": 60}}, now=0)
    assert state == "training" and "step 0" in lines[0]


def test_a_memory_starved_stall_names_the_cause():
    """The +/-256 run held its process runnable at high CPU while it swap-thrashed
    for 101 minutes.  Every signal said 'training' except the absence of output,
    so the watchdog must say which of the two it is."""
    state, lines = gpu_train.watchdog_verdict(
        ["suvi-run-1"], runs(v7_s256=(41000, 47.47, 101)), now=0,
        busy_age=1545 * 60, free_gb=0.2)
    assert state == "stalled"
    assert "STALLED" in lines[0]
    assert "thrashing" in lines[1] and "0.2 GB" in lines[1]


def test_low_memory_warns_before_it_becomes_a_stall():
    """Memory runs out gradually; the warning is only useful while there is still
    time to react."""
    state, lines = gpu_train.watchdog_verdict(
        ["suvi-run-1"], runs(v7=(1000, 46.0, 3)), now=0, busy_age=600,
        free_gb=1.5)
    assert state == "training"
    assert any("1.5 GB" in line for line in lines)


def test_healthy_memory_adds_no_noise():
    state, lines = gpu_train.watchdog_verdict(
        ["suvi-run-1"], runs(v7=(1000, 46.0, 3)), now=0, busy_age=600,
        free_gb=65.0)
    assert state == "training" and len(lines) == 1


def test_memory_unknown_is_not_treated_as_starved():
    """An older host, or a failed read, must not manufacture a false alarm."""
    state, lines = gpu_train.watchdog_verdict(
        ["suvi-run-1"], runs(v7=(1000, 46.0, 3)), now=0, busy_age=600,
        free_gb=None)
    assert state == "training" and len(lines) == 1


def test_progress_reports_epochs_and_steps_separately():
    """Epochs gauge real progress; steps are the short-term measure.  Both, always."""
    row = {"step": 27000, "epoch": 3, "steps_per_epoch": 8759, "total_steps": 57000}
    assert gpu_train.progress(row) == "epoch 3.08 of 6.51, step 27000 of 57000"


def test_progress_falls_back_to_the_epoch_index_for_older_histories():
    assert gpu_train.progress({"step": 27000, "epoch": 3}) == \
        "epoch 4 (in progress), step 27000"
    assert gpu_train.progress({"step": 0}) == "step 0"


def test_the_watchdog_line_carries_epochs_when_the_history_has_them():
    entry = {"step": 9000, "psnr": 47.1, "age": 60, "epoch": 1,
             "steps_per_epoch": 8759, "total_steps": 57000}
    state, lines = gpu_train.watchdog_verdict(["suvi-run-1"], {"v8": entry}, now=0)
    assert state == "training"
    assert "epoch 1.03 of 6.51" in lines[0] and "step 9000 of 57000" in lines[0]