build.sh 10.1 KB
Newer Older
1
#!/bin/bash -e
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
TAG=
RUN_PREFIX=
PLATFORM=linux/amd64
21
22
23
24

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

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

28
29
# Get latest TAG and add COMMIT_ID for dev
latest_tag=$(git describe --tags --abbrev=0 $(git rev-list --tags --max-count=1 main) | sed 's/^v//') || true
30
31
32
33
if [[ -z ${latest_tag} ]]; then
    latest_tag="0.0.1"
    echo "No git release tag found, setting to unknown version: ${latest_tag}"
fi
34

35
36
37
38
# 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}
39
40
41
42
43
44

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

ptarasiewiczNV's avatar
ptarasiewiczNV committed
46
declare -A FRAMEWORKS=(["STANDARD"]=1 ["TENSORRTLLM"]=2 ["VLLM"]=3 ["VLLM_NIXL"]=4)
47
48
49
50
51
52
53
54
DEFAULT_FRAMEWORK=STANDARD

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

# Base Images

55
STANDARD_BASE_VERSION=25.01
56
57
58
STANDARD_BASE_IMAGE=nvcr.io/nvidia/tritonserver
STANDARD_BASE_IMAGE_TAG=${STANDARD_BASE_VERSION}-py3

59
TENSORRTLLM_BASE_VERSION=25.01
60
61
TENSORRTLLM_BASE_IMAGE=nvcr.io/nvidia/tritonserver
TENSORRTLLM_BASE_IMAGE_TAG=${TENSORRTLLM_BASE_VERSION}-trtllm-python-py3
62
63
# IMPORTANT NOTE: Ensure the repo tag complies with the TRTLLM backend version
# used in the base image above.
64
TENSORRTLLM_BACKEND_REPO_TAG=triton-llm/v0.17.0
65
66
67
# Set this as 1 to rebuild and replace trtllm backend bits in the container.
# This will allow building triton distributed container image with custom
# trt-llm backend repo branch.
68
69
70
71
72
TENSORRTLLM_BACKEND_REBUILD=0
# Set this as 1 to skip cloning the trt-llm backend repo. If cloning is skipped, trt-llm
# backend repo tag and rebuild flag will be ignored. Use this option if you are using
# trtllm llmapi worker.
TENSORRTLLM_SKIP_CLONE=0
73

74
75
VLLM_BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base"
VLLM_BASE_IMAGE_TAG="25.01-cuda12.8-devel-ubuntu24.04"
76

ptarasiewiczNV's avatar
ptarasiewiczNV committed
77
78
79
VLLM_NIXL_BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base"
VLLM_NIXL_BASE_IMAGE_TAG="25.01-cuda12.8-devel-ubuntu24.04"

80
81
82
83
84
85
86
get_options() {
    while :; do
        case $1 in
        -h | -\? | --help)
            show_help
            exit
            ;;
87
        --platform)
88
89
90
91
            if [ "$2" ]; then
                PLATFORM=$2
                shift
            else
92
                missing_requirement $1
93
94
            fi
            ;;
95
        --framework)
96
97
98
99
            if [ "$2" ]; then
                FRAMEWORK=$2
                shift
            else
100
                missing_requirement $1
101
102
            fi
            ;;
103
        --tensorrtllm-backend-repo-tag)
104
105
106
107
            if [ "$2" ]; then
                TRTLLM_BACKEND_COMMIT=$2
                shift
            else
108
                missing_requirement $1
109
110
            fi
            ;;
111
112
113
114
115
        --tensorrtllm-backend-rebuild)
            if [ "$2" ]; then
                TRTLLM_BACKEND_REBUILD=$2
                shift
            else
116
                missing_requirement $1
117
118
            fi
            ;;
119
120
121
122
123
124
125
126
        --skip-clone-tensorrtllm)
            if [ "$2" ]; then
                TENSORRTLLM_SKIP_CLONE=$2
                shift
            else
                missing_requirement $1
            fi
            ;;
127
128
129
130
131
        --base-image)
            if [ "$2" ]; then
                BASE_IMAGE=$2
                shift
            else
132
                missing_requirement $1
133
134
            fi
            ;;
135
        --base-image-tag)
136
137
138
139
            if [ "$2" ]; then
                BASE_IMAGE_TAG=$2
                shift
            else
140
141
142
143
144
145
146
147
148
                missing_requirement $1
            fi
            ;;
        --target)
            if [ "$2" ]; then
                TARGET=$2
                shift
            else
                missing_requirement $1
149
150
151
152
153
154
155
            fi
            ;;
        --build-arg)
            if [ "$2" ]; then
                BUILD_ARGS+="--build-arg $2 "
                shift
            else
156
                missing_requirement $1
157
158
159
160
            fi
            ;;
        --tag)
            if [ "$2" ]; then
161
                TAG="--tag $2"
162
163
                shift
            else
164
                missing_requirement $1
165
166
167
168
169
170
171
172
173
174
            fi
            ;;
        --dry-run)
            RUN_PREFIX="echo"
            echo ""
            echo "=============================="
            echo "DRY RUN: COMMANDS PRINTED ONLY"
            echo "=============================="
            echo ""
            ;;
175
176
        --no-cache)
            NO_CACHE=" --no-cache"
177
            ;;
178
179
        --cache-from)
            if [ "$2" ]; then
180
181
182
                CACHE_FROM="--cache-from $2"
                shift
            else
183
                missing_requirement $1
184
185
            fi
            ;;
ptarasiewiczNV's avatar
ptarasiewiczNV committed
186
187
188
189
190
191
192
193
        --build-context)
            if [ "$2" ]; then
                BUILD_CONTEXT_ARG="--build-context $2"
                shift
            else
                missing_requirement $1
            fi
            ;;
194
195
196
197
198
        --)
            shift
            break
            ;;
         -?*)
199
            error 'ERROR: Unknown option: ' $1
200
            ;;
201
202
         ?*)
            error 'ERROR: Unknown option: ' $1
203
204
205
206
207
208
209
210
211
            ;;
        *)
            break
            ;;
        esac
        shift
    done

    if [ -z "$FRAMEWORK" ]; then
212
        FRAMEWORK=$DEFAULT_FRAMEWORK
213
214
215
    fi

    if [ ! -z "$FRAMEWORK" ]; then
216
        FRAMEWORK=${FRAMEWORK^^}
217

218
219
220
        if [[ ! -n "${FRAMEWORKS[$FRAMEWORK]}" ]]; then
            error 'ERROR: Unknown framework: ' $FRAMEWORK
        fi
221

222
223
224
225
        if [ -z $BASE_IMAGE_TAG ]; then
            BASE_IMAGE_TAG=${FRAMEWORK}_BASE_IMAGE_TAG
            BASE_IMAGE_TAG=${!BASE_IMAGE_TAG}
        fi
226

227
228
229
230
        if [ -z $BASE_IMAGE ]; then
            BASE_IMAGE=${FRAMEWORK}_BASE_IMAGE
            BASE_IMAGE=${!BASE_IMAGE}
        fi
231

232
233
234
        if [ -z $BASE_IMAGE ]; then
            error "ERROR: Framework $FRAMEWORK without BASE_IMAGE"
        fi
235

236
237
        BASE_VERSION=${FRAMEWORK}_BASE_VERSION
        BASE_VERSION=${!BASE_VERSION}
238
239
240
241

    fi

    if [ -z "$TAG" ]; then
242
        TAG="--tag triton-distributed:${VERSION}-${FRAMEWORK,,}"
243
244
245
        if [ ! -z ${TARGET} ]; then
            TAG="${TAG}-${TARGET}"
        fi
246
247
248
249
250
251
    fi

    if [ ! -z "$PLATFORM" ]; then
        PLATFORM="--platform ${PLATFORM}"
    fi

252
253
254
    if [ ! -z "$TARGET" ]; then
        TARGET_STR="--target ${TARGET}"
    fi
255
256
257
258
259
260
261
262
263
264
}


show_image_options() {
    echo ""
    echo "Building Triton Distributed Image: '${TAG}'"
    echo ""
    echo "   Base: '${BASE_IMAGE}'"
    echo "   Base_Image_Tag: '${BASE_IMAGE_TAG}'"
    if [[ $FRAMEWORK == "TENSORRTLLM" ]]; then
265
        echo "   Tensorrtllm Backend Repo Tag: '${TENSORRTLLM_BACKEND_REPO_TAG}'"
266
        echo "   Tensorrtllm Backend Rebuild: '${TENSORRTLLM_BACKEND_REBUILD}'"
267
        echo "   Tensorrtllm Skip Clone: '${TENSORRTLLM_SKIP_CLONE}'"
268
269
270
271
272
273
274
275
276
277
278
279
280
    fi
    echo "   Build Context: '${BUILD_CONTEXT}'"
    echo "   Build Arguments: '${BUILD_ARGS}'"
    echo "   Framework: '${FRAMEWORK}'"
    echo ""
}

show_help() {
    echo "usage: build.sh"
    echo "  [--base base image]"
    echo "  [--base-imge-tag base image tag]"
    echo "  [--platform platform for docker build"
    echo "  [--framework framework one of ${!FRAMEWORKS[@]}]"
281
282
    echo "  [--tensorrtllm-backend-repo-tag commit or tag]"
    echo "  [--tensorrtllm-backend-rebuild whether or not to rebuild the backend]"
283
    echo "  [--skip-clone-tensorrtllm whether or not to skip cloning the trt-llm backend repo]"
284
285
286
287
    echo "  [--build-arg additional build args to pass to docker build]"
    echo "  [--tag tag for image]"
    echo "  [--no-cache disable docker build cache]"
    echo "  [--dry-run print docker commands without running]"
ptarasiewiczNV's avatar
ptarasiewiczNV committed
288
    echo "  [--build-context name=path to add build context]"
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
    exit 0
}

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

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

get_options "$@"


304
305
306
# Update DOCKERFILE if framework is VLLM
if [[ $FRAMEWORK == "VLLM" ]]; then
    DOCKERFILE=${SOURCE_DIR}/Dockerfile.vllm
ptarasiewiczNV's avatar
ptarasiewiczNV committed
307
308
elif [[ $FRAMEWORK == "VLLM_NIXL" ]]; then
    DOCKERFILE=${SOURCE_DIR}/Dockerfile.vllm_nixl
309
310
fi

311
312
# BUILD DEV IMAGE

313
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"
314
315
316
317
318
319
320
321
322

if [ ! -z ${GITHUB_TOKEN} ]; then
    BUILD_ARGS+=" --build-arg GITHUB_TOKEN=${GITHUB_TOKEN} "
fi

if [ ! -z ${GITLAB_TOKEN} ]; then
    BUILD_ARGS+=" --build-arg GITLAB_TOKEN=${GITLAB_TOKEN} "
fi

323
324
325
if [[ $FRAMEWORK == "TENSORRTLLM" ]] && [ ! -z ${TENSORRTLLM_BACKEND_REPO_TAG} ]; then
    BUILD_ARGS+=" --build-arg TENSORRTLLM_BACKEND_REPO_TAG=${TENSORRTLLM_BACKEND_REPO_TAG} "
    BUILD_ARGS+=" --build-arg TENSORRTLLM_BACKEND_REBUILD=${TENSORRTLLM_BACKEND_REBUILD} "
326
    BUILD_ARGS+=" --build-arg TENSORRTLLM_SKIP_CLONE=${TENSORRTLLM_SKIP_CLONE} "
327
328
329
330
331
332
fi

if [ ! -z ${HF_TOKEN} ]; then
    BUILD_ARGS+=" --build-arg HF_TOKEN=${HF_TOKEN} "
fi

333
LATEST_TAG="--tag triton-distributed:latest-${FRAMEWORK,,}"
334
335
336
if [ ! -z ${TARGET} ]; then
    LATEST_TAG="${LATEST_TAG}-${TARGET}"
fi
337

338
339
340
341
342
343
show_image_options

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

ptarasiewiczNV's avatar
ptarasiewiczNV committed
344
$RUN_PREFIX docker build -f $DOCKERFILE $TARGET_STR $PLATFORM $BUILD_ARGS $CACHE_FROM $TAG $LATEST_TAG $BUILD_CONTEXT_ARG $BUILD_CONTEXT $NO_CACHE
345
346
347
348
349
350
351
352

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

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

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