FROM nvidia/cuda:12.1.1-devel-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV OMP_NUM_THREADS=1
ENV PYTORCH_ALLOC_CONF=expandable_segments:True,max_split_size_mb=128
ENV ENABLE_CUDNN_BENCHMARK=true
ENV ENABLE_TORCH_COMPILE=false
ENV INFERENCE_TIMEOUT=1800
ENV MAX_GRADIO_CONCURRENCY=1
ENV TORCH_CUDA_ARCH_LIST=8.9
ENV CUDA_HOME=/usr/local/cuda
ENV MPLCONFIGDIR=/tmp/matplotlib

# Cap CUDA-extension build parallelism. mamba_ssm hardcodes 8+ GPU arches
# in setup.py (sm_53..sm_90) which makes the build both memory-heavy AND
# wall-clock heavy on HF Spaces (~16 GB RAM, ~60 min build cap).
# We further patch mamba_ssm below to compile for only sm_89 (L40S target),
# so MAX_JOBS=2 is now safe and ~6x faster than the all-arch serial build.
ENV MAX_JOBS=2
ENV NVCC_THREADS=2

RUN apt-get update && apt-get install -y \
    python3.10 \
    python3.10-venv \
    python3-pip \
    git \
    curl \
    build-essential \
    ninja-build \
    libgl1 \
    libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

RUN nvcc --version

ENV VENV=/opt/venv
RUN python3.10 -m venv $VENV
ENV PATH="$VENV/bin:$PATH"

RUN pip install --upgrade pip wheel setuptools packaging ninja

RUN pip install --index-url https://download.pytorch.org/whl/cu121 \
    torch torchvision torchaudio

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY . .

# Download mamba-ssm source, patch its hardcoded cc_flag list to only build
# for sm_89 (L40S / Ada / RTX 40-series), then install. This avoids HF
# Spaces 60-minute build timeouts when compiling for all 8+ archs serially.
#
# Notes:
# - --no-build-isolation on `pip download` is critical: without it, pip
#   spins up a clean build env that tries to fetch torch from PyPI, which
#   fails because the cu121 torch wheel only lives on the pytorch.org index.
# - Falls back to a direct PyPI sdist download (curl) if `pip download`
#   still misbehaves on this pip/setuptools combination.
COPY scripts/patch_mamba_ssm.py /tmp/patch_mamba_ssm.py
RUN set -eux; \
    mkdir -p /tmp/mamba-src; \
    cd /tmp/mamba-src; \
    if ! pip download "mamba-ssm>=2.2.2" \
            --no-deps --no-binary=:all: --no-build-isolation \
            -d /tmp/mamba-src; then \
        echo "pip download failed; falling back to direct PyPI download"; \
        VER=$(pip index versions mamba-ssm 2>/dev/null \
              | sed -n 's/.*Available versions: \([0-9.]*\).*/\1/p' | head -1 \
              || echo "2.2.2"); \
        curl -fsSLO "https://files.pythonhosted.org/packages/source/m/mamba-ssm/mamba_ssm-${VER}.tar.gz" \
          || curl -fsSLO "https://files.pythonhosted.org/packages/source/m/mamba_ssm/mamba_ssm-${VER}.tar.gz"; \
    fi; \
    SDIST=$(ls mamba_ssm*.tar.gz mamba-ssm*.tar.gz 2>/dev/null | head -1); \
    test -n "$SDIST"; \
    tar -xzf "$SDIST"; \
    SRCDIR=$(tar -tzf "$SDIST" | head -1 | sed 's:/.*::'); \
    cd "$SRCDIR"; \
    python /tmp/patch_mamba_ssm.py setup.py; \
    pip install . --no-build-isolation -v

RUN python - <<'PY'
import torch, sys
print("Torch:", torch.__version__, "CUDA:", torch.version.cuda, "Avail:", torch.cuda.is_available())
PY

RUN python - <<'PY'
import os
setup_path = 'SRMA-Mamba/selective_scan/setup.py'
with open(setup_path, 'r') as f:
    content = f.read()

old_func = '''def get_compute_capability():
    device = torch.device("cuda")
    capability = torch.cuda.get_device_capability(device)
    return int(str(capability[0]) + str(capability[1]))'''

new_func = '''def get_compute_capability():
    if not torch.cuda.is_available():
        arch_list = os.getenv("TORCH_CUDA_ARCH_LIST", "8.9")
        arch = arch_list.split(";")[0].split(",")[0].strip()
        if "." in arch:
            major, minor = arch.split(".")
            return int(major + minor)
        else:
            return int(arch)
    device = torch.device("cuda")
    capability = torch.cuda.get_device_capability(device)
    return int(str(capability[0]) + str(capability[1]))'''

if old_func in content:
    content = content.replace(old_func, new_func)
    with open(setup_path, 'w') as f:
        f.write(content)
    print("Patched setup.py to use TORCH_CUDA_ARCH_LIST when CUDA not available")
else:
    print("WARNING: Could not find function to patch")
PY

RUN cd SRMA-Mamba/selective_scan && pip install --no-build-isolation -e . -v

RUN python - <<'PY'
import torch, sys
print("Torch:", torch.__version__, "CUDA:", torch.version.cuda, "Avail:", torch.cuda.is_available())
import mamba_ssm
print("mamba_ssm OK:", mamba_ssm.__file__)
import selective_scan_cuda_oflex as s
print("selective_scan_cuda_oflex OK:", s.__file__)
print("All CUDA extensions verified successfully.")
PY

EXPOSE 7860

ENV APP_FILE=app.py
CMD ["sh", "-c", "exec python ${APP_FILE}"]