FROM registry.fedoraproject.org/fedora:43

# Proxy configuration ARGs
ARG HTTP_PROXY=""
ARG HTTPS_PROXY=""
ARG SOCKS_PROXY=""

# Set environment variables if proxies are provided
ENV http_proxy=$HTTP_PROXY
ENV https_proxy=$HTTPS_PROXY
ENV no_proxy="localhost,127.0.0.1,::1"

# Configure DNF proxy if provided
RUN if [ -n "$HTTP_PROXY" ]; then \
  echo "proxy=$HTTP_PROXY" >> /etc/dnf/dnf.conf; \
  fi

# Copy proxychains configuration
COPY proxychains4.conf /etc/proxychains.conf

# Update proxychains4.conf if SOCKS_PROXY is provided (format: host:port)
RUN if [ -n "$SOCKS_PROXY" ]; then \
  HOST=$(echo $SOCKS_PROXY | cut -d: -f1); \
  PORT=$(echo $SOCKS_PROXY | cut -d: -f2); \
  sed -i "s/socks5.*/socks5 $HOST $PORT/g" /etc/proxychains.conf; \
  fi

# Define a proxy wrapper
RUN echo '#!/bin/bash' > /usr/local/bin/maybe_proxy && \
    echo 'if grep -q "^socks5 [0-9]" /etc/proxychains.conf; then if command -v proxychains4 >/dev/null 2>&1; then proxychains4 "$@"; else proxychains "$@"; fi; else "$@"; fi' >> /usr/local/bin/maybe_proxy && \
    chmod +x /usr/local/bin/maybe_proxy

# 1. System Base & Build Tools
RUN dnf -y install --setopt=install_weak_deps=False --nodocs \
  python3.13 python3.13-devel git rsync libatomic bash ca-certificates curl \
  gcc gcc-c++ binutils make ffmpeg-free \
  cmake ninja-build aria2c tar xz vim nano proxychains-ng \
  libdrm-devel zlib-devel openssl-devel jq \
  numactl-devel gperftools-libs procps-ng \
  && dnf clean all && rm -rf /var/cache/dnf/*

# 2. Install "TheRock" ROCm SDK (Tarball Method)
WORKDIR /tmp
ARG ROCM_MAJOR_VER=7
ARG GFX=gfx120X-all

# 智能处理：如果 build context 里有 therock.tar.gz 就用本地的，没有就从网络拉取
# 注意：Dockerfile 的 COPY 无法直接做条件判断，我们通过 shell 逻辑处理
RUN set -euo pipefail; \
    if [ -f "therock.tar.gz" ]; then \
        echo "Using local therock.tar.gz from build context."; \
    else \
        echo "Local SDK not found, downloading from remote..."; \
        BASE="https://therock-nightly-tarball.s3.amazonaws.com"; \
        PREFIX="therock-dist-linux-${GFX}-${ROCM_MAJOR_VER}"; \
        KEY="$(maybe_proxy curl -s "${BASE}?list-type=2&prefix=${PREFIX}" \
        | tr '<' '\n' \
        | grep -o "therock-dist-linux-${GFX}-${ROCM_MAJOR_VER}\..*\.tar\.gz" \
        | sort -V | tail -n1)"; \
        echo "Downloading Latest Tarball: ${KEY}"; \
        maybe_proxy aria2c -x 16 -s 16 -k 1M --file-allocation=none "${BASE}/${KEY}" -o therock.tar.gz; \
    fi && \
    mkdir -p /opt/rocm && \
    tar xzf therock.tar.gz -C /opt/rocm --strip-components=1 && \
    rm therock.tar.gz

# 3. Configure Global ROCm Environment
RUN export ROCM_PATH=/opt/rocm && \
  BITCODE_PATH=$(find /opt/rocm -type d -name bitcode -print -quit) && \
  printf '%s\n' \
  "export ROCM_PATH=/opt/rocm" \
  "export HIP_PLATFORM=amd" \
  "export HIP_PATH=/opt/rocm" \
  "export HIP_CLANG_PATH=/opt/rocm/llvm/bin" \
  "export HIP_DEVICE_LIB_PATH=$BITCODE_PATH" \
  "export PATH=$ROCM_PATH/bin:$ROCM_PATH/llvm/bin:\$PATH" \
  "export LD_LIBRARY_PATH=$ROCM_PATH/lib:$ROCM_PATH/lib64:$ROCM_PATH/llvm/lib:\$LD_LIBRARY_PATH" \
  "export ROCBLAS_USE_HIPBLASLT=1" \
  "export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1" \
  "export VLLM_TARGET_DEVICE=rocm" \
  "export HIP_FORCE_DEV_KERNARG=1" \
  "export RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1" \
  "export LD_PRELOAD=/usr/lib64/libtcmalloc_minimal.so.4" \
  > /etc/profile.d/rocm-sdk.sh && \
  chmod 0644 /etc/profile.d/rocm-sdk.sh

# 4. Python Venv Setup
RUN /usr/bin/python3.13 -m venv /opt/venv
ENV VIRTUAL_ENV=/opt/venv
ENV PATH=/opt/venv/bin:$PATH
ENV PIP_NO_CACHE_DIR=1
RUN printf 'source /opt/venv/bin/activate\n' > /etc/profile.d/venv.sh
RUN python -m pip install --upgrade pip wheel packaging "setuptools<80.0.0"

# 5. Install PyTorch (TheRock Nightly)
RUN maybe_proxy python -m pip install \
  --index-url https://rocm.nightlies.amd.com/v2-staging/gfx120X-all/ \
  --pre torch torchaudio torchvision && \
  maybe_proxy python -m pip install pyyaml

# Flash-Attention
WORKDIR /opt
ENV FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE"
ENV ROCM_HOME="/opt/rocm"
ENV HIP_PATH="/opt/rocm"
ENV PATH="/opt/rocm/bin:/opt/rocm/llvm/bin:$PATH"
ENV LD_LIBRARY_PATH="/opt/rocm/lib:/opt/rocm/lib64:/opt/rocm/llvm/lib"

RUN set -euo pipefail; \
  git config --global http.postBuffer 524288000; \
  git config --global http.lowSpeedLimit 0; \
  git config --global http.lowSpeedTime 999999; \
  git config --global http.version HTTP/1.1; \
  if [ -n "$SOCKS_PROXY" ]; then \
    git config --global http.proxy "socks5://$SOCKS_PROXY"; \
    git config --global https.proxy "socks5://$SOCKS_PROXY"; \
  fi; \
  CLONED=0; \
  for i in 1 2 3 4 5; do \
    rm -rf /opt/flash-attention; \
    if maybe_proxy git clone --recursive --branch main_perf --depth 1 --shallow-submodules https://github.com/ROCm/flash-attention.git /opt/flash-attention; then \
      CLONED=1; \
      break; \
    fi; \
    echo "flash-attention clone failed (attempt ${i}/5), retrying..."; \
    sleep $((i * 5)); \
  done; \
  if [ "${CLONED}" -ne 1 ]; then \
    echo "flash-attention clone failed after retries"; \
    exit 1; \
  fi; \
  cd /opt/flash-attention; \
  python setup.py install; \
  cd /opt && rm -rf /opt/flash-attention

# 6. Clone vLLM
RUN maybe_proxy git clone https://github.com/vllm-project/vllm.git /opt/vllm
WORKDIR /opt/vllm

# --- PATCHING ---
RUN echo "import sys, re" > patch_vllm.py && \
  echo "from pathlib import Path" >> patch_vllm.py && \
  echo "p = Path('vllm/platforms/__init__.py')" >> patch_vllm.py && \
  echo "txt = p.read_text()" >> patch_vllm.py && \
  echo "txt = txt.replace('import amdsmi', '# import amdsmi')" >> patch_vllm.py && \
  echo "txt = re.sub(r'is_rocm = .*', 'is_rocm = True', txt)" >> patch_vllm.py && \
  echo "txt = re.sub(r'if len\(amdsmi\.amdsmi_get_processor_handles\(\)\) > 0:', 'if True:', txt)" >> patch_vllm.py && \
  echo "txt = txt.replace('amdsmi.amdsmi_init()', 'pass')" >> patch_vllm.py && \
  echo "txt = txt.replace('amdsmi.amdsmi_shut_down()', 'pass')" >> patch_vllm.py && \
  echo "p.write_text(txt)" >> patch_vllm.py && \
  echo "p = Path('vllm/platforms/rocm.py')" >> patch_vllm.py && \
  echo "txt = p.read_text()" >> patch_vllm.py && \
  echo "header = 'import sys\nfrom unittest.mock import MagicMock\nsys.modules[\"amdsmi\"] = MagicMock()\n'" >> patch_vllm.py && \
  echo "txt = header + txt" >> patch_vllm.py && \
  echo "txt = re.sub(r'device_type = .*', 'device_type = \"rocm\"', txt)" >> patch_vllm.py && \
  echo "txt = re.sub(r'device_name = .*', 'device_name = \"gfx1201\"', txt)" >> patch_vllm.py && \
  echo "txt += '\n    def get_device_name(self, device_id: int = 0) -> str:\n        return \"AMD-gfx1201\"\n'" >> patch_vllm.py && \
  echo "p.write_text(txt)" >> patch_vllm.py && \
  echo "print('Successfully patched vLLM for R9700')" >> patch_vllm.py && \
  python patch_vllm.py

# 7. Build vLLM (Wheel Method) with CLANG Host Compiler
RUN python -m pip install --upgrade cmake ninja packaging wheel numpy "setuptools-scm>=8" "setuptools<80.0.0" scikit-build-core pybind11
ENV ROCM_HOME="/opt/rocm"
ENV HIP_PATH="/opt/rocm"
ENV VLLM_TARGET_DEVICE="rocm"
ENV PYTORCH_ROCM_ARCH="gfx1201"
ENV HIP_ARCHITECTURES="gfx1201"
ENV AMDGPU_TARGETS="gfx1201"
ENV MAX_JOBS="32"

ENV CC="/opt/rocm/llvm/bin/clang"
ENV CXX="/opt/rocm/llvm/bin/clang++"

RUN export HIP_DEVICE_LIB_PATH=$(find /opt/rocm -type d -name bitcode -print -quit) && \
  echo "Compiling with Bitcode: $HIP_DEVICE_LIB_PATH" && \
  export CMAKE_PREFIX_PATH="/opt/venv/lib64/python3.13/site-packages/torch/share/cmake:/opt/rocm" && \
  export CMAKE_ARGS="-DROCM_PATH=/opt/rocm -DHIP_PATH=/opt/rocm -DAMDGPU_TARGETS=gfx1201 -DHIP_ARCHITECTURES=gfx1201 -DCMAKE_PREFIX_PATH=/opt/venv/lib64/python3.13/site-packages/torch/share/cmake:/opt/rocm" && \
  python -m pip wheel --no-build-isolation --no-deps -w /tmp/dist -v . && \
  python -m pip install /tmp/dist/*.whl

# bitsandbytes
WORKDIR /opt
RUN maybe_proxy git clone -b rocm_enabled_multi_backend https://github.com/ROCm/bitsandbytes.git
WORKDIR /opt/bitsandbytes

ENV HIP_PLATFORM="amd"
ENV CMAKE_PREFIX_PATH="/opt/rocm"

RUN cmake -S . \
  -DGPU_TARGETS="gfx1201" \
  -DBNB_ROCM_ARCH="gfx1201" \
  -DCOMPUTE_BACKEND=hip \
  -DCMAKE_HIP_COMPILER=/opt/rocm/llvm/bin/clang++ \
  -DCMAKE_CXX_COMPILER=/opt/rocm/llvm/bin/clang++ \
  && \
  make -j$(nproc) && \
  python -m pip install --no-cache-dir . --no-build-isolation --no-deps

# Cleanup
WORKDIR /opt
RUN chmod -R a+rwX /opt && \
  find /opt/venv -type f -name "*.so" -exec strip -s {} + 2>/dev/null || true && \
  find /opt/venv -type d -name "__pycache__" -prune -exec rm -rf {} + && \
  rm -rf /root/.cache/pip || true && \
  dnf clean all && rm -rf /var/cache/dnf/*

RUN mkdir -p /etc/vllm && mkdir -p /model

COPY scripts/01-rocm-envs.sh /etc/profile.d/01-rocm-envs.sh
COPY scripts/99-toolbox-banner.sh /etc/profile.d/99-toolbox-banner.sh
COPY scripts/zz-venv-last.sh /etc/profile.d/zz-venv-last.sh
COPY scripts/start_vllm.py /usr/local/bin/start-vllm
RUN chmod 0644 /etc/profile.d/*.sh && chmod +x /usr/local/bin/start-vllm
RUN printf 'ulimit -S -c 0\n' > /etc/profile.d/90-nocoredump.sh && chmod 0644 /etc/profile.d/90-nocoredump.sh

CMD ["/bin/bash"]
