unsloth/tests/test_import_fixes_drift.py

# Unsloth - 2x faster, 60% less VRAM LLM training and finetuning
# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.

"""Drift detectors for the upstream pathologies ``unsloth/import_fixes.py``
works around; one test per ``fix_*`` / ``patch_*``, each fails (never skips)
when the pathology is active. Runs under the GPU-free ``tests/conftest.py``."""

from __future__ import annotations

import importlib
import importlib.util
import inspect
import os
import re
import sys
from pathlib import Path
from importlib.metadata import version as importlib_version

import pytest


# Mirrors import_fixes.py's local Version(): strip dev/alpha/beta/rc/local suffixes.
from packaging.version import Version as _PkgVersion


def _safe_version(raw):
    raw_str = str(raw)
    base = raw_str.split("+", 1)[0]
    try:
        return _PkgVersion(base)
    except Exception:
        match = re.match(r"[0-9]+(?:\.[0-9]+)*", base)
        if not match:
            raise
        return _PkgVersion(match.group(0))


# protobuf


def test_protobuf_message_factory_get_prototype_or_get_message_class_present():
    """``fix_message_factory_issue``."""
    mf = pytest.importorskip("google.protobuf.message_factory")
    has_mf_class = hasattr(mf, "MessageFactory")
    has_get_prototype = has_mf_class and hasattr(mf.MessageFactory, "GetPrototype")
    has_get_message_class = hasattr(mf, "GetMessageClass")
    if not has_mf_class:
        pytest.fail(
            "DRIFT DETECTED: google.protobuf.message_factory.MessageFactory is "
            "missing entirely -- fix_message_factory_issue would inject a stub."
        )
    if not (has_get_prototype or has_get_message_class):
        pytest.fail(
            "DRIFT DETECTED: neither MessageFactory.GetPrototype nor "
            "module-level GetMessageClass is present; fix_message_factory_issue "
            "would inject the GetPrototype/GetMessageClass shim."
        )
    assert has_get_prototype or has_get_message_class


# datasets


def test_datasets_version_not_in_broken_recursion_range():
    """``patch_datasets``: datasets 4.4.0-4.5.0 hit RLock recursion in the Arrow loader."""
    pytest.importorskip("datasets")
    ds_v = _safe_version(importlib_version("datasets"))
    lo = _PkgVersion("4.4.0")
    hi = _PkgVersion("4.5.0")
    assert not (lo <= ds_v <= hi), (
        f"datasets=={ds_v} lies in the 4.4.0-4.5.0 recursion-error "
        f"range that patch_datasets explicitly forbids. Downgrade to "
        f"datasets==4.3.0 or upgrade past 4.5.0."
    )


# trl


def test_trl_is_x_available_returns_bool_not_tuple():
    """``fix_trl_vllm_ascend``: TRL's ``is_*_available`` must still return bools
    after transformers >=4.48 made ``_is_package_available`` return a tuple."""
    pytest.importorskip("trl")
    try:
        import trl.import_utils as tiu
    except Exception as exc:
        pytest.skip(f"trl.import_utils not importable: {exc!r}")

    accessor_names = [
        n
        for n in dir(tiu)
        if n.startswith("is_") and n.endswith("_available") and callable(getattr(tiu, n, None))
    ]
    assert accessor_names, "trl.import_utils has no is_*_available accessors"

    bad = {}
    for name in accessor_names:
        accessor = getattr(tiu, name)
        try:
            sig = inspect.signature(accessor)
            required = [
                p
                for p in sig.parameters.values()
                if p.default is inspect.Parameter.empty
                and p.kind
                in (
                    inspect.Parameter.POSITIONAL_ONLY,
                    inspect.Parameter.POSITIONAL_OR_KEYWORD,
                )
            ]
            if required:
                continue
            result = accessor()
        except Exception:
            continue
        if not isinstance(result, bool):
            bad[name] = (type(result).__name__, result)

    if bad:
        pytest.fail(
            "DRIFT DETECTED: fix_trl_vllm_ascend coerces these accessors "
            f"from tuple-cached values to bool: {bad}"
        )


def test_trl_cached_available_flags_are_not_tuples():
    """``fix_trl_vllm_ascend``: same drift on the module-level cached ``_*_available`` attrs."""
    pytest.importorskip("trl")
    try:
        import trl.import_utils as tiu
    except Exception as exc:
        pytest.skip(f"trl.import_utils not importable: {exc!r}")

    tuple_flags = {
        name: value
        for name, value in vars(tiu).items()
        if name.startswith("_") and name.endswith("_available") and isinstance(value, tuple)
    }
    if tuple_flags:
        pytest.fail(
            "DRIFT DETECTED: fix_trl_vllm_ascend needs to coerce these tuple-"
            f"cached flags to bool: {sorted(tuple_flags)}"
        )


# transformers


def test_pretrained_model_enable_input_require_grads_uses_old_pattern():
    """``patch_enable_input_require_grads``: HF PR #41993 made
    enable_input_require_grads iterate ``self.modules()``, so vision submodules
    raise NotImplementedError unless the tolerant replacement is installed."""
    pytest.importorskip("transformers")
    from transformers import PreTrainedModel

    try:
        src = inspect.getsource(PreTrainedModel.enable_input_require_grads)
    except Exception as exc:
        pytest.skip(f"could not getsource(enable_input_require_grads): {exc!r}")

    if "for module in self.modules()" not in src:
        return  # pre-HF#41993 shape
    if "NotImplementedError" in src:
        return  # tolerant replacement installed

    pytest.fail(
        "DRIFT DETECTED: PreTrainedModel.enable_input_require_grads now "
        "iterates self.modules() (post HF#41993) and has NOT been "
        "wrapped by patch_enable_input_require_grads; vision submodules "
        "(e.g. GLM V4.6's self.visual) will raise NotImplementedError "
        "from get_input_embeddings and crash the whole call."
    )


def test_transformers_torchcodec_available_flag_is_present():
    """``disable_torchcodec_if_broken``: needs the pre-5.x ``_torchcodec_available``
    flag or 5.x ``is_torchcodec_available`` as its patch site when FFmpeg is missing."""
    tf_iu = pytest.importorskip("transformers.utils.import_utils")
    has_flag = hasattr(tf_iu, "_torchcodec_available")
    has_func = callable(getattr(tf_iu, "is_torchcodec_available", None))
    assert has_flag or has_func, (
        "transformers.utils.import_utils dropped both "
        "``_torchcodec_available`` (pre-5.x) AND "
        "``is_torchcodec_available`` (>=5.x); "
        "disable_torchcodec_if_broken can no longer disable a broken "
        "torchcodec install."
    )


def test_transformers_is_causal_conv1d_available_symbol_present():
    """``_disable_transformers_causal_conv1d``: needs a causal_conv1d availability hook."""
    tf_iu = pytest.importorskip("transformers.utils.import_utils")
    candidates = [
        "is_causal_conv1d_available",
        "_causal_conv1d_available",
        "_is_causal_conv1d_available",
    ]
    present = [name for name in candidates if hasattr(tf_iu, name)]
    if not present:
        pytest.fail(
            "DRIFT DETECTED: transformers.utils.import_utils dropped every "
            f"hook in {candidates}; _disable_transformers_causal_conv1d "
            "can no longer mask a broken causal_conv1d binary."
        )


# transformers + accelerate (wandb checkers)


def test_transformers_and_accelerate_is_wandb_available_callable():
    """``disable_broken_wandb``: patches is_wandb_available in three modules
    (transformers integration_utils + accelerate imports/utils); all must exist."""
    pytest.importorskip("transformers")
    pytest.importorskip("accelerate")
    from transformers.integrations import integration_utils as tf_integration
    import accelerate.utils.imports as acc_imports
    import accelerate.utils as acc_utils

    assert callable(getattr(tf_integration, "is_wandb_available", None)), (
        "transformers.integrations.integration_utils.is_wandb_available "
        "was removed/renamed; disable_broken_wandb can no longer mask a "
        "broken wandb install for trl trainers."
    )
    assert callable(getattr(acc_imports, "is_wandb_available", None)), (
        "accelerate.utils.imports.is_wandb_available removed; "
        "disable_broken_wandb cannot patch the source module."
    )
    assert callable(getattr(acc_utils, "is_wandb_available", None)), (
        "accelerate.utils.is_wandb_available removed; "
        "disable_broken_wandb cannot patch the re-export namespace "
        "consulted by trl/trainer/callbacks.py."
    )


# peft


def test_peft_transformers_weight_conversion_importable_and_signature():
    """``patch_peft_weight_converter_compatibility``: wraps build_peft_weight_mapping;
    silently no-ops if the module is unimportable."""
    pytest.importorskip("peft")
    try:
        from peft.utils import transformers_weight_conversion as twc
    except Exception as exc:
        pytest.fail(
            "DRIFT DETECTED: peft.utils.transformers_weight_conversion "
            f"is unimportable on this stack ({exc!r}). "
            "patch_peft_weight_converter_compatibility will silently no-op."
        )

    assert hasattr(
        twc, "build_peft_weight_mapping"
    ), "build_peft_weight_mapping vanished from peft.utils.transformers_weight_conversion."
    sig = inspect.signature(twc.build_peft_weight_mapping)
    expected_params = {"weight_conversions", "adapter_name"}
    actual_params = set(sig.parameters)
    assert expected_params.issubset(actual_params), (
        f"build_peft_weight_mapping signature drifted: expected at "
        f"least {sorted(expected_params)}, got {sorted(actual_params)}."
    )


# triton


def test_triton_compiled_kernel_has_num_ctas_and_cluster_dims():
    """``fix_triton_compiled_kernel_missing_attrs``: triton 3.6+ dropped
    num_ctas/cluster_dims on CompiledKernel, but Inductor's make_launcher needs them."""
    pytest.importorskip("torch")
    triton_mod = pytest.importorskip("triton")  # noqa: F841
    tc = pytest.importorskip("triton.compiler.compiler")

    ck_cls = tc.CompiledKernel
    # Healthy if pre-3.6 class attr present, or __init__ wrapped to install
    # num_ctas + cluster_dims per instance (the post-3.6 fix).
    if hasattr(ck_cls, "num_ctas"):
        return
    init = getattr(ck_cls, "__init__", None)
    if init is not None:
        code = getattr(init, "__code__", None)
        freevars = set(getattr(code, "co_freevars", ()) or ())
        co_names = set(getattr(code, "co_names", ()) or ())
        if "_orig_init" in freevars or {"num_ctas", "cluster_dims"}.issubset(co_names):
            return

    pytest.fail(
        "DRIFT DETECTED: triton.CompiledKernel lacks the `num_ctas` "
        "class attribute AND ``__init__`` has not been wrapped by "
        "fix_triton_compiled_kernel_missing_attrs; torch Inductor's "
        "``make_launcher`` will crash on the eager "
        "``binary.metadata.num_ctas, *binary.metadata.cluster_dims`` "
        "unpack under torch.compile."
    )


# torch + torchvision pairing table


# Mirrors TORCH_TORCHVISION_COMPAT in torchvision_compatibility_check.
_TORCH_TORCHVISION_COMPAT = {
    (2, 9): (0, 24),
    (2, 8): (0, 23),
    (2, 7): (0, 22),
    (2, 6): (0, 21),
    (2, 5): (0, 20),
    (2, 4): (0, 19),
}


def _is_custom_torch_build(raw_version_str):
    if "+" not in raw_version_str:
        return False
    local = raw_version_str.split("+", 1)[1]
    if not local:
        return False
    return not re.fullmatch(r"cu\d[\d.]*|rocm\d[\d.]*|cpu|xpu", local, re.IGNORECASE)


def test_installed_torch_torchvision_pair_is_compatible():
    """``torchvision_compatibility_check``: raises when the (torch, torchvision)
    pair fails the pinned table; custom/prerelease builds are warning-only."""
    pytest.importorskip("torch")
    pytest.importorskip("torchvision")

    torch_raw = importlib_version("torch")
    tv_raw = importlib_version("torchvision")
    torch_v = _safe_version(torch_raw)
    tv_v = _safe_version(tv_raw)

    torch_major = torch_v.release[0]
    torch_minor = torch_v.release[1] if len(torch_v.release) > 1 else 0

    required = _TORCH_TORCHVISION_COMPAT.get((torch_major, torch_minor))
    if required is None:
        pytest.skip(
            f"torch=={torch_raw} is outside the pinned compatibility "
            f"table (entries cover 2.4-2.9). The formula fallback "
            f"in _infer_required_torchvision handles it at runtime."
        )

    pre_tags = (".dev", "a0", "b0", "rc", "alpha", "beta", "nightly")
    is_prerelease = any(t in torch_raw for t in pre_tags) or any(t in tv_raw for t in pre_tags)
    is_custom = _is_custom_torch_build(torch_raw) or _is_custom_torch_build(tv_raw)
    if is_prerelease or is_custom:
        pytest.skip(
            f"torch=={torch_raw} torchvision=={tv_raw} is a custom/"
            f"prerelease build; the runtime check downgrades to warning."
        )

    required_str = f"{required[0]}.{required[1]}.0"
    assert tv_v >= _PkgVersion(required_str), (
        f"DRIFT DETECTED: torch=={torch_raw} requires "
        f"torchvision>={required_str}, but torchvision=={tv_raw} is "
        f"installed. torchvision_compatibility_check would raise."
    )


# vllm


def test_vllm_guided_decoding_params_or_structured_outputs_present():
    """``fix_vllm_guided_decoding_params``: vLLM PR #22772 renamed
    GuidedDecodingParams -> StructuredOutputsParams; the fix re-aliases for trl."""
    pytest.importorskip("vllm")
    try:
        sp = importlib.import_module("vllm.sampling_params")
    except Exception as exc:
        pytest.skip(f"vllm.sampling_params unimportable: {exc!r}")

    has_guided = hasattr(sp, "GuidedDecodingParams")
    has_structured = hasattr(sp, "StructuredOutputsParams")
    assert has_guided or has_structured, (
        "vllm.sampling_params has neither GuidedDecodingParams nor "
        "StructuredOutputsParams; fix_vllm_guided_decoding_params "
        "cannot re-alias. trl import path will break."
    )
    if not has_guided:
        pytest.fail(
            "DRIFT DETECTED: vllm.sampling_params only exposes "
            "StructuredOutputsParams (post PR #22772); "
            "fix_vllm_guided_decoding_params injects a GuidedDecodingParams "
            "alias so trl keeps importing."
        )


def test_vllm_aimv2_ovis_config_is_past_fix_version():
    """``fix_vllm_aimv2_issue``: vLLM <0.10.1 double-registers ``aimv2`` (duplicate-key
    ValueError); the fix only touches old versions."""
    pytest.importorskip("vllm")
    vllm_v = _safe_version(importlib_version("vllm"))
    cutoff = _PkgVersion("0.10.1")
    if vllm_v < cutoff:
        pytest.fail(
            f"DRIFT DETECTED: vllm=={vllm_v} < {cutoff}; "
            "fix_vllm_aimv2_issue rewrites ovis.py to skip the duplicate "
            'AutoConfig.register("aimv2", ...) call.'
        )


# huggingface_hub


def test_huggingface_hub_is_offline_mode_or_hf_hub_offline_present():
    """``fix_huggingface_hub``: re-injects top-level ``is_offline_mode`` from
    ``constants.HF_HUB_OFFLINE`` after huggingface_hub dropped it."""
    hub = pytest.importorskip("huggingface_hub")
    has_top_level = False
    try:
        has_top_level = callable(getattr(hub, "is_offline_mode", None))
    except Exception:
        has_top_level = False

    has_constant = False
    try:
        constants_mod = importlib.import_module("huggingface_hub.constants")
        has_constant = hasattr(constants_mod, "HF_HUB_OFFLINE")
    except Exception:
        has_constant = False

    assert has_top_level or has_constant, (
        "huggingface_hub dropped both ``is_offline_mode`` AND "
        "``huggingface_hub.constants.HF_HUB_OFFLINE``; "
        "fix_huggingface_hub can no longer re-inject the helper."
    )


# torch


def test_torch_nn_init_trunc_normal_exists():
    """``patch_trunc_normal_precision_issue``: fp16/bf16 wrapper monkey-patches
    torch.nn.init.trunc_normal_, which must still exist."""
    pytest.importorskip("torch")
    import torch.nn.init as init_mod

    assert callable(getattr(init_mod, "trunc_normal_", None)), (
        "torch.nn.init.trunc_normal_ removed/renamed; "
        "patch_trunc_normal_precision_issue cannot wrap it."
    )


# xformers


def test_xformers_is_post_num_splits_key_fix_or_not_installed():
    """``fix_xformers_performance_issue``: xformers <0.0.29 has the
    ``num_splits_key=-1`` perf bug Unsloth rewrites at install time."""
    if importlib.util.find_spec("xformers") is None:
        pytest.skip("xformers not installed -- nothing to drift-check.")
    x_v = _safe_version(importlib_version("xformers"))
    cutoff = _PkgVersion("0.0.29")
    if x_v < cutoff:
        pytest.fail(
            f"DRIFT DETECTED: xformers=={x_v} < {cutoff}; "
            "fix_xformers_performance_issue rewrites "
            "ops/fmha/cutlass.py num_splits_key=-1 -> None."
        )


# transformers (PreTrainedModel base import sanity)


def test_transformers_pretrained_model_has_get_input_embeddings():
    """``patch_enable_input_require_grads``: its replacement calls
    ``get_input_embeddings`` per submodule, so the accessor must still exist."""
    pytest.importorskip("transformers")
    from transformers import PreTrainedModel

    assert hasattr(PreTrainedModel, "get_input_embeddings"), (
        "PreTrainedModel.get_input_embeddings was renamed or removed; "
        "patch_enable_input_require_grads's replacement no longer compiles."
    )


# accelerate -- ``is_X_available`` API stability used across the fixes


# Regression for https://github.com/unslothai/unsloth/issues/4188:
# Qwen3_5ForConditionalGeneration uses loss_type='ForConditionalGeneration', a
# separate LOSS_MAPPING key left unpatched, falling back to stock ForCausalLMLoss
# whose logits.float() OOMs on <=24 GB GPUs.


def _reset_loss_mapping(mapping, saved):
    mapping.clear()
    mapping.update(saved)


def test_patch_loss_functions_covers_conditional_generation():
    """patch_loss_functions() must repoint every ForCausalLMLoss alias to the
    Unsloth kernel, not just LOSS_MAPPING['ForCausalLM']."""
    lu = pytest.importorskip("transformers.loss.loss_utils")
    cel = pytest.importorskip("unsloth.kernels.cross_entropy_loss")

    saved = dict(lu.LOSS_MAPPING)
    try:
        cel.patch_loss_functions(torch_compile = False)

        unsloth_loss = lu.LOSS_MAPPING.get("ForCausalLM")
        assert unsloth_loss is not None
        assert "Unsloth" in str(
            unsloth_loss
        ), f"LOSS_MAPPING['ForCausalLM'] was not replaced: {unsloth_loss}"

        cg_loss = lu.LOSS_MAPPING.get("ForConditionalGeneration")
        assert cg_loss is unsloth_loss, (
            f"LOSS_MAPPING['ForConditionalGeneration'] not patched: {cg_loss}. "
            f"Qwen3_5ForConditionalGeneration will silently use the stock "
            f"ForCausalLMLoss and OOM at large sequence lengths."
        )
    finally:
        _reset_loss_mapping(lu.LOSS_MAPPING, saved)


def test_patch_loss_functions_does_not_touch_other_loss_types():
    """patch_loss_functions() must not overwrite unrelated loss types with the causal-LM kernel."""
    lu = pytest.importorskip("transformers.loss.loss_utils")
    cel = pytest.importorskip("unsloth.kernels.cross_entropy_loss")

    non_causal_keys = {
        k for k, v in lu.LOSS_MAPPING.items() if getattr(v, "__name__", "") != "ForCausalLMLoss"
    }

    saved = dict(lu.LOSS_MAPPING)
    try:
        cel.patch_loss_functions(torch_compile = False)

        unsloth_loss = lu.LOSS_MAPPING.get("ForCausalLM")
        for key in non_causal_keys:
            assert lu.LOSS_MAPPING.get(key) is not unsloth_loss, (
                f"patch_loss_functions() incorrectly overwrote "
                f"LOSS_MAPPING['{key}'] with the Unsloth ForCausalLM kernel."
            )
    finally:
        _reset_loss_mapping(lu.LOSS_MAPPING, saved)


def test_accelerate_utils_imports_module_present():
    """``disable_broken_wandb`` + ``fix_trl_vllm_ascend`` both reach into
    accelerate.utils.imports."""
    pytest.importorskip("accelerate")
    mod = pytest.importorskip("accelerate.utils.imports")
    # is_wandb_available is the canonical target of disable_broken_wandb.
    assert hasattr(mod, "is_wandb_available"), (
        "accelerate.utils.imports.is_wandb_available is gone; "
        "disable_broken_wandb cannot patch the source module."
    )


def test_accelerate_recursively_apply_empty_logits_patch():
    """patch_accelerate_recursively_apply overrides recursively_apply to bypass EmptyLogits."""
    pytest.importorskip("accelerate")

    import accelerate.utils.operations as acc_ops
    from unsloth.import_fixes import patch_accelerate_recursively_apply

    class EmptyLogits:
        pass

    e = EmptyLogits()
    patch_accelerate_recursively_apply()

    res = acc_ops.recursively_apply(lambda x: x, e, error_on_other_type = True)
    assert res is e


def test_accelerate_gather_empty_logits_debug_mode_patch():
    """gather and broadcast bypass EmptyLogits when debug mode is enabled."""
    pytest.importorskip("accelerate")
    from accelerate.state import PartialState, DistributedType
    import accelerate.utils.operations as acc_ops
    from unsloth.import_fixes import patch_accelerate_recursively_apply
    import unittest.mock as mock
    import torch

    class EmptyLogits:
        pass

    e = EmptyLogits()
    patch_accelerate_recursively_apply()

    # Enable debug mode and mock a 2-process distributed state
    state = PartialState()
    orig_debug = state.debug
    orig_dist_type = state.distributed_type
    orig_num_processes = state.num_processes
    orig_device = state.device

    state.debug = True
    state.distributed_type = DistributedType.MULTI_GPU
    state.num_processes = 2

    def mock_gather_object(obj, *args, **kwargs):
        return [obj] * state.num_processes

    def mock_gpu_gather(tensor, *args, **kwargs):
        def _gather_one(t):
            if t.ndim == 0:
                t = t.clone()[None]
            return torch.cat([t] * state.num_processes, dim = 0)

        return acc_ops.recursively_apply(_gather_one, tensor, error_on_other_type = True)

    def mock_gpu_broadcast(data, *args, **kwargs):
        return data

    try:
        with (
            mock.patch(
                "accelerate.utils.operations.gather_object",
                side_effect = mock_gather_object,
            ),
            mock.patch("accelerate.utils.operations._gpu_gather", side_effect = mock_gpu_gather),
            mock.patch(
                "accelerate.utils.operations._gpu_broadcast",
                side_effect = mock_gpu_broadcast,
            ),
        ):
            state.device = torch.device("cpu")

            # Top-level EmptyLogits gathers to itself
            res = acc_ops.gather(e)
            assert res is e

            # Nested EmptyLogits
            res_nested = acc_ops.gather([e])
            assert isinstance(res_nested, list) and res_nested[0] is e

            # Mixed payload: real tensor gets gathered, EmptyLogits passes through.
            # Tensor must live on state.device or debug-mode device check fails on GPUs.
            real_tensor = torch.tensor([42], device = state.device)
            payload = {"labels": real_tensor, "logits": e}
            res_mixed = acc_ops.gather(payload)

            assert isinstance(res_mixed, dict)
            assert res_mixed["logits"] is e
            # num_processes = 2 -> gathered to [42, 42]
            assert torch.equal(res_mixed["labels"], torch.tensor([42, 42], device = state.device))

            # Broadcast with EmptyLogits
            res_broadcast = acc_ops.broadcast(e)
            assert res_broadcast is e

            # Mixed payload broadcast
            res_broadcast_mixed = acc_ops.broadcast(payload)
            assert isinstance(res_broadcast_mixed, dict)
            assert res_broadcast_mixed["logits"] is e
            assert torch.equal(res_broadcast_mixed["labels"], real_tensor)
    finally:
        state.debug = orig_debug
        state.distributed_type = orig_dist_type
        state.num_processes = orig_num_processes
        state.device = orig_device


def test_accelerate_patch_is_idempotent():
    """Calling patch_accelerate_recursively_apply twice must not stack wrappers."""
    pytest.importorskip("accelerate")
    import accelerate.utils.operations as acc_ops
    from unsloth.import_fixes import patch_accelerate_recursively_apply

    patch_accelerate_recursively_apply()
    recursively_apply = acc_ops.recursively_apply
    find_device = acc_ops.find_device
    patch_accelerate_recursively_apply()
    assert (
        acc_ops.recursively_apply is recursively_apply
    ), "DRIFT DETECTED: recursively_apply was wrapped twice."
    assert acc_ops.find_device is find_device, "DRIFT DETECTED: find_device was wrapped twice."


def test_accelerate_find_device_skips_empty_logits():
    """find_device must search past EmptyLogits and keep None for tensor-free data."""
    pytest.importorskip("accelerate")
    import torch
    import accelerate.utils.operations as acc_ops
    from accelerate.state import PartialState
    from unsloth.import_fixes import patch_accelerate_recursively_apply

    class EmptyLogits:
        pass

    patch_accelerate_recursively_apply()
    tensor = torch.tensor([1.0])
    # Leading sentinel must not stop the search before the real tensor
    assert acc_ops.find_device({"logits": EmptyLogits(), "labels": tensor}) == tensor.device
    # Tensor-free payloads keep returning None (AlignDevicesHook needs it to skip moves)
    assert acc_ops.find_device({"a": 1}) is None
    # Sentinel-only payloads fall back to current device so debug-mode
    # find_device(...).type doesn't raise AttributeError
    assert acc_ops.find_device(EmptyLogits()) == PartialState().device


def test_accelerate_patch_wired_into_gpu_init():
    """The patch must be installed at startup, not only importable."""
    source = Path(__file__).resolve().parent.parent / "unsloth" / "_gpu_init.py"
    source = source.read_text()
    assert "patch_accelerate_recursively_apply()" in source, (
        "DRIFT DETECTED: patch_accelerate_recursively_apply is defined but "
        "never called in _gpu_init.py, so real imports never install it."
    )


# ===========================================================================
# bitsandbytes -- ROCm arch / warp-size detection shape
# ===========================================================================


def test_bitsandbytes_rocm_detection_helpers_recognizable():
    """``fix_bitsandbytes_rocm_arch_detection``: the source sniff only patches
    bnb's ROCm helpers in recognized shapes; fail (don't import) when it drifts."""
    spec = importlib.util.find_spec("bitsandbytes")
    if spec is None:
        pytest.skip("bitsandbytes not installed -- nothing to drift-check.")
    cuda_specs_path = None
    for location in spec.submodule_search_locations or []:
        candidate = os.path.join(location, "cuda_specs.py")
        if os.path.isfile(candidate):
            cuda_specs_path = candidate
            break
    if cuda_specs_path is None:
        pytest.skip("bitsandbytes has no cuda_specs.py (pre-ROCm version).")

    import ast

    with open(cuda_specs_path, "r", encoding = "utf-8") as f:
        source = f.read()
    helpers = [
        node
        for node in ast.walk(ast.parse(source))
        if isinstance(node, ast.FunctionDef)
        and node.name in ("get_rocm_gpu_arch", "get_rocm_warpsize")
    ]
    if not helpers:
        pytest.skip("bitsandbytes cuda_specs has no ROCm detection helpers.")
    for node in helpers:
        segment = ast.get_source_segment(source, node) or ""
        recognized = (
            "subprocess" in segment
            or "get_device_properties" in segment
            or "gcnArchName" in segment
        )
        if not recognized:
            pytest.fail(
                f"DRIFT DETECTED: bitsandbytes.cuda_specs.{node.name} uses "
                "neither subprocess nor torch device properties; "
                "fix_bitsandbytes_rocm_arch_detection's shape sniff will "
                "decline to patch it and Windows ROCm import-time noise / "
                "wrong ROCM_GPU_ARCH may return."
            )


# ===========================================================================
# torchao Windows-ROCm import shim -- fix_torchao_windows_rocm_import
# ===========================================================================
# The shim FRAGMENT-registers the `_c10d_functional` op schemas so real torchao
# imports on a distributed-less Windows ROCm wheel. These verify the schema table
# tracks the installed torch, the FRAGMENT (not DEF) collision semantics, the
# strict no-op / capability gating, and that it is wired into startup. Windows ROCm
# itself cannot be reproduced here, so the transactional acceptance-import + rollback
# is what guarantees no regression on the real device.


def _torch_minor_tuple():
    import torch
    base = torch.__version__.split("+", 1)[0].split(".")
    return (int(base[0]), int(base[1]))


def _live_c10d_functional_ops():
    import torch

    get_ops = getattr(torch._C, "_dispatch_get_all_op_names", None)
    if not callable(get_ops):
        pytest.skip("dispatcher op enumeration unavailable")
    return sorted({n.split("::", 1)[1] for n in get_ops() if n.startswith("_c10d_functional::")})


def _live_dtensor_ops():
    import torch

    get_ops = getattr(torch._C, "_dispatch_get_all_op_names", None)
    if not callable(get_ops):
        pytest.skip("dispatcher op enumeration unavailable")
    return sorted({n.split("::", 1)[1] for n in get_ops() if n.startswith("_dtensor::")})


def test_torchao_rocm_shim_schema_table_matches_installed_torch():
    """The `_c10d_functional` schema table must exactly match the ops the installed
    torch registers (op set + canonical schema strings). A minor with no row means the
    shim fail-closes there (safe, no coverage) -> skip; a present row must be exact."""
    from unsloth.import_fixes import _C10D_FUNCTIONAL_SCHEMAS, _schema_op_name

    import torch

    native = _live_c10d_functional_ops()
    if not native:
        pytest.skip("no native _c10d_functional ops (distributed-less torch build).")

    minor = _torch_minor_tuple()
    schemas = _C10D_FUNCTIONAL_SCHEMAS.get(minor)
    if schemas is None:
        pytest.skip(
            f"no shim schema row for torch {minor}; fix_torchao_windows_rocm_import "
            f"fail-closes here (safe). Add a reviewed tuple to enable it (ops: {native})."
        )

    table_ops = sorted(_schema_op_name(s) for s in schemas)
    assert table_ops == native, (
        f"DRIFT DETECTED: torchao shim _c10d_functional table for torch {minor} lists "
        f"{table_ops} but the installed torch registers {native}. Update "
        f"_C10D_FUNCTIONAL_SCHEMAS."
    )

    parse = getattr(torch._C, "parse_schema", None)
    if not callable(parse):
        return
    real = {}
    for op in native:
        packet = getattr(torch.ops._c10d_functional, op)
        overload = packet.overloads()[0]
        real[op] = str(getattr(packet, overload)._schema)
    for s in schemas:
        parsed = parse(f"_c10d_functional::{s}")  # must not raise
        name = _schema_op_name(s)
        assert str(parsed) == real[name], (
            f"DRIFT DETECTED: torchao shim schema for _c10d_functional::{name}\n"
            f"  shim:  {parsed}\n  torch: {real[name]}"
        )


def test_torchao_rocm_shim_dtensor_schema_matches_installed_torch():
    """The `_dtensor` schema table must exactly match the ops the installed torch registers.
    torchao's `from torch.distributed._tensor import DTensor` runs
    `register_fake("_dtensor::shard_dim_alltoall")` at import, which raises unless the op is
    defined, so the shim must define this namespace too (not only _c10d_functional)."""
    from unsloth.import_fixes import _DTENSOR_SCHEMAS, _schema_op_name

    import torch

    native = _live_dtensor_ops()
    if not native:
        pytest.skip("no native _dtensor ops (distributed-less torch build).")

    minor = _torch_minor_tuple()
    schemas = _DTENSOR_SCHEMAS.get(minor)
    if schemas is None:
        pytest.skip(
            f"no shim _dtensor row for torch {minor}; fix_torchao_windows_rocm_import "
            f"fail-closes here (safe). Add a reviewed tuple to enable it (ops: {native})."
        )

    table_ops = sorted(_schema_op_name(s) for s in schemas)
    assert table_ops == native, (
        f"DRIFT DETECTED: torchao shim _dtensor table for torch {minor} lists {table_ops} "
        f"but the installed torch registers {native}. Update _DTENSOR_SCHEMAS."
    )

    parse = getattr(torch._C, "parse_schema", None)
    if not callable(parse):
        return
    real = {}
    for op in native:
        packet = getattr(torch.ops._dtensor, op)
        overload = packet.overloads()[0]
        real[op] = str(getattr(packet, overload)._schema)
    for s in schemas:
        parsed = parse(f"_dtensor::{s}")  # must not raise
        name = _schema_op_name(s)
        assert str(parsed) == real[name], (
            f"DRIFT DETECTED: torchao shim schema for _dtensor::{name}\n"
            f"  shim:  {parsed}\n  torch: {real[name]}"
        )


def test_torchao_rocm_shim_strict_noop_on_non_windows():
    """On a non-Windows / distributed-present box the shim must not touch sys.modules,
    torch.ops, or torch.distributed.is_available()."""
    from unsloth.import_fixes import fix_torchao_windows_rocm_import

    import torch

    assert sys.platform != "win32"
    before_ext = "torch._C._distributed_c10d" in sys.modules
    before_avail = torch.distributed.is_available()
    before_ops = set(_live_c10d_functional_ops())

    fix_torchao_windows_rocm_import()

    assert ("torch._C._distributed_c10d" in sys.modules) == before_ext
    assert torch.distributed.is_available() == before_avail
    assert set(_live_c10d_functional_ops()) == before_ops


def test_torchao_rocm_shim_native_present_builds_no_library(monkeypatch):
    """Even with the platform gates spoofed to look like Windows ROCm, a box that already
    has real distributed (native _c10d_functional ops / is_available) must trip a guard
    before any torch.library.Library is constructed."""
    from unsloth.import_fixes import (
        fix_torchao_windows_rocm_import,
        _native_c10d_functional_present,
    )

    import torch

    monkeypatch.setattr(sys, "platform", "win32")
    if not getattr(getattr(torch, "version", None), "hip", None):
        monkeypatch.setattr(torch.version, "hip", "6.4.0", raising = False)

    assert _native_c10d_functional_present(torch) is True

    calls = {"n": 0}
    real_library = torch.library.Library

    def _tripwire(*args, **kwargs):
        calls["n"] += 1
        return real_library(*args, **kwargs)

    monkeypatch.setattr(torch.library, "Library", _tripwire)
    fix_torchao_windows_rocm_import()
    assert (
        calls["n"] == 0
    ), "torchao shim constructed a torch.library.Library despite a real distributed build."


_TORCHAO_ROCM_FRAGMENT_PROBE = """
import torch
ns = "_unsloth_torchao_shim_probe_ns"
# A second DEF on a namespace raises; FRAGMENT must not -- that is why the shim uses
# FRAGMENT (no fatal collision with a native C++ TORCH_LIBRARY). NB: the first DEF must be
# held by a strong ref, else CPython GCs it (its __del__ calls _destroy) and releases the
# namespace before the second call -- the same reason the shim keeps a strong ref to its
# FRAGMENT Library so its registered schemas are not dropped.
_hold = torch.library.Library(ns, "DEF")
raised = False
try:
    torch.library.Library(ns, "DEF")
except Exception:
    raised = True
assert raised, "second DEF unexpectedly did not raise"
frag = torch.library.Library(ns, "FRAGMENT")  # must not raise
frag.define("myop(Tensor x) -> Tensor")
assert hasattr(torch.ops, ns) and hasattr(getattr(torch.ops, ns), "myop")
frag._destroy()
print("FRAGMENT_OK")
"""


def test_torchao_rocm_shim_fragment_semantics_subprocess():
    """FRAGMENT-define + resolve + _destroy work and FRAGMENT (unlike a second DEF) never
    collides -- run in a subprocess since dispatcher registration is process-global."""
    import subprocess

    result = subprocess.run(
        [sys.executable, "-c", _TORCHAO_ROCM_FRAGMENT_PROBE],
        capture_output = True,
        text = True,
        timeout = 300,
    )
    assert (
        "FRAGMENT_OK" in result.stdout
    ), f"FRAGMENT probe failed:\nSTDOUT:{result.stdout}\nSTDERR:{result.stderr}"


def test_torchao_rocm_shim_source_has_guards_fragment_and_rollback():
    """The shim source must keep its win32 + HIP + is_available() gates, use FRAGMENT (not a
    DEF on _c10d_functional), and roll back via Library._destroy."""
    import inspect

    from unsloth import import_fixes

    src = inspect.getsource(import_fixes.fix_torchao_windows_rocm_import)
    assert "win32" in src, "missing win32 guard"
    assert "hip" in src, "missing torch.version.hip guard"
    assert '"rocm" in' in src, (
        "shim ROCm detection must also accept a 'rocm'-tagged __version__ wheel (parity with "
        "the Studio is_win32_rocm() helper), not gate on torch.version.hip alone"
    )
    assert "is_available()" in src, "missing is_available() guard"
    assert '"FRAGMENT"' in src, "shim must register with FRAGMENT, not DEF"
    assert '"_c10d_functional", "DEF"' not in src, "shim must never DEF _c10d_functional"
    assert '"_dtensor", "DEF"' not in src, "shim must never DEF _dtensor"
    assert '"_dtensor"' in src, "shim must also register the _dtensor namespace"
    assert "_destroy" in src, "missing rollback via Library._destroy"


def test_torchao_rocm_shim_wired_into_gpu_init():
    """The shim must be called at startup (before `import unsloth_zoo`), not merely
    importable (mirrors test_accelerate_patch_wired_into_gpu_init)."""
    source = Path(__file__).resolve().parent.parent / "unsloth" / "_gpu_init.py"
    text = source.read_text()
    assert "fix_torchao_windows_rocm_import()" in text, (
        "DRIFT DETECTED: fix_torchao_windows_rocm_import is defined but never called in "
        "_gpu_init.py, so real imports never install it."
    )