run.sh 6.66 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 ["TENSORRTLLM"]=2)
28
DEFAULT_FRAMEWORK=VLLM
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

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

IMAGE=
HF_CACHE=
DEFAULT_HF_CACHE=${SOURCE_DIR}/.cache/huggingface
GPUS="all"
PRIVILEGED=
VOLUME_MOUNTS=
MOUNT_WORKSPACE=
ENVIRONMENT_VARIABLES=
REMAINING_ARGS=
INTERACTIVE=

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

	--gpus)
            if [ "$2" ]; then
                GPUS=$2
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
96
		missing_requirement "$1"
97
98
            fi
            ;;
99
	--entrypoint)
100
            if [ "$2" ]; then
101
                ENTRYPOINT=$2
102
103
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
104
		missing_requirement "$1"
105
106
107
108
109
110
111
            fi
            ;;
	--privileged)
            if [ "$2" ]; then
                PRIVILEGED=$2
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
112
		missing_requirement "$1"
113
114
            fi
            ;;
115
116
117
118
119
	--rm)
            if [ "$2" ]; then
                RM=$2
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
120
		missing_requirement "$1"
121
122
            fi
            ;;
123
124
125
126
127
	-v)
            if [ "$2" ]; then
                VOLUME_MOUNTS+=" -v $2 "
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
128
		missing_requirement "$1"
129
130
131
132
133
134
135
            fi
            ;;
	-e)
            if [ "$2" ]; then
                ENVIRONMENT_VARIABLES+=" -e $2 "
                shift
            else
Carsten Csiky's avatar
Carsten Csiky committed
136
		missing_requirement "$1"
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
            fi
            ;;
	-it)
	    INTERACTIVE=" -it "
	    ;;
	--mount-workspace)
	    MOUNT_WORKSPACE=TRUE
	    ;;
        --dry-run)
            RUN_PREFIX="echo"
            echo ""
            echo "=============================="
            echo "DRY RUN: COMMANDS PRINTED ONLY"
            echo "=============================="
            echo ""
            ;;
        --)
            shift
            break
            ;;
         -?*)
Carsten Csiky's avatar
Carsten Csiky committed
158
	    error 'ERROR: Unknown option: ' "$1"
159
160
            ;;
	 ?*)
Carsten Csiky's avatar
Carsten Csiky committed
161
	    error 'ERROR: Unknown option: ' "$1"
162
163
164
165
166
167
168
169
170
171
172
173
174
            ;;
        *)
            break
            ;;
        esac

        shift
    done

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

Carsten Csiky's avatar
Carsten Csiky committed
175
    if [ -n "$FRAMEWORK" ]; then
176
	FRAMEWORK=${FRAMEWORK^^}
Carsten Csiky's avatar
Carsten Csiky committed
177
178
	if [[ -n "${FRAMEWORKS[$FRAMEWORK]}" ]]; then
	    error 'ERROR: Unknown framework: ' "$FRAMEWORK"
179
180
181
182
	fi
    fi

    if [ -z "$IMAGE" ]; then
183
        IMAGE="dynamo:latest-${FRAMEWORK,,}"
Carsten Csiky's avatar
Carsten Csiky committed
184
        if [ -n "${TARGET}" ]; then
185
186
            IMAGE="${IMAGE}-${TARGET}"
        fi
187
188
189
190
191
192
193
194
195
196
197
198
199
200
    fi

    if [[ ${GPUS^^} == "NONE" ]]; then
	GPU_STRING=""
    else
	GPU_STRING="--gpus ${GPUS}"
    fi

    if [[ ${NAME^^} == "" ]]; then
	NAME_STRING=""
    else
	NAME_STRING="--name ${NAME}"
    fi

201
202
203
204
205
206
    if [[ ${ENTRYPOINT^^} == "" ]]; then
	ENTRYPOINT_STRING=""
    else
	ENTRYPOINT_STRING="--entrypoint ${ENTRYPOINT}"
    fi

Carsten Csiky's avatar
Carsten Csiky committed
207
    if [ -n "$MOUNT_WORKSPACE" ]; then
208
209
210
211
212
213
214
215
	VOLUME_MOUNTS+=" -v ${SOURCE_DIR}/..:/workspace "
	VOLUME_MOUNTS+=" -v /tmp:/tmp "
	VOLUME_MOUNTS+=" -v /mnt/:/mnt "

	if [ -z "$HF_CACHE" ]; then
	    HF_CACHE=$DEFAULT_HF_CACHE
	fi

Carsten Csiky's avatar
Carsten Csiky committed
216
	if [ -z "${PRIVILEGED}" ]; then
217
218
219
220
221
222
223
224
225
226
227
228
	    PRIVILEGED="TRUE"
	fi

	ENVIRONMENT_VARIABLES+=" -e HF_TOKEN"

	INTERACTIVE=" -it "
    fi

    if [[ ${HF_CACHE^^} == "NONE" ]]; then
	HF_CACHE=
    fi

Carsten Csiky's avatar
Carsten Csiky committed
229
230
    if [ -n "$HF_CACHE" ]; then
	mkdir -p "$HF_CACHE"
231
232
233
	VOLUME_MOUNTS+=" -v $HF_CACHE:/root/.cache/huggingface"
    fi

Carsten Csiky's avatar
Carsten Csiky committed
234
    if [ -z "${PRIVILEGED}" ]; then
235
236
237
	PRIVILEGED="FALSE"
    fi

Carsten Csiky's avatar
Carsten Csiky committed
238
    if [ -z "${RM}" ]; then
239
240
241
	RM="TRUE"
    fi

242
243
244
245
246
247
    if [[ ${PRIVILEGED^^} == "FALSE" ]]; then
	PRIVILEGED_STRING=""
    else
	PRIVILEGED_STRING="--privileged"
    fi

248
249
250
251
252
253
254
    if [[ ${RM^^} == "FALSE" ]]; then
	RM_STRING=""
    else
	RM_STRING=" --rm "
    fi


255
256
257
258
259
260
    REMAINING_ARGS=("$@")
}

show_help() {
    echo "usage: run.sh"
    echo "  [--image image]"
Carsten Csiky's avatar
Carsten Csiky committed
261
    echo "  [--framework framework one of ${!FRAMEWORKS[*]}]"
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
    echo "  [--name name for launched container, default NONE] "
    echo "  [--privileged whether to launch in privileged mode, default FALSE unless mounting workspace]"
    echo "  [--dry-run print docker commands without running]"
    echo "  [--hf-cache directory to volume mount as the hf cache, default is NONE unless mounting workspace]"
    echo "  [--gpus gpus to enable, default is 'all', 'none' disables gpu support]"
    echo "  [-v add volume mount]"
    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]"
    exit 0
}

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

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

get_options "$@"

# RUN the image

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

291
${RUN_PREFIX} docker run ${GPU_STRING} ${INTERACTIVE} ${RM_STRING} --network host --shm-size=10G --ulimit memlock=-1 --ulimit stack=67108864 ${ENVIRONMENT_VARIABLES} ${VOLUME_MOUNTS} -w /workspace --cap-add CAP_SYS_PTRACE --ipc host ${PRIVILEGED_STRING} ${NAME_STRING} ${ENTRYPOINT_STRING} ${IMAGE} "${REMAINING_ARGS[@]}"
292
293

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