code/conftest.py

# SPDX-FileCopyrightText: (c) 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0

"""Minimal pytest fixtures for tt-MoGe.

Provides a single-device ``device`` fixture so the tt_moge tests run without the
tt-metal monorepo's own conftest.py. The device is selected with ``--device-id N``
on the pytest CLI or the ``MOGE_DEVICE`` environment variable (default 0). The
fixture opens with a trace region + 2 command queues so the perf tests exercise
the same deployment config as the benchmark. Set ``TT_VISIBLE_DEVICES`` to
constrain which PCIe chips ttnn enumerates at all (e.g. to pin one p150a).
"""
from __future__ import annotations

import gc
import os

import pytest


def pytest_addoption(parser):
    parser.addoption("--device-id", action="store", default=None, help="ttnn device id")


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


@pytest.fixture(scope="session")
def device(request):
    import ttnn

    opt = request.config.getoption("--device-id")
    device_id = int(opt) if opt is not None else int(os.environ.get("MOGE_DEVICE", "0"))
    two_cq = bool(int(os.environ.get("MOGE_2CQ", "0")))
    kw = dict(device_id=device_id, l1_small_size=32768, trace_region_size=1500000000,
              num_command_queues=2 if two_cq else 1)
    # Default = the p150 configuration of the published numbers: ETH dispatch, 1 CQ, 12x10 compute grid
    # (tt_moge.device.open_device; falls back to worker dispatch with a warning without the ETH-dispatch patch).
    # MOGE_DISPATCH=worker [MOGE_2CQ=1] = the Galaxy-only worker-dispatch mode; MOGE_GRID=12x10|11x10|full.
    from tt_moge.device import open_device
    dispatch = os.environ.get("MOGE_DISPATCH", "eth")
    kw.pop("device_id")
    dev, _info = open_device(device_id, grid=os.environ.get("MOGE_GRID", "12x10"), dispatch=dispatch, **kw)
    dispatch = _info["dispatch"]
    g = dev.compute_with_storage_grid_size()
    print(f"\n[conftest] dispatch={dispatch or 'default'} grid={g.x}x{g.y}")
    yield dev
    ttnn.close_device(dev)