build.sh 22.1 KB
Newer Older
Carsten Csiky's avatar
Carsten Csiky committed
1
#!/usr/bin/env bash
Neelay Shah's avatar
Neelay Shah committed
2
# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3
# SPDX-License-Identifier: Apache-2.0
4
5
6
7
8
9
10
11
12
13
14
15
#
# 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.
16

17
18
19
20
21
if [ "${BASH_VERSINFO[0]}" -lt 4 ]; then
    echo "Error: Bash version 4.0 or higher is required. Current version: ${BASH_VERSINFO[0]}.${BASH_VERSINFO[1]}"
    exit 1
fi

Carsten Csiky's avatar
Carsten Csiky committed
22
set -e
23

24
25
26
TAG=
RUN_PREFIX=
PLATFORM=linux/amd64
27
28
USER_UID=
USER_GID=
29
30
31
32

# Get short commit hash
commit_id=$(git rev-parse --short HEAD)

33
34
# if COMMIT_ID matches a TAG use that
current_tag=$(git describe --tags --exact-match 2>/dev/null | sed 's/^v//') || true
35

36
# Get latest TAG and add COMMIT_ID for dev
Carsten Csiky's avatar
Carsten Csiky committed
37
latest_tag=$(git describe --tags --abbrev=0 "$(git rev-list --tags --max-count=1 main)" | sed 's/^v//') || true
38
39
40
41
if [[ -z ${latest_tag} ]]; then
    latest_tag="0.0.1"
    echo "No git release tag found, setting to unknown version: ${latest_tag}"
fi
42

43
44
45
46
# Use tag if available, otherwise use latest_tag.dev.commit_id
VERSION=v${current_tag:-$latest_tag.dev.$commit_id}

PYTHON_PACKAGE_VERSION=${current_tag:-$latest_tag.dev+$commit_id}
47
48
49
50
51
52

# Frameworks
#
# Each framework has a corresponding base image.  Additional
# dependencies are specified in the /container/deps folder and
# installed within framework specific sections of the Dockerfile.
53

54
declare -A FRAMEWORKS=(["VLLM"]=1 ["TRTLLM"]=2 ["NONE"]=3 ["SGLANG"]=4)
Ryan Olson's avatar
Ryan Olson committed
55

56
DEFAULT_FRAMEWORK=VLLM
57
58
59
60
61
62

SOURCE_DIR=$(dirname "$(readlink -f "$0")")
DOCKERFILE=${SOURCE_DIR}/Dockerfile
BUILD_CONTEXT=$(dirname "$(readlink -f "$SOURCE_DIR")")

# Base Images
63
TRTLLM_BASE_IMAGE=nvcr.io/nvidia/pytorch
64
TRTLLM_BASE_IMAGE_TAG=25.06-py3
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

# Important Note: Because of ABI compatibility issues between TensorRT-LLM and NGC PyTorch,
# we need to build the TensorRT-LLM wheel from source.
#
# There are two ways to build the dynamo image with TensorRT-LLM.
# 1. Use the local TensorRT-LLM wheel directory.
# 2. Use the TensorRT-LLM wheel on artifactory.
#
# If using option 1, the TENSORRTLLM_PIP_WHEEL_DIR must be a path to a directory
# containing TensorRT-LLM wheel file along with commit.txt file with the
# <arch>_<commit ID> as contents. If no valid trtllm wheel is found, the script
# will attempt to build the wheel from source and store the built wheel in the
# specified directory. TRTLLM_COMMIT from the TensorRT-LLM main branch will be
# used to build the wheel.
#
# If using option 2, the TENSORRTLLM_PIP_WHEEL must be the TensorRT-LLM wheel
# package that will be installed from the specified TensorRT-LLM PyPI Index URL.
# This option will ignore the TRTLLM_COMMIT option. As the TensorRT-LLM wheel from PyPI
# is not ABI compatible with NGC PyTorch, you can use TENSORRTLLM_INDEX_URL to specify
# a private PyPI index URL which has your pre-built TensorRT-LLM wheel.
#
# By default, we will use option 1. If you want to use option 2, you can set
# TENSORRTLLM_PIP_WHEEL to the TensorRT-LLM wheel on artifactory.
#
# Path to the local TensorRT-LLM wheel directory or the wheel on artifactory.
TENSORRTLLM_PIP_WHEEL_DIR="/tmp/trtllm_wheel/"
# TensorRT-LLM commit to use for building the trtllm wheel if not provided.
# Important Note: This commit is not used in our CI pipeline. See the CI
# variables to learn how to run a pipeline with a specific commit.
94
DEFAULT_EXPERIMENTAL_TRTLLM_COMMIT="a16ba6445c61ed70e7aadfe787d6f316bb422652"
95
TRTLLM_COMMIT=""
96
TRTLLM_USE_NIXL_KVCACHE_EXPERIMENTAL="0"
97
TRTLLM_GIT_URL=""
98

99
100
# TensorRT-LLM PyPI index URL
TENSORRTLLM_INDEX_URL="https://pypi.python.org/simple"
101
# TODO: Remove the version specification from here and use the ai-dynamo[trtllm] package.
102
# Need to update the Dockerfile.trtllm to use the ai-dynamo[trtllm] package.
103
DEFAULT_TENSORRTLLM_PIP_WHEEL="tensorrt-llm==1.0.0rc6"
104
105
106
TENSORRTLLM_PIP_WHEEL=""


107
VLLM_BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base"
108
109
110
111
112
# FIXME: NCCL will hang with 25.03, so use 25.01 for now
# Please check https://github.com/ai-dynamo/dynamo/pull/1065
# for details and reproducer to manually test if the image
# can be updated to later versions.
VLLM_BASE_IMAGE_TAG="25.01-cuda12.8-devel-ubuntu24.04"
113

114
115
NONE_BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base"
NONE_BASE_IMAGE_TAG="25.01-cuda12.8-devel-ubuntu24.04"
116

117
118
119
SGLANG_BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base"
SGLANG_BASE_IMAGE_TAG="25.01-cuda12.8-devel-ubuntu24.04"

120
NIXL_REF=0.4.1
121
122
NIXL_UCX_REF=v1.19.0
NIXL_UCX_EFA_REF=9d2b88a1f67faf9876f267658bd077b379b8bb76
123

124
125
NO_CACHE=""

126
127
128
129
130
# sccache configuration for S3
USE_SCCACHE=""
SCCACHE_BUCKET=""
SCCACHE_REGION=""

131
132
133
134
135
136
137
get_options() {
    while :; do
        case $1 in
        -h | -\? | --help)
            show_help
            exit
            ;;
138
        --platform)
139
140
141
142
            if [ "$2" ]; then
                PLATFORM=$2
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
143
                missing_requirement "$1"
144
145
            fi
            ;;
146
        --framework)
147
148
149
150
            if [ "$2" ]; then
                FRAMEWORK=$2
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
151
                missing_requirement "$1"
152
153
            fi
            ;;
154
        --tensorrtllm-pip-wheel-dir)
155
            if [ "$2" ]; then
156
157
158
159
160
161
162
163
164
165
166
167
168
169
                TENSORRTLLM_PIP_WHEEL_DIR=$2
                shift
            else
                missing_requirement "$1"
            fi
            ;;
        --tensorrtllm-commit)
            if [ "$2" ]; then
                TRTLLM_COMMIT=$2
                shift
            else
                missing_requirement "$1"
            fi
            ;;
170
171
172
173
174
175
176
        --use-default-experimental-tensorrtllm-commit)
            if [ -n "$2" ] && [[ "$2" != --* ]]; then
                echo "ERROR: --use-default-experimental-tensorrtllm-commit does not take any argument"
                exit 1
            fi
            USE_DEFAULT_EXPERIMENTAL_TRTLLM_COMMIT=true
            ;;
177
178
179
180
181
182
183
        --trtllm-use-nixl-kvcache-experimental)
            if [ -n "$2" ] && [[ "$2" != --* ]]; then
                echo "ERROR: --trtllm-use-nixl-kvcache-experimental does not take any argument"
                exit 1
            fi
            TRTLLM_USE_NIXL_KVCACHE_EXPERIMENTAL="1"
            ;;
184
185
186
187
188
189
190
191
192
193
194
        --tensorrtllm-pip-wheel)
            if [ "$2" ]; then
                TENSORRTLLM_PIP_WHEEL=$2
                shift
            else
                missing_requirement "$1"
            fi
            ;;
        --tensorrtllm-index-url)
            if [ "$2" ]; then
                TENSORRTLLM_INDEX_URL=$2
195
196
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
197
                missing_requirement "$1"
198
199
            fi
            ;;
200
201
202
203
204
205
206
207
        --tensorrtllm-git-url)
            if [ "$2" ]; then
                TRTLLM_GIT_URL=$2
                shift
            else
                missing_requirement "$1"
            fi
            ;;
208
209
210
211
212
        --base-image)
            if [ "$2" ]; then
                BASE_IMAGE=$2
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
213
                missing_requirement "$1"
214
215
            fi
            ;;
216
        --base-image-tag)
217
218
219
220
            if [ "$2" ]; then
                BASE_IMAGE_TAG=$2
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
221
                missing_requirement "$1"
222
223
224
225
226
227
228
            fi
            ;;
        --target)
            if [ "$2" ]; then
                TARGET=$2
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
229
                missing_requirement "$1"
230
231
232
233
234
235
236
            fi
            ;;
        --build-arg)
            if [ "$2" ]; then
                BUILD_ARGS+="--build-arg $2 "
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
237
                missing_requirement "$1"
238
239
240
241
            fi
            ;;
        --tag)
            if [ "$2" ]; then
242
                TAG="--tag $2"
243
244
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
245
                missing_requirement "$1"
246
247
            fi
            ;;
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
        --uid)
            if [ "$2" ]; then
                USER_UID="$2"
                shift
            else
                missing_requirement "$1"
            fi
            ;;
        --gid)
            if [ "$2" ]; then
                USER_GID="$2"
                shift
            else
                missing_requirement "$1"
            fi
            ;;
264
265
266
267
268
269
270
271
        --dry-run)
            RUN_PREFIX="echo"
            echo ""
            echo "=============================="
            echo "DRY RUN: COMMANDS PRINTED ONLY"
            echo "=============================="
            echo ""
            ;;
272
273
        --no-cache)
            NO_CACHE=" --no-cache"
274
            ;;
275
276
        --cache-from)
            if [ "$2" ]; then
277
278
279
                CACHE_FROM="--cache-from $2"
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
280
                missing_requirement "$1"
281
282
            fi
            ;;
283
284
285
286
287
        --cache-to)
            if [ "$2" ]; then
                CACHE_TO="--cache-to $2"
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
288
                missing_requirement "$1"
289
290
            fi
            ;;
ptarasiewiczNV's avatar
ptarasiewiczNV committed
291
292
293
294
295
        --build-context)
            if [ "$2" ]; then
                BUILD_CONTEXT_ARG="--build-context $2"
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
296
                missing_requirement "$1"
ptarasiewiczNV's avatar
ptarasiewiczNV committed
297
298
            fi
            ;;
299
300
301
        --release-build)
            RELEASE_BUILD=true
            ;;
302
303
304
        --enable-kvbm)
            ENABLE_KVBM=true
            ;;
305
306
307
        --make-efa)
            NIXL_UCX_REF=$NIXL_UCX_EFA_REF
            ;;
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
        --use-sccache)
            USE_SCCACHE=true
            ;;
        --sccache-bucket)
            if [ "$2" ]; then
                SCCACHE_BUCKET=$2
                shift
            else
                missing_requirement "$1"
            fi
            ;;

        --sccache-region)
            if [ "$2" ]; then
                SCCACHE_REGION=$2
                shift
            else
                missing_requirement "$1"
            fi
327
328
            ;;
         -?*)
Carsten Csiky's avatar
Carsten Csiky committed
329
            error 'ERROR: Unknown option: ' "$1"
330
            ;;
331
         ?*)
Carsten Csiky's avatar
Carsten Csiky committed
332
            error 'ERROR: Unknown option: ' "$1"
333
334
335
336
337
338
339
340
341
            ;;
        *)
            break
            ;;
        esac
        shift
    done

    if [ -z "$FRAMEWORK" ]; then
342
        FRAMEWORK=$DEFAULT_FRAMEWORK
343
344
    fi

Carsten Csiky's avatar
Carsten Csiky committed
345
    if [ -n "$FRAMEWORK" ]; then
346
        FRAMEWORK=${FRAMEWORK^^}
347

Carsten Csiky's avatar
Carsten Csiky committed
348
349
        if [[ -z "${FRAMEWORKS[$FRAMEWORK]}" ]]; then
            error 'ERROR: Unknown framework: ' "$FRAMEWORK"
350
        fi
351

Carsten Csiky's avatar
Carsten Csiky committed
352
        if [ -z "$BASE_IMAGE_TAG" ]; then
353
354
355
            BASE_IMAGE_TAG=${FRAMEWORK}_BASE_IMAGE_TAG
            BASE_IMAGE_TAG=${!BASE_IMAGE_TAG}
        fi
356

Carsten Csiky's avatar
Carsten Csiky committed
357
        if [ -z "$BASE_IMAGE" ]; then
358
359
360
            BASE_IMAGE=${FRAMEWORK}_BASE_IMAGE
            BASE_IMAGE=${!BASE_IMAGE}
        fi
361

Carsten Csiky's avatar
Carsten Csiky committed
362
        if [ -z "$BASE_IMAGE" ]; then
363
364
            error "ERROR: Framework $FRAMEWORK without BASE_IMAGE"
        fi
365

366
367
        BASE_VERSION=${FRAMEWORK}_BASE_VERSION
        BASE_VERSION=${!BASE_VERSION}
368
369
370
371

    fi

    if [ -z "$TAG" ]; then
372
        TAG="--tag dynamo:${VERSION}-${FRAMEWORK,,}"
Carsten Csiky's avatar
Carsten Csiky committed
373
        if [ -n "${TARGET}" ]; then
374
375
            TAG="${TAG}-${TARGET}"
        fi
376
377
    fi

Carsten Csiky's avatar
Carsten Csiky committed
378
    if [ -n "$PLATFORM" ]; then
379
380
381
        PLATFORM="--platform ${PLATFORM}"
    fi

Carsten Csiky's avatar
Carsten Csiky committed
382
    if [ -n "$TARGET" ]; then
383
        TARGET_STR="--target ${TARGET}"
384
    else
385
        TARGET_STR="--target dev"
386
    fi
387
388
389
390
391
392
393
394
395
396

    # Validate sccache configuration
    if [ "$USE_SCCACHE" = true ]; then
        if [ -z "$SCCACHE_BUCKET" ]; then
            error "ERROR: --sccache-bucket is required when --use-sccache is specified"
        fi
        if [ -z "$SCCACHE_REGION" ]; then
            error "ERROR: --sccache-region is required when --use-sccache is specified"
        fi
    fi
397
398
399
400
401
}


show_image_options() {
    echo ""
402
    echo "Building Dynamo Image: '${TAG}'"
403
404
405
    echo ""
    echo "   Base: '${BASE_IMAGE}'"
    echo "   Base_Image_Tag: '${BASE_IMAGE_TAG}'"
406
    if [[ $FRAMEWORK == "TRTLLM" ]]; then
407
        echo "   Tensorrtllm_Pip_Wheel: '${TENSORRTLLM_PIP_WHEEL}'"
408
409
410
411
    fi
    echo "   Build Context: '${BUILD_CONTEXT}'"
    echo "   Build Arguments: '${BUILD_ARGS}'"
    echo "   Framework: '${FRAMEWORK}'"
412
413
414
415
416
417
418
419
420
    if [ "$USE_SCCACHE" = true ]; then
        echo "   sccache: Enabled"
        echo "   sccache Bucket: '${SCCACHE_BUCKET}'"
        echo "   sccache Region: '${SCCACHE_REGION}'"

        if [ -n "$SCCACHE_S3_KEY_PREFIX" ]; then
            echo "   sccache S3 Key Prefix: '${SCCACHE_S3_KEY_PREFIX}'"
        fi
    fi
421
422
423
424
425
    echo ""
}

show_help() {
    echo "usage: build.sh"
426
    echo "  [--base-image base image]"
Carsten Csiky's avatar
Carsten Csiky committed
427
    echo "  [--base-image-tag base image tag]"
428
    echo "  [--platform platform for docker build]"
Carsten Csiky's avatar
Carsten Csiky committed
429
    echo "  [--framework framework one of ${!FRAMEWORKS[*]}]"
430
431
    echo "  [--tensorrtllm-pip-wheel-dir path to tensorrtllm pip wheel directory]"
    echo "  [--tensorrtllm-commit tensorrtllm commit to use for building the trtllm wheel if the wheel is not provided]"
432
    echo "  [--use-default-experimental-tensorrtllm-commit] Use the default experimental commit (${DEFAULT_EXPERIMENTAL_TRTLLM_COMMIT}) to build TensorRT-LLM. This is a flag (no argument). Do not combine with --tensorrtllm-commit or --tensorrtllm-pip-wheel."
433
434
    echo "  [--tensorrtllm-pip-wheel tensorrtllm pip wheel on artifactory]"
    echo "  [--tensorrtllm-index-url tensorrtllm PyPI index URL if providing the wheel from artifactory]"
435
    echo "  [--tensorrtllm-git-url tensorrtllm git repository URL for cloning]"
436
    echo "  [--build-arg additional build args to pass to docker build]"
437
438
    echo "  [--cache-from cache location to start from]"
    echo "  [--cache-to location where to cache the build output]"
439
    echo "  [--tag tag for image]"
440
441
    echo "  [--uid user ID for dev target (default: current user)]"
    echo "  [--gid group ID for dev target (default: current group)]"
442
443
    echo "  [--no-cache disable docker build cache]"
    echo "  [--dry-run print docker commands without running]"
ptarasiewiczNV's avatar
ptarasiewiczNV committed
444
    echo "  [--build-context name=path to add build context]"
445
446
    echo "  [--release-build perform a release build]"
    echo "  [--make-efa Enables EFA support for NIXL]"
447
    echo "  [--enable-kvbm Enables KVBM support in Python 3.12]"
448
    echo "  [--trtllm-use-nixl-kvcache-experimental Enables NIXL KVCACHE experimental support for TensorRT-LLM]"
449
450
451
    echo "  [--use-sccache enable sccache for Rust/C/C++ compilation caching]"
    echo "  [--sccache-bucket S3 bucket name for sccache (required with --use-sccache)]"
    echo "  [--sccache-region S3 region for sccache (required with --use-sccache)]"
452
453
454
455
456
457
458
459
460
461
462
463
464
465
    exit 0
}

missing_requirement() {
    error "ERROR: $1 requires an argument."
}

error() {
    printf '%s %s\n' "$1" "$2" >&2
    exit 1
}

get_options "$@"

466
467
468
469
470
471
472
# Validate UID/GID flags are only used with dev target
if [ -n "$USER_UID" ] || [ -n "$USER_GID" ]; then
    if [[ "$TARGET" != "dev" ]]; then
        echo "⚠️  Warning: --uid and --gid flags are only effective with --target dev"
        echo "   Current target: ${TARGET:-}"
    fi
fi
473

474
# Automatically set ARCH and ARCH_ALT if PLATFORM is linux/arm64
475
ARCH="amd64"
476
if [[ "$PLATFORM" == *"linux/arm64"* ]]; then
477
    ARCH="arm64"
478
479
480
    BUILD_ARGS+=" --build-arg ARCH=arm64 --build-arg ARCH_ALT=aarch64 "
fi

481
482
483
# Update DOCKERFILE if framework is VLLM
if [[ $FRAMEWORK == "VLLM" ]]; then
    DOCKERFILE=${SOURCE_DIR}/Dockerfile.vllm
484
485
elif [[ $FRAMEWORK == "TRTLLM" ]]; then
    DOCKERFILE=${SOURCE_DIR}/Dockerfile.trtllm
486
elif [[ $FRAMEWORK == "NONE" ]]; then
487
    DOCKERFILE=${SOURCE_DIR}/Dockerfile
488
489
elif [[ $FRAMEWORK == "SGLANG" ]]; then
    DOCKERFILE=${SOURCE_DIR}/Dockerfile.sglang
490
491
fi

492
493
# Add NIXL_REF as a build argument
BUILD_ARGS+=" --build-arg NIXL_REF=${NIXL_REF} "
494

495
if [[ $TARGET == "dev" ]]; then
496
497
498
499
500
501
502
503
    # Use provided UID/GID or default to current user
    if [ -z "$USER_UID" ]; then
        USER_UID=$(id -u)
    fi
    if [ -z "$USER_GID" ]; then
        USER_GID=$(id -g)
    fi
    BUILD_ARGS+=" --build-arg USER_UID=$USER_UID --build-arg USER_GID=$USER_GID "
504
505
fi

506
507
# BUILD DEV IMAGE

508
BUILD_ARGS+=" --build-arg BASE_IMAGE=$BASE_IMAGE --build-arg BASE_IMAGE_TAG=$BASE_IMAGE_TAG --build-arg FRAMEWORK=$FRAMEWORK --build-arg ${FRAMEWORK}_FRAMEWORK=1 --build-arg VERSION=$VERSION --build-arg PYTHON_PACKAGE_VERSION=$PYTHON_PACKAGE_VERSION"
509

Carsten Csiky's avatar
Carsten Csiky committed
510
if [ -n "${GITHUB_TOKEN}" ]; then
511
512
513
    BUILD_ARGS+=" --build-arg GITHUB_TOKEN=${GITHUB_TOKEN} "
fi

Carsten Csiky's avatar
Carsten Csiky committed
514
if [ -n "${GITLAB_TOKEN}" ]; then
515
516
517
    BUILD_ARGS+=" --build-arg GITLAB_TOKEN=${GITLAB_TOKEN} "
fi

518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557

check_wheel_file() {
    local wheel_dir="$1"
    # Check if directory exists
    if [ ! -d "$wheel_dir" ]; then
        echo "Error: Directory '$wheel_dir' does not exist"
        return 1
    fi

    # Look for .whl files
    wheel_count=$(find "$wheel_dir" -name "*.whl" | wc -l)

    if [ "$wheel_count" -eq 0 ]; then
        echo "WARN: No .whl files found in '$wheel_dir'"
        return 1
    elif [ "$wheel_count" -gt 1 ]; then
        echo "Warning: Multiple wheel files found in '$wheel_dir'. Will use first one found."
        find "$wheel_dir" -name "*.whl" | head -n 1
        return 0
    else
        echo "Found $wheel_count wheel files in '$wheel_dir'"
        # Check if commit file exists
        commit_file="$wheel_dir/commit.txt"
        if [ ! -f "$commit_file" ]; then
            echo "Error: Commit file '$commit_file' does not exist"
            return 1
        fi

        # Check if commit ID matches, otherwise re-build the wheel
        # Commit ID is of the form <arch>_<commit_id>
        commit_id=$(cat "$commit_file")
        if [ "$commit_id" != "$2" ]; then
            echo "Error: Commit ID mismatch. Expected '$2', got '$commit_id'"
            rm -rf $wheel_dir/*.whl
            return 1
        fi
        return 0
    fi
}

558
if [[ $FRAMEWORK == "TRTLLM" ]]; then
559
560
561
562
563
564
565
566
    if [ "$USE_DEFAULT_EXPERIMENTAL_TRTLLM_COMMIT" = true ]; then
        if [ -n "$TRTLLM_COMMIT" ] || [ -n "$TENSORRTLLM_PIP_WHEEL" ]; then
            echo "ERROR: When using --use-default-experimental-trtllm-commit, do not set --tensorrtllm-commit or --tensorrtllm-pip-wheel."
            exit 1
        fi
        TRTLLM_COMMIT="$DEFAULT_EXPERIMENTAL_TRTLLM_COMMIT"
    fi

567
568
569
570
    if [ -n "${TRTLLM_USE_NIXL_KVCACHE_EXPERIMENTAL}" ]; then
        BUILD_ARGS+=" --build-arg TRTLLM_USE_NIXL_KVCACHE_EXPERIMENTAL=${TRTLLM_USE_NIXL_KVCACHE_EXPERIMENTAL} "
    fi

571
572
573
574
575
    # If user didn't set both wheel and commit, use default tensorrt_llm pip wheel
    if [ -z "$TENSORRTLLM_PIP_WHEEL" ] && [ -z "$TRTLLM_COMMIT" ]; then
        TENSORRTLLM_PIP_WHEEL="$DEFAULT_TENSORRTLLM_PIP_WHEEL"
    fi

576
577
578
579
580
581
582
583
584
585
    if [ -z "${TENSORRTLLM_PIP_WHEEL}" ]; then
        # Use option 1
        if [ ! -d "${TENSORRTLLM_PIP_WHEEL_DIR}" ]; then
            # Create the directory if it doesn't exist
            mkdir -p ${TENSORRTLLM_PIP_WHEEL_DIR}
        fi
        BUILD_ARGS+=" --build-arg HAS_TRTLLM_CONTEXT=1"
        echo "Checking for TensorRT-LLM wheel in ${TENSORRTLLM_PIP_WHEEL_DIR}"
        if ! check_wheel_file "${TENSORRTLLM_PIP_WHEEL_DIR}" "${ARCH}_${TRTLLM_COMMIT}"; then
            echo "WARN: Valid trtllm wheel file not found in ${TENSORRTLLM_PIP_WHEEL_DIR}, attempting to build from source"
586
587
588
589
590
            GIT_URL_ARG=""
            if [ -n "${TRTLLM_GIT_URL}" ]; then
                GIT_URL_ARG="-u ${TRTLLM_GIT_URL}"
            fi
            if ! env -i ${SOURCE_DIR}/build_trtllm_wheel.sh -o ${TENSORRTLLM_PIP_WHEEL_DIR} -c ${TRTLLM_COMMIT} -a ${ARCH} -n ${NIXL_REF} ${GIT_URL_ARG}; then
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
                error "ERROR: Failed to build TensorRT-LLM wheel"
            fi
        fi
        echo "Installing TensorRT-LLM from local wheel directory"
        BUILD_CONTEXT_ARG+=" --build-context trtllm_wheel=${TENSORRTLLM_PIP_WHEEL_DIR}"

    else
        BUILD_ARGS+=" --build-arg HAS_TRTLLM_CONTEXT=0"
        BUILD_ARGS+=" --build-arg TENSORRTLLM_PIP_WHEEL=${TENSORRTLLM_PIP_WHEEL}"
        BUILD_ARGS+=" --build-arg TENSORRTLLM_INDEX_URL=${TENSORRTLLM_INDEX_URL}"

        # Create a dummy directory to satisfy the build context requirement
        # There is no way to conditionally copy the build context in dockerfile.
        mkdir -p /tmp/dummy_dir
        BUILD_CONTEXT_ARG+=" --build-context trtllm_wheel=/tmp/dummy_dir"
606
    fi
607
608
fi

Carsten Csiky's avatar
Carsten Csiky committed
609
if [ -n "${HF_TOKEN}" ]; then
610
611
    BUILD_ARGS+=" --build-arg HF_TOKEN=${HF_TOKEN} "
fi
612
613
614
615
if [  ! -z ${RELEASE_BUILD} ]; then
    echo "Performing a release build!"
    BUILD_ARGS+=" --build-arg RELEASE_BUILD=${RELEASE_BUILD} "
fi
616

617
618
619
620
621
if [[ $FRAMEWORK == "VLLM" ]]; then
    echo "Forcing enable_kvbm to true in vLLM image build"
    ENABLE_KVBM=true
fi

622
623
624
625
626
if [  ! -z ${ENABLE_KVBM} ]; then
    echo "Enabling the KVBM in the ai-dynamo-runtime"
    BUILD_ARGS+=" --build-arg ENABLE_KVBM=${ENABLE_KVBM} "
fi

627
628
629
630
if [ -n "${NIXL_UCX_REF}" ]; then
    BUILD_ARGS+=" --build-arg NIXL_UCX_REF=${NIXL_UCX_REF} "
fi

631
632
633
634
635
636
637
# Add sccache build arguments
if [ "$USE_SCCACHE" = true ]; then
    BUILD_ARGS+=" --build-arg USE_SCCACHE=true"
    BUILD_ARGS+=" --build-arg SCCACHE_BUCKET=${SCCACHE_BUCKET}"
    BUILD_ARGS+=" --build-arg SCCACHE_REGION=${SCCACHE_REGION}"
fi

638
LATEST_TAG="--tag dynamo:latest-${FRAMEWORK,,}"
Carsten Csiky's avatar
Carsten Csiky committed
639
if [ -n "${TARGET}" ]; then
640
641
    LATEST_TAG="${LATEST_TAG}-${TARGET}"
fi
642

643
644
645
646
647
648
show_image_options

if [ -z "$RUN_PREFIX" ]; then
    set -x
fi

649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
# TODO: Follow 2-step build process for all frameworks once necessary changes are made to the sglang and TRT-LLM backend Dockerfiles.
if [[ $FRAMEWORK == "VLLM" ]]; then
    # Define base image tag before using it
    DYNAMO_BASE_IMAGE="dynamo-base:${VERSION}"
    # Start base image build
    echo "======================================"
    echo "Starting Build 1: Base Image"
    echo "======================================"
    $RUN_PREFIX docker build -f "${SOURCE_DIR}/Dockerfile" --target dev $PLATFORM $BUILD_ARGS $CACHE_FROM $CACHE_TO --tag $DYNAMO_BASE_IMAGE $BUILD_CONTEXT_ARG $BUILD_CONTEXT $NO_CACHE
    # Start framework build
    echo "======================================"
    echo "Starting Build 2: Framework Image"
    echo "======================================"
    BUILD_ARGS+=" --build-arg DYNAMO_BASE_IMAGE=${DYNAMO_BASE_IMAGE}"
    $RUN_PREFIX docker build -f $DOCKERFILE $TARGET_STR $PLATFORM $BUILD_ARGS $CACHE_FROM $CACHE_TO $TAG $LATEST_TAG $BUILD_CONTEXT_ARG $BUILD_CONTEXT $NO_CACHE
else
    $RUN_PREFIX docker build -f $DOCKERFILE $TARGET_STR $PLATFORM $BUILD_ARGS $CACHE_FROM $CACHE_TO $TAG $LATEST_TAG $BUILD_CONTEXT_ARG $BUILD_CONTEXT $NO_CACHE
fi

668
669

{ set +x; } 2>/dev/null