scripts/patch_mamba_ssm.py

"""Patch mamba-ssm's setup.py to compile CUDA kernels for only sm_89.

Mamba-ssm's upstream setup.py builds for 8+ GPU architectures (sm_53,
sm_60, sm_70, sm_75, sm_80, sm_86, sm_89, sm_90). On the HF Spaces
build VM (~16 GB RAM, ~60 min wall-clock cap) compiling all of them
either OOMKills or hits the build timeout. This patch reduces the list
to only sm_89 (the L40S target the production VoxPixel and bench
Spaces both run on).

Strategy (line-based, robust to upstream formatting changes):

    1. Find the line ``cc_flag = []``.
    2. Replace it with a list literal pinned to sm_89.
    3. Comment out every subsequent line containing
       ``cc_flag.append(`` (these would otherwise re-inject the other
       architectures, including via conditional ``if`` blocks).

Idempotent: the patch leaves a sentinel marker and is a no-op on
already-patched files.

Usage:
    python scripts/patch_mamba_ssm.py path/to/setup.py
"""

import re
import sys
from pathlib import Path


SENTINEL = "# PATCHED: cc_flag limited to sm_89 by patch_mamba_ssm.py"
TARGET_ARCH = "89"


def _leading_indent(line: str) -> str:
    return line[: len(line) - len(line.lstrip())]


def patch(path: Path) -> bool:
    src = path.read_text()
    if SENTINEL in src:
        print(f"{path}: already patched (sentinel present), skipping.")
        return False
    if "cc_flag" not in src:
        raise SystemExit(
            f"ERROR: 'cc_flag' not found in {path}. "
            "Upstream mamba-ssm setup.py may have changed."
        )

    init_pattern = re.compile(r"^cc_flag\s*=\s*\[\]\s*$")
    out_lines: list[str] = []
    init_patched = False
    appends_neutralized = 0

    for line in src.splitlines(keepends=True):
        stripped = line.strip()

        if not init_patched and init_pattern.match(stripped):
            indent = _leading_indent(line)
            out_lines.append(f"{indent}{SENTINEL}\n")
            out_lines.append(
                f'{indent}cc_flag = ["-gencode", '
                f'"arch=compute_{TARGET_ARCH},code=sm_{TARGET_ARCH}"]\n'
            )
            init_patched = True
            continue

        if init_patched and "cc_flag.append(" in stripped:
            # Replace with `pass` (not a comment) so any enclosing
            # `if:` / `else:` block keeps a syntactically valid body.
            indent = _leading_indent(line)
            out_lines.append(
                f"{indent}pass  # (was: {stripped})\n"
            )
            appends_neutralized += 1
            continue

        out_lines.append(line)

    if not init_patched:
        raise SystemExit(
            f"ERROR: line 'cc_flag = []' not found in {path}. "
            "Upstream mamba-ssm setup.py format may have changed."
        )

    path.write_text("".join(out_lines))
    print(
        f"{path}: patched. cc_flag pinned to sm_{TARGET_ARCH}; "
        f"{appends_neutralized} subsequent cc_flag.append calls neutralized."
    )
    return True


def main() -> None:
    if len(sys.argv) != 2:
        print("usage: patch_mamba_ssm.py <path/to/setup.py>", file=sys.stderr)
        sys.exit(2)
    path = Path(sys.argv[1])
    if not path.is_file():
        raise SystemExit(f"ERROR: {path} not found.")
    patch(path)


if __name__ == "__main__":
    main()