ai-lab/m01_autograd/check.py

"""Run the AI01 lab tests one lab step at a time.

    python check.py            every step (1-8)
    python check.py 3          steps 1-3 only; use this while you work
    python check.py --stretch  every step plus the optional stretch test

A step "passes" only when all of its tests pass. When one fails, read the
test's name, re-read that lab step in AI01, and climb the step's hint ladder
one rung at a time.
"""
import argparse
import os
import sys
import unittest
from pathlib import Path

HERE = Path(__file__).resolve().parent
STEP_NAMES = {
    1: "forward pass", 2: "local backward", 3: "backward()", 4: "accumulation",
    5: "more operations", 6: "gradient checking", 7: "neural net", 8: "training",
}
# Which test module holds which steps.
MODULE_STEPS = {
    "test_engine": {1, 2, 3, 4, 5, 6},
    "test_gradcheck": {6},
    "test_nn": {7},
    "test_train": {8},
}


def step_of(test):
    """Step number for a test, from its class name (Step3Backward -> 3)."""
    test = getattr(test, "test_case", test)      # unwrap a failing subTest
    name = type(test).__name__
    if name.startswith("Step") and name[4:5].isdigit():
        return int(name[4])
    if name.startswith("Stretch"):
        return "stretch"
    return None


def flatten(suite):
    for item in suite:
        if isinstance(item, unittest.TestSuite):
            yield from flatten(item)
        else:
            yield item


def main():
    parser = argparse.ArgumentParser(description="Run the AI01 lab tests by step.")
    parser.add_argument("step", nargs="?", type=int, default=8, help="run steps 1..STEP (default 8)")
    parser.add_argument("--stretch", action="store_true", help="also run the stretch test")
    args = parser.parse_args()
    if not 1 <= args.step <= 8:
        parser.error("step must be between 1 and 8")
    if args.stretch:
        os.environ["AILAB_STRETCH"] = "1"

    wanted = set(range(1, args.step + 1))
    sys.path.insert(0, str(HERE / "tests"))
    loader = unittest.TestLoader()
    suite = unittest.TestSuite()
    for module, steps in MODULE_STEPS.items():
        if not steps & wanted:
            continue
        for test in flatten(loader.loadTestsFromName(module)):
            step = step_of(test)
            if step is None:                       # an import error: always show it
                suite.addTest(test)
            elif step == "stretch" and args.stretch:
                suite.addTest(test)
            elif step in wanted:
                suite.addTest(test)

    print(f"AI01 lab: running steps 1-{args.step}" + (" plus stretch" if args.stretch else "") + "\n")
    result = unittest.TextTestRunner(verbosity=2, stream=sys.stdout).run(suite)

    broken = {}
    for test, trace in result.failures + result.errors:
        step = step_of(test) or "import"
        not_written = "NotImplementedError" in trace
        broken.setdefault(step, []).append((test, not_written))

    print("\nSummary")
    for step in sorted(wanted):
        if step in broken:
            tests = broken[step]
            if all(nw for _, nw in tests):
                status = "not written yet"
            else:
                first = next(t for t, nw in tests if not nw)
                first = getattr(first, "test_case", first)
                status = f"FAIL  first failing test: {type(first).__name__}.{first._testMethodName}"
        else:
            status = "pass"
        print(f"  step {step} ({STEP_NAMES[step]}): {status}")
    if args.stretch:
        print(f"  stretch: {'FAIL' if 'stretch' in broken else 'pass'}")
    if "import" in broken:
        print("  A test file could not import your code. Read the traceback above first.")
    print(f"\n{result.testsRun} tests run, {len(result.failures)} failures, "
          f"{len(result.errors)} errors, {len(result.skipped)} skipped")
    sys.exit(0 if result.wasSuccessful() else 1)


if __name__ == "__main__":
    main()