code/conftest.py

# SPDX-License-Identifier: Apache-2.0
"""pytest fixtures of bevformer-p150 (no dependency on tt-metal's own conftest).

- ``device`` (session): one chip opened like the published numbers (``tt_bevformer.device.open_device``:
  ETH dispatch, 12x10 grid, the port's validated sizes and CQs). ``--device-id N`` or ``TT_DEVICE_ID`` selects
  the chip (default 0); ``BEVFORMER_DISPATCH=worker`` is the A/B opt-in, and the other ``BEVFORMER_*``
  variables apply as for the server. A failing ETH open is an error here, never a silent WORKER fallback: gates
  are only valid on the published setup. The test chip gets a 1 GiB trace region (:data:`TEST_TRACE_REGION`): the
  debug variants of ``test_pcc_device.py`` (taps, teacher forcing) capture far more than the production pair, whose
  ~12 MiB fit the 32 MiB default of ``tt_bevformer.device.DEVICE_DEFAULTS``; ``BEVFORMER_TRACE_REGION`` overrides it.
- Tests marked ``device`` are skipped when ttnn is missing or ``TT_VISIBLE_DEVICES=none`` (host-only runs).
  On the shared workspace box run them through the lock and name the test files:
  ``bin/devrun -t 1800 -- python -m pytest -q -s code/tt_bevformer/tests/test_pcc_device.py``.
"""
from __future__ import annotations

import gc
import importlib.util
import os

import pytest


def pytest_addoption(parser):
    parser.addoption("--device-id", action="store", default=None, help="chip id (default $TT_DEVICE_ID or 0)")


def pytest_configure(config):
    config.addinivalue_line("markers", "device: needs a Tenstorrent chip (skipped on host-only runs)")


def _no_device_reason():
    if os.environ.get("TT_VISIBLE_DEVICES", "").lower() == "none":
        return "TT_VISIBLE_DEVICES=none (host-only run)"
    if importlib.util.find_spec("ttnn") is None:
        return "ttnn is not installed"
    return None


def pytest_collection_modifyitems(config, items):
    reason = _no_device_reason()
    if reason:
        skip = pytest.mark.skip(reason=reason)
        for item in items:
            if "device" in item.keywords:
                item.add_marker(skip)


@pytest.fixture(autouse=True)
def _gc_between_tests():
    gc.collect()


TEST_TRACE_REGION = 1 << 30


@pytest.fixture(scope="session")
def device(request):
    from tt_bevformer.device import close_device, open_device

    cli = request.config.getoption("--device-id")
    region = None if os.environ.get("BEVFORMER_TRACE_REGION") else TEST_TRACE_REGION
    dev = open_device(int(cli) if cli is not None else None, allow_fallback=False, trace_region_size=region)
    yield dev
    close_device(dev)