ARG BASE_DOCKER=rocm/vllm-dev:nightly_main_20250914
FROM $BASE_DOCKER

WORKDIR /unsloth-workspace

# Single RUN so the apt index can't go stale; cleanup keeps the layer small.
RUN apt-get update \
    && apt-get install -y --no-install-recommends zip unzip wget build-essential cmake \
    && rm -rf /var/lib/apt/lists/*

ARG COMPUTE_BACKEND="hip"
ARG BNB_ROCM_ARCHS="gfx942"
# bnb PR #1887 floor (0.50.0.dev0); earlier builds NaN at 4-bit GEMV decode on AMD.
ARG BNB_COMMIT="713a3b83e405d32449e1cb392e5e25985ab927c6"
ARG BNB_REPO="https://github.com/bitsandbytes-foundation/bitsandbytes.git"

# Plain install, not `-e`: upstream reserves editable for bnb development, and
# scikit-build-core 1.0.0 rejects it unless pyproject opts into
# editable.mode = "inplace", which this pin predates. cmake already dropped the
# .so into the package dir, so setuptools package-data ships it as built.
RUN mkdir bitsandbytes \
    && cd bitsandbytes \
    && git init \
    && git remote add origin ${BNB_REPO} \
    && git fetch --depth 1 origin ${BNB_COMMIT} \
    && git checkout FETCH_HEAD \
    && rm -rf .git \
    && cmake -S . -B build \
         -DCOMPUTE_BACKEND=${COMPUTE_BACKEND} \
         -DBNB_ROCM_ARCH="${BNB_ROCM_ARCHS}" \
    && cmake --build build -j"$(nproc)" \
    && pip install . \
    && rm -rf build

ARG UNSLOTH_REPO="https://github.com/unslothai/unsloth.git"
ARG UNSLOTH_BRANCH="main"


# huggingfacenotorch skips torch/torchvision so the base image's ROCm torch survives; bare `-e .` can't train.
RUN git clone --branch ${UNSLOTH_BRANCH} --depth 1 ${UNSLOTH_REPO} unsloth \
    && cd unsloth \
    && rm -rf .git \
    && pip install -e '.[huggingfacenotorch]'


ARG UNSLOTH_ZOO_REPO="https://github.com/unslothai/unsloth-zoo.git"
ARG UNSLOTH_ZOO_BRANCH="main"

RUN git clone --branch ${UNSLOTH_ZOO_BRANCH} --depth 1 ${UNSLOTH_ZOO_REPO} unsloth-zoo \
    && cd unsloth-zoo \
    && rm -rf .git \
    && pip install -e .


# pip show reads dist-info only, so it stayed green on an image whose freshly
# compiled HIP library was missing. bitsandbytes is no help either: with no AMD
# GPU on the builder it falls back to libbitsandbytes_cpu.so and reports
# success, so load the ROCm library explicitly. unsloth_zoo refuses to import
# without a torch accelerator, so only resolve it and unsloth.
# Still uncaught: a wrong-arch build loads fine here and only fails once a
# 4-bit GEMV runs on real AMD silicon.
RUN python3 -c "import ctypes, importlib.util, pathlib; \
    import bitsandbytes; \
    pkg = pathlib.Path(bitsandbytes.__file__).parent; \
    libs = sorted(pkg.glob('libbitsandbytes*.so')); \
    assert any('_rocm' in so.name for so in libs), f'no libbitsandbytes_rocm*.so in {pkg}'; \
    [ctypes.CDLL(str(so)) for so in libs]; \
    missing = [m for m in ('unsloth', 'unsloth_zoo') if importlib.util.find_spec(m) is None]; \
    assert not missing, f'not importable: {missing}'"