Dockerfile.amd

# Full Hailo Dataflow Compiler toolchain on an AMD GPU (ROCm).
#
# The stock DFC optimization engine is TensorFlow-based and effectively
# CUDA-only. This project instead drives DFC's PyTorch-based optimization
# engine (use_saitama=True in the pipeline), which makes AMD GPUs usable for
# quantization. The base image therefore provides ROCm PyTorch, and the DFC
# is installed around it without letting dependency resolution replace torch.
#
# IMPORTANT — the DFC wheel is NOT redistributable and is NOT part of this
# repository. Obtain it from the Hailo Developer Zone (free account) and
# copy it next to this file before building. See README.md in this directory.
#
# Build:
#   cp hailo_dataflow_compiler-*.whl docker/
#   docker build -f docker/Dockerfile.amd -t dfc-amd:5.3.0 docker/
#
# Run (AMD GPUs need more passthrough flags than NVIDIA):
#   docker run --rm -it \
#       --device=/dev/kfd --device=/dev/dri \
#       --group-add <render-gid> --group-add <video-gid> \
#       --security-opt seccomp=unconfined --ipc=host \
#       dfc-amd:5.3.0 bash

# Community ROCm PyTorch image; validated with tag v2.11.0-rocm-7.2.1 for
# gfx906-class GPUs. Substitute the base matching your GPU generation.
ARG BASE_IMAGE=mixa3607/pytorch-gfx906:v2.11.0-rocm-7.2.1
FROM ${BASE_IMAGE}

ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && apt-get install -y --no-install-recommends \
        build-essential \
        graphviz graphviz-dev \
        git curl ca-certificates \
    && rm -rf /var/lib/apt/lists/*

# --- Version alignment: the DFC quantization stage is strict about NumPy
# --- semantics, and the PyTorch-ROCm stack pins scipy/onnxruntime ranges.
# --- These three were validated together; change them as a set.
RUN pip install --no-cache-dir \
        "numpy==2.1.3" \
        "scipy==1.13.1" \
        "onnxruntime==1.20.1"

# --- DFC wheel (bring your own — see header comment) ---
COPY hailo_dataflow_compiler-*.whl /tmp/
RUN WHEEL=$(ls /tmp/hailo_dataflow_compiler-*.whl) && \
    # Bare install: the wheel's requirements would replace the ROCm torch
    # build from the base image with a CUDA one.
    pip install --no-cache-dir --no-deps "$WHEEL"

# --- Remaining DFC dependencies, again protecting the preinstalled torch.
# --- matplotlib >= 3.9 required: older releases use np.Inf, removed in
# --- NumPy 2 (plotting would crash inside notebooks).
RUN pip install --no-cache-dir \
        protobuf graphviz tqdm pillow pandas "matplotlib>=3.9" tensorflow-cpu \
    && pip install --no-cache-dir --no-deps tensordict torchvision

# Two environment requirements, both mandatory:
# - legacy NumPy promotion: float64/float32 distinction matters inside
#   create_hw_params(); NumPy >= 2 default semantics break it.
# - CUDA_VISIBLE_DEVICES looks like a typo on ROCm but is not: PyTorch's
#   ROCm build consumes it, and if it is unset at run time the DFC sets it
#   to "99", hiding every GPU. It must be provided by `docker run -e`.
ENV NPY_PROMOTION_STATE=legacy

WORKDIR /workdir
CMD ["bash"]