train.sh

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

ROOT_DIR="$(cd -- "$(dirname -- "$0")" && pwd)"
SCRIPT_PATH="$ROOT_DIR/$(basename -- "$0")"
ENTRY="${TRAIN_ENTRY:-train_hope.py}"
PYTHON_BIN="${TRAIN_PYTHON:-/opt/conda/bin/python}"
ENV_FILE="${TRAIN_ENV_FILE:-/root/.config/infimory.ai/.env}"
LOG_FILE="${TRAIN_LOG_FILE:-$ROOT_DIR/nohup.out}"
PID_FILE="${TRAIN_PID_FILE:-$ROOT_DIR/train_hope.pid}"
SESSION_NAME="${TRAIN_SESSION_NAME:-hope_train}"
KILL_PATTERN="${TRAIN_KILL_PATTERN:-train_hope.py|train_hope_debug.py}"
REQUIRED_ENV_VARS="${TRAIN_REQUIRED_ENV_VARS:-INFIMORY_API_KEY}"

usage() {
  cat <<USAGE
Usage: $(basename "$0") [start|stop|restart|status|logs [N]|follow]

Env overrides:
  TRAIN_ENTRY=train_hope.py
  TRAIN_PYTHON=/opt/conda/bin/python
  TRAIN_ENV_FILE=/root/.config/infimory.ai/.env
  TRAIN_LOG_FILE=$ROOT_DIR/nohup.out
  TRAIN_PID_FILE=$ROOT_DIR/train_hope.pid
  TRAIN_SESSION_NAME=hope_train
  TRAIN_KILL_PATTERN='train_hope.py|train_hope_debug.py'
  TRAIN_REQUIRED_ENV_VARS=INFIMORY_API_KEY
USAGE
}

have_tmux() {
  command -v tmux >/dev/null 2>&1
}

have_setsid() {
  command -v setsid >/dev/null 2>&1
}

find_pids() {
  pgrep -f "$KILL_PATTERN" || true
}

write_pid_file() {
  local pid="$1"
  printf '%s\n' "$pid" > "$PID_FILE"
}

clear_pid_file() {
  rm -f "$PID_FILE"
}

run_job() {
  cd "$ROOT_DIR"

  if [[ -f "$ENV_FILE" ]]; then
    set -a
    # shellcheck disable=SC1090
    . "$ENV_FILE"
    set +a
  fi

  if [[ -n "$REQUIRED_ENV_VARS" ]]; then
    local var_name
    for var_name in $REQUIRED_ENV_VARS; do
      if [[ -z "${!var_name:-}" ]]; then
        echo "Missing required env var: $var_name" >&2
        exit 1
      fi
    done
  fi

  exec "$PYTHON_BIN" -u "$ENTRY" >> "$LOG_FILE" 2>&1
}

stop_job() {
  local pids

  if have_tmux && tmux has-session -t "$SESSION_NAME" 2>/dev/null; then
    tmux kill-session -t "$SESSION_NAME"
  fi

  pids="$(find_pids)"
  if [[ -n "$pids" ]]; then
    kill $pids 2>/dev/null || true
    sleep 3
    pids="$(find_pids)"
    if [[ -n "$pids" ]]; then
      kill -9 $pids 2>/dev/null || true
    fi
  fi

  clear_pid_file
}

start_job() {
  local launcher pid pids

  pids="$(find_pids)"
  if [[ -n "$pids" ]]; then
    echo "Training already running:" >&2
    ps -fp "$(echo "$pids" | paste -sd, -)"
    return 0
  fi

  mkdir -p "$(dirname -- "$LOG_FILE")"
  : > /dev/null

  if have_tmux; then
    tmux new-session -d -s "$SESSION_NAME" "$SCRIPT_PATH __run"
    launcher="tmux:$SESSION_NAME"
    sleep 2
    pid="$(pgrep -nf "$ENTRY" || true)"
  elif have_setsid; then
    nohup setsid "$SCRIPT_PATH" __run >/dev/null 2>&1 < /dev/null &
    pid="$!"
    launcher="setsid+nohup"
  else
    nohup "$SCRIPT_PATH" __run >/dev/null 2>&1 < /dev/null &
    pid="$!"
    launcher="nohup"
  fi

  if [[ -z "$pid" ]]; then
    echo "Failed to determine training PID" >&2
    exit 1
  fi

  write_pid_file "$pid"
  sleep 3

  if ! ps -p "$pid" >/dev/null 2>&1; then
    echo "Training failed to stay alive; recent log:" >&2
    tail -n 80 "$LOG_FILE" >&2 || true
    exit 1
  fi

  echo "Started $ENTRY with PID $pid via $launcher"
}

status_job() {
  local pids pid_csv

  if have_tmux && tmux has-session -t "$SESSION_NAME" 2>/dev/null; then
    echo "tmux session: $SESSION_NAME"
  else
    echo "tmux session: not present"
  fi

  pids="$(find_pids)"
  if [[ -z "$pids" ]]; then
    echo "training: not running"
    return 1
  fi

  pid_csv="$(echo "$pids" | paste -sd, -)"
  echo "training: running"
  ps -o pid=,ppid=,etimes=,cmd= -p "$pid_csv"
  echo "log: $LOG_FILE"
}

logs_job() {
  local lines="${1:-80}"
  tail -n "$lines" "$LOG_FILE"
}

follow_job() {
  tail -f "$LOG_FILE"
}

action="${1:-restart}"

case "$action" in
  __run)
    run_job
    ;;
  start)
    start_job
    ;;
  stop)
    stop_job
    ;;
  restart)
    stop_job
    start_job
    ;;
  status)
    status_job
    ;;
  logs)
    shift || true
    logs_job "${1:-80}"
    ;;
  follow)
    follow_job
    ;;
  -h|--help|help)
    usage
    ;;
  *)
    usage >&2
    exit 1
    ;;
esac