# syntax=docker/dockerfile:1.7
# AMD Strix Halo Ray Worker for khelben (gfx1151 / RDNA 3.5)
# Used for: vLLM (Llama 3.1 70B)
#
# Build:
#   docker build -t git.daviestechlabs.io/daviestechlabs/ray-worker-strixhalo:latest \
#     -f dockerfiles/Dockerfile.ray-worker-strixhalo .
#
# Multi-stage build: Extract ROCm 7.1 from vendor image, use Ray base for Python 3.11
# Note: Uses TheRock gfx110X wheels due to ROCm/ROCm#5853 segfault issue

# Stage 1: ROCm 7.1 libraries from AMD vendor image
FROM docker.io/rocm/pytorch:rocm7.1_ubuntu24.04_py3.12_pytorch_release_2.9.1 AS rocm-source

# Stage 2: Production image
FROM rayproject/ray:2.53.0-py311 AS production

# OCI Image Spec labels
LABEL org.opencontainers.image.title="Ray Worker - AMD Strix Halo"
LABEL org.opencontainers.image.description="Ray Serve worker for AMD Strix Halo (vLLM LLM inference)"
LABEL org.opencontainers.image.vendor="DaviesTechLabs"
LABEL org.opencontainers.image.source="https://git.daviestechlabs.io/daviestechlabs/kuberay-images"
LABEL org.opencontainers.image.licenses="MIT"
LABEL gpu.target="amd-rocm-7.1-gfx1151"
LABEL ray.version="2.53.0"

WORKDIR /app

# Copy ROCm stack from vendor image
COPY --from=rocm-source /opt/rocm /opt/rocm

# ROCm environment variables - split to ensure ROCM_HOME is set first
ENV ROCM_HOME=/opt/rocm
ENV PATH="/opt/rocm/bin:/opt/rocm/llvm/bin:/home/ray/anaconda3/bin:/home/ray/.local/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin" \
    LD_LIBRARY_PATH="/opt/rocm/lib:/opt/rocm/lib64" \
    HSA_PATH="/opt/rocm/hsa" \
    HIP_PATH="/opt/rocm/hip" \
    # Strix Halo (gfx1151) specific settings
    HIP_VISIBLE_DEVICES=0 \
    HSA_ENABLE_SDMA=0 \
    PYTORCH_HIP_ALLOC_CONF="expandable_segments:True,max_split_size_mb:512" \
    HSA_OVERRIDE_GFX_VERSION="11.0.0" \
    ROCM_TARGET_LST="gfx1151,gfx1100"

# Install system dependencies
USER root
RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
    --mount=type=cache,target=/var/lib/apt,sharing=locked \
    apt-get update && apt-get install -y --no-install-recommends \
        libelf1 \
        libnuma1 \
        libdrm2 \
        libdrm-amdgpu1 \
        kmod \
    && rm -rf /var/lib/apt/lists/*

# Install uv for fast Python package management (ADR-0014)
COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv

USER ray

# WORKAROUND: ROCm/ROCm#5853 - Standard PyTorch ROCm wheels cause segfault
# in libhsa-runtime64.so during VRAM allocation on gfx1151 (Strix Halo).
# TheRock gfx110X-all packages provide compatible Python 3.11 wheels.
RUN --mount=type=cache,target=/home/ray/.cache/uv,uid=1000,gid=1000 \
    uv pip install --system \
        --index-url https://rocm.nightlies.amd.com/v2/gfx110X-all/ \
        torch torchaudio torchvision

# Install vLLM and inference dependencies (uv is 10-100x faster than pip)
RUN --mount=type=cache,target=/home/ray/.cache/uv,uid=1000,gid=1000 \
    uv pip install --system \
        'vllm>=0.5.0' \
        'transformers>=4.35.0,<5.0' \
        'accelerate>=0.25.0,<1.0' \
        'sentence-transformers>=2.3.0,<3.0' \
        'httpx>=0.27.0,<1.0' \
        'scipy>=1.11.0,<2.0'

# FIX: Ray base image has pandas 1.5.3 which is incompatible with numpy 2.x
# The TheRock PyTorch wheels require numpy 2.x, so upgrade pandas to match
RUN --mount=type=cache,target=/home/ray/.cache/uv,uid=1000,gid=1000 \
    uv pip install --system 'pandas>=2.0.0,<3.0'

# Pre-download common models for faster cold starts (optional, increases image size)
# RUN python3 -c "from sentence_transformers import SentenceTransformer; SentenceTransformer('BAAI/bge-large-en-v1.5')"

# Copy entrypoint script (ray-serve-apps is installed from PyPI at runtime)
COPY --chown=ray:ray --chmod=755 dockerfiles/ray-entrypoint.sh /app/ray-entrypoint.sh

# Environment configuration
ENV PYTHONPATH=/app \
    PYTHONUNBUFFERED=1 \
    PYTHONDONTWRITEBYTECODE=1 \
    RAY_HEAD_SVC="ai-inference-raycluster-head-svc" \
    GPU_RESOURCE="gpu_amd_strixhalo" \
    NUM_GPUS="1"

# Health check
HEALTHCHECK --interval=30s --timeout=10s --start-period=120s --retries=3 \
    CMD ray status --address=localhost:6379 || exit 1

ENTRYPOINT ["/app/ray-entrypoint.sh"]