28 lines
717 B
Bash
Executable File
28 lines
717 B
Bash
Executable File
#!/bin/bash
|
|
set -e
|
|
|
|
MODEL_REPLICA=${MODEL_REPLICA:-1}
|
|
|
|
if [ -d "$MODEL_NAME" ]; then
|
|
MODEL_DIR="$MODEL_NAME"
|
|
else
|
|
# Get model dir.
|
|
MODEL_DIR=$(python3 <<EOF
|
|
from huggingface_hub import snapshot_download
|
|
|
|
print(snapshot_download(repo_id='$MODEL_NAME', allow_patterns='triton/**/*', local_files_only=True))
|
|
EOF
|
|
)
|
|
fi
|
|
|
|
# Set model dir in triton config.
|
|
sed -i 's@${MODEL_DIR}@'$MODEL_DIR'@g' $MODEL_DIR/triton/fastertransformer/config.pbtxt
|
|
|
|
# SET model replica in triton config.
|
|
sed -i "s/count: [[:digit:]]\+/count: $MODEL_REPLICA/g" $MODEL_DIR/triton/fastertransformer/config.pbtxt
|
|
|
|
# Start triton server.
|
|
mpirun -n 1 \
|
|
--allow-run-as-root /opt/tritonserver/bin/tritonserver \
|
|
--model-repository=$MODEL_DIR/triton
|