test_smoke.py

"""
Structural / integration smoke test.

This sandbox has no GPU and no network access to huggingface.co, so we
can't download a real checkpoint or launch the Triton kernel. To still
exercise the real code paths (not just "does it import"), this script:

  1. Builds a tiny, locally-constructed GPT2LMHeadModel + a minimal fake
     tokenizer (no network) and monkeypatches transformers' from_pretrained
     calls inside benchmark.py to return them regardless of model_name.
  2. Runs run_benchmark() for a handful of real configs (eager, compiled,
     fp16-on-cpu-fallback) and checks the results are well-formed.
  3. Exercises ExperimentAgent's heuristic proposal loop + select_best.
  4. Exercises kernels.benchmark_softmax_variants on a non-cuda device
  5. Exercises main.py's table formatting on the results produced above.

Run from the project root:
    python3 tests/test_smoke.py
"""
from __future__ import annotations

import os
import sys
import traceback

import torch
from transformers import GPT2Config, GPT2LMHeadModel

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import benchmark
import main as main_module
from agent import ExperimentAgent, CANDIDATE_GRID
from config import CompileMode, DType, ExperimentConfig
from kernels import benchmark_softmax_variants

FAILURES = []


def check(name, cond, detail=""):
    status = "PASS" if cond else "FAIL"
    print(f"[{status}] {name}" + (f" -- {detail}" if detail and not cond else ""))
    if not cond:
        FAILURES.append(name)


# 1. Fake HF backend (no network): tiny GPT-2 + minimal tokenizer stand-in
class FakeTokenizer:
    vocab_size = 200

    def __init__(self):
        self.pad_token = None
        self.eos_token = "<eos>"

_tiny_cfg = GPT2Config(
    vocab_size=200, n_positions=512, n_embd=32, n_layer=2, n_head=2, n_inner=64,
    bos_token_id=0, eos_token_id=0,
)
torch.manual_seed(1234)
_tiny_state_dict = GPT2LMHeadModel(_tiny_cfg).state_dict()


def _fake_model_from_pretrained(*args, **kwargs):
    m = GPT2LMHeadModel(_tiny_cfg)
    m.load_state_dict(_tiny_state_dict)
    return m

def _fake_tokenizer_from_pretrained(*args, **kwargs):
    return FakeTokenizer()

benchmark.AutoModelForCausalLM.from_pretrained = _fake_model_from_pretrained
benchmark.AutoTokenizer.from_pretrained = _fake_tokenizer_from_pretrained
benchmark._get_tokenizer.cache_clear()
benchmark._reference_cache.clear()

# 2. run_benchmark over a handful of real configs
print("\n=== benchmark.run_benchmark ===")
test_configs = [
    ExperimentConfig(batch_size=1, seq_len=16, dtype=DType.FP32, compile_mode=CompileMode.EAGER),
    ExperimentConfig(batch_size=2, seq_len=16, dtype=DType.FP16, compile_mode=CompileMode.EAGER),  # fp16 -> fallback to fp32 on CPU
    ExperimentConfig(batch_size=1, seq_len=16, dtype=DType.FP32, compile_mode=CompileMode.DEFAULT),
]

bench_results = []
for cfg in test_configs:
    try:
        r = benchmark.run_benchmark(cfg, model_name="tiny-fake", device="cpu", n_warmup=2, n_iters=3)
        bench_results.append(r)
        check(f"run_benchmark ok: {cfg.key()}", r.ok, r.error or "")
        if r.ok:
            check(f"  latency >= 0: {cfg.key()}", r.mean_latency_ms >= 0)
            check(f"  correctness computed: {cfg.key()}", r.top1_agreement is not None)
    except Exception:
        check(f"run_benchmark raised: {cfg.key()}", False, traceback.format_exc())

# 3. ExperimentAgent heuristic loop (no HF_TOKEN -> forced heuristic)
print("\n=== agent.ExperimentAgent (heuristic) ===")
import asyncio

async def run_agent_test():
    agent = ExperimentAgent(constraint="Minimize latency.", trial_budget=5, use_llm=False)
    check("agent.use_llm is False without HF_TOKEN", agent.use_llm is False)

    proposed_keys = set()
    for i in range(5):
        cfg = await agent.propose_next()
        check(f"propose_next returns a config (trial {i+1})", cfg is not None)
        if cfg is None:
            continue
        check(f"propose_next returns a NEW config (trial {i+1})", cfg.key() not in proposed_keys)
        proposed_keys.add(cfg.key())
        fake_result = benchmark.run_benchmark(cfg, model_name="tiny-fake", device="cpu", n_warmup=1, n_iters=2)
        agent.record(fake_result)

    check("candidate grid larger than budget", len(CANDIDATE_GRID) > 5)
    best = agent.select_best()
    check("select_best returns a result", best is not None)
    if best is not None:
        check("select_best result is ok", best.ok)
    return agent

agent = asyncio.run(run_agent_test())

# 4. kernels.benchmark_softmax_variants on CPU (should skip cleanly)
print("\n=== kernels.benchmark_softmax_variants (cpu) ===")
rows = benchmark_softmax_variants(device="cpu")
check("softmax variants returns skip message on cpu", rows and "error" in rows[0])

# 5. main.py table formatting on real results produced above
print("\n=== main.format_results_table / format_softmax_table ===")
try:
    table = main_module.format_results_table(bench_results + agent.results)
    check("format_results_table produces non-empty text", len(table) > 0)
    print(table)
except Exception:
    check("format_results_table raised", False, traceback.format_exc())

try:
    sm_table = main_module.format_softmax_table(rows)
    check("format_softmax_table produces non-empty text", len(sm_table) > 0)
    print(sm_table)
except Exception:
    check("format_softmax_table raised", False, traceback.format_exc())
 

# 6. config.py validation (pydantic constraints actually enforced)
print("\n=== config.ExperimentConfig validation ===")
try:
    ExperimentConfig(batch_size=0, seq_len=16, dtype=DType.FP32, compile_mode=CompileMode.EAGER)
    check("rejects batch_size=0", False)
except Exception:
    check("rejects batch_size=0", True)

try:
    ExperimentConfig(batch_size=1, seq_len=16, dtype="not-a-dtype", compile_mode=CompileMode.EAGER)
    check("rejects invalid dtype", False)
except Exception:
    check("rejects invalid dtype", True)

# ---------------------------------------------------------------------------
print("\n" + "=" * 60)
if FAILURES:
    print(f"{len(FAILURES)} CHECK(S) FAILED:")
    for f in FAILURES:
        print(f"  - {f}")
    sys.exit(1)
else:
    print("ALL CHECKS PASSED")