# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Note this container is built from a local dockerfile
# Please see instructions in examples/sglang/README.md
FROM deepep:latest

# Add NIXL build dependencies
RUN apt-get update -y && \
    apt-get install -y \
    cmake \
    meson \
    ninja-build \
    pybind11-dev \
    patchelf

# Install Python build dependencies
RUN pip install --break-system-packages meson-python wheel build

# Add architecture args for NIXL build
ARG ARCH=amd64
ARG ARCH_ALT=x86_64

WORKDIR /sgl-workspace

# Install UCX dependencies
RUN apt-get update -y && \
    apt-get install -y --no-install-recommends \
    --reinstall libibverbs-dev rdma-core ibverbs-utils libibumad-dev \
    libnuma-dev librdmacm-dev ibverbs-providers \
    autoconf libtool

# Build UCX from source
ARG NIXL_UCX_REF=v1.19.x
RUN rm -rf /opt/hpcx/ucx && \
    rm -rf /usr/local/ucx && \
    cd /usr/local/src && \
    git clone https://github.com/openucx/ucx.git && \
    cd ucx && \
    git checkout $NIXL_UCX_REF && \
    ./autogen.sh && ./configure \
    --prefix=/usr/local/ucx \
    --enable-shared \
    --disable-static \
    --disable-doxygen-doc \
    --enable-optimizations \
    --enable-cma \
    --enable-devel-headers \
    --with-cuda=/usr/local/cuda \
    --with-verbs \
    --with-efa \
    --with-dm \
    --with-gdrcopy=/usr/local \
    --enable-mt && \
    make -j && \
    make -j install-strip && \
    ldconfig

ENV LD_LIBRARY_PATH=/usr/lib:/usr/local/ucx/lib:$LD_LIBRARY_PATH

# Pinning to NIXL 0.2.1 right now
# TODO: investigate pip install failure with 0.3.0 release
ARG NIXL_COMMIT="5e4c179ee850d482a83cb2a211e0947e46281060"
RUN git clone https://github.com/ai-dynamo/nixl.git && cd nixl && git checkout ${NIXL_COMMIT} && pip install --break-system-packages . --config-settings=setup-args="-Ducx_path=/usr/local/ucx"

WORKDIR /sgl-workspace

RUN pip uninstall --break-system-packages -y sglang
RUN rm -rf sglang
# 0.4.8 has a bug with CUDA graphs and decode worker
# https://github.com/sgl-project/sglang/issues/7511
RUN pip install --break-system-packages "sglang==0.4.7.post1"

# Allow forceful shutdown of inflight requests
ENV SGL_FORCE_SHUTDOWN=1

WORKDIR /sgl-workspace
# https://github.com/ai-dynamo/dynamo/pull/1510
ARG DYNAMO_COMMIT="382e3aedc421b3b3abc338062b332b54b5aa8529"
ARG DYNAMO_BRANCH="ishan/cmpl-token-id"
RUN git clone https://github.com/ai-dynamo/dynamo.git && cd dynamo && git checkout ${DYNAMO_BRANCH}

# install dynamo in editable mode
WORKDIR /sgl-workspace/dynamo
# Rust build/dev dependencies
RUN apt update -y && \
    apt install --no-install-recommends -y \
    build-essential \
    protobuf-compiler \
    cmake \
    libssl-dev \
    pkg-config \
    clang \
    libclang-dev \
    git

# Define Rust target based on ARCH_ALT ARG
ARG RUSTARCH=${ARCH_ALT}-unknown-linux-gnu

ENV RUSTUP_HOME=/usr/local/rustup \
    CARGO_HOME=/usr/local/cargo \
    PATH=/usr/local/cargo/bin:$PATH \
    RUST_VERSION=1.86.0

# Install Rust using RUSTARCH derived from ARCH_ALT
RUN wget --tries=3 --waitretry=5 "https://static.rust-lang.org/rustup/archive/1.28.1/${RUSTARCH}/rustup-init" && \
    # TODO: Add SHA check back based on RUSTARCH
    chmod +x rustup-init && \
    ./rustup-init -y --no-modify-path --profile minimal --default-toolchain $RUST_VERSION --default-host ${RUSTARCH} && \
    rm rustup-init && \
    chmod -R a+w $RUSTUP_HOME $CARGO_HOME

ARG CARGO_BUILD_JOBS
# Set CARGO_BUILD_JOBS to 16 if not provided
# This is to prevent cargo from building $(nproc) jobs in parallel,
# which might exceed the number of opened files limit.
ENV CARGO_BUILD_JOBS=${CARGO_BUILD_JOBS:-16}

RUN cargo build --release
RUN mkdir -p deploy/sdk/src/dynamo/sdk/cli/bin
RUN cp target/release/http deploy/sdk/src/dynamo/sdk/cli/bin
RUN cp target/release/llmctl deploy/sdk/src/dynamo/sdk/cli/bin
RUN cp target/release/dynamo-run deploy/sdk/src/dynamo/sdk/cli/bin

RUN cd lib/bindings/python && pip install --break-system-packages -e . && cd ../../..
RUN pip install --break-system-packages -e .

ENV PYTHONPATH=/sgl-workspace/dynamo/components/planner/src

RUN wget --tries=3 --waitretry=5 https://github.com/nats-io/nats-server/releases/download/v2.10.24/nats-server-v2.10.24-${ARCH}.deb && \
    dpkg -i nats-server-v2.10.24-${ARCH}.deb && rm nats-server-v2.10.24-${ARCH}.deb

ENV ETCD_VERSION="v3.5.18"
RUN wget --tries=3 --waitretry=5 https://github.com/etcd-io/etcd/releases/download/$ETCD_VERSION/etcd-$ETCD_VERSION-linux-${ARCH}.tar.gz -O /tmp/etcd.tar.gz && \
    mkdir -p /usr/local/bin/etcd && \
    tar -xvf /tmp/etcd.tar.gz -C /usr/local/bin/etcd --strip-components=1 && \
    rm /tmp/etcd.tar.gz
ENV PATH=/usr/local/bin/etcd/:$PATH

COPY examples/sglang/configs/deepep/* /sgl-workspace/dynamo/examples/sglang/configs/

WORKDIR /sgl-workspace/dynamo/examples/sglang