unsloth/tests/test_new_mapper_fetched_fp8.py

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.

"""Regression tests for what ``_get_new_mapper`` hands back to the upgrade probe.

``test_new_mapper_no_global_leak.py`` serves the repo's own ``mapper.py`` as both installed
and fetched source, so it cannot tell a fetched table from a fresh copy of the installed one.
Two gaps it misses:

1. The probe must answer for an fp8 repo only the FETCHED mapper knows, so an extra ``"8"``
   entry is spliced into the fetched source only. Isolating the exec without returning the
   fetched fp8 tables would silently drop the fp8 half of the upgrade check.
2. The probe must survive a fetched ``mapper.py`` with no fp8 tables (anything older, or a
   future rename): reading them with ``[]`` raises ``KeyError`` into the bare ``except``,
   taking the 4bit half, the probe's whole purpose, down with it.

``loader_utils`` imports torch, so ast-extract the resolvers and run them against a stubbed
``requests``, as in ``tests/test_bad_mappings_redirect.py``.
"""

import ast
import os
import sys
import types

_MODELS = os.path.join(os.path.dirname(__file__), os.pardir, "unsloth", "models")

_WANTED = {"__get_model_name", "_resolve_with_mappers", "_get_new_mapper", "get_model_name"}

# An fp8 ("8") model, spliced into the FETCHED mapper only.
_NEW_KEY = "unsloth/Zeta-9B-Only-On-Main"
_NEW_OFFICIAL = "zeta-org/Zeta-9B-Only-On-Main-FP8"
_NEW_BLOCK = "unsloth/Zeta-9B-Only-On-Main-FP8-Block"
_NEW_ROW = "unsloth/Zeta-9B-Only-On-Main-FP8-Row"
_ANCHOR = '    "unsloth/Kimi-K2-Instruct-BF16" : ('


def _mapper_source():
    with open(os.path.join(_MODELS, "mapper.py"), encoding = "utf-8") as f:
        return f.read()


def _with_extra_fp8_model(source):
    assert _ANCHOR in source, "anchor moved; update this test"
    entry = (
        f'    "{_NEW_KEY}" : {{\n'
        f'        "16" : ("{_NEW_KEY}", "zeta-org/Zeta-9B-Only-On-Main"),\n'
        f'        "8"  : ("{_NEW_OFFICIAL}", "{_NEW_BLOCK}", "{_NEW_ROW}"),\n'
        f"    }},\n"
    )
    return source.replace(_ANCHOR, entry + _ANCHOR, 1)


def _without_fp8_tables(source):
    """A mapper.py from before the fp8 tables existed."""
    return source.replace("FLOAT_TO_FP8_BLOCK_MAPPER", "SOME_OTHER_BLOCK_TABLE").replace(
        "FLOAT_TO_FP8_ROW_MAPPER", "SOME_OTHER_ROW_TABLE"
    )


class _FakeResponse:
    def __init__(self, text):
        self.text = text

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        return False


def _install_fake_requests(monkeypatch, text):
    module = types.ModuleType("requests")
    module.get = lambda url, timeout = None: _FakeResponse(text)
    monkeypatch.setitem(sys.modules, "requests", module)


def _install_fake_vllm_absent(monkeypatch, namespace):
    """vllm >= 0.12.0 returns early from __get_model_name, leaving the probe unreachable."""
    monkeypatch.delitem(sys.modules, "vllm", raising = False)
    fake = types.ModuleType("importlib")
    fake.util = types.SimpleNamespace(find_spec = lambda name: None)
    namespace["importlib"] = fake


def _load_resolver(installed_source):
    """Stand-in for loader_utils' module globals, built from `installed_source`."""
    from unsloth_zoo.utils import Version

    mapper_ns = {}
    exec(compile(installed_source, "mapper.py", "exec"), mapper_ns)

    namespace = {
        "INT_TO_FLOAT_MAPPER": mapper_ns["INT_TO_FLOAT_MAPPER"],
        "FLOAT_TO_INT_MAPPER": mapper_ns["FLOAT_TO_INT_MAPPER"],
        "MAP_TO_UNSLOTH_16bit": mapper_ns["MAP_TO_UNSLOTH_16bit"],
        "FLOAT_TO_FP8_BLOCK_MAPPER": mapper_ns["FLOAT_TO_FP8_BLOCK_MAPPER"],
        "FLOAT_TO_FP8_ROW_MAPPER": mapper_ns["FLOAT_TO_FP8_ROW_MAPPER"],
        "SUPPORTS_FOURBIT": True,
        "transformers_version": Version("4.57.6"),
        "Version": Version,
        "os": os,
    }
    with open(os.path.join(_MODELS, "loader_utils.py"), encoding = "utf-8") as f:
        tree = ast.parse(f.read())
    for node in tree.body:
        if isinstance(node, ast.Assign) and any(
            getattr(t, "id", None) in ("BAD_MAPPINGS", "_OFFLINE_ENV_VALUES", "_OFFLINE_ENV_KEYS")
            for t in node.targets
        ):
            exec(compile(ast.Module([node], []), "<assign>", "exec"), namespace)
        elif isinstance(node, ast.FunctionDef) and (
            node.name in _WANTED or node.name == "_env_says_offline"
        ):
            exec(compile(ast.Module([node], []), node.name, "exec"), namespace)
    return namespace


def test_probe_answers_for_an_fp8_repo_only_the_fetched_mapper_knows(monkeypatch):
    installed = _mapper_source()
    namespace = _load_resolver(installed)
    installed_block = namespace["FLOAT_TO_FP8_BLOCK_MAPPER"]
    installed_row = namespace["FLOAT_TO_FP8_ROW_MAPPER"]
    assert _NEW_OFFICIAL.lower() not in installed_block, "the installed table must not know it"

    _install_fake_requests(monkeypatch, _with_extra_fp8_model(installed))
    _install_fake_vllm_absent(monkeypatch, namespace)

    try:
        resolved = namespace["get_model_name"](
            _NEW_OFFICIAL, load_in_4bit = False, load_in_fp8 = "block"
        )
    except NotImplementedError as error:
        assert "not supported in your current Unsloth version" in str(error)
    else:
        raise AssertionError(
            f"a fetched-only fp8 repo must raise the upgrade error, got {resolved!r}"
        )

    # Answering must not have adopted the fetched tables.
    assert namespace["FLOAT_TO_FP8_BLOCK_MAPPER"] is installed_block
    assert namespace["FLOAT_TO_FP8_ROW_MAPPER"] is installed_row
    assert _NEW_OFFICIAL.lower() not in namespace["FLOAT_TO_FP8_BLOCK_MAPPER"]


def test_probe_survives_a_fetched_mapper_without_the_fp8_tables(monkeypatch):
    installed = _mapper_source()
    namespace = _load_resolver(installed)
    _install_fake_requests(monkeypatch, _without_fp8_tables(installed))

    int_to_float, float_to_int, map_to_16bit = namespace["_get_new_mapper"]()[:3]

    assert (
        int_to_float and float_to_int and map_to_16bit
    ), "a fetched mapper.py without the fp8 tables must not take the 4bit upgrade check down"