build.sh 27.6 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
29
30

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

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

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

41
42
43
44
# 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}
45
46
47
48
49
50

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

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

54
DEFAULT_FRAMEWORK=VLLM
55
56
57
58
59
60

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

# Base Images
61
TRTLLM_BASE_IMAGE=nvcr.io/nvidia/pytorch
62
TRTLLM_BASE_IMAGE_TAG=25.06-py3
63
64
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

# 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.
92
DEFAULT_EXPERIMENTAL_TRTLLM_COMMIT="0c9430e5a530ba958fc9dca561a3ad865ad9f492"
93
TRTLLM_COMMIT=""
94
TRTLLM_USE_NIXL_KVCACHE_EXPERIMENTAL="0"
95
TRTLLM_GIT_URL=""
96

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


105
VLLM_BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base"
106
107
108
109
110
# 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"
111

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

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

118
NIXL_REF=0.6.0
119
120
NIXL_UCX_REF=v1.19.0
NIXL_UCX_EFA_REF=9d2b88a1f67faf9876f267658bd077b379b8bb76
121

122
123
NO_CACHE=""

124
125
126
127
128
# sccache configuration for S3
USE_SCCACHE=""
SCCACHE_BUCKET=""
SCCACHE_REGION=""

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

        --vllm-max-jobs)
334
            # Set MAX_JOBS for vLLM compilation (only used by Dockerfile.vllm)
335
336
337
338
339
340
341
            if [ "$2" ]; then
                MAX_JOBS=$2
                shift
            else
                missing_requirement "$1"
            fi
            ;;
342
         -?*)
Carsten Csiky's avatar
Carsten Csiky committed
343
            error 'ERROR: Unknown option: ' "$1"
344
            ;;
345
         ?*)
Carsten Csiky's avatar
Carsten Csiky committed
346
            error 'ERROR: Unknown option: ' "$1"
347
348
349
350
351
352
353
354
            ;;
        *)
            break
            ;;
        esac
        shift
    done

355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
    # Validate argument combinations
    if [[ -n "${DEV_IMAGE_INPUT:-}" && -n "${BASE_IMAGE:-}" ]]; then
        error "ERROR: --dev-image cannot be used with --base-image. Use --dev-image to build from existing images or --base-image to build new images."
    fi

    # Validate that --target and --dev-image cannot be used together
    if [[ -n "${DEV_IMAGE_INPUT:-}" && -n "${TARGET:-}" ]]; then
        error "ERROR: --target cannot be used with --dev-image. Use --target to build from scratch or --dev-image to build from existing images."
    fi

    # Validate that --uid and --gid are only used with local-dev related options
    if [[ -n "${CUSTOM_UID:-}" || -n "${CUSTOM_GID:-}" ]]; then
        if [[ -z "${DEV_IMAGE_INPUT:-}" && "${TARGET:-}" != "local-dev" ]]; then
            error "ERROR: --uid and --gid can only be used with --dev-image or --target local-dev"
        fi
    fi

372
    if [ -z "$FRAMEWORK" ]; then
373
        FRAMEWORK=$DEFAULT_FRAMEWORK
374
375
    fi

Carsten Csiky's avatar
Carsten Csiky committed
376
    if [ -n "$FRAMEWORK" ]; then
377
        FRAMEWORK=${FRAMEWORK^^}
378

Carsten Csiky's avatar
Carsten Csiky committed
379
380
        if [[ -z "${FRAMEWORKS[$FRAMEWORK]}" ]]; then
            error 'ERROR: Unknown framework: ' "$FRAMEWORK"
381
        fi
382

Carsten Csiky's avatar
Carsten Csiky committed
383
        if [ -z "$BASE_IMAGE_TAG" ]; then
384
385
386
            BASE_IMAGE_TAG=${FRAMEWORK}_BASE_IMAGE_TAG
            BASE_IMAGE_TAG=${!BASE_IMAGE_TAG}
        fi
387

Carsten Csiky's avatar
Carsten Csiky committed
388
        if [ -z "$BASE_IMAGE" ]; then
389
390
391
            BASE_IMAGE=${FRAMEWORK}_BASE_IMAGE
            BASE_IMAGE=${!BASE_IMAGE}
        fi
392

Carsten Csiky's avatar
Carsten Csiky committed
393
        if [ -z "$BASE_IMAGE" ]; then
394
395
            error "ERROR: Framework $FRAMEWORK without BASE_IMAGE"
        fi
396

397
398
        BASE_VERSION=${FRAMEWORK}_BASE_VERSION
        BASE_VERSION=${!BASE_VERSION}
399
400
401
402

    fi

    if [ -z "$TAG" ]; then
403
        TAG="--tag dynamo:${VERSION}-${FRAMEWORK,,}"
404
        if [ -n "${TARGET}" ] && [ "${TARGET}" != "local-dev" ]; then
405
406
            TAG="${TAG}-${TARGET}"
        fi
407
408
    fi

Carsten Csiky's avatar
Carsten Csiky committed
409
    if [ -n "$PLATFORM" ]; then
410
411
412
        PLATFORM="--platform ${PLATFORM}"
    fi

Carsten Csiky's avatar
Carsten Csiky committed
413
    if [ -n "$TARGET" ]; then
414
        TARGET_STR="--target ${TARGET}"
415
    else
416
        TARGET_STR="--target dev"
417
    fi
418
419
420
421
422
423
424
425
426
427

    # 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
428
429
430
431
432
}


show_image_options() {
    echo ""
433
    echo "Building Dynamo Image: '${TAG}'"
434
435
436
    echo ""
    echo "   Base: '${BASE_IMAGE}'"
    echo "   Base_Image_Tag: '${BASE_IMAGE_TAG}'"
437
    if [[ $FRAMEWORK == "TRTLLM" ]]; then
438
        echo "   Tensorrtllm_Pip_Wheel: '${TENSORRTLLM_PIP_WHEEL}'"
439
440
441
442
    fi
    echo "   Build Context: '${BUILD_CONTEXT}'"
    echo "   Build Arguments: '${BUILD_ARGS}'"
    echo "   Framework: '${FRAMEWORK}'"
443
444
445
446
447
448
449
450
451
    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
452
453
454
455
456
    echo ""
}

show_help() {
    echo "usage: build.sh"
457
    echo "  [--base-image base image]"
Carsten Csiky's avatar
Carsten Csiky committed
458
    echo "  [--base-image-tag base image tag]"
459
    echo "  [--platform platform for docker build]"
Carsten Csiky's avatar
Carsten Csiky committed
460
    echo "  [--framework framework one of ${!FRAMEWORKS[*]}]"
461
462
    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]"
463
    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."
464
465
    echo "  [--tensorrtllm-pip-wheel tensorrtllm pip wheel on artifactory]"
    echo "  [--tensorrtllm-index-url tensorrtllm PyPI index URL if providing the wheel from artifactory]"
466
    echo "  [--tensorrtllm-git-url tensorrtllm git repository URL for cloning]"
467
    echo "  [--build-arg additional build args to pass to docker build]"
468
469
    echo "  [--cache-from cache location to start from]"
    echo "  [--cache-to location where to cache the build output]"
470
    echo "  [--tag tag for image]"
471
472
473
    echo "  [--dev-image dev image to build local-dev from]"
    echo "  [--uid user ID for local-dev images (only with --dev-image or --target local-dev)]"
    echo "  [--gid group ID for local-dev images (only with --dev-image or --target local-dev)]"
474
475
    echo "  [--no-cache disable docker build cache]"
    echo "  [--dry-run print docker commands without running]"
ptarasiewiczNV's avatar
ptarasiewiczNV committed
476
    echo "  [--build-context name=path to add build context]"
477
478
    echo "  [--release-build perform a release build]"
    echo "  [--make-efa Enables EFA support for NIXL]"
479
    echo "  [--enable-kvbm Enables KVBM support in Python 3.12]"
480
    echo "  [--trtllm-use-nixl-kvcache-experimental Enables NIXL KVCACHE experimental support for TensorRT-LLM]"
481
482
483
    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)]"
484
    echo "  [--vllm-max-jobs number of parallel jobs for compilation (only used by vLLM framework)]"
485
486
487
488
    echo ""
    echo "  Note: When using --use-sccache, AWS credentials must be set:"
    echo "        export AWS_ACCESS_KEY_ID=your_access_key"
    echo "        export AWS_SECRET_ACCESS_KEY=your_secret_key"
489
490
491
492
493
494
495
496
497
498
499
500
501
502
    exit 0
}

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

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

get_options "$@"

503
# Automatically set ARCH and ARCH_ALT if PLATFORM is linux/arm64
504
ARCH="amd64"
505
if [[ "$PLATFORM" == *"linux/arm64"* ]]; then
506
    ARCH="arm64"
507
508
509
    BUILD_ARGS+=" --build-arg ARCH=arm64 --build-arg ARCH_ALT=aarch64 "
fi

510
511
512
# Update DOCKERFILE if framework is VLLM
if [[ $FRAMEWORK == "VLLM" ]]; then
    DOCKERFILE=${SOURCE_DIR}/Dockerfile.vllm
513
514
elif [[ $FRAMEWORK == "TRTLLM" ]]; then
    DOCKERFILE=${SOURCE_DIR}/Dockerfile.trtllm
515
elif [[ $FRAMEWORK == "NONE" ]]; then
516
    DOCKERFILE=${SOURCE_DIR}/Dockerfile
517
518
elif [[ $FRAMEWORK == "SGLANG" ]]; then
    DOCKERFILE=${SOURCE_DIR}/Dockerfile.sglang
519
520
fi

521
522
# Add NIXL_REF as a build argument
BUILD_ARGS+=" --build-arg NIXL_REF=${NIXL_REF} "
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
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
# Function to build local-dev image with header
build_local_dev_with_header() {
    local dev_base_image="$1"
    local tags="$2"
    local success_msg="$3"
    local header_title="$4"

    echo "======================================"
    echo "$header_title"
    echo "======================================"

    # Get user info right before using it
    USER_UID=${CUSTOM_UID:-$(id -u)}
    USER_GID=${CUSTOM_GID:-$(id -g)}

    # Set up dockerfile path
    DOCKERFILE_LOCAL_DEV="${SOURCE_DIR}/Dockerfile.local_dev"

    if [[ ! -f "$DOCKERFILE_LOCAL_DEV" ]]; then
        echo "ERROR: Dockerfile.local_dev not found at: $DOCKERFILE_LOCAL_DEV"
        exit 1
    fi

    echo "Building new local-dev image from: $dev_base_image"
    echo "User 'ubuntu' will have UID: $USER_UID, GID: $USER_GID"

    # Show the docker command being executed if not in dry-run mode
    if [ -z "$RUN_PREFIX" ]; then
        set -x
    fi

    $RUN_PREFIX docker build \
        --build-arg DEV_BASE="$dev_base_image" \
        --build-arg USER_UID="$USER_UID" \
        --build-arg USER_GID="$USER_GID" \
        --build-arg ARCH="$ARCH" \
        --file "$DOCKERFILE_LOCAL_DEV" \
        $tags \
        "$SOURCE_DIR" || {
        { set +x; } 2>/dev/null
        echo "ERROR: Failed to build local_dev image"
        exit 1
    }

    { set +x; } 2>/dev/null
    echo "$success_msg"

    # Show usage instructions
    echo ""
    echo "To run the local-dev image as the local user ($USER_UID/$USER_GID):"
    # Extract the last tag from the tags string
    last_tag=$(echo "$tags" | grep -o -- '--tag [^ ]*' | tail -1 | cut -d' ' -f2)
    # Calculate relative path to run.sh from current working directory
    # Get the directory where build.sh is located
    build_dir="$(dirname "${BASH_SOURCE[0]}")"
    # Get the absolute path to run.sh (in the same directory as build.sh)
    run_abs_path="$(realpath "$build_dir/run.sh")"
    # Calculate relative path from current PWD to run.sh
    run_path="$(python3 -c "import os; print(os.path.relpath('$run_abs_path', '$PWD'))")"
    echo "  $run_path --image $last_tag --mount-workspace ..."
}


# Handle local-dev target
588
if [[ $TARGET == "local-dev" ]]; then
589
590
    LOCAL_DEV_BUILD=true
    TARGET_STR="--target dev"
591
592
fi

593
594
# BUILD DEV IMAGE

595
BUILD_ARGS+=" --build-arg BASE_IMAGE=$BASE_IMAGE --build-arg BASE_IMAGE_TAG=$BASE_IMAGE_TAG"
596

Carsten Csiky's avatar
Carsten Csiky committed
597
if [ -n "${GITHUB_TOKEN}" ]; then
598
599
600
    BUILD_ARGS+=" --build-arg GITHUB_TOKEN=${GITHUB_TOKEN} "
fi

Carsten Csiky's avatar
Carsten Csiky committed
601
if [ -n "${GITLAB_TOKEN}" ]; then
602
603
604
    BUILD_ARGS+=" --build-arg GITLAB_TOKEN=${GITLAB_TOKEN} "
fi

605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644

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
}

645
if [[ $FRAMEWORK == "TRTLLM" ]]; then
646
    BUILD_ARGS+=" --build-arg GITHUB_TRTLLM_COMMIT=${TRTLLM_COMMIT}"
647
648
649
650
651
652
653
654
    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

655
656
657
658
    if [ -n "${TRTLLM_USE_NIXL_KVCACHE_EXPERIMENTAL}" ]; then
        BUILD_ARGS+=" --build-arg TRTLLM_USE_NIXL_KVCACHE_EXPERIMENTAL=${TRTLLM_USE_NIXL_KVCACHE_EXPERIMENTAL} "
    fi

659
660
661
662
663
    # 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

664
665
666
667
668
669
670
671
672
673
    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"
674
675
676
677
678
            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
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
                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"
694
    fi
695
696
fi

697
698
# ENABLE_KVBM: Used in base Dockerfile for block-manager feature.
#              Declared but not currently used in Dockerfile.{vllm,trtllm}.
699
700
if [[ $FRAMEWORK == "VLLM" ]] || [[ $FRAMEWORK == "TRTLLM" ]]; then
    echo "Forcing enable_kvbm to true in ${FRAMEWORK} image build"
701
702
703
    ENABLE_KVBM=true
fi

704
705
706
707
708
if [  ! -z ${ENABLE_KVBM} ]; then
    echo "Enabling the KVBM in the ai-dynamo-runtime"
    BUILD_ARGS+=" --build-arg ENABLE_KVBM=${ENABLE_KVBM} "
fi

709
710
# NIXL_UCX_REF: Used in base Dockerfile only.
#               Passed to framework Dockerfile.{vllm,sglang,...} where it's NOT used.
711
712
713
714
if [ -n "${NIXL_UCX_REF}" ]; then
    BUILD_ARGS+=" --build-arg NIXL_UCX_REF=${NIXL_UCX_REF} "
fi

715
# MAX_JOBS is only used by Dockerfile.vllm
716
717
718
719
if [ -n "${MAX_JOBS}" ]; then
    BUILD_ARGS+=" --build-arg MAX_JOBS=${MAX_JOBS} "
fi

720
721
722
723
724
# 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}"
725
726
    BUILD_ARGS+=" --secret id=aws-key-id,env=AWS_ACCESS_KEY_ID"
    BUILD_ARGS+=" --secret id=aws-secret-id,env=AWS_SECRET_ACCESS_KEY"
727
728
fi

729
LATEST_TAG="--tag dynamo:latest-${FRAMEWORK,,}"
730
if [ -n "${TARGET}" ] && [ "${TARGET}" != "local-dev" ]; then
731
732
    LATEST_TAG="${LATEST_TAG}-${TARGET}"
fi
733

734
735
736
737
738
739
show_image_options

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

740
741
742

# Skip Build 1 and Build 2 if DEV_IMAGE_INPUT is set (we'll handle it at the bottom)
if [[ -z "${DEV_IMAGE_INPUT:-}" ]]; then
743
744
    # Follow 2-step build process for all frameworks
    if [[ $FRAMEWORK != "NONE" ]]; then
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
        # 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
761
762
fi

763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
# Handle --dev-image option (build local-dev from existing dev image)
if [[ -n "${DEV_IMAGE_INPUT:-}" ]]; then
    # Validate that the dev image is not already a local-dev image
    if [[ "$DEV_IMAGE_INPUT" == *"-local-dev" ]]; then
        echo "ERROR: Cannot use local-dev image as dev image input: '$DEV_IMAGE_INPUT'"
        exit 1
    fi

    # Build tag arguments - always add -local-dev suffix for --dev-image
    # Generate local-dev tag from input image
    if [[ "$DEV_IMAGE_INPUT" == *:* ]]; then
        LOCAL_DEV_TAG="--tag ${DEV_IMAGE_INPUT}-local-dev"
    else
        LOCAL_DEV_TAG="--tag ${DEV_IMAGE_INPUT}:latest-local-dev"
    fi

    build_local_dev_with_header "$DEV_IMAGE_INPUT" "$LOCAL_DEV_TAG" "Successfully built local-dev image: ${LOCAL_DEV_TAG#--tag }" "Building Local-Dev Image"
elif [[ "${LOCAL_DEV_BUILD:-}" == "true" ]]; then
    # Use the first tag name (TAG) if available, otherwise use latest
    if [[ -n "$TAG" ]]; then
        DEV_IMAGE=$(echo "$TAG" | sed 's/--tag //' | sed 's/-local-dev$//')
    else
        DEV_IMAGE="dynamo:latest-${FRAMEWORK,,}"
    fi

    # Build local-dev tags from existing tags
    LOCAL_DEV_TAGS=""
    if [[ -n "$TAG" ]]; then
        # Extract tag name, remove any existing -local-dev suffix, then add -local-dev
        TAG_NAME=$(echo "$TAG" | sed 's/--tag //' | sed 's/-local-dev$//')
        LOCAL_DEV_TAGS+=" --tag ${TAG_NAME}-local-dev"
    fi

    if [[ -n "$LATEST_TAG" ]]; then
        # Extract tag name, remove any existing -local-dev suffix, then add -local-dev
        LATEST_TAG_NAME=$(echo "$LATEST_TAG" | sed 's/--tag //' | sed 's/-local-dev$//')
        LOCAL_DEV_TAGS+=" --tag ${LATEST_TAG_NAME}-local-dev"
    fi

    build_local_dev_with_header "$DEV_IMAGE" "$LOCAL_DEV_TAGS" "Successfully built local-dev images" "Starting Build 3: Local-Dev Image"
fi
804

805

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