#!/usr/bin/env bash
# Run a smoke test under Nsight Systems profiling on L40S via Argo Workflows.
#
# Output: .nsys-rep uploaded to MinIO at
#   foxhunt-training-results/profiles/smoke/<short-sha>/profile-<pod>.nsys-rep
#
# Download + open:
#   mc cp foxhunt/foxhunt-training-results/profiles/smoke/<sha>/profile.nsys-rep .
#   nsys-ui profile.nsys-rep
#
# Why L40S not local: the laptop GPU has no nsys-ui counters bound and
# limited VRAM. L40S has all the counters and ample headroom; nsys overhead
# is typically ~5-10% so runtime stays close to native.
#
# Usage:
#   ./scripts/argo-nsys.sh                                   # performance::test_real_data_single_epoch
#   ./scripts/argo-nsys.sh --test magnitude_distribution     # different test
#   ./scripts/argo-nsys.sh --multi-fold                       # exercise fold-boundary code
#   ./scripts/argo-nsys.sh --watch
#
# Requires: argo CLI

set -euo pipefail

TEST_NAME="performance::test_real_data_single_epoch"
COMMIT_REF="HEAD"
GPU_POOL=""
EXTRA_ARGS=""
WATCH=false

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

Wraps a smoke test under nsys profile on L40S. Output uploaded to MinIO.

Options:
  --test <name>        Smoke test name (default: $TEST_NAME)
                       Common values:
                         performance::test_real_data_single_epoch
                         iqn_quantile_monotonicity::iqn_multi_quantile_heads_produce_monotonic_estimates
                         multi_fold_convergence::test_multi_fold_convergence
                         magnitude_distribution::test_magnitude_distribution
  --multi-fold         Shortcut: --test multi_fold_convergence::test_multi_fold_convergence
                       Exercises fold-boundary code paths (IQN sync, aux Adam reset).
  --ref <ref>          Git ref to test (default: HEAD)
  --gpu-pool <pool>    GPU node pool (default: ci-training-l40s)
  --extra-args <args>  Pass-through args to nsys profile.
                       Default: "--trace=cuda,nvtx,osrt --gpu-metrics-devices=0
                                 --gpu-metrics-set=ga10x"
                       For range capture: "--capture-range=cudaProfilerApi"
                       For more detail: "--cuda-memory-usage=true --cudabacktrace=all"
  --watch              Follow workflow logs
  -h, --help           Show this help
EOF
  exit 0
}

while [[ $# -gt 0 ]]; do
  case $1 in
    --test)        TEST_NAME="$2"; shift 2 ;;
    --multi-fold)  TEST_NAME="multi_fold_convergence::test_multi_fold_convergence"; shift ;;
    --ref)         COMMIT_REF="$2"; shift 2 ;;
    --gpu-pool)    GPU_POOL="$2"; shift 2 ;;
    --extra-args)  EXTRA_ARGS="$2"; shift 2 ;;
    --watch)       WATCH=true; shift ;;
    -h|--help)     usage ;;
    *)             echo "Unknown option: $1"; usage ;;
  esac
done

if [ "$COMMIT_REF" = "HEAD" ]; then
  COMMIT_REF=$(git rev-parse HEAD)
fi

CMD="argo submit -n foxhunt --from=wftmpl/nsys-test"
CMD="$CMD -p commit-ref=$COMMIT_REF"
CMD="$CMD -p test-name=$TEST_NAME"
[[ -n "$GPU_POOL" ]]     && CMD="$CMD -p gpu-pool=$GPU_POOL"
[[ -n "$EXTRA_ARGS" ]]   && CMD="$CMD -p nsys-extra-args=\"$EXTRA_ARGS\""

if $WATCH; then
  CMD="$CMD --watch"
fi

echo "Submitting nsys workflow..."
echo "  commit: $(echo "$COMMIT_REF" | cut -c1-8)"
echo "  test:   $TEST_NAME"
[[ -n "$EXTRA_ARGS" ]] && echo "  extra:  $EXTRA_ARGS"
echo ""
eval "$CMD"