# 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.

# This should be pinned to the sglang version that is installed with Dynamo
# in the pyproject.toml
FROM lmsysorg/sglang:v0.4.9.post6-cu126

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

# 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

ARG NIXL_TAG=0.3.1
RUN git clone https://github.com/ai-dynamo/nixl.git && cd nixl && git checkout ${NIXL_TAG} && pip install --break-system-packages . --config-settings=setup-args="-Ducx_path=/usr/local/ucx"

WORKDIR /sgl-workspace

# Allow forceful shutdown of inflight requests
ENV SGL_FORCE_SHUTDOWN=1

WORKDIR /sgl-workspace
RUN git clone https://github.com/ai-dynamo/dynamo.git

# 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 cd lib/bindings/python && pip install --break-system-packages -e . && cd ../../..
RUN pip install --break-system-packages -e .

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

ENV ETCD_VERSION="v3.5.21"
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

ARG CMAKE_VERSION=3.31.8
RUN mkdir /sgl-workspace/cmake_build
WORKDIR /sgl-workspace/cmake_build

# uninstall CMake
RUN apt-get purge -y cmake
# download newer version of CMake
RUN wget https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}-linux-$(uname -m).tar.gz && \
    tar -xvzf cmake-${CMAKE_VERSION}-linux-$(uname -m).tar.gz && \
    mv cmake-${CMAKE_VERSION}-linux-$(uname -m) custom_cmake
ENV PATH=/sgl-workspace/cmake_build/custom_cmake/bin:$PATH

# should be 3.31.8
RUN cmake --version

# Install perf_analyzer and genai-perf
RUN apt-get update -y && \
    apt-get install -y --no-install-recommends \
    rapidjson-dev \
    zlib1g-dev

RUN git clone --depth=1 https://github.com/triton-inference-server/perf_analyzer.git && \
    mkdir perf_analyzer/build && \
    cmake -B perf_analyzer/build -S perf_analyzer && \
    cmake --build perf_analyzer/build -- -j8

ENV PATH=/sgl-workspace/perf_analyzer/build/perf_analyzer/src/perf-analyzer-build:$PATH

RUN pip install --break-system-packages genai-perf

# https://pypi.org/project/sglang-router/0.1.5 is latest
RUN pip install sglang-router==0.1.5

WORKDIR /sgl-workspace/dynamo/components/backends/sglang