code/conftest.py

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

- ``device`` (session): one chip opened like the published numbers (``tt_diffusion_planner.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); ``DIFFUSION_PLANNER_DISPATCH=worker`` is the A/B opt-in, and the other ``DIFFUSION_PLANNER_*``
  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.
- 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_diffusion_planner/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()


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

    cli = request.config.getoption("--device-id")
    dev = open_device(int(cli) if cli is not None else None, allow_fallback=False)
    yield dev
    close_device(dev)