code/tests/test_api_device.py · changh95/hamer-p150 at main

hamer-p150 / code /tests /test_api_device.py

Python API (2026-10-04): pip install -e code/, from_pretrained() + model(...), Python-first quickstart

1b17257 verified about 1 month ago

8.11 kB

# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""Device tests of the Python API (needs one Blackhole chip as device 0, the released weights in the HF cache).
1. ``test_api_equals_server``: boots the HTTP server (``uvicorn hamer.server.app:app`` subprocess, the ``tt-model.yaml``
   serve env), takes ``/predict`` for the 8 demo hands of ``media/``, stops it, then
   builds ``HamerModel.from_pretrained(device_id=0)`` and requires ``hand.to_dict()`` == the server response (all keys
   except ``timing_ms`` / ``trace_active``) for every hand: the same numbers at the server's rounding.
2. ``test_api_call_time``: warm repeated calls, fresh crop every call (the 8 demo crops in turn), median / min of
   - ``model.tt(u8)``            the bare ``TtHamer.__call__`` (= the published "model call" / served call),
   - ``model.run_crop(u8)``      the same call through the API (lock + inference_mode),
   - ``model(frame, bbox)``      the full API call (host crop + resize, model call, post-processing),
   and requires ``run_crop`` not slower than the bare call (+3 %).  Results go to ``.work/api_timing.json``.
    source tools/chipenv.sh $CHIP models/hamer-p150/.venv
    cd models/hamer-p150 && timeout -s INT 1200 python -m pytest code/tests/test_api_device.py -q -s
"""
from __future__ import annotations

import base64
import json
import os
import statistics
import time
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
import torch  # noqa: E402
from PIL import Image  # noqa: E402

REPO = Path(__file__).resolve().parents[2]
MEDIA = REPO / "media"
PUBLISHED_MODEL_CALL_MS = (3.480, 3.499)  # OPT_REPORT.md round 15 headline, chip 5: served TtHamer.__call__ median
ITERS = int(os.environ.get("HAMER_API_ITERS", "200"))


def _hands():
    boxes = json.loads((MEDIA / "bboxes.json").read_text())
    return [(name, e["bbox"], bool(e.get("is_right", 1))) for name in sorted(boxes) for e in boxes[name]]


def _serve_env():
    from hamer.served_config import SERVED_DISPATCH, SERVED_ENV, WEIGHTS_REVISION

    env = dict(SERVED_ENV, HAMER_DISPATCH=SERVED_DISPATCH, HAMER_WARMUP_ITERS="3", TT_MESH_SHAPE="1x1",
               TT_DEVICE_ID="0", TT_WEIGHTS_REVISION=WEIGHTS_REVISION)
    for k, v in env.items():
        os.environ[k] = v
    return env


def _free_port() -> int:
    import socket

    with socket.socket() as so:
        so.bind(("127.0.0.1", 0))
        return so.getsockname()[1]


@pytest.fixture(scope="module")
def server_outputs():
    """The HTTP server as deployed (uvicorn subprocess, serve env), queried for every demo hand, then stopped."""
    import signal
    import subprocess
    import sys
    import urllib.request

    pytest.importorskip("fastapi")
    pytest.importorskip("uvicorn")
    env = dict(os.environ, **_serve_env())
    port = _free_port()
    log = REPO / ".work" / "api_test_server.log"
    log.parent.mkdir(exist_ok=True)
    with open(log, "w") as fh:
        proc = subprocess.Popen([sys.executable, "-m", "uvicorn", "--host", "127.0.0.1", "--port", str(port),
                                 "--lifespan", "on", "hamer.server.app:app"], env=env, stdout=fh, stderr=subprocess.STDOUT)
    url = f"http://127.0.0.1:{port}"

    def get(path, body=None):
        req = urllib.request.Request(url + path, data=None if body is None else json.dumps(body).encode(),
                                     headers={"Content-Type": "application/json"})
        return json.loads(urllib.request.urlopen(req, timeout=60).read())

    try:
        for _ in range(600):
            if proc.poll() is not None:
                raise RuntimeError(f"server exited ({proc.returncode}); see {log}")
            if "Application startup complete" in log.read_text():
                break
            time.sleep(1)
        else:
            raise RuntimeError(f"server did not become ready; see {log}")
        outs = {}
        for name, bbox, right in _hands():
            img = base64.b64encode((MEDIA / name).read_bytes()).decode()
            out = get("/predict", {"image": img, "bbox": bbox, "is_right": right})
            assert out["trace_active"] is True
            out.pop("timing_ms")
            out.pop("trace_active")
            outs[(name, tuple(bbox), right)] = out
        info = get("/info")
    finally:
        proc.send_signal(signal.SIGINT)
        try:
            proc.wait(timeout=120)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait()
    return outs, info


@pytest.fixture(scope="module")
def model(server_outputs):
    from hamer import HamerModel

    m = HamerModel.from_pretrained(device_id=0, verbose=True)
    yield m
    m.close()


def test_api_equals_server(server_outputs, model):
    outs, info = server_outputs
    assert model.trace_active and model.info["grid"] == "12x10" and model.info["dispatch"] == "eth"
    print(f"\n[api] info: {json.dumps({k: v for k, v in model.info.items() if k != 'fused'})}")
    print(f"[api] server warm-up: {info['warmup']}")
    n = 0
    for (name, bbox, right), want in outs.items():
        hand = model(MEDIA / name, bbox=list(bbox), is_right=right)
        got = hand.to_dict()
        diff = [k for k in want if got.get(k) != want[k]]
        assert not diff and set(got) == set(want), f"{name} {bbox}: keys differ from the server: {diff}"
        n += 1
    # the same hands through the list form (one decode per frame)
    by_frame = {}
    for (name, bbox, right) in outs:
        by_frame.setdefault(name, []).append((list(bbox), right))
    for name, items in by_frame.items():
        hands = model(MEDIA / name, bbox=[b for b, _ in items], is_right=[r for _, r in items])
        for (b, r), h in zip(items, hands):
            assert h.to_dict() == outs[(name, tuple(b), r)]
    print(f"[api] {n} demo hands: HamerModel output == server /predict output (all keys, server rounding)")


def _stats(v):
    return {"median_ms": round(statistics.median(v), 4), "min_ms": round(min(v), 4), "n": len(v)}


def test_api_call_time(model):
    from hamer.preprocess import preprocess_pil

    hands = _hands()
    frames = {name: Image.open(MEDIA / name).convert("RGB") for name, _, _ in hands}
    crops = [preprocess_pil(frames[n], b, is_right=r, raw=True)[0] for n, b, r in hands]
    tt = model.tt

    def loop(fn, iters=ITERS):
        ts = []
        for i in range(iters):
            t0 = time.perf_counter()
            fn(i)
            ts.append((time.perf_counter() - t0) * 1e3)
        return ts

    with torch.inference_mode():
        loop(lambda i: tt(crops[i % len(crops)]), 20)  # warm
    res = {}
    # alternate the two measurements (3 rounds each) so host-load drift hits both equally
    bare, api = [], []
    for _ in range(3):
        with torch.inference_mode():
            bare += loop(lambda i: tt(crops[i % len(crops)]))
        api += loop(lambda i: model.run_crop(crops[i % len(crops)]))
    res["bare_TtHamer_call"] = _stats(bare)
    res["api_run_crop"] = _stats(api)
    full, inner = [], []
    for i in range(ITERS):
        n, b, r = hands[i % len(hands)]
        t0 = time.perf_counter()
        h = model(frames[n], bbox=b, is_right=r)
        full.append((time.perf_counter() - t0) * 1e3)
        inner.append(h.timing_ms["model"])
    res["api_call_full_frame"] = _stats(full)
    res["api_call_full_frame_model_stage"] = _stats(inner)
    res["published_model_call_ms"] = PUBLISHED_MODEL_CALL_MS
    res["grid"], res["dispatch"] = model.info["grid"], model.info["dispatch"]
    res["chip"] = os.environ.get("TT_CHIP")
    res["host_load"] = [round(v, 1) for v in os.getloadavg()]
    print("\n[api] timing " + json.dumps(res))
    out = REPO / ".work"
    out.mkdir(exist_ok=True)
    (out / "api_timing.json").write_text(json.dumps(res, indent=2))
    assert res["api_run_crop"]["median_ms"] <= res["bare_TtHamer_call"]["median_ms"] * 1.03

	# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
	# SPDX-License-Identifier: Apache-2.0
	"""Device tests of the Python API (needs one Blackhole chip as device 0, the released weights in the HF cache).

	1. ``test_api_equals_server``: boots the HTTP server (``uvicorn hamer.server.app:app`` subprocess, the ``tt-model.yaml``
	serve env), takes ``/predict`` for the 8 demo hands of ``media/``, stops it, then
	builds ``HamerModel.from_pretrained(device_id=0)`` and requires ``hand.to_dict()`` == the server response (all keys
	except ``timing_ms`` / ``trace_active``) for every hand: the same numbers at the server's rounding.
	2. ``test_api_call_time``: warm repeated calls, fresh crop every call (the 8 demo crops in turn), median / min of
	- ``model.tt(u8)`` the bare ``TtHamer.__call__`` (= the published "model call" / served call),
	- ``model.run_crop(u8)`` the same call through the API (lock + inference_mode),
	- ``model(frame, bbox)`` the full API call (host crop + resize, model call, post-processing),
	and requires ``run_crop`` not slower than the bare call (+3 %). Results go to ``.work/api_timing.json``.

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

	import base64
	import json
	import os
	import statistics
	import time
	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
	import torch # noqa: E402
	from PIL import Image # noqa: E402

	REPO = Path(__file__).resolve().parents[2]
	MEDIA = REPO / "media"
	PUBLISHED_MODEL_CALL_MS = (3.480, 3.499) # OPT_REPORT.md round 15 headline, chip 5: served TtHamer.__call__ median
	ITERS = int(os.environ.get("HAMER_API_ITERS", "200"))


	def _hands():
	boxes = json.loads((MEDIA / "bboxes.json").read_text())
	return [(name, e["bbox"], bool(e.get("is_right", 1))) for name in sorted(boxes) for e in boxes[name]]


	def _serve_env():
	from hamer.served_config import SERVED_DISPATCH, SERVED_ENV, WEIGHTS_REVISION

	env = dict(SERVED_ENV, HAMER_DISPATCH=SERVED_DISPATCH, HAMER_WARMUP_ITERS="3", TT_MESH_SHAPE="1x1",
	TT_DEVICE_ID="0", TT_WEIGHTS_REVISION=WEIGHTS_REVISION)
	for k, v in env.items():
	os.environ[k] = v
	return env


	def _free_port() -> int:
	import socket

	with socket.socket() as so:
	so.bind(("127.0.0.1", 0))
	return so.getsockname()[1]


	@pytest.fixture(scope="module")
	def server_outputs():
	"""The HTTP server as deployed (uvicorn subprocess, serve env), queried for every demo hand, then stopped."""
	import signal
	import subprocess
	import sys
	import urllib.request

	pytest.importorskip("fastapi")
	pytest.importorskip("uvicorn")
	env = dict(os.environ, **_serve_env())
	port = _free_port()
	log = REPO / ".work" / "api_test_server.log"
	log.parent.mkdir(exist_ok=True)
	with open(log, "w") as fh:
	proc = subprocess.Popen([sys.executable, "-m", "uvicorn", "--host", "127.0.0.1", "--port", str(port),
	"--lifespan", "on", "hamer.server.app:app"], env=env, stdout=fh, stderr=subprocess.STDOUT)
	url = f"http://127.0.0.1:{port}"

	def get(path, body=None):
	req = urllib.request.Request(url + path, data=None if body is None else json.dumps(body).encode(),
	headers={"Content-Type": "application/json"})
	return json.loads(urllib.request.urlopen(req, timeout=60).read())

	try:
	for _ in range(600):
	if proc.poll() is not None:
	raise RuntimeError(f"server exited ({proc.returncode}); see {log}")
	if "Application startup complete" in log.read_text():
	break
	time.sleep(1)
	else:
	raise RuntimeError(f"server did not become ready; see {log}")
	outs = {}
	for name, bbox, right in _hands():
	img = base64.b64encode((MEDIA / name).read_bytes()).decode()
	out = get("/predict", {"image": img, "bbox": bbox, "is_right": right})
	assert out["trace_active"] is True
	out.pop("timing_ms")
	out.pop("trace_active")
	outs[(name, tuple(bbox), right)] = out
	info = get("/info")
	finally:
	proc.send_signal(signal.SIGINT)
	try:
	proc.wait(timeout=120)
	except subprocess.TimeoutExpired:
	proc.kill()
	proc.wait()
	return outs, info


	@pytest.fixture(scope="module")
	def model(server_outputs):
	from hamer import HamerModel

	m = HamerModel.from_pretrained(device_id=0, verbose=True)
	yield m
	m.close()


	def test_api_equals_server(server_outputs, model):
	outs, info = server_outputs
	assert model.trace_active and model.info["grid"] == "12x10" and model.info["dispatch"] == "eth"
	print(f"\n[api] info: {json.dumps({k: v for k, v in model.info.items() if k != 'fused'})}")
	print(f"[api] server warm-up: {info['warmup']}")
	n = 0
	for (name, bbox, right), want in outs.items():
	hand = model(MEDIA / name, bbox=list(bbox), is_right=right)
	got = hand.to_dict()
	diff = [k for k in want if got.get(k) != want[k]]
	assert not diff and set(got) == set(want), f"{name} {bbox}: keys differ from the server: {diff}"
	n += 1
	# the same hands through the list form (one decode per frame)
	by_frame = {}
	for (name, bbox, right) in outs:
	by_frame.setdefault(name, []).append((list(bbox), right))
	for name, items in by_frame.items():
	hands = model(MEDIA / name, bbox=[b for b, _ in items], is_right=[r for _, r in items])
	for (b, r), h in zip(items, hands):
	assert h.to_dict() == outs[(name, tuple(b), r)]
	print(f"[api] {n} demo hands: HamerModel output == server /predict output (all keys, server rounding)")


	def _stats(v):
	return {"median_ms": round(statistics.median(v), 4), "min_ms": round(min(v), 4), "n": len(v)}


	def test_api_call_time(model):
	from hamer.preprocess import preprocess_pil

	hands = _hands()
	frames = {name: Image.open(MEDIA / name).convert("RGB") for name, _, _ in hands}
	crops = [preprocess_pil(frames[n], b, is_right=r, raw=True)[0] for n, b, r in hands]
	tt = model.tt

	def loop(fn, iters=ITERS):
	ts = []
	for i in range(iters):
	t0 = time.perf_counter()
	fn(i)
	ts.append((time.perf_counter() - t0) * 1e3)
	return ts

	with torch.inference_mode():
	loop(lambda i: tt(crops[i % len(crops)]), 20) # warm
	res = {}
	# alternate the two measurements (3 rounds each) so host-load drift hits both equally
	bare, api = [], []
	for _ in range(3):
	with torch.inference_mode():
	bare += loop(lambda i: tt(crops[i % len(crops)]))
	api += loop(lambda i: model.run_crop(crops[i % len(crops)]))
	res["bare_TtHamer_call"] = _stats(bare)
	res["api_run_crop"] = _stats(api)
	full, inner = [], []
	for i in range(ITERS):
	n, b, r = hands[i % len(hands)]
	t0 = time.perf_counter()
	h = model(frames[n], bbox=b, is_right=r)
	full.append((time.perf_counter() - t0) * 1e3)
	inner.append(h.timing_ms["model"])
	res["api_call_full_frame"] = _stats(full)
	res["api_call_full_frame_model_stage"] = _stats(inner)
	res["published_model_call_ms"] = PUBLISHED_MODEL_CALL_MS
	res["grid"], res["dispatch"] = model.info["grid"], model.info["dispatch"]
	res["chip"] = os.environ.get("TT_CHIP")
	res["host_load"] = [round(v, 1) for v in os.getloadavg()]
	print("\n[api] timing " + json.dumps(res))
	out = REPO / ".work"
	out.mkdir(exist_ok=True)
	(out / "api_timing.json").write_text(json.dumps(res, indent=2))
	assert res["api_run_crop"]["median_ms"] <= res["bare_TtHamer_call"]["median_ms"] * 1.03