mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 07:17:08 +08:00
114 lines
4.6 KiB
Python
114 lines
4.6 KiB
Python
# conftest.py for tilelang kernel tests.
|
|
#
|
|
# Loads the tile_kernels pytest plugins (CLI options, markers, fixtures) for
|
|
# tests in deepinfer/model_executor/layers/tilelang_ops/, replacing the former
|
|
# local pytest_benchmark_plugin.py.
|
|
#
|
|
# The tile_kernels plugins read their data directories from environment
|
|
# variables, so we set them here (with ``setdefault`` so external overrides
|
|
# win), defaulting to the benchmark_baselines/ and gpu_mem_profiles/
|
|
# directories under this tilelang_ops/ package.
|
|
#
|
|
# GPU binding for xdist workers happens at module import time, before the
|
|
# ``pytest_plugins`` declaration below: the plugins import torch/tilelang
|
|
# which may initialize CUDA, and each worker must be bound to its GPU first.
|
|
# We only use the lightweight ``deep_cluster.hardware`` module, which never
|
|
# initializes CUDA.
|
|
#
|
|
# This file intentionally does not import anything heavy from `deepinfer` or
|
|
# `torch` to avoid early initialization.
|
|
|
|
import os
|
|
from collections.abc import Generator
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from deep_cluster.hardware import accelerator_kind, get_actual_num_gpus
|
|
|
|
|
|
def _bind_worker_gpu() -> None:
|
|
"""Bind each xdist worker to a GPU via CUDA_VISIBLE_DEVICES.
|
|
|
|
Mirrors the old pytest_benchmark_plugin.py logic but runs at module
|
|
import time, using only ``deep_cluster.hardware`` (no CUDA initialization).
|
|
"""
|
|
worker_id = os.environ.get("PYTEST_XDIST_WORKER", None)
|
|
if worker_id is None:
|
|
return
|
|
gpu_id = int(worker_id.replace("gw", ""))
|
|
num_gpus = get_actual_num_gpus()
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id % num_gpus)
|
|
if accelerator_kind() == "ascend":
|
|
os.environ["ASCEND_RT_VISIBLE_DEVICES"] = os.environ["CUDA_VISIBLE_DEVICES"]
|
|
|
|
|
|
_bind_worker_gpu()
|
|
|
|
# ---- Environment variables consumed by the tile_kernels pytest plugins ----
|
|
# benchmark 与 gpu_mem 必须指向两个不同目录
|
|
_OPS_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
# Common
|
|
os.environ.setdefault("TK_TEST_MERGE_MODE", "overwrite")
|
|
os.environ.setdefault("TK_TEST_PREFIX", "deepinfer/model_executor/layers/tilelang_ops/")
|
|
# Benchmark plugin
|
|
os.environ.setdefault("TK_BENCHMARK_TESTS_DIR", _OPS_DIR)
|
|
os.environ.setdefault("TK_BENCHMARK_BASELINES_DIR", os.path.join(_OPS_DIR, "benchmark_baselines"))
|
|
os.environ.setdefault("TK_BENCHMARK_FAIL_MODE", "none")
|
|
os.environ.setdefault("TK_BENCHMARK_REGRESSION_THRESHOLD", "0.15")
|
|
os.environ.setdefault("TK_BENCHMARK_MIN_DELTA_US", "3.0")
|
|
# GPU mem plugin
|
|
os.environ.setdefault("TK_GPU_MEM_PROFILES_DIR", os.path.join(_OPS_DIR, "gpu_mem_profiles"))
|
|
|
|
|
|
pytest_plugins = [
|
|
"tile_kernels.testing.pytest.benchmark",
|
|
"tile_kernels.testing.pytest.gpu_mem",
|
|
"tile_kernels.testing.pytest.random",
|
|
"tile_kernels.testing.pytest.tilelang_warning",
|
|
"tile_kernels.testing.pytest.xdist_failfast",
|
|
]
|
|
|
|
|
|
def pytest_collect_file(file_path: Path, parent: pytest.Collector) -> pytest.Collector | None:
|
|
"""Collect test files directly under tilelang_ops/ so the directory can be
|
|
passed to pytest directly.
|
|
|
|
Their names don't match pytest's default ``python_files`` pattern
|
|
(``test_*.py`` / ``*_test.py``), so without this hook a bare
|
|
``python -m pytest deepinfer/model_executor/layers/tilelang_ops`` collects 0 items.
|
|
Collects any top-level ``*.py`` except ``__init__.py`` / ``conftest.py`` that
|
|
contains a ``def test_``.
|
|
"""
|
|
ops_dir = Path(__file__).parent
|
|
if file_path.parent != ops_dir or file_path.suffix != ".py":
|
|
return None
|
|
if file_path.name in ("__init__.py", "conftest.py"):
|
|
return None
|
|
if parent.session.isinitpath(file_path):
|
|
# Explicitly passed file args: the core pytest_collect_file already
|
|
# collects them (isinitpath bypasses python_files), so skip here to
|
|
# avoid double collection.
|
|
return None
|
|
if "def test_" not in file_path.read_text(encoding="utf-8"):
|
|
return None
|
|
ihook = parent.session.gethookproxy(file_path)
|
|
return ihook.pytest_pycollect_makemodule(module_path=file_path, parent=parent)
|
|
|
|
|
|
@pytest.hookimpl(wrapper=True)
|
|
def pytest_runtest_call(item: pytest.Item) -> Generator[None, None, None]:
|
|
"""Skip tests hitting kernels the selected backend has not implemented.
|
|
|
|
ops.KernelUnavailableError is matched by name to avoid importing ops
|
|
(heavy deepinfer.model_executor deps, see module docstring). On cuda the
|
|
backend is complete, so this only ever fires on a typo — keep it fatal.
|
|
"""
|
|
try:
|
|
return (yield)
|
|
except AttributeError as e:
|
|
if type(e).__name__ != "KernelUnavailableError" or accelerator_kind() == "nvidia":
|
|
raise
|
|
pytest.skip(str(e))
|