unsloth/tests/python/test_ssm_kernels_install_python_stack.py

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

"""Setup installs flash-attn but not causal-conv1d or mamba-ssm: those (~550 MB) install when a
model that needs them is trained or loaded (core/training/worker.py, utils/ssm_runtime.py), and
`unsloth install-kernels` still provisions them ahead of time for offline images. After torch
moves, setup still removes an installed copy whose compiled extension no longer imports."""

from __future__ import annotations

import ast
import contextlib
import re
import subprocess
import sys
from pathlib import Path

import pytest

STUDIO_DIR = Path(__file__).resolve().parents[2] / "studio"
sys.path.insert(0, str(STUDIO_DIR))
sys.path.insert(0, str(STUDIO_DIR / "backend"))

import install_python_stack as ips  # noqa: E402

SOURCE = (STUDIO_DIR / "install_python_stack.py").read_text(encoding = "utf-8")
TREE = ast.parse(SOURCE)


def _called_names() -> set[str]:
    return {
        node.func.id
        for node in ast.walk(TREE)
        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name)
    }


def test_setup_has_no_ssm_kernel_step():
    assert not hasattr(ips, "_ensure_ssm_kernels")
    assert "install_kernel" not in _called_names()
    assert '_progress("mamba kernels")' not in SOURCE
    assert "causal-conv1d + mamba-ssm" not in SOURCE


def test_flash_attn_and_the_stale_kernel_check_run_at_all_three_sites():
    calls = [m.start() for m in re.finditer(r"\n\s+_ensure_flash_attn\(\)\n", SOURCE)]
    assert len(calls) == 3
    for start in calls:
        following = SOURCE[start : start + 120]
        assert "_remove_stale_ssm_kernels()" in following, following
    assert "base_total += 1  # flash-attn, Linux and Windows" in SOURCE


class _Kernels:
    """Which SSM distributions are installed and how their extension import ends: True loads,
    False raises ImportError, None hangs, an int is a raw exit code (negative = signal)."""

    def __init__(
        self,
        installed: dict,
        torch_imports = True,
    ):
        self.installed = dict(installed)
        self.torch_imports = torch_imports
        self.probes: list[str] = []
        self.removed: list[str] = []

    def version(self, distribution):
        return "1.0" if distribution in self.installed else None

    def run(self, cmd, *args, **kwargs):
        code = cmd[-1]
        self.probes.append(code)
        if code == "import torch":
            if self.torch_imports is None:
                raise subprocess.TimeoutExpired(cmd, 1)
            if self.torch_imports:
                return subprocess.CompletedProcess(cmd, 0, stderr = "")
            return subprocess.CompletedProcess(
                cmd, 1, stderr = "OSError: libtorch_cuda.so: cannot open shared object file"
            )
        assert "import torch" in code  # the extension probe proves torch imports too
        extension = "causal_conv1d_cuda" if "causal_conv1d_cuda" in code else "selective_scan_cuda"
        distribution = "causal-conv1d" if extension == "causal_conv1d_cuda" else "mamba-ssm"
        state = self.installed[distribution]
        if state is None:
            raise subprocess.TimeoutExpired(cmd, 1)
        if state is True:
            return subprocess.CompletedProcess(cmd, 0, stderr = "")
        if state is False:
            return subprocess.CompletedProcess(
                cmd,
                ips._SSM_EXTENSION_BROKEN_EXIT,
                stderr = f"{extension}: ImportError: {extension}.so: undefined symbol: _ZN3c10",
            )
        return subprocess.CompletedProcess(cmd, state, stderr = "Killed")

    def uninstall(self, distribution):
        self.removed.append(distribution)
        self.installed.pop(distribution, None)
        return True


@contextlib.contextmanager
def _free_lock(name, **kwargs):
    yield True


def _run_stale_check(monkeypatch, host: _Kernels, **flags):
    monkeypatch.setattr(ips, "NO_TORCH", flags.get("no_torch", False))
    monkeypatch.setattr(ips, "IS_MACOS", flags.get("macos", False))
    monkeypatch.setattr(ips, "_installed_distribution_version", host.version)
    monkeypatch.setattr(ips, "_uninstall_distribution", host.uninstall)
    monkeypatch.setattr(ips, "_count_install_action", lambda: None)
    monkeypatch.setattr(ips, "kernel_install_lock", flags.get("lock", _free_lock))
    warnings = []
    monkeypatch.setattr(ips, "_step", lambda label, value, color_fn = None: warnings.append(value))
    monkeypatch.setattr(ips, "_note", lambda message, color_fn = None: warnings.append(message))
    monkeypatch.setattr(ips.subprocess, "run", host.run)
    installs = []
    monkeypatch.setattr(ips, "install_prebuilt", lambda *a, **k: installs.append(a) or "installed")
    ips._remove_stale_ssm_kernels()
    assert installs == []  # never installs anything
    return warnings


def test_stale_kernels_that_no_longer_import_are_removed(monkeypatch):
    host = _Kernels({"causal-conv1d": False, "mamba-ssm": False})
    warnings = _run_stale_check(monkeypatch, host)
    assert host.removed == ["causal-conv1d", "mamba-ssm"]
    assert len(warnings) == 2
    # The probe's stderr tail says why.
    assert "undefined symbol" in warnings[0] and "selective_scan_cuda" in warnings[1]


def test_kernels_that_still_import_are_kept(monkeypatch):
    host = _Kernels({"causal-conv1d": True, "mamba-ssm": False})
    _run_stale_check(monkeypatch, host)
    assert host.removed == ["mamba-ssm"]


def test_absent_kernels_are_not_probed_or_installed(monkeypatch):
    host = _Kernels({})
    _run_stale_check(monkeypatch, host)
    assert host.probes == [] and host.removed == []


def test_a_hung_probe_is_inconclusive_and_keeps_the_kernel(monkeypatch):
    host = _Kernels({"causal-conv1d": None})
    _run_stale_check(monkeypatch, host)
    assert host.removed == []


def test_no_torch_or_macos_skips_the_check(monkeypatch):
    for flags in ({"no_torch": True}, {"macos": True}):
        host = _Kernels({"causal-conv1d": False})
        _run_stale_check(monkeypatch, host, **flags)
        assert host.probes == []


@pytest.mark.parametrize("torch_imports", [False, None])
def test_a_torch_that_does_not_import_keeps_every_kernel(monkeypatch, torch_imports):
    # A half-finished repair or a transient failure in torch says nothing about the kernels, and an
    # offline machine could not reinstall them.
    host = _Kernels({"causal-conv1d": False, "mamba-ssm": False}, torch_imports = torch_imports)
    warnings = _run_stale_check(monkeypatch, host)
    assert host.removed == []
    assert host.probes == ["import torch"]
    if torch_imports is False:
        assert "libtorch_cuda.so" in warnings[0]


@pytest.mark.parametrize("exit_code", [-9, -11, 1, 2])
def test_a_killed_or_unrelated_failure_is_inconclusive(monkeypatch, exit_code):
    # A signal (OOM killer, segfault) or an exit that is not the extension's ImportError / OSError.
    host = _Kernels({"causal-conv1d": exit_code, "mamba-ssm": exit_code})
    warnings = _run_stale_check(monkeypatch, host)
    assert host.removed == []
    assert all("kept it" in message for message in warnings)


def test_the_probed_extensions_match_the_kernel_checks():
    for name, extension in ips._SSM_KERNEL_EXTENSIONS.items():
        assert ips.KERNELS[name][1] == f"import torch, {extension}"


def test_the_kernel_cli_still_provisions_both_ahead_of_time():
    from utils import kernel_install
    assert {"causal_conv1d", "mamba_ssm"} <= set(kernel_install.KERNELS)
    assert callable(kernel_install.install_kernel)


def test_a_busy_install_lock_keeps_the_kernel(monkeypatch):
    # A Studio load verifying a fresh install holds this lock: never remove it underneath that.
    @contextlib.contextmanager
    def _busy(name, **kwargs):
        yield False

    host = _Kernels({"causal-conv1d": False})
    warnings = _run_stale_check(monkeypatch, host, lock = _busy)
    assert host.removed == []
    assert any("install lock is busy" in message for message in warnings)


def test_a_kernel_reinstalled_while_waiting_on_the_lock_is_kept(monkeypatch):
    # Broken when first probed, but a Studio load / `unsloth install-kernels` held the lock and
    # reinstalled it: the re-probe under the lock loads, so nothing is removed.
    host = _Kernels({"causal-conv1d": False})

    @contextlib.contextmanager
    def _lock_after_reinstall(name, **kwargs):
        host.installed["causal-conv1d"] = True
        yield True

    warnings = _run_stale_check(monkeypatch, host, lock = _lock_after_reinstall)
    assert host.removed == []
    assert any("now loads" in message and "kept it" in message for message in warnings)
    extension_probes = [code for code in host.probes if code != "import torch"]
    assert len(extension_probes) == 2  # probed again under the lock


def test_an_inconclusive_re_probe_under_the_lock_keeps_the_kernel(monkeypatch):
    host = _Kernels({"causal-conv1d": False})

    @contextlib.contextmanager
    def _lock_then_hang(name, **kwargs):
        host.installed["causal-conv1d"] = None
        yield True

    warnings = _run_stale_check(monkeypatch, host, lock = _lock_then_hang)
    assert host.removed == []
    assert any("kept it" in message for message in warnings)