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

# Host-side launcher for the Megatron-LM text generation server.
# Run on g0050 from /ssd/workspace/yi/laoyao_2b_moe.

REPO_ROOT=${REPO_ROOT:-/ssd/workspace/yi/laoyao_2b_moe}
IMAGE=${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}
SERVER_PORT=${SERVER_PORT:-5000}
MASTER_PORT=${MASTER_PORT:-6000}
CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0}
CONTAINER_NAME=${CONTAINER_NAME:-laoyao_megatron_server}
DETACH=${DETACH:-0}
LOG_DIR=${LOG_DIR:-${REPO_ROOT}/runs/megatron_inference}

mkdir -p "${LOG_DIR}"

DOCKER_ARGS=(
  --gpus all
  --network host
  --ipc host
  --shm-size 64g
  --ulimit memlock=-1
  --ulimit stack=67108864
  --name "${CONTAINER_NAME}"
  -e CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}"
  -e SERVER_PORT="${SERVER_PORT}"
  -e MASTER_PORT="${MASTER_PORT}"
  -v "${REPO_ROOT}:/work"
  -w /work
)

docker rm -f "${CONTAINER_NAME}" >/dev/null 2>&1 || true

if [[ "${DETACH}" == "1" ]]; then
  docker run -d "${DOCKER_ARGS[@]}" "${IMAGE}" \
    bash scripts/serve_laoyao_megatron_inner.sh
  echo "Started ${CONTAINER_NAME}"
  echo "Logs: docker logs -f ${CONTAINER_NAME}"
else
  docker run --rm "${DOCKER_ARGS[@]}" "${IMAGE}" \
    bash scripts/serve_laoyao_megatron_inner.sh 2>&1 | tee "${LOG_DIR}/server_$(date +%Y%m%d_%H%M%S).log"
fi
