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

ARG SGLANG_IMAGE_TAG="v0.4.10-cu126"

FROM lmsysorg/sglang:${SGLANG_IMAGE_TAG}

ARG MODE="hopper"
ARG ARCH="amd64"
ARG ARCH_ALT="x86_64"
ARG NIXL_UCX_REF="v1.19.x"
ARG NIXL_TAG="0.4.1"
ARG CMAKE_VERSION="3.31.8"
ARG RUST_VERSION="1.87.0"
ARG CARGO_BUILD_JOBS="16"

RUN apt-get update -y && \
    apt-get install -y \
      cmake meson ninja-build pybind11-dev patchelf net-tools \
      build-essential protobuf-compiler libssl-dev pkg-config \
      clang libclang-dev git rapidjson-dev zlib1g-dev && \
    pip install --break-system-packages meson-python wheel build

# Build UCX + NIXL for x86/hopper until its fully tested on GB200
RUN if [ "$MODE" = "hopper" ]; then \
      apt-get install -y --no-install-recommends \
        libibverbs-dev rdma-core ibverbs-utils libibumad-dev \
        libnuma-dev librdmacm-dev ibverbs-providers autoconf libtool && \
      # UCX from source
      rm -rf /opt/hpcx/ucx /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 install-strip && ldconfig && \
      # NIXL
      git clone https://github.com/ai-dynamo/nixl.git /opt/nixl && \
      cd /opt/nixl && git checkout $NIXL_TAG && \
      pip install --break-system-packages . \
        --config-settings="setup-args=-Ducx_path=/usr/local/ucx"; \
    fi

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

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

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

RUN wget --tries=3 --waitretry=5 \
    "https://static.rust-lang.org/rustup/archive/1.28.1/${ARCH_ALT}-unknown-linux-gnu/rustup-init" && \
    chmod +x rustup-init && \
    ./rustup-init -y \
      --no-modify-path \
      --profile minimal \
      --default-toolchain $RUST_VERSION \
      --default-host ${ARCH_ALT}-unknown-linux-gnu && \
    rm rustup-init && \
    chmod -R a+w $RUSTUP_HOME $CARGO_HOME

ARG CARGO_BUILD_JOBS
ENV CARGO_BUILD_JOBS=${CARGO_BUILD_JOBS}

RUN cd dynamo && cargo build --release

RUN cd dynamo/lib/bindings/python && \
    pip install --break-system-packages -e . && \
    cd /sgl-workspace/dynamo && \
    pip install --break-system-packages .

RUN pip install --break-system-packages sglang-router==0.1.5

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 -xzf /tmp/etcd.tar.gz \
        -C /usr/local/bin/etcd --strip-components=1 && \
    rm /tmp/etcd.tar.gz

ENV PATH=/usr/local/bin/etcd:$PATH

# GenAI Perf
RUN apt-get purge -y cmake

RUN mkdir /sgl-workspace/cmake_build && \
    cd /sgl-workspace/cmake_build && \
    wget https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/\
cmake-${CMAKE_VERSION}-linux-$(uname -m).tar.gz && \
    tar -xzf cmake-${CMAKE_VERSION}-linux-$(uname -m).tar.gz && \
    mv cmake-${CMAKE_VERSION}-linux-$(uname -m) custom_cmake && \
    rm cmake-${CMAKE_VERSION}-linux-$(uname -m).tar.gz

ENV PATH=/sgl-workspace/cmake_build/custom_cmake/bin:$PATH
RUN cmake --version

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 -- -j$(nproc)

ENV PATH=/sgl-workspace/perf_analyzer/build/perf_analyzer/src/perf-analyzer-build:$PATH
RUN pip install --break-system-packages genai-perf

# Enable forceful shutdown of inflight requests
ENV SGL_FORCE_SHUTDOWN=1

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