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 _dump(path: Path, src: str) -> None:
    """Print every line of `src` numbered, plus matches for `cc_flag`,
    so we can adapt the patcher regex offline without re-running the
    Docker build to debug."""
    print(f"\n========== BEGIN DUMP OF {path} ==========", flush=True)
    for i, line in enumerate(src.splitlines(), 1):
        print(f"{i:4d} | {line}", flush=True)
    print(f"========== END DUMP OF {path} ==========\n", flush=True)
    print("Lines mentioning 'cc_flag':", flush=True)
    for i, line in enumerate(src.splitlines(), 1):
        if "cc_flag" in line:
            print(f"  {i:4d}: {line}", flush=True)


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:
        _dump(path, src)
        raise SystemExit(
            f"ERROR: 'cc_flag' not found in {path}. "
            "Upstream mamba-ssm setup.py may have changed. "
            "Full file contents printed above."
        )

    # Match `cc_flag = []`, `cc_flag = list()`, and typed variants like
    # `cc_flag: list[str] = []` / `cc_flag: List[str] = list()`.
    init_pattern = re.compile(
        r"""
        ^cc_flag                 # variable name
        (?:\s*:\s*[^=]+?)?        # optional type annotation
        \s*=\s*
        (?:\[\s*\]|list\(\s*\))   # [] or list()
        \s*(?:\#.*)?$             # optional trailing comment
        """,
        re.VERBOSE,
    )
    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:
        _dump(path, src)
        raise SystemExit(
            f"ERROR: line 'cc_flag = []' not found in {path}. "
            "Upstream mamba-ssm setup.py format may have changed. "
            "Full file contents printed above so the patcher regex "
            "can be adapted."
        )

    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()