gistfile1.txt

#!/bin/bash

# Usage:
# bash run_tpu_benchmark_client.sh --model Qwen/Qwen2.5-1.5B-Instruct --tp 1

LONGOPTS=model:,tp:,profile
# Parse arguments
PARSED=$(getopt --options=$OPTIONS --longoptions=$LONGOPTS --name "$0" -- "$@")
if [[ $? -ne 0 ]]; then
  exit 2
fi
eval set -- "$PARSED"
profile=false
# Option parsing
while true; do
  case "$1" in
    --model)
      model="$2"
      shift 2
      ;;
    --tp)
      tp=$2
      shift 2
      ;;
    --profile)
      profile=true
      shift
      ;;
    --)
      shift
      break
      ;;
    *)
      echo "Unknown option: $1"
      exit 3
      ;;
  esac
done
if [[ -z "$model" ]]; then
  echo "Error: --model is required"
  exit 1
fi
if [[ -z "$tp" ]]; then
  echo "Error: --tp is required"
  exit 1
fi
if $profile; then
  rm -rfv /tmp/tpu-profile/ && mkdir /tmp/tpu-profile
  export VLLM_TORCH_PROFILER_DIR=/tmp/tpu-profile
fi

if [ ! -e oldjax/vllm/benchmarks/sonnet_4x.txt ]; then
  echo "Creating the data file sonnet_4x.txt in old jax"
  touch oldjax/vllm/benchmarks/sonnet_4x.txt
  for _ in {1..4}
  do
  cat oldjax/vllm/benchmarks/sonnet.txt >> oldjax/vllm/benchmarks/sonnet_4x.txt
  done
else
  echo "oldjax/vllm/bencharmks/sonnet_4x.txt exists. Skip creating the file."
fi

if [ ! -e newjax/vllm/benchmarks/sonnet_4x.txt ]; then
  echo "Creating the data file sonnet_4x.txt in new jax"
  touch newjax/vllm/benchmarks/sonnet_4x.txt
  for _ in {1..4}
  do
  cat newjax/vllm/benchmarks/sonnet.txt >> newjax/vllm/benchmarks/sonnet_4x.txt
  done
else
  echo "newjax/vllm/benchmarks/sonnet_4x.txt exists. Skip creating the file."
fi

DEFAULT_HOST=127.0.0.1
DEFAULT_PORT=8000

printf $profile
if ! $profile; then
  echo "Profiling is disabled"
else
  echo "Profiling is enabled"
fi


nc -zv $DEFAULT_HOST $DEFAULT_PORT
while [ $? -ne 0 ]; do
  echo "Waiting for the server to start..."
  sleep 15
  nc -zv $DEFAULT_HOST $DEFAULT_PORT
done

echo "Server is up and running"
if ! $profile; then
  echo "Running the benchmark..."
  vllm bench serve --model $model --dataset-name sonnet --dataset-path oldjax/vllm/benchmarks/sonnet_4x.txt --sonnet-input-len 1800 --sonnet-output-len 128 --num-prompts 1000 --percentile-metrics "ttft,tpot,itl,e2el" --ignore_eos 2>&1 | tee vllm_benchmark_out.txt
else
  echo "Running the benchmark with profiling..."
  vllm bench serve --model $model --dataset-name sonnet --dataset-path oldjax/vllm/benchmarks/sonnet_4x.txt --sonnet-input-len 1800 --sonnet-output-len 128 --num-prompts 1000 --percentile-metrics "ttft,tpot,itl,e2el" --ignore_eos --profile
fi

echo "All done. Killing the server..."
kill %1