deepspeedai--deepspeed/tests/unit/ops/test_op_builder.py

# Copyright (c) DeepSpeed Team.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import os
import sys
import subprocess
import importlib.util
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest

BUILDER_PATH = Path(__file__).resolve().parents[3] / "op_builder" / "builder.py"
BUILDER_SPEC = importlib.util.spec_from_file_location("test_op_builder_module", BUILDER_PATH)
builder_module = importlib.util.module_from_spec(BUILDER_SPEC)
BUILDER_SPEC.loader.exec_module(builder_module)
CUDAOpBuilder = builder_module.CUDAOpBuilder

BUILDER_MODULE = builder_module
CUDA_API = BUILDER_MODULE.torch.cuda  #ignore-cuda


class _StubCUDAOpBuilder(CUDAOpBuilder):
    BUILD_VAR = "STUB_BUILDER"
    NAME = "stub"

    def __init__(self):
        super().__init__(name="stub")

    def absolute_name(self):
        return "deepspeed.ops.stub"

    def sources(self):
        return []

    def include_paths(self):
        return []


def make_builder(**overrides):
    builder = _StubCUDAOpBuilder()
    for key, value in overrides.items():
        setattr(builder, key, value)
    return builder


def assert_jit_uses_explicit_arch_list(builder, expected_arch_list, env_updates=None):
    env_updates = env_updates or {}

    with patch.dict(os.environ, env_updates, clear=False):
        if "TORCH_CUDA_ARCH_LIST" not in env_updates:
            os.environ.pop("TORCH_CUDA_ARCH_LIST", None)
        with patch.object(CUDA_API, "device_count",
                          side_effect=AssertionError("probe should not be called")) as device_count:
            with patch.object(CUDA_API,
                              "get_device_capability",
                              side_effect=AssertionError("probe should not be called")) as get_device_capability:
                assert builder.compute_capability_args() == []
                assert os.environ["TORCH_CUDA_ARCH_LIST"] == expected_arch_list

    device_count.assert_not_called()
    get_device_capability.assert_not_called()


def test_jit_mode_prefers_explicit_arch_lists_before_cuda_probe():
    assert_jit_uses_explicit_arch_list(make_builder(jit_mode=True, _jit_arch_list="8.0;8.9"), "8.0;8.9+PTX")
    assert_jit_uses_explicit_arch_list(make_builder(jit_mode=True), "8.0;8.9+PTX", {"TORCH_CUDA_ARCH_LIST": "8.0 8.9"})


def test_bad_fork_jit_without_arch_list_raises_actionable_error():
    builder = make_builder(jit_mode=True)

    with patch.dict(os.environ, {}, clear=False):
        os.environ.pop("TORCH_CUDA_ARCH_LIST", None)
        with patch.object(CUDA_API, "_is_in_bad_fork", return_value=True):
            with patch.object(CUDA_API, "device_count",
                              side_effect=AssertionError("probe should not be called")) as device_count:
                with pytest.raises(RuntimeError, match="TORCH_CUDA_ARCH_LIST"):
                    builder.compute_capability_args()

    device_count.assert_not_called()


def test_jit_mode_probes_devices_when_safe_and_errors_without_visible_gpus():
    builder = make_builder(jit_mode=True)

    with patch.dict(os.environ, {}, clear=False):
        os.environ.pop("TORCH_CUDA_ARCH_LIST", None)
        with patch.object(CUDA_API, "_is_in_bad_fork", return_value=False):
            with patch.object(CUDA_API, "device_count", return_value=2) as device_count:
                with patch.object(CUDA_API, "get_device_capability", side_effect=[(7, 0),
                                                                                  (8, 9)]) as get_device_capability:
                    assert builder.compute_capability_args() == []
                    assert os.environ["TORCH_CUDA_ARCH_LIST"] == "7.0;8.9+PTX"
                    assert builder.enable_bf16 is False

    device_count.assert_called_once_with()
    assert get_device_capability.call_count == 2

    builder = make_builder(jit_mode=True)
    with patch.dict(os.environ, {}, clear=False):
        os.environ.pop("TORCH_CUDA_ARCH_LIST", None)
        with patch.object(CUDA_API, "_is_in_bad_fork", return_value=False):
            with patch.object(CUDA_API, "device_count", return_value=0):
                with pytest.raises(RuntimeError, match="no CUDA devices"):
                    builder.compute_capability_args()


def test_jit_load_restores_env_and_state_after_failure():
    builder = make_builder()

    def fail_nvcc_args():
        assert getattr(builder, "_jit_arch_list", None) == "8.9"
        os.environ["TORCH_CUDA_ARCH_LIST"] = "8.9+PTX"
        raise RuntimeError("build failed")

    with patch.dict(os.environ, {"TORCH_CUDA_ARCH_LIST": "8.9"}, clear=False):
        with patch.object(builder, "is_compatible", return_value=True):
            with patch.object(CUDAOpBuilder, "is_rocm_pytorch", return_value=False):
                with patch.object(CUDA_API, "is_available", return_value=True):
                    with patch("torch.utils.cpp_extension.verify_ninja_availability", return_value=None):
                        with patch.object(builder, "nvcc_args", side_effect=fail_nvcc_args):
                            with pytest.raises(RuntimeError, match="build failed"):
                                builder.jit_load(verbose=False)

        assert getattr(builder, "_jit_arch_list", None) is None
        assert builder.jit_mode is False
        assert os.environ["TORCH_CUDA_ARCH_LIST"] == "8.9"


def test_jit_load_restores_state_after_success():
    builder = make_builder()
    op_module = MagicMock()

    def successful_nvcc_args():
        assert builder._jit_arch_list == "8.9"
        os.environ["TORCH_CUDA_ARCH_LIST"] = "8.9+PTX"
        return []

    with patch.dict(os.environ, {"TORCH_CUDA_ARCH_LIST": "8.9"}, clear=False):
        with patch.object(builder, "is_compatible", return_value=True):
            with patch.object(CUDAOpBuilder, "is_rocm_pytorch", return_value=False):
                with patch.object(CUDA_API, "is_available", return_value=True):
                    with patch("torch.utils.cpp_extension.verify_ninja_availability", return_value=None):
                        with patch.object(builder, "nvcc_args", side_effect=successful_nvcc_args):
                            with patch.object(builder, "cxx_args", return_value=[]):
                                with patch("torch.utils.cpp_extension.load", return_value=op_module):
                                    assert builder.jit_load(verbose=False) is op_module

        assert os.environ["TORCH_CUDA_ARCH_LIST"] == "8.9"
        assert getattr(builder, "_jit_arch_list", None) is None
        assert builder.jit_mode is False


def test_non_jit_branch_unchanged():
    builder = make_builder(jit_mode=False)

    with patch.dict(os.environ, {"TORCH_CUDA_ARCH_LIST": "8.0;8.9+PTX"}, clear=False):
        args = builder.compute_capability_args()

    assert args == [
        "-gencode=arch=compute_80,code=sm_80",
        "-gencode=arch=compute_89,code=sm_89",
        "-gencode=arch=compute_89,code=compute_89",
    ]


def test_non_jit_branch_sorts_and_dedupes_gencode_flags():
    builder = make_builder(jit_mode=False)

    with patch.dict(os.environ, {"TORCH_CUDA_ARCH_LIST": "8.0;7.5;8.0;7.0"}, clear=False):
        args = builder.compute_capability_args()
        assert os.environ["TORCH_CUDA_ARCH_LIST"] == "7.0;7.5;8.0"

    assert args == [
        "-gencode=arch=compute_70,code=sm_70",
        "-gencode=arch=compute_75,code=sm_75",
        "-gencode=arch=compute_80,code=sm_80",
    ]


def test_non_jit_branch_canonicalizes_mixed_ptx_variants_to_one_sm_and_one_ptx():
    # For mixed inputs such as "8.0;8.0+PTX" or "8.0+PTX;8.0", PyTorch
    # canonicalizes the architecture to one sm_80 entry plus one compute_80
    # PTX entry. Dedupe by the canonical numeric arch so we match.
    expected_arch_list = "7.5;8.0+PTX"
    expected_args = [
        "-gencode=arch=compute_75,code=sm_75",
        "-gencode=arch=compute_80,code=sm_80",
        "-gencode=arch=compute_80,code=compute_80",
    ]

    for arch_input in ("8.0;8.0+PTX;7.5", "7.5;8.0+PTX;8.0", "8.0+PTX;7.5;8.0", "8.0;7.5;8.0+PTX"):
        builder = make_builder(jit_mode=False)
        with patch.dict(os.environ, {"TORCH_CUDA_ARCH_LIST": arch_input}, clear=False):
            args = builder.compute_capability_args()
            assert os.environ["TORCH_CUDA_ARCH_LIST"] == expected_arch_list, arch_input
        assert args == expected_args, arch_input


def test_non_jit_branch_canonical_dedupe_mixed_ptx_combinations():
    # Lock in the four mixed-PTX combinations for a single arch so the dedupe
    # behavior cannot regress on either ordering or duplication.
    builder = make_builder(jit_mode=False)
    cases = [
        ("8.0;8.0+PTX", "8.0+PTX", ["-gencode=arch=compute_80,code=sm_80",
                                    "-gencode=arch=compute_80,code=compute_80"]),
        ("8.0+PTX;8.0", "8.0+PTX", ["-gencode=arch=compute_80,code=sm_80",
                                    "-gencode=arch=compute_80,code=compute_80"]),
        ("8.0;8.0", "8.0", ["-gencode=arch=compute_80,code=sm_80"]),
        ("8.0+PTX;8.0+PTX", "8.0+PTX",
         ["-gencode=arch=compute_80,code=sm_80", "-gencode=arch=compute_80,code=compute_80"]),
    ]
    for arch_input, expected_arch_list, expected_args in cases:
        with patch.dict(os.environ, {"TORCH_CUDA_ARCH_LIST": arch_input}, clear=False):
            args = builder.compute_capability_args()
            assert os.environ["TORCH_CUDA_ARCH_LIST"] == expected_arch_list, arch_input
        assert args == expected_args, arch_input


def test_cuda_capability_major_skips_probe_when_context_not_initialized():
    # Probing device properties forces a lazy CUDA-context init, which creates a
    # CUDA context. Doing that while checking op compatibility at "import deepspeed"
    # time poisons fork()-based multiprocessing (issue #7918): a forked child cannot
    # reuse the parent's context. With no context yet, the probe must be skipped.
    builder = make_builder()
    with patch.object(CUDA_API, "is_initialized", return_value=False):
        with patch.object(
                CUDA_API, "get_device_properties",
                side_effect=AssertionError("must not initialize CUDA / poison fork")) as get_device_properties:
            assert builder.cuda_capability_major() is None
    get_device_properties.assert_not_called()


def test_cuda_capability_major_probes_when_context_already_initialized():
    # When a CUDA context already exists (e.g. at op load time), probing is safe
    # and must report the real compute-capability major.
    builder = make_builder()
    device_properties = MagicMock(major=8)
    with patch.object(CUDA_API, "is_initialized", return_value=True):
        with patch.object(CUDA_API, "_is_in_bad_fork", return_value=False):
            with patch.object(CUDA_API, "get_device_properties",
                              return_value=device_properties) as get_device_properties:
                assert builder.cuda_capability_major() == 8
    get_device_properties.assert_called_once_with(0)


def test_cuda_capability_major_skips_probe_in_bad_fork():
    # Inside a forked child that inherited an initialized-but-invalid context,
    # probing would raise "Cannot re-initialize CUDA in forked subprocess", so it
    # must be skipped there as well.
    builder = make_builder()
    with patch.object(CUDA_API, "is_initialized", return_value=True):
        with patch.object(CUDA_API, "_is_in_bad_fork", return_value=True):
            with patch.object(CUDA_API,
                              "get_device_properties",
                              side_effect=AssertionError("must not probe in a forked child")) as get_device_properties:
                assert builder.cuda_capability_major() is None
    get_device_properties.assert_not_called()


def test_forked_child_can_use_cuda_after_importing_deepspeed():
    # Core contract of issue #7918: after the parent process runs
    # ``import deepspeed``, a forked child must still be able to initialize and
    # use CUDA. If import created a CUDA context in the parent, the child fails
    # with "Cannot re-initialize CUDA in forked subprocess". Everything runs in a
    # dedicated subprocess so a poisoned parent cannot leak into the pytest worker
    # or other tests.
    program = "\n".join([
        "import os, sys",
        "import torch",
        "import deepspeed  # must not create a CUDA context in the parent",
        # device_count() is NVML-based and never initializes a context, so it is
        # a fork-safe way to check for a GPU before forking.
        "if torch.cuda.device_count() == 0:",  #ignore-cuda
        "    print('NO_CUDA'); sys.exit(0)",
        "pid = os.fork()",
        "if pid == 0:",
        "    try:",
        "        torch.ones(1, device='cuda')",
        "        os._exit(0)",
        "    except Exception as exc:",
        "        sys.stderr.write(repr(exc))",
        "        os._exit(1)",
        "_, status = os.waitpid(pid, 0)",
        "sys.exit(os.waitstatus_to_exitcode(status))",
    ])
    env = os.environ.copy()
    repo_root = str(Path(__file__).resolve().parents[3])
    env["PYTHONPATH"] = repo_root + (os.pathsep + env["PYTHONPATH"] if env.get("PYTHONPATH") else "")
    result = subprocess.run([sys.executable, "-c", program], capture_output=True, text=True, env=env, timeout=300)
    if result.returncode != 0 and ("No module named 'deepspeed'" in result.stderr
                                   or "No module named 'torch'" in result.stderr):
        pytest.skip("deepspeed/torch not importable in a subprocess in this environment")
    if result.stdout.strip() == "NO_CUDA":
        pytest.skip("no CUDA device available")
    assert result.returncode == 0, ("forked child could not use CUDA after 'import deepspeed' "
                                    "(a CUDA context was created during import, issue #7918):\n" + result.stderr)