File size: 1,555 Bytes
629b123 | 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 | # SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import os
import shutil
from typing import TYPE_CHECKING
import torch
from loguru import logger
if TYPE_CHECKING:
import pytest
def pytest_configure(config: pytest.Config) -> None:
"""Register the log-start plugin only when output capture is disabled."""
if config.option.capture == "no":
config.pluginmanager.register(_LogStartPlugin(), "tt_dit_logstart")
class _LogStartPlugin:
@staticmethod
def pytest_runtest_logstart(nodeid: str, location: tuple[str, int | None, str]) -> None: # noqa: ARG004
parts = nodeid.split("::")
filename = parts[0].rsplit("/", 1)[-1]
rest = "::".join(parts[1:])
# split off params from last part
params = []
if "[" in rest:
rest, param_str = rest.split("[", 1)
params = param_str.rstrip("]").split("-")
dim = "\033[2m"
bold = "\033[1m"
reset = "\033[0m"
label = f"{dim}{filename}{reset} {bold}{rest}{reset}"
if params:
label += f" {dim}({', '.join(params)}){reset}"
raw = f"{filename} {rest}"
if params:
raw += f" ({', '.join(params)})"
width = shutil.get_terminal_size()[0]
print(f"\n\n{'━' * width}\n{label}\n") # noqa: T201
num_torch_threads = max(1, os.cpu_count())
logger.info(f"Setting torch num_threads to {num_torch_threads}")
torch.set_num_threads(num_torch_threads)
|