| |
|
|
| |
|
|
| """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) |
| |
| 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 |
| 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 |
|
|
|
|
| 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) |
| 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 |
|
|
|
|
| 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) |
| 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) |
| 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" |
|
|
