tests/test_progressive_training.py

from __future__ import annotations

import pytest

from adam.executor import ToolContext, ToolExecutionError
from adam.progressive_training import parse_stages, stage_batch_settings, suggested_stages
from adam.registry import ToolSpec
from adam.tools import ddpm_adapter


def _context(root):
    import threading

    running = threading.Event()
    running.set()
    return ToolContext(
        root=root, job_id="PROGRESSIVE", tool=ToolSpec("ddpm_trainer", "DDPM", "", "Training", "train_ddpm"),
        cancel_event=threading.Event(), run_event=running,
        progress_callback=lambda *_args, **_kwargs: None, log_callback=lambda _message: None,
    )


def test_progressive_stages_require_an_increasing_complete_budget() -> None:
    stages = parse_stages(
        [{"resolution": 64, "epochs": 30}, {"resolution": 128, "epochs": 10}],
        trainer="ddpm", total_epochs=40,
    )
    assert [(stage.resolution, stage.epochs) for stage in stages] == [(64, 30), (128, 10)]

    with pytest.raises(ToolExecutionError, match="total"):
        parse_stages(
            [{"resolution": 64, "epochs": 30}, {"resolution": 128, "epochs": 10}],
            trainer="ddpm", total_epochs=41,
        )
    with pytest.raises(ToolExecutionError, match="increase"):
        parse_stages(
            [{"resolution": 128, "epochs": 10}, {"resolution": 64, "epochs": 10}],
            trainer="ddpm", total_epochs=20,
        )


def test_auto_stage_batch_never_keeps_an_unsafe_final_batch() -> None:
    batch, accumulation = stage_batch_settings(
        trainer="flow", stage_resolution=512, final_resolution=512,
        final_batch_size=12, base_accumulation=1, auto_batch=True,
    )
    assert (batch, accumulation) == (1, 12)

    low_batch, low_accumulation = stage_batch_settings(
        trainer="flow", stage_resolution=64, final_resolution=512,
        final_batch_size=1, base_accumulation=12, auto_batch=True,
    )
    assert low_batch == 12
    assert low_accumulation == 1


def test_suggested_schedule_uses_the_entire_epoch_budget() -> None:
    stages = suggested_stages(512, 485)
    assert stages[0].resolution == 64
    assert stages[-1].resolution == 512
    assert sum(stage.epochs for stage in stages) == 485


def test_ddpm_curriculum_hands_weights_forward_and_publishes_only_final_model(tmp_path, monkeypatch) -> None:
    calls: list[tuple[int, str, str]] = []

    def fake_stage(_context, _dataset, _name, stage_epochs, stage_output, resume_from, *args):
        path = __import__("pathlib").Path(stage_output)
        path.mkdir(parents=True)
        (path / "model_index.json").write_text("{}", encoding="utf-8")
        calls.append((stage_epochs, str(path), str(resume_from)))
        return {}

    monkeypatch.setattr(ddpm_adapter, "_train_ddpm_stage", fake_stage)
    final_output = tmp_path / "output" / "Example"
    result = ddpm_adapter.train_ddpm(
        _context(tmp_path), dataset_dir=str(tmp_path), model_name="Example", epochs=5,
        output_dir=str(final_output), resolution=128, batch_size=1,
        progressive_stages=[{"resolution": 64, "epochs": 3}, {"resolution": 128, "epochs": 2}],
    )

    assert len(calls) == 2
    assert calls[0][2] == ""
    assert calls[1][2] == calls[0][1]
    assert (final_output / "model_index.json").is_file()
    assert result["output_folder"] == str(final_output)