conftest.py

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import subprocess
from dataclasses import dataclass
from functools import cache
from pathlib import Path

import pytest
from typing_extensions import Self

_ALLOWED_GPUS_BY_LEVEL = {
    0: [0, 1],
    1: [0, 1, 8],
    2: [0, 1, 8],
}


@pytest.fixture(scope="module")
def original_datadir(request: pytest.FixtureRequest) -> Path:
    root_dir = request.config.rootpath
    relative_path = request.path.with_suffix("").relative_to(root_dir)
    return root_dir / "tests/data" / relative_path


@cache
def _get_available_gpus() -> int:
    try:
        return len(subprocess.check_output(["nvidia-smi", "--list-gpus"], text=True).splitlines())
    except Exception as e:
        print(f"WARNING: Failed to get available GPUs: {e}")
        return 0


@dataclass(frozen=True)
class _Args:
    worker_id: str
    worker_index: int

    enable_manual: bool
    num_gpus: int | None
    levels: list[int] | None

    @classmethod
    def from_config(cls, config: pytest.Config) -> Self:
        worker_id = os.environ.get("PYTEST_XDIST_WORKER", "master")
        if worker_id == "master":
            worker_index = 0
        else:
            worker_index = int(worker_id.removeprefix("gw"))

        if config.option.levels:
            # Requires enforcing level order: https://pytest-dev.github.io/pytest-order/stable/other_plugins.html
            levels = [int(config.option.levels)]
        else:
            levels = None

        return cls(
            worker_id=worker_id,
            worker_index=worker_index,
            enable_manual=config.option.manual,
            num_gpus=config.option.num_gpus,
            levels=levels,
        )


_ARGS: _Args = None  # type: ignore


def pytest_addoption(parser: pytest.Parser):
    parser.addoption("--manual", action="store_true", default=False, help="Run manual tests")
    parser.addoption("--num-gpus", default=None, type=int, help="Run tests with the specified number of GPUs")
    parser.addoption("--levels", default=None, help="Run tests with the specified level")


def pytest_xdist_auto_num_workers(config: pytest.Config) -> int | None:
    num_gpus: int | None = config.option.num_gpus
    if num_gpus is None:
        return 1
    if num_gpus == 0:
        # CPU
        return None

    available_gpus = _get_available_gpus()
    if available_gpus < num_gpus:
        raise ValueError(f"Not enough GPUs available. Required: {num_gpus}, Available: {available_gpus}")
    return available_gpus // num_gpus


def pytest_configure(config: pytest.Config):
    global _ARGS
    _ARGS = _Args.from_config(config)

    if _ARGS.worker_id == "master":
        return

    if _ARGS.worker_index > 1:
        if _ARGS.num_gpus is None:
            raise NotImplementedError(f"Running parallel tests requires --num-gpus to be set.")

    # Check if there are enough GPUs available.
    if _ARGS.num_gpus is not None and _ARGS.num_gpus > 0:
        required_gpus = _ARGS.num_gpus * _ARGS.worker_index
        available_gpus = _get_available_gpus()
        if available_gpus < required_gpus:
            raise ValueError(f"Not enough GPUs available. Required: {required_gpus}, Available: {available_gpus}")


def _get_marker(item: pytest.Item, name: str) -> pytest.Mark | None:
    markers = list(item.iter_markers(name=name))
    if not markers:
        return None
    if len(markers) != 1:
        raise ValueError(f"Multiple markers found for {name}: {markers}")
    return markers[0]


def _parse_level_marker(mark: pytest.Mark) -> int:
    if len(mark.args) != 1:
        raise ValueError(f"Invalid arguments: {mark.args}")
    if mark.kwargs:
        raise ValueError(f"Invalid keyword arguments: {mark.kwargs}")
    level = int(mark.args[0])
    if level not in [0, 1, 2]:
        raise ValueError(f"Invalid level: {level}")
    return level


def _parse_gpus_marker(mark: pytest.Mark) -> int:
    if len(mark.args) != 1:
        raise ValueError(f"Invalid arguments: {mark.args}")
    if mark.kwargs:
        raise ValueError(f"Invalid keyword arguments: {mark.kwargs}")
    required_gpus = int(mark.args[0])
    if required_gpus not in [0, 1, 8]:
        raise ValueError(f"Invalid number of GPUs: {required_gpus}")
    return required_gpus


def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]):
    for item in items:
        manual_mark = _get_marker(item, "manual")
        level_mark = _get_marker(item, "level")
        gpus_mark = _get_marker(item, "gpus")
        try:
            level = _parse_level_marker(level_mark) if level_mark else 0
            gpus = _parse_gpus_marker(gpus_mark) if gpus_mark else 0
        except ValueError as e:
            pytest.fail(f"Invalid marker on test {item.name}: {e}")
            assert False, "unreachable"

        allowed_gpus = _ALLOWED_GPUS_BY_LEVEL[level]
        if gpus not in allowed_gpus:
            pytest.fail(f"Level {level} tests must have {allowed_gpus} GPUs, but {item.name} has {gpus} GPUs")

        # Check if the test should be skipped
        if not _ARGS.enable_manual and manual_mark is not None:
            item.add_marker(pytest.mark.skip(reason="test requires --manual"))
        if _ARGS.levels is not None and level not in _ARGS.levels:
            item.add_marker(pytest.mark.skip(reason=f"test requires --levels={level}"))
        if _ARGS.num_gpus is not None and gpus != _ARGS.num_gpus:
            item.add_marker(pytest.mark.skip(reason=f"test requires --num-gpus={gpus}"))
        available_gpus = _get_available_gpus()
        if gpus > available_gpus:
            item.add_marker(
                pytest.mark.skip(reason=f"test requires {gpus} GPUs, but only {available_gpus} are available")
            )

    # Exclude skipped tests
    selected_items = []
    deselected_items = []
    for item in items:
        if item.get_closest_marker("skip"):
            deselected_items.append(item)
            continue
        selected_items.append(item)
    items[:] = selected_items
    config.hook.pytest_deselected(items=deselected_items)


def pytest_runtest_setup(item: pytest.Item):
    gpus_mark = item.get_closest_marker(name="gpus")
    try:
        gpus = _parse_gpus_marker(gpus_mark) if gpus_mark else 0
    except ValueError as e:
        pytest.fail(f"Invalid marker on test {item.name}: {e}")
        assert False, "unreachable"

    # Limit the number of GPUs used by the test
    if gpus > 0:
        device_start = _ARGS.worker_index * gpus
        device_end = device_start + gpus
        os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, range(device_start, device_end)))
    else:
        os.environ["CUDA_VISIBLE_DEVICES"] = ""
    os.environ["NUM_GPUS"] = str(gpus)

    # Set master port to a unique port for each worker.
    os.environ["MASTER_PORT"] = str(12341 + _ARGS.worker_index)


def pytest_sessionfinish(session: pytest.Session, exitstatus: int):
    """Combine coverage data files after all tests complete.

    This hook runs after all tests finish. It combines all .coverage.* files
    created by parallel coverage runs into a single .coverage file that
    pytest-cov can then use to generate reports.
    """
    # Only run on the master worker
    if _ARGS and _ARGS.worker_id != "master":
        return

    # Check if coverage is enabled
    if not session.config.pluginmanager.has_plugin("pytest_cov"):
        return

    # Check if there are any .coverage.* files to combine
    coverage_files = list(Path.cwd().glob(".coverage.*"))
    if not coverage_files:
        return

    try:
        # Combine all coverage data files
        result = subprocess.run(
            ["coverage", "combine"],
            capture_output=True,
            text=True,
            timeout=60,
        )
        if result.returncode == 0:
            print(f"\n✓ Combined {len(coverage_files)} coverage data file(s)")
        else:
            print(f"\n⚠ Warning: coverage combine failed: {result.stderr}")
    except (subprocess.TimeoutExpired, FileNotFoundError) as e:
        print(f"\n⚠ Warning: Could not combine coverage data: {e}")