# DeepInfer on AWS Trainium (Neuron basic mode), on top of the native PyTorch Neuron DLC.
#
# The base image ships torch 2.14 + torch-neuronx 2.14 (no torch-xla), neuronx-cc, nki and the Neuron tools on
# Ubuntu 26.04 / Python 3.14 (Neuron DLCs are Ubuntu 26.04-only from SDK 2.33; the torch-neuronx wheel needs
# glibc >= 2.43). The Neuron driver (aws-neuronx-dkms) must be installed on the host.
#
#   docker build -f docker/neuron/Dockerfile -t deepinfer-neuron .
#   docker run --rm -it --privileged --ipc=host --network=host \
#       -v /mnt/nvme/models:/models deepinfer-neuron \
#       python examples/neuron/offline_inference_neuron.py --model_name /models/DeepSeek-V4.1-Flash-1017 --ep 64
ARG BASE_IMAGE=public.ecr.aws/neuron/torch-neuronx:2.14.0-neuronx-py314-sdk2.33.0-ubuntu26.04
FROM ${BASE_IMAGE}

ENV DEEPINFER_ACCELERATOR=neuron \
    PIP_NO_CACHE_DIR=1 \
    PIP_ROOT_USER_ACTION=ignore

WORKDIR /opt/deepinfer
COPY . .

# requirements-neuron.txt (torch / torch-neuronx are already satisfied by the base image), no vendor build
# (Neuron uses the pure-torch drop-ins in deepinfer/neuron_ops/vendor), then the host-only C++ extension.
# No NeuronCore is visible during `docker build`: keep torch from auto-loading torch_neuronx (it queries the device
# when it initializes); compiling the extension only needs the torch headers. The self-check runs in Neuron CPU mode.
RUN TORCH_DEVICE_BACKEND_AUTOLOAD=0 MAX_JOBS=16 ./build.sh \
 && TORCH_DEVICE_BACKEND_AUTOLOAD=0 DEEPINFER_NEURON_DEVICE=cpu python -c "import deepinfer, torch; deepinfer.PLATFORM; \
from deepinfer._C import engine; from deepinfer.model_executor.layers.tilelang_ops import ops; \
print('deepinfer-neuron OK:', type(deepinfer.PLATFORM).__name__, 'torch', torch.__version__)"

ENV PYTHONPATH=/opt/deepinfer
CMD ["/bin/bash"]