#!/usr/bin/env bash
set -Eeuo pipefail

readonly IMAGE="${B12X_TEST_IMAGE:-ai01/vllm-v9-b12x-tp3:20260710}"
readonly SERVICE="deepseek-v4-flash.service"
readonly EXTENSION_CACHE="/srv/ai/cache/b12x-v9-tp3-torch-extensions"
readonly TEST_FILE="${B12X_TEST_FILE:-/srv/ai/src/b12x-tp3/tests/distributed/test_pcie_oneshot_torture.py}"

restart_service() {
  docker rm -f b12x-tp3-test >/dev/null 2>&1 || true
  systemctl start "${SERVICE}"
}

mkdir -p "${EXTENSION_CACHE}"
systemctl stop "${SERVICE}"
trap restart_service EXIT

timeout --signal=TERM --kill-after=10 "${B12X_TEST_TIMEOUT:-300}" docker run \
  --rm \
  --name b12x-tp3-test \
  --gpus all \
  --ipc=host \
  -e TORCH_EXTENSIONS_DIR=/torch-cache \
  -e B12X_RUN_PCIE_ONESHOT_TORTURE=1 \
  -e B12X_PCIE_ONESHOT_TORTURE_WORLD_SIZE=3 \
  -e B12X_PCIE_ONESHOT_TORTURE_EAGER_ITERS="${B12X_TORTURE_ITERS:-256}" \
  -e B12X_PCIE_ONESHOT_TORTURE_GRAPH_REPLAYS="${B12X_TORTURE_ITERS:-256}" \
  -e B12X_PCIE_ONESHOT_TORTURE_MULTISTREAM_ITERS="${B12X_TORTURE_ITERS:-256}" \
  -v "${EXTENSION_CACHE}:/torch-cache" \
  -v "${TEST_FILE}:/workspace/test_pcie_oneshot_torture.py:ro" \
  --entrypoint bash \
  "${IMAGE}" -lc \
  'python -m pip install -q pytest && python -m pytest /workspace/test_pcie_oneshot_torture.py -q -s'

restart_service
trap - EXIT