Blame for code/tests/test_first_call_device.py

# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""Device test: the FIRST real call after ``HamerModel.from_pretrained`` runs at the steady speed (warm-up), and the
warm-up does not change any output.

Each case runs ``code/scripts/first_call.py`` in a FRESH process (one chip as device 0, the released weights in the HF
cache): ``from_pretrained``, then 30 real calls with new inputs (the demo hands; never the warm-up input).

- ``test_first_call_fast[variant]``: call 1 <= FIRST_CALL_TOL x the median of the later calls on the same demo hand
  (calls 9, 17, 25; the crop size, and so the host time, depends on the hand).  Before the warm-up API call 1 was
  1.29-1.55x for a JPEG file (right / left hand), after it 0.99-1.06x; see PYTHON.md "Warm-up".
- ``test_warmup_does_not_change_outputs``: the 30 outputs of ``warmup_variants="default"`` and ``"none"`` are bit-identical.

    source tools/chipenv.sh $CHIP models/hamer-p150/.venv
    cd models/hamer-p150 && timeout -s INT 900 python -m pytest code/tests/test_first_call_device.py -q -s
"""
from __future__ import annotations

import json
import os
import subprocess
import sys
from pathlib import Path

import pytest

if os.environ.get("TT_VISIBLE_DEVICES") == "none" or os.environ.get("HAMER_USE_TT") == "0":
    pytest.skip("device test (TT_VISIBLE_DEVICES=none / HAMER_USE_TT=0)", allow_module_level=True)

import numpy as np  # noqa: E402

REPO = Path(__file__).resolve().parents[2]
SCRIPT = REPO / "code" / "scripts" / "first_call.py"
# The residual is the cost of the first call after ANY idle gap (cold CPU caches), measured at 1.04-1.13x on the same
# hosts in a warm process (first_call.py "after_idle"); the tolerance leaves room for host-load noise on a shared host.
FIRST_CALL_TOL = float(os.environ.get("HAMER_FIRST_CALL_TOL", "1.20"))


def _run(variant: str, out: Path, warmup="default") -> dict:
    cmd = [sys.executable, str(SCRIPT), "--variant", variant, "--out", str(out),
           "--kwargs", json.dumps({"warmup_variants": warmup})]
    r = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
    assert r.returncode == 0, r.stdout[-3000:] + r.stderr[-3000:]
    return json.loads((out / f"{variant}.json").read_text())


@pytest.mark.parametrize("variant", ["right", "left", "run_crop"])
def test_first_call_fast(variant, tmp_path):
    res = _run(variant, tmp_path)
    print(f"\n[first_call] {variant}: from_pretrained {res['from_pretrained_s']} s, call1 {res['call1_ms']} ms, "
          f"call2 {res['call2_ms']} ms, call10 {res['call10_ms']} ms, steady {res['steady_median_ms']} ms, "
          f"same-hand {res['same_input_median_ms']} ms, after idle {res['after_idle_ms']} ms")
    assert res["warmup_variants"], "from_pretrained did not record its warm-up variants"
    assert res["call1_over_same_input"] <= FIRST_CALL_TOL, res


def test_warmup_does_not_change_outputs(tmp_path):
    a, b = tmp_path / "default", tmp_path / "none"
    _run("right", a, "default")
    res = _run("right", b, None)
    assert res["warmup_variants"] == []
    za, zb = np.load(a / "right.npz"), np.load(b / "right.npz")
    assert za.files == zb.files and len(za.files) == 30
    for k in za.files:
        assert np.array_equal(za[k], zb[k]), k