ARG BASE_IMAGE=vllm/vllm-openai@sha256:c4392d76e3eec8983fa152651365158cb062e348fd40398963f499d5867b9e28
FROM ${BASE_IMAGE}
COPY vendor/vllm_ple_mmap.py patches/spark_ngram_adapter.py /usr/local/lib/python3.12/dist-packages/
RUN python3 - <<'PY'
from pathlib import Path
import ast
root = Path('/usr/local/lib/python3.12/dist-packages/vllm')
p = root / 'models/qwen4_exp/nvidia/ngram_embedding.py'
s = p.read_text()
assert 'class Qwen4ExpPLEPinnedHostEmbedding' in s
s += '\nfrom spark_ngram_adapter import apply as _spark_disk_apply\n_spark_disk_apply(Qwen4ExpNGramEmbedding)\n'
ast.parse(s)
p.write_text(s)
for file, before, after in [
    ('third_party/flash_linear_attention/ops/utils.py', 'DEFAULT = 102400', 'DEFAULT = 101376'),
    ('third_party/flash_linear_attention/ops/chunk_delta_h.py', 'for num_warps in [2, 4]', 'for num_warps in [2]'),
]:
    p = root / file
    s = p.read_text()
    assert before in s, file
    p.write_text(s.replace(before, after))
PY

COPY patches/fix_mamba_prefix_blocks.py /tmp/fix_mamba_prefix_blocks.py
RUN python3 /tmp/fix_mamba_prefix_blocks.py /usr/local/lib/python3.12/dist-packages
COPY vendor/LICENSE vendor/Apache-2.0.txt /usr/local/share/licenses/spark-ple/