# syntax=docker/dockerfile:1.3-labs
#
# Ray Image Builder
# ==============================
# Installs the Ray wheel into a base image (CPU or CUDA), includes
# pip freeze output for reproducibility.
#
ARG PYTHON_VERSION=3.10
ARG PLATFORM=cpu
ARG ARCH_SUFFIX=
ARG IMAGE_TYPE=ray
ARG BASE_VARIANT=base
ARG BASE_IMAGE=cr.ray.io/rayproject/${IMAGE_TYPE}-py${PYTHON_VERSION}-${PLATFORM}-${BASE_VARIANT}${ARCH_SUFFIX}
ARG RAY_WHEEL_IMAGE=cr.ray.io/rayproject/ray-wheel-py${PYTHON_VERSION}${ARCH_SUFFIX}

FROM ${RAY_WHEEL_IMAGE} AS wheel-source
FROM ${BASE_IMAGE}

ARG IMAGE_TYPE=ray
ARG PLATFORM=cpu
ARG RAY_COMMIT=unknown-commit
ARG RAY_VERSION=3.0.0.dev0

LABEL io.ray.ray-commit="${RAY_COMMIT}"
LABEL io.ray.ray-version="${RAY_VERSION}"

COPY --from=wheel-source /opt/artifacts/*.whl /home/ray/

# Install Ray wheel with all extras
# Uses requirements_compiled.txt from base image (already at /home/ray/)
RUN <<EOF
#!/bin/bash
set -euo pipefail

WHEEL_FILES=(/home/ray/ray-*.whl)
if [[ ${#WHEEL_FILES[@]} -ne 1 ]]; then
    echo "Error: Expected 1 ray wheel file, but found ${#WHEEL_FILES[@]} in /home/ray/." >&2
    ls -l /home/ray/*.whl >&2
    exit 1
fi
WHEEL_FILE="${WHEEL_FILES[0]}"

echo "Installing wheel: $WHEEL_FILE"

if [[ "${IMAGE_TYPE}" == "ray-llm" ]]; then
  RAY_EXTRAS="default,data,serve"
else
  RAY_EXTRAS="all"
fi

# TODO(cu130): ray[all]'s cgraph extra hard-pins cupy-cuda12x, so this install
# always pulls the CUDA-12 build even on cu130 images (no PEP 508 marker exists
# to select cupy by CUDA version). Until the cgraph extra can resolve cupy per
# CUDA runtime (or cupy ships a unified package), we patch it up with the
# uninstall/reinstall swap below. Drop that swap once this install can pick the
# right cupy directly.
$HOME/anaconda3/bin/pip --no-cache-dir install \
    -c /home/ray/requirements_compiled.txt \
    "${WHEEL_FILE}[${RAY_EXTRAS}]"

# ray[all]'s cgraph extra hard-pins cupy-cuda12x (a CUDA-12 build), but cu130
# images ship a CUDA-13 runtime where that build is broken. Swap it for the
# matching CUDA-13 build. cupy-cuda12x and cupy-cuda13x both own the top-level
# `cupy` package and cannot coexist, so this is an uninstall-then-install.
# Scoped to IMAGE_TYPE=ray (covers ray + ray-extra); ray-llm flows through this
# same Dockerfile but manages cupy via its own llm locks, so leave it untouched.
if [[ "${IMAGE_TYPE}" == "ray" && "${PLATFORM}" == cu13* ]]; then
    $HOME/anaconda3/bin/pip --no-cache-dir uninstall -y cupy-cuda12x
    $HOME/anaconda3/bin/pip --no-cache-dir install "cupy-cuda13x==13.6.0"
fi

$HOME/anaconda3/bin/pip freeze > /home/ray/pip-freeze.txt

echo "Ray version: $($HOME/anaconda3/bin/python -c 'import ray; print(ray.__version__)')"
EOF

CMD ["/bin/bash"]