Chat with imposter-72b on kurtz (Qwen2.5-72B + LoRA via vLLM)

pi-kurtz

#!/usr/bin/env bash
# Chat with imposter-72b on kurtz (Qwen2.5-72B + LoRA via vLLM)
# Sets up SSH tunnel automatically, tears it down on exit

set -euo pipefail

LOCAL_PORT=8000
TUNNEL_PID=""

cleanup() {
  if [[ -n "$TUNNEL_PID" ]] && kill -0 "$TUNNEL_PID" 2>/dev/null; then
    kill "$TUNNEL_PID" 2>/dev/null
    wait "$TUNNEL_PID" 2>/dev/null || true
  fi
}
trap cleanup EXIT INT TERM

# Check if tunnel/port is already active
if lsof -i :"$LOCAL_PORT" -sTCP:LISTEN >/dev/null 2>&1; then
  echo "Port $LOCAL_PORT already in use — assuming tunnel is up."
else
  echo "Starting SSH tunnel to kurtz via gcloud..."
  gcloud compute ssh kurtz --zone=us-east1-b --project=volta-489906 \
    -- -f -N -L "$LOCAL_PORT":localhost:"$LOCAL_PORT"
  TUNNEL_PID=$(lsof -ti :"$LOCAL_PORT" -sTCP:LISTEN 2>/dev/null | head -1)

  for i in {1..10}; do
    if curl -s --max-time 2 http://localhost:"$LOCAL_PORT"/v1/models >/dev/null 2>&1; then
      echo "Tunnel ready."
      break
    fi
    if [[ $i -eq 10 ]]; then
      echo "Warning: tunnel may not be ready yet, proceeding anyway..."
    fi
    sleep 0.5
  done
fi

BASE_URL="http://localhost:$LOCAL_PORT/v1"
MODEL="/scratch/imposter_merged_model"

# Non-interactive: single prompt from args or stdin
if [[ $# -gt 0 ]]; then
  exec python3 -c "
import sys
from openai import OpenAI
client = OpenAI(base_url='$BASE_URL', api_key='not-needed')
r = client.chat.completions.create(
    model='$MODEL',
    messages=[{'role':'user','content':' '.join(sys.argv[1:])}],
    max_tokens=1024, stream=True)
for chunk in r:
    if chunk.choices and chunk.choices[0].delta.content:
        print(chunk.choices[0].delta.content, end='', flush=True)
print()
" "$@"
fi

# Interactive REPL
exec python3 -c "
from openai import OpenAI
client = OpenAI(base_url='$BASE_URL', api_key='not-needed')
messages = []
print('imposter-72b on kurtz — Ctrl-C to quit\n')
try:
    while True:
        user = input('\033[1mYou:\033[0m ')
        if not user.strip():
            continue
        messages.append({'role': 'user', 'content': user})
        r = client.chat.completions.create(
            model='$MODEL', messages=messages,
            max_tokens=1024, stream=True)
        print('\033[1mImposter:\033[0m ', end='', flush=True)
        full = []
        for chunk in r:
            if chunk.choices and chunk.choices[0].delta.content:
                t = chunk.choices[0].delta.content
                full.append(t)
                print(t, end='', flush=True)
        print('\n')
        messages.append({'role': 'assistant', 'content': ''.join(full)})
except (KeyboardInterrupt, EOFError):
    print('\nBye.')
"