Benchmark-Triton-vLLM/build.sbatch at main · cyfronet/Benchmark-Triton-vLLM · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
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
55
56
57
58
59
60
61
62
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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
#!/bin/bash -l
#SBATCH -J TRT_BUILD
#SBATCH -N 1
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=72
#SBATCH --mem=120G
#SBATCH -p plgrid-gpu-gh200
#SBATCH -A plggdvps01-gpu-gh200
#SBATCH --time=00:30:00
#SBATCH --output="build_output.out"
#SBATCH --error="build_error.err"


# Cleaning host environment
ml purge
unset LD_LIBRARY_PATH
unset PYTHONPATH
unset CUDA_HOME

# Loading configuration
if [ -f ".env" ]; then
    set -a
    source .env

    set +a
else
    echo "ERROR: No .env file found"
    exit 1
fi

export MODEL_NAME=$(python3 utils/config_parser.py --config configs/common.yaml --mode var --key global.model_name)
export HF_REPO=$(python3 utils/config_parser.py --config configs/common.yaml --mode var --key global.hf_repo)
export NUM_GPUS=$(python3 utils/config_parser.py --config configs/common.yaml --mode var --key global.num_gpus)
export TARGET_MAX_BATCH_SIZE=$(python3 utils/config_parser.py --config configs/common.yaml --mode var --key global.target_max_batch_size)
export TARGET_MAX_INPUT_LEN=$(python3 utils/config_parser.py --config configs/common.yaml --mode var --key global.target_max_input_len)
export TARGET_MAX_SEQUENCE_LEN=$(python3 utils/config_parser.py --config configs/common.yaml --mode var --key global.target_max_seq_len)
export CONVERT_SCRIPT=$(python3 utils/config_parser.py --config configs/common.yaml --mode var --key global.convert_script)

mkdir -p $PROJECT_DIR
cd $PROJECT_DIR

#Bash scripts for build triton
BUILD_SCRIPT=steps_build.sh
PREPARE_SCRIPT=prepare_repo.sh

# Apptainer cache on Scratch
mkdir -p $SCRATCH/.apptainer_cache $SCRATCH/.apptainer_tmp
export APPTAINER_CACHEDIR=$SCRATCH/.apptainer_cache
export APPTAINER_TMPDIR=$SCRATCH/.apptainer_tmp

# Fake home setup
export FAKE_HOME="$PROJECT_DIR/fake_home"
mkdir -p "$FAKE_HOME"

# Cache variables pointing to Fake Home
export HF_HOME="$FAKE_HOME/.cache/huggingface"
export PIP_CACHE_DIR="$FAKE_HOME/.cache/pip"
export TORCH_EXTENSIONS_DIR="$FAKE_HOME/.cache/torch_extensions"
export FLASHINFER_WORKSPACE_DIR="$FAKE_HOME/.cache/flashinfer"
export XDG_CACHE_HOME="$FAKE_HOME/.cache"

mkdir -p "$HF_HOME" "$PIP_CACHE_DIR" "$TORCH_EXTENSIONS_DIR" "$FLASHINFER_WORKSPACE_DIR" "$XDG_CACHE_HOME"

# Downloading TensorRT-LLM source code via ZIP
if [ "$SHOULD_CLONE_TRTLLM" = "true" ]; then
    VERSION_NUM=$(echo "$TRTLLM_REPO_TAG" | sed 's/^v//')
    rm -rf TensorRT-LLM "TensorRT-LLM-$VERSION_NUM" trtllm.zip

    echo "Downloading TensorRT-LLM $TRTLLM_REPO_TAG zip file..."
    wget --no-check-certificate "${TRTLLM_REPO_URL}/archive/refs/tags/${TRTLLM_REPO_TAG}.zip" -O trtllm.zip
    unzip -q trtllm.zip

    mv "TensorRT-LLM-$VERSION_NUM" TensorRT-LLM
    rm trtllm.zip
fi

chmod +x $BUILD_SCRIPT
chmod +x $PREPARE_SCRIPT

# Running build script inside container
echo "Starting engine build inside container..."

# --- Directory Path Definitions ---
WEIGHTS_DIR="model_weights/$MODEL_NAME"
CKPT_NAME="${MODEL_NAME}_ckpt_gpu${NUM_GPUS}"

BACKUP_BASE="backup_engines"
ACTIVE_DIR="engines/$MODEL_NAME"
MODEL_TOKENIZER_PATH="$WEIGHTS_DIR"
rm -rf engines/*
# Using --home "$FAKE_HOME" instead of --env HOME=...
# Deleted --no-home (--home overwrites it)
apptainer exec --nv --cleanenv \
    --home "$FAKE_HOME" \
    --bind $SCRATCH:$SCRATCH \
    --env HF_TOKEN="$HF_TOKEN" \
    --env PROJECT_DIR="$PROJECT_DIR" \
    --env HF_HOME="$HF_HOME" \
    --env PIP_CACHE_DIR="$PIP_CACHE_DIR" \
    --env TORCH_EXTENSIONS_DIR="$TORCH_EXTENSIONS_DIR" \
    --env FLASHINFER_WORKSPACE_DIR="$FLASHINFER_WORKSPACE_DIR" \
    --env XDG_CACHE_HOME="$XDG_CACHE_HOME" \
    --env PYTHONUSERBASE="$PYTHONUSERBASE" \
    --env PATH="$PATH" \
    --env WEIGHTS_DIR="$WEIGHTS_DIR" \
    --env CKPT_NAME="$CKPT_NAME" \
    --env BACKUP_BASE="$BACKUP_BASE" \
    --env ACTIVE_DIR="$ACTIVE_DIR" \
    --env MODEL_NAME="$MODEL_NAME" \
    --env NUM_GPUS="$NUM_GPUS" \
    --env HF_REPO="$HF_REPO" \
    --env TARGET_MAX_BATCH_SIZE="$TARGET_MAX_BATCH_SIZE" \
    --env TARGET_MAX_INPUT_LEN="$TARGET_MAX_INPUT_LEN" \
    --env TARGET_MAX_SEQUENCE_LEN="$TARGET_MAX_SEQUENCE_LEN" \
    --env CONVERT_SCRIPT="$CONVERT_SCRIPT" \
    $IMAGE \
    ./$BUILD_SCRIPT

# Preparing Triton model repository
echo "Preparing Triton model repository..."
apptainer exec --nv --cleanenv \
    --home "$FAKE_HOME" \
    --bind $SCRATCH:$SCRATCH \
    --env PROJECT_DIR="$PROJECT_DIR" \
    --env MODEL_TOKENIZER_PATH="$MODEL_TOKENIZER_PATH" \
    --env TRTLLM_REPO_TAG="$TRTLLM_REPO_TAG" \
    --env NUM_GPUS="$NUM_GPUS" \
    --env TARGET_MAX_BATCH_SIZE="$TARGET_MAX_BATCH_SIZE" \
    --env MODEL_NAME="$MODEL_NAME" \
    $IMAGE \
    ./$PREPARE_SCRIPT