File size: 3,964 Bytes
33d23e1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | # SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0
"""Device-free coverage for the weight-load / JIT-compile observability primitives:
the wall-time ledger, the stall Watchdog, and the per-kernel-compile env opt-in."""
from __future__ import annotations
import importlib
import os
import sys
import time
import pytest
from loguru import logger
from ...utils import progress, walltime
from ...utils.progress import Watchdog
@pytest.fixture(autouse=True)
def _fresh_ledger():
"""The ledger is process-global; isolate each test."""
walltime._ledger = walltime._Ledger()
yield
def test_ledger_reconciles_to_wall():
walltime.record("weight_load", "vae", 2.0, cached=True)
walltime.record("gen", "denoise", 3.0)
out = walltime.render("t", wall=10.0)
# tracked (5.0) + untracked (5.0) must reconcile to the 10.0 wall we passed.
assert "TOTAL (tracked)" in out and "5.0" in out
assert "untracked" in out and "TOTAL (wall)" in out and "10.0" in out
assert "weight_load" in out and "gen" in out
def test_anomalies_flag_only_cache_misses():
walltime.record("weight_load", "hit_tensor", 1.0, cached=True)
assert "ANOMALIES: none" in walltime.render("t")
walltime.record("weight_load", "miss_tensor", 4.0, cached=False, detail="TT_DIT_CACHE_DIR unset")
with_miss = walltime.render("t")
assert "miss_tensor" in with_miss and "TT_DIT_CACHE_DIR unset" in with_miss
def test_atexit_emits_under_pytest(capsys):
"""Regression guard: the ledger must surface at teardown even though pytest is imported (the bug
was an unconditional `"pytest" in sys.modules` suppression that dropped it on the only LTX path)."""
assert "pytest" in sys.modules # we ARE under pytest — the exact suppressed condition
walltime.record("gen", "denoise", 1.0)
walltime._atexit()
assert "WALL-TIME LEDGER" in capsys.readouterr().out
def test_disabled_is_a_noop(monkeypatch):
monkeypatch.setattr(walltime, "_ENABLED", False)
walltime.record("gen", "denoise", 9.9)
assert not walltime._ledger.cats # nothing recorded when disabled
def test_watchdog_heartbeats_and_records_phase():
msgs: list[str] = []
sink = logger.add(lambda m: msgs.append(str(m)), level="INFO")
try:
with Watchdog("unit phase", interval=0.05):
time.sleep(0.14) # long enough for >=1 heartbeat tick
finally:
logger.remove(sink)
assert any("still working" in m for m in msgs), "watchdog emitted no heartbeat"
assert any("done in" in m for m in msgs), "watchdog logged no completion"
assert "phase" in walltime._ledger.cats # non-cache-load label is recorded on exit
def test_dit_model_import_opts_in(monkeypatch):
"""Importing a DiT model package runs models/tt_dit/models/__init__, which opts the process in."""
monkeypatch.delenv("TT_METAL_LOG_KERNEL_COMPILE", raising=False)
import models.tt_dit.models as dit_models
importlib.reload(dit_models) # re-run the package __init__ side effect deterministically
assert os.environ.get("TT_METAL_LOG_KERNEL_COMPILE") == "1"
def test_bare_util_import_stays_quiet(monkeypatch):
"""The gate must NOT flip on for code that only borrows a tt_dit util (e.g. Qwen36 imports
tt_dit.utils.tensor) — otherwise it defeats the build.cpp default-off gate. Only models/ opts in."""
monkeypatch.delenv("TT_METAL_LOG_KERNEL_COMPILE", raising=False)
importlib.reload(progress) # a representative tt_dit util — re-running it must not set the flag
assert "TT_METAL_LOG_KERNEL_COMPILE" not in os.environ
def test_explicit_opt_out_wins(monkeypatch):
"""An explicit =0 survives the model-import setdefault (a DiT run can still be silenced)."""
monkeypatch.setenv("TT_METAL_LOG_KERNEL_COMPILE", "0")
import models.tt_dit.models as dit_models
importlib.reload(dit_models)
assert os.environ.get("TT_METAL_LOG_KERNEL_COMPILE") == "0"
|