run.sh 10.3 KB
Newer Older
Carsten Csiky's avatar
Carsten Csiky committed
1
#!/usr/bin/env bash
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# SPDX-FileCopyrightText: Copyright (c) 2024-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.

Carsten Csiky's avatar
Carsten Csiky committed
17
18
set -e

19
20
21
22
23
24
25
26
RUN_PREFIX=

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

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

29
DEFAULT_FRAMEWORK=VLLM
30
31
32
33

SOURCE_DIR=$(dirname "$(readlink -f "$0")")

IMAGE=
34
35
HF_HOME=${HF_HOME:-}
DEFAULT_HF_HOME=${SOURCE_DIR}/.cache/huggingface
36
37
38
GPUS="all"
PRIVILEGED=
VOLUME_MOUNTS=
39
PORT_MAPPINGS=
40
41
42
43
MOUNT_WORKSPACE=
ENVIRONMENT_VARIABLES=
REMAINING_ARGS=
INTERACTIVE=
44
USE_NIXL_GDS=
45
46
RUNTIME=nvidia
WORKDIR=/workspace
47
NETWORK=host
48
USER=
49
50
51
52
53
54
55
56

get_options() {
    while :; do
        case $1 in
        -h | -\? | --help)
            show_help
            exit
            ;;
57
        --framework)
58
59
60
61
            if [ "$2" ]; then
                FRAMEWORK=$2
                shift
            else
62
                missing_requirement "$1"
63
64
65
66
67
68
69
            fi
            ;;
        --image)
            if [ "$2" ]; then
                IMAGE=$2
                shift
            else
70
                missing_requirement "$1"
71
72
            fi
            ;;
73
74
75
76
77
        --target)
            if [ "$2" ]; then
                TARGET=$2
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
78
                missing_requirement "$1"
79
80
            fi
            ;;
81
        --name)
82
83
84
85
            if [ "$2" ]; then
                NAME=$2
                shift
            else
86
                missing_requirement "$1"
87
88
            fi
            ;;
89
        --hf-cache|--hf-home)
90
            if [ "$2" ]; then
91
                HF_HOME=$2
92
93
                shift
            else
94
                missing_requirement "$1"
95
96
97
            fi
            ;;

98
        --gpus)
99
100
101
102
            if [ "$2" ]; then
                GPUS=$2
                shift
            else
103
                missing_requirement "$1"
104
105
            fi
            ;;
106
        --runtime)
107
108
109
110
            if [ "$2" ]; then
                RUNTIME=$2
                shift
            else
111
                missing_requirement "$1"
112
113
            fi
            ;;
114
        --entrypoint)
115
            if [ "$2" ]; then
116
                ENTRYPOINT=$2
117
118
                shift
            else
119
120
121
122
123
124
125
126
127
                missing_requirement "$1"
            fi
            ;;
        --workdir)
            if [ "$2" ]; then
                WORKDIR="$2"
                shift
            else
                missing_requirement "$1"
128
129
            fi
            ;;
130
        --privileged)
131
132
133
134
            if [ "$2" ]; then
                PRIVILEGED=$2
                shift
            else
135
                missing_requirement "$1"
136
137
            fi
            ;;
138
        --rm)
139
140
141
142
            if [ "$2" ]; then
                RM=$2
                shift
            else
143
                missing_requirement "$1"
144
145
            fi
            ;;
146
        -v)
147
148
149
150
            if [ "$2" ]; then
                VOLUME_MOUNTS+=" -v $2 "
                shift
            else
151
                missing_requirement "$1"
152
153
            fi
            ;;
154
155
156
157
158
159
160
161
        -p|--port)
            if [ "$2" ]; then
                PORT_MAPPINGS+=" -p $2 "
                shift
            else
                missing_requirement "$1"
            fi
            ;;
162
        -e)
163
164
165
166
            if [ "$2" ]; then
                ENVIRONMENT_VARIABLES+=" -e $2 "
                shift
            else
167
                missing_requirement "$1"
168
169
            fi
            ;;
170
171
172
173
174
175
        -it)
            INTERACTIVE=" -it "
            ;;
        --mount-workspace)
            MOUNT_WORKSPACE=TRUE
            ;;
176
177
178
        --use-nixl-gds)
            USE_NIXL_GDS=TRUE
            ;;
179
180
181
182
183
184
185
186
        --network)
            if [ "$2" ]; then
                NETWORK=$2
                shift
            else
                missing_requirement "$1"
            fi
            ;;
187
188
189
190
191
192
193
194
        --user)
            if [ "$2" ]; then
                USER=$2
                shift
            else
                missing_requirement "$1"
            fi
            ;;
195
196
197
198
199
200
201
202
203
204
205
206
207
        --dry-run)
            RUN_PREFIX="echo"
            echo ""
            echo "=============================="
            echo "DRY RUN: COMMANDS PRINTED ONLY"
            echo "=============================="
            echo ""
            ;;
        --)
            shift
            break
            ;;
         -?*)
208
            error 'ERROR: Unknown option: ' "$1"
209
            ;;
210
211
         ?*)
            error 'ERROR: Unknown option: ' "$1"
212
213
214
215
216
217
218
219
220
221
            ;;
        *)
            break
            ;;
        esac

        shift
    done

    if [ -z "$FRAMEWORK" ]; then
222
        FRAMEWORK=$DEFAULT_FRAMEWORK
223
224
    fi

Carsten Csiky's avatar
Carsten Csiky committed
225
    if [ -n "$FRAMEWORK" ]; then
226
227
228
229
        FRAMEWORK=${FRAMEWORK^^}
        if [[ -z "${FRAMEWORKS[$FRAMEWORK]}" ]]; then
            error 'ERROR: Unknown framework: ' "$FRAMEWORK"
        fi
230
231
232
    fi

    if [ -z "$IMAGE" ]; then
233
        IMAGE="dynamo:latest-${FRAMEWORK,,}"
Carsten Csiky's avatar
Carsten Csiky committed
234
        if [ -n "${TARGET}" ]; then
235
236
            IMAGE="${IMAGE}-${TARGET}"
        fi
237
238
239
    fi

    if [[ ${GPUS^^} == "NONE" ]]; then
240
        GPU_STRING=""
241
    else
242
        GPU_STRING="--gpus ${GPUS}"
243
244
245
    fi

    if [[ ${NAME^^} == "" ]]; then
246
        NAME_STRING=""
247
    else
248
        NAME_STRING="--name ${NAME}"
249
250
    fi

251
    if [[ ${ENTRYPOINT^^} == "" ]]; then
252
        ENTRYPOINT_STRING=""
253
    else
254
        ENTRYPOINT_STRING="--entrypoint ${ENTRYPOINT}"
255
256
    fi

257
258
259
260
261
    if [ -n "$MOUNT_WORKSPACE" ]; then
        VOLUME_MOUNTS+=" -v ${SOURCE_DIR}/..:/workspace "
        VOLUME_MOUNTS+=" -v /tmp:/tmp "
        VOLUME_MOUNTS+=" -v /mnt/:/mnt "

262
263
        if [ -z "$HF_HOME" ]; then
            HF_HOME=$DEFAULT_HF_HOME
264
265
266
267
268
269
270
271
272
        fi

        if [ -z "${PRIVILEGED}" ]; then
            PRIVILEGED="TRUE"
        fi

        ENVIRONMENT_VARIABLES+=" -e HF_TOKEN"
    fi

273
274
    if [[ ${HF_HOME^^} == "NONE" ]]; then
        HF_HOME=
275
276
    fi

277
278
    if [ -n "$HF_HOME" ]; then
        mkdir -p "$HF_HOME"
279
        if [[ ${USER} == "root" ]] || [[ ${USER} == "0" ]]; then
280
            HF_HOME_TARGET="/root/.cache/huggingface"
281
282
        else
            HF_HOME_TARGET="/home/dynamo/.cache/huggingface"
283
        fi
284
        VOLUME_MOUNTS+=" -v $HF_HOME:$HF_HOME_TARGET"
285
    fi
286

Carsten Csiky's avatar
Carsten Csiky committed
287
    if [ -z "${PRIVILEGED}" ]; then
288
        PRIVILEGED="FALSE"
289
290
    fi

Carsten Csiky's avatar
Carsten Csiky committed
291
    if [ -z "${RM}" ]; then
292
        RM="TRUE"
293
294
    fi

295
296
297
    if [[ ${PRIVILEGED^^} == "FALSE" ]]; then
        PRIVILEGED_STRING=""
    else
298
        PRIVILEGED_STRING="--privileged"
299
300
    fi

301
    if [[ ${RM^^} == "FALSE" ]]; then
302
        RM_STRING=""
303
    else
304
        RM_STRING=" --rm "
305
306
    fi

307
308
309
    if [ -n "$USE_NIXL_GDS" ]; then
        VOLUME_MOUNTS+=" -v /run/udev:/run/udev:ro "
        NIXL_GDS_CAPS="--cap-add=IPC_LOCK"
Ryan Olson's avatar
Ryan Olson committed
310
311
312
313
314
315
316
        # NOTE(jthomson04): In the KVBM disk pools, we currently allocate our files in /tmp.
        # For some arcane reason, GDS requires that /tmp be mounted.
        # This is already handled for us if we set --mount-workspace
        # If we aren't mounting our workspace but need GDS, we need to mount /tmp.
        if [ -z "$MOUNT_WORKSPACE" ]; then
            VOLUME_MOUNTS+=" -v /tmp:/tmp "
        fi
317
318
319
    else
        NIXL_GDS_CAPS=""
    fi
320
    if [[ "$GPUS" == "none" || "$GPUS" == "NONE" ]]; then
321
            RUNTIME=""
322
    fi
323

324
325
326
327
328
329
    if [[ ${USER} == "" ]]; then
        USER_STRING=""
    else
        USER_STRING="--user ${USER}"
    fi

330
331
332
333
334
335
    REMAINING_ARGS=("$@")
}

show_help() {
    echo "usage: run.sh"
    echo "  [--image image]"
Carsten Csiky's avatar
Carsten Csiky committed
336
    echo "  [--framework framework one of ${!FRAMEWORKS[*]}]"
337
    echo "  [--name name for launched container, default NONE]"
338
339
    echo "  [--privileged whether to launch in privileged mode, default FALSE unless mounting workspace]"
    echo "  [--dry-run print docker commands without running]"
340
    echo "  [--hf-home|--hf-cache directory to volume mount as the hf home, default is NONE unless mounting workspace]"
341
    echo "  [--gpus gpus to enable, default is 'all', 'none' disables gpu support]"
342
    echo "  [--use-nixl-gds add volume mounts and capabilities needed for NVIDIA GPUDirect Storage]"
343
344
345
346
    echo "  [--network network mode for container, default is 'host']"
    echo "           Options: 'host' (default), 'bridge', 'none', 'container:name'"
    echo "           Examples: --network bridge (isolated), --network none (no network - WARNING: breaks most functionality)"
    echo "                    --network container:redis (share network with 'redis' container)"
347
348
    echo "  [--user <name|uid>[:<group|gid>] specify user to run container as]"
    echo "           Format: username or numeric UID, optionally with group/GID (e.g., 'root', '0', '1000:0')"
349
    echo "  [-v add volume mount]"
350
    echo "  [-p|--port add port mapping (host_port:container_port)]"
351
352
353
    echo "  [-e add environment variable]"
    echo "  [--mount-workspace set up for local development]"
    echo "  [-- stop processing and pass remaining args as command to docker run]"
354
355
    echo "  [--workdir set the working directory inside the container]"
    echo "  [--runtime add runtime variables]"
356
357
    echo "  [--entrypoint override container entrypoint]"
    echo "  [-h, --help show this help]"
358
359
360
361
362
363
364
365
366
367
368
369
370
371
    exit 0
}

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

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

get_options "$@"

372
# RUN the image
373
374
375
376
if [ -z "$RUN_PREFIX" ]; then
    set -x
fi

377
378
379
380
${RUN_PREFIX} docker run \
    ${GPU_STRING} \
    ${INTERACTIVE} \
    ${RM_STRING} \
381
    --network "$NETWORK" \
382
    ${RUNTIME:+--runtime "$RUNTIME"} \
383
384
385
386
387
388
    --shm-size=10G \
    --ulimit memlock=-1 \
    --ulimit stack=67108864 \
    --ulimit nofile=65536:65536 \
    ${ENVIRONMENT_VARIABLES} \
    ${VOLUME_MOUNTS} \
389
    ${PORT_MAPPINGS} \
390
    -w "$WORKDIR" \
391
    --cap-add CAP_SYS_PTRACE \
392
    ${NIXL_GDS_CAPS} \
393
394
    --ipc host \
    ${PRIVILEGED_STRING} \
395
    ${USER_STRING} \
396
397
398
399
    ${NAME_STRING} \
    ${ENTRYPOINT_STRING} \
    ${IMAGE} \
    "${REMAINING_ARGS[@]}"
400
401

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