#!/usr/bin/env bash
# Документальный профиль. Редакцией на GPU не запускался.
set -euo pipefail
HERE="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
SLUG="qwen38"
LOCK="${LOCK_DIR:-$HERE/locks}/$SLUG.json"
ROOT_MODELS="${ROOT_MODELS:-/srv/ai/models}"
test -f "$LOCK" || { echo "Сначала prepare.py $SLUG"; exit 1; }
IMAGE="$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1]))["image_digest"])' "$LOCK")"
REVISION="$(python3 -c 'import json,sys; print(json.load(open(sys.argv[1]))["revision"])' "$LOCK")"
MODEL_DIR="$ROOT_MODELS/$SLUG/$REVISION"
test -f "$MODEL_DIR/config.json"
test -f "$MODEL_DIR/.telecobra-download-complete.json" || { echo "Загрузка не завершена: повторите prepare.py $SLUG"; exit 1; }
EXTRA=()
DOCKER_EXTRA=()
# На первом узле rank=0; остальные последовательно 1,2,...
: "${NODE_RANK:?Укажите номер узла}"
: "${HEAD_ADDR:?Укажите приватный IP узла 0}"
: "${NODE_IP:?Укажите приватный IP текущего узла}"
[[ "$NODE_RANK" =~ ^[0-9]+$ ]] && (( NODE_RANK < 4 )) || { echo "Неверный NODE_RANK"; exit 1; }
DOCKER_EXTRA+=(-e "VLLM_HOST_IP=$NODE_IP" -e "NCCL_DEBUG=INFO" --cap-add IPC_LOCK)
if [ -d /dev/infiniband ]; then DOCKER_EXTRA+=(--device /dev/infiniband); fi
if [ -n "${NCCL_SOCKET_IFNAME:-}" ]; then DOCKER_EXTRA+=(-e "NCCL_SOCKET_IFNAME=$NCCL_SOCKET_IFNAME"); fi
if [ -n "${GLOO_SOCKET_IFNAME:-}" ]; then DOCKER_EXTRA+=(-e "GLOO_SOCKET_IFNAME=$GLOO_SOCKET_IFNAME"); fi
if [ -n "${NCCL_IB_HCA:-}" ]; then DOCKER_EXTRA+=(-e "NCCL_IB_HCA=$NCCL_IB_HCA"); fi
if [ "$NODE_RANK" -gt 0 ]; then EXTRA+=(--headless); fi
exec docker run --rm --name "tc-$SLUG" --gpus all \
  --network host --ipc host --ulimit memlock=-1 --ulimit stack=67108864 \
  -v "$MODEL_DIR:/model:ro" \
  -e VLLM_ENGINE_READY_TIMEOUT_S=7200 \
  -e VLLM_WORKER_MULTIPROC_METHOD=spawn \
  "${DOCKER_EXTRA[@]}" \
  --entrypoint vllm "$IMAGE" serve /model \
  --host 127.0.0.1 \
  --port 8000 \
  --served-model-name telecobra-model \
  --tensor-parallel-size 32 \
  --max-model-len 32768 \
  --max-num-seqs 4 \
  --gpu-memory-utilization 0.90 \
  --reasoning-parser qwen3 \
  --kv-cache-dtype fp8 \
  --trust-remote-code \
  --tool-call-parser qwen3_coder \
  --enable-auto-tool-choice \
  --distributed-executor-backend mp \
  --nnodes 4 \
  --node-rank "$NODE_RANK" \
  --master-addr "$HEAD_ADDR" "${EXTRA[@]}"