foxhunt/scripts/argo-train.sh

#!/usr/bin/env bash
# Train a model via Argo Workflows.
#
# Usage:
#   ./scripts/argo-train.sh dqn                          # defaults: HEAD, H100, 50 epochs
#   ./scripts/argo-train.sh dqn --sha abc1234            # specific commit
#   ./scripts/argo-train.sh dqn --epochs 100 --trials 40 # override training params
#   ./scripts/argo-train.sh ppo --gpu-pool ci-training    # L40S instead of H100
#   ./scripts/argo-train.sh dqn --baseline                # skip hyperopt
#   ./scripts/argo-train.sh dqn --watch                  # follow logs
#
# Supported models:
#   RL:          dqn, ppo
#   Supervised:  tft, mamba2, tggn, tlob, liquid, kan, xlstm, diffusion
#
# Requires: argo CLI
set -euo pipefail

SHA="HEAD"
BRANCH="main"
TRIALS=""
EPOCHS=""
GPU_POOL=""
SYMBOL=""
WATCH=false
BASELINE=false
CAPITAL=""
SANITIZER="none"
MULTI_SEED=1
FOLDS=1
DRY_RUN=false
TAG=""

usage() {
  cat <<EOF
Usage: $(basename "$0") <model> [OPTIONS]

Models:
  dqn, ppo                                          (RL)
  tft, mamba2, tggn, tlob, liquid, kan, xlstm       (supervised)

Options:
  --sha <commit>       Git commit SHA (default: HEAD)
  --branch <branch>    Git branch (default: main)
  --trials <n>         Hyperopt trials (default: 20)
  --epochs <n>         Training epochs (default: 50)
  --gpu-pool <pool>    GPU node pool (default: ci-training-h100)
  --symbol <sym>       Trading symbol (default: ES.FUT)
  --capital <n>        Initial capital (default: 35000)
  --baseline           Skip hyperopt (trials=0)
  --sanitizer <tool>   Run under compute-sanitizer (memcheck|racecheck|synccheck)
  --watch              Follow workflow logs
  --multi-seed <n>     Run N seeds in parallel (default: 1, fans out via DAG when >1)
  --folds <k>          Walk-forward fold count (default: 1, fans out via DAG when >1)
  --tag <t>            Label workflow with foxhunt-tag=<t> for log aggregation
  --dry-run            Print rendered workflow YAML to stdout, do not submit
  -h, --help           Show this help
EOF
  exit 0
}

[[ $# -eq 0 ]] && { echo "Error: model argument required"; usage; }

MODEL="$1"; shift

case "$MODEL" in
  dqn|ppo|tft|mamba2|tggn|tlob|liquid|kan|xlstm|diffusion) ;;
  *) echo "Error: unknown model '$MODEL'"; usage ;;
esac

while [[ $# -gt 0 ]]; do
  case $1 in
    --sha)         SHA="$2"; shift 2 ;;
    --branch)      BRANCH="$2"; shift 2 ;;
    --trials)      TRIALS="$2"; shift 2 ;;
    --epochs)      EPOCHS="$2"; shift 2 ;;
    --gpu-pool)    GPU_POOL="$2"; shift 2 ;;
    --symbol)      SYMBOL="$2"; shift 2 ;;
    --capital)     CAPITAL="$2"; shift 2 ;;
    --baseline)    BASELINE=true; shift ;;
    --sanitizer)   SANITIZER="$2"; shift 2 ;;
    --watch)       WATCH=true; shift ;;
    --multi-seed)  MULTI_SEED="$2"; shift 2 ;;
    --folds)       FOLDS="$2"; shift 2 ;;
    --tag)         TAG="$2"; shift 2 ;;
    --dry-run)     DRY_RUN=true; shift ;;
    -h|--help)     usage ;;
    *)             echo "Unknown option: $1"; usage ;;
  esac
done

# Validate seed/fold counts are positive integers.
if ! [[ "$MULTI_SEED" =~ ^[0-9]+$ ]] || [[ "$MULTI_SEED" -lt 1 ]]; then
  echo "Error: --multi-seed must be a positive integer, got '$MULTI_SEED'"
  exit 1
fi
if ! [[ "$FOLDS" =~ ^[0-9]+$ ]] || [[ "$FOLDS" -lt 1 ]]; then
  echo "Error: --folds must be a positive integer, got '$FOLDS'"
  exit 1
fi

# Auto-derive cuda-compute-cap from GPU pool — cubins must match device sm_XX.
# Default pool is ci-training-h100 (sm_90). Override for other architectures:
#   ci-training-h100*  → sm_90 (Hopper)
#   ci-training-l40s   → sm_89 (Ada Lovelace)
case "${GPU_POOL:-ci-training-h100}" in
  *l40s*)     CUDA_COMPUTE_CAP="89" ;;
  *h100*|*)   CUDA_COMPUTE_CAP="90" ;;  # default Hopper
esac

# ── Route: single-job (existing template) vs multi-seed DAG (new template) ──
# Backward compat: --multi-seed 1 --folds 1 keeps the existing single-template
# call path verbatim. Only when N>1 OR K>1 do we render the matrix DAG.
USE_MULTI_SEED=false
if [[ "$MULTI_SEED" -gt 1 || "$FOLDS" -gt 1 ]]; then
  USE_MULTI_SEED=true
fi

if [[ "$USE_MULTI_SEED" == "false" ]]; then
  CMD="argo submit -n foxhunt --from=wftmpl/train"
  CMD="$CMD -p commit-sha=$SHA"
  CMD="$CMD -p git-branch=$BRANCH"
  CMD="$CMD -p model=$MODEL"
  CMD="$CMD -p cuda-compute-cap=$CUDA_COMPUTE_CAP"

  [[ -n "$TRIALS" ]]    && CMD="$CMD -p hyperopt-trials=$TRIALS"
  [[ -n "$EPOCHS" ]]    && CMD="$CMD -p train-epochs=$EPOCHS"
  [[ -n "$GPU_POOL" ]]  && CMD="$CMD -p gpu-pool=$GPU_POOL"
  [[ -n "$SYMBOL" ]]    && CMD="$CMD -p symbol=$SYMBOL"
  [[ -n "$CAPITAL" ]]   && CMD="$CMD -p initial-capital=$CAPITAL"
  [[ "$SANITIZER" != "none" ]] && CMD="$CMD -p sanitizer=$SANITIZER"

  $BASELINE && CMD="$CMD -p hyperopt-trials=0"
  $WATCH    && CMD="$CMD --watch"

  if [[ "$DRY_RUN" == "true" ]]; then
    # `argo submit --dry-run -o yaml` renders the resolved Workflow without
    # contacting the cluster. Single-template path emits one Workflow object
    # (no `kind: WorkflowTask` lines — that string only appears in the
    # multi-seed DAG matrix).
    eval "$CMD --dry-run -o yaml"
    exit 0
  fi

  echo "Submitting $MODEL training workflow..."
  echo "  sha:    $SHA"
  echo "  branch: $BRANCH"
  echo "  model:  $MODEL"
  $BASELINE              && echo "  mode:   baseline (no hyperopt)"
  [[ -n "$TRIALS" ]]     && echo "  trials: $TRIALS"
  [[ -n "$EPOCHS" ]]     && echo "  epochs: $EPOCHS"
  [[ -n "$GPU_POOL" ]]   && echo "  gpu:    $GPU_POOL"
  echo "  sm:     $CUDA_COMPUTE_CAP"
  echo ""
  eval "$CMD"
  exit 0
fi

# ── Multi-seed × multi-fold path ──
# Render the train-multi-seed-template.yaml with the (seed, fold) matrix
# expanded inline. The base template ships with a placeholder marker
# (`# __MATRIX_TASKS__`) which we replace with N*K generated WorkflowTask
# stanzas. This avoids hand-writing a 30-task matrix and keeps the source
# template human-readable.
TEMPLATE_SRC="infra/k8s/argo/train-multi-seed-template.yaml"
if [[ ! -f "$TEMPLATE_SRC" ]]; then
  echo "Error: multi-seed template not found at $TEMPLATE_SRC"
  exit 1
fi

# Build matrix YAML. Each task is a dag.tasks[] entry that targets the
# `train-single` template with seed/fold parameters bound from inputs.
build_matrix_tasks() {
  local seeds="$1"
  local folds="$2"
  # 10 spaces — matches the sibling `- name: ensure-binary` list-item indent
  # under `dag.tasks:` (which is itself at 8 spaces). Wrong indent here
  # produces a YAML parse error in the rendered template.
  local indent="          "
  local s
  local f
  for ((s=0; s<seeds; s++)); do
    for ((f=0; f<folds; f++)); do
      cat <<EOF
${indent}- name: train-s${s}-f${f}
${indent}  template: train-single
${indent}  dependencies: [ensure-fxcache, gpu-warmup]
${indent}  arguments:
${indent}    parameters:
${indent}      - name: seed
${indent}        value: "${s}"
${indent}      - name: fold
${indent}        value: "${f}"
EOF
    done
  done
}

MATRIX_TASKS=$(build_matrix_tasks "$MULTI_SEED" "$FOLDS")

# Substitute the placeholder. Match the *exact* placeholder line (leading
# whitespace + the marker as the only content on the line) — the marker
# string also appears in doc comments above and we must not replace those.
# Use awk (not sed) — multi-line replacement with sed is fragile across
# platforms.
RENDERED=$(awk -v repl="$MATRIX_TASKS" '
  /^[[:space:]]*# __MATRIX_TASKS__[[:space:]]*$/ { print repl; next }
  { print }
' "$TEMPLATE_SRC")

if [[ "$DRY_RUN" == "true" ]]; then
  # Emit the rendered template + a synthetic per-task marker line so test
  # harnesses can count generated jobs without piping through `argo submit`
  # (which would require a live cluster). The `kind: WorkflowTask` marker
  # is what the test_multi_seed_harness.sh asserts against.
  echo "$RENDERED"
  for ((s=0; s<MULTI_SEED; s++)); do
    for ((f=0; f<FOLDS; f++)); do
      echo "# kind: WorkflowTask  seed=${s} fold=${f}"
    done
  done
  exit 0
fi

# Submit: write the rendered template to a temp file, apply it as a
# WorkflowTemplate, then submit a Workflow that references it.
TMP_TEMPLATE=$(mktemp -t train-multi-seed.XXXXXX.yaml)
trap 'rm -f "$TMP_TEMPLATE"' EXIT
echo "$RENDERED" > "$TMP_TEMPLATE"

echo "Submitting multi-seed $MODEL training workflow..."
echo "  sha:        $SHA"
echo "  branch:     $BRANCH"
echo "  model:      $MODEL"
echo "  multi-seed: $MULTI_SEED"
echo "  folds:      $FOLDS"
echo "  total jobs: $((MULTI_SEED * FOLDS))"
[[ -n "$TAG" ]] && echo "  tag:        $TAG"
[[ -n "$EPOCHS" ]] && echo "  epochs:     $EPOCHS"
[[ -n "$GPU_POOL" ]] && echo "  gpu:        $GPU_POOL"
echo "  sm:         $CUDA_COMPUTE_CAP"
echo ""

# Apply the rendered WorkflowTemplate, then submit a workflow from it.
kubectl apply -n foxhunt -f "$TMP_TEMPLATE"

CMD="argo submit -n foxhunt --from=wftmpl/train-multi-seed"
CMD="$CMD -p commit-sha=$SHA"
CMD="$CMD -p git-branch=$BRANCH"
CMD="$CMD -p model=$MODEL"
CMD="$CMD -p cuda-compute-cap=$CUDA_COMPUTE_CAP"
CMD="$CMD -p multi-seed=$MULTI_SEED"
CMD="$CMD -p folds=$FOLDS"

[[ -n "$TRIALS" ]]   && CMD="$CMD -p hyperopt-trials=$TRIALS"
[[ -n "$EPOCHS" ]]   && CMD="$CMD -p train-epochs=$EPOCHS"
[[ -n "$GPU_POOL" ]] && CMD="$CMD -p gpu-pool=$GPU_POOL"
[[ -n "$SYMBOL" ]]   && CMD="$CMD -p symbol=$SYMBOL"
[[ -n "$CAPITAL" ]]  && CMD="$CMD -p initial-capital=$CAPITAL"
[[ -n "$TAG" ]]      && CMD="$CMD --labels foxhunt-tag=$TAG"
[[ "$SANITIZER" != "none" ]] && CMD="$CMD -p sanitizer=$SANITIZER"

$BASELINE && CMD="$CMD -p hyperopt-trials=0"
$WATCH    && CMD="$CMD --watch"

eval "$CMD"