# infra/k8s/argo/gpu-test-pipeline-template.yaml
apiVersion: argoproj.io/v1alpha1
kind: WorkflowTemplate
metadata:
  name: gpu-test-pipeline
  namespace: foxhunt
  labels:
    app.kubernetes.io/name: gpu-test-pipeline
    app.kubernetes.io/part-of: foxhunt
spec:
  entrypoint: pipeline
  onExit: notify-result
  serviceAccountName: argo-workflow
  podMetadata:
    labels:
      app.kubernetes.io/part-of: foxhunt
      app.kubernetes.io/component: gpu-test
  securityContext:
    fsGroup: 0
  podGC:
    strategy: OnPodCompletion
  ttlStrategy:
    secondsAfterCompletion: 3600
  activeDeadlineSeconds: 7200
  arguments:
    parameters:
      - name: commit-ref
        value: HEAD
      - name: models
        value: "dqn,ppo,tft"
      - name: test-scope
        value: all
      - name: gpu-pool
        value: ci-training-h100
      - name: cuda-compute-cap
        value: "90"
      - name: notify
        value: "true"
  volumes:
    - name: git-ssh-key
      secret:
        secretName: argo-git-ssh-key
        defaultMode: 256
    - name: cargo-target
      persistentVolumeClaim:
        claimName: cargo-target-cuda-test
    - name: test-data
      persistentVolumeClaim:
        claimName: test-data-pvc
        readOnly: true

  templates:
    # ── pipeline: DAG entrypoint ──
    # compile-and-test runs on GPU node (RWO PVC can't be shared cross-node).
    # gpu-warmup ensures H100 is scaled up before compile starts.
    - name: pipeline
      dag:
        tasks:
          - name: gpu-warmup
            template: gpu-warmup
          - name: compile-and-test
            template: compile-and-test
            dependencies: [gpu-warmup]
          - name: perf-benchmark
            template: perf-benchmark
            dependencies: [compile-and-test]

    # ── gpu-warmup: trigger H100 autoscale ──
    # Requests GPU to force autoscaler to provision node, then releases it.
    - name: gpu-warmup
      nodeSelector:
        k8s.scaleway.com/pool-name: "{{workflow.parameters.gpu-pool}}"
      tolerations:
        - key: nvidia.com/gpu
          operator: Exists
          effect: NoSchedule
        - key: node.cilium.io/agent-not-ready
          operator: Exists
          effect: NoSchedule
      container:
        image: alpine:3.21
        command: ["/bin/sh", "-c"]
        args:
          - |
            echo "GPU warmup: triggering node autoscale..."
            nvidia-smi --query-gpu=name,memory.total --format=csv,noheader || true
            echo "GPU node ready, releasing for compile-and-test"
        resources:
          requests:
            nvidia.com/gpu: "1"
            cpu: 100m
            memory: 64Mi
          limits:
            nvidia.com/gpu: "1"
            cpu: 200m
            memory: 128Mi

    # ── compile-and-test: compile + run GPU tests in single H100 pod ──
    - name: compile-and-test
      nodeSelector:
        k8s.scaleway.com/pool-name: "{{workflow.parameters.gpu-pool}}"
      tolerations:
        - key: nvidia.com/gpu
          operator: Exists
          effect: NoSchedule
        - key: node.cilium.io/agent-not-ready
          operator: Exists
          effect: NoSchedule
      outputs:
        parameters:
          - name: results
            valueFrom:
              path: /tmp/outputs/results
              default: "unknown:FAIL"
          - name: failures
            valueFrom:
              path: /tmp/outputs/failures
              default: "1"
      container:
        image: gitlab-registry.foxhunt.svc.cluster.local:5000/root/foxhunt/ci-builder:latest
        command: ["/bin/bash", "-c"]
        env:
          - name: SQLX_OFFLINE
            value: "true"
          - name: CARGO_TERM_COLOR
            value: always
          - name: CARGO_TARGET_DIR
            value: /cargo-target
          - name: CARGO_HOME
            value: /cargo-target/cargo-home
          - name: CUDA_VISIBLE_DEVICES
            value: "0"
          - name: CUBLAS_WORKSPACE_CONFIG
            value: ":4096:8"
          - name: TEST_DATA_DIR
            value: /data/test-data
          - name: RUST_LOG
            value: info
          - name: LD_LIBRARY_PATH
            value: /usr/local/nvidia/lib64:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64
          - name: CARGO_PROFILE_TEST_OPT_LEVEL
            value: "2"
        volumeMounts:
          - name: git-ssh-key
            mountPath: /etc/git-ssh
            readOnly: true
          - name: cargo-target
            mountPath: /cargo-target
          - name: test-data
            mountPath: /data/test-data
            readOnly: true
        resources:
          requests:
            nvidia.com/gpu: "1"
            cpu: "4"
            memory: 32Gi
          limits:
            nvidia.com/gpu: "1"
            cpu: "8"
            memory: 64Gi
        args:
          - |
            set -e

            REF="{{workflow.parameters.commit-ref}}"
            MODELS="{{workflow.parameters.models}}"
            SCOPE="{{workflow.parameters.test-scope}}"

            # --- SSH setup (same as compile-and-train) ---
            mkdir -p ~/.ssh
            cp /etc/git-ssh/ssh-privatekey ~/.ssh/id_ed25519
            chmod 600 ~/.ssh/id_ed25519
            printf 'Host *\n  StrictHostKeyChecking no\n  UserKnownHostsFile /dev/null\n' > ~/.ssh/config
            chmod 600 ~/.ssh/config

            REPO="ssh://git@gitlab-gitlab-shell.foxhunt.svc.cluster.local:2222/root/foxhunt.git"
            BUILD="/cargo-target/src"
            git config --global --add safe.directory "$BUILD"

            # --- Persistent checkout on PVC ---
            if [ -d "$BUILD/.git" ]; then
              cd "$BUILD"
              echo "=== Fetching latest refs ==="
              git fetch origin
              # Resolve REF after fetch so we always get the latest commit.
              # Try origin/$REF (branch), then $REF directly (tag or SHA).
              if git rev-parse --verify "origin/$REF" >/dev/null 2>&1; then
                TARGET=$(git rev-parse "origin/$REF")
              else
                TARGET=$(git rev-parse --verify "$REF" 2>/dev/null || echo "$REF")
              fi
              CURRENT=$(git rev-parse HEAD 2>/dev/null || echo "none")
              if [ "$CURRENT" = "$TARGET" ]; then
                echo "=== Already at $REF ($TARGET) ==="
              else
                echo "=== Updating checkout to $REF ($TARGET) ==="
                git checkout --force --detach "$TARGET"
                git clean -fd
              fi
            else
              echo "=== Initial clone ==="
              git clone --filter=blob:none "$REPO" "$BUILD"
              cd "$BUILD"
              git fetch origin
              if git rev-parse --verify "origin/$REF" >/dev/null 2>&1; then
                git checkout --force --detach "origin/$REF"
              else
                git checkout "$REF"
              fi
            fi
            echo "Checked out $(git rev-parse --short HEAD)"

            export PATH="${CARGO_HOME}/bin:${PATH}"
            export CUDA_COMPUTE_CAP={{workflow.parameters.cuda-compute-cap}}

            # --- PTX cache invalidation ---
            # Purge stale cached PTX if any CUDA kernel source changed since last run.
            bash scripts/ptx-cache-invalidate.sh "${CARGO_TARGET_DIR}/.ptx_cache"

            # --- PVC size guard ---
            TARGET_SIZE_MB=$(du -sm "$CARGO_TARGET_DIR" 2>/dev/null | cut -f1 || echo 0)
            echo "PVC usage: ${TARGET_SIZE_MB}MB"
            if [ "$TARGET_SIZE_MB" -gt 25000 ]; then
              echo "PVC exceeds 25GB, pruning..."
              rm -rf "$CARGO_TARGET_DIR/release" "$CARGO_TARGET_DIR/debug"
            fi

            # --- Compile ---
            echo "=== Compiling test binaries (--features cuda) ==="
            cargo test -p ml -p ml-dqn -p ml-core --features cuda --no-run 2>&1 | tee /cargo-target/gpu-compile.log
            echo "=== Compilation done ==="

            # --- Expand "all" ---
            if [ "$MODELS" = "all" ]; then
              MODELS="dqn,ppo,tft,mamba2,tggn,tlob,liquid,kan,xlstm,diffusion"
            fi

            # --- Test runner ---
            RESULTS=""
            FAILURES=0

            # Reset CUDA context between test binaries to prevent cuBLAS
            # CUBLAS_STATUS_NOT_INITIALIZED cascades. Each test binary creates
            # and destroys hundreds of cuBLAS handles; without a reset, the
            # driver fails to re-init for the next binary.
            gpu_context_drain() {
              nvidia-smi -rgc >/dev/null 2>&1 || true
              sleep 1
            }

            run_tests() {
              local NAME="$1"; shift
              echo ""
              echo "========================================"
              echo "  TEST: $NAME"
              echo "========================================"
              # --nocapture is a test-binary flag, must come after --
              # If args already contain --, append after it; otherwise add -- first
              local HAS_SEP=false
              for arg in "$@"; do
                [ "$arg" = "--" ] && HAS_SEP=true && break
              done
              set +e
              if $HAS_SEP; then
                "$@" --nocapture 2>&1
              else
                "$@" -- --nocapture 2>&1
              fi
              EXIT=$?
              set -e
              if [ $EXIT -eq 0 ]; then
                RESULTS="${RESULTS}${NAME}:PASS\n"
              else
                RESULTS="${RESULTS}${NAME}:FAIL\n"
                FAILURES=$((FAILURES + 1))
              fi
              # Drain CUDA context after each GPU test suite
              gpu_context_drain
            }

            # --- Core tests (always run) ---
            # --test-threads=1 for all GPU lib tests: prevents concurrent cuBLAS
            # handle creation that causes CUBLAS_STATUS_NOT_INITIALIZED cascades.
            run_tests "core-lib" cargo test -p ml-core --features cuda --lib -- --test-threads=1
            run_tests "bayesian" cargo test -p ml --features cuda --test bayesian_changepoint_test -- --test-threads=1

            # --- Per-model tests ---
            IFS=',' read -ra MODEL_LIST <<< "$MODELS"
            for MODEL in "${MODEL_LIST[@]}"; do
              case "$MODEL" in
                dqn)
                  if [ "$SCOPE" = "lib" ] || [ "$SCOPE" = "all" ]; then
                    # --test-threads=1: GPU lib tests must run serially — each test
                    # creates a cuBLAS handle via Device::new_cuda(0). Under parallel
                    # execution, concurrent cuBLAS init races cause
                    # CUBLAS_STATUS_NOT_INITIALIZED failures (51 test cascade).
                    run_tests "dqn-lib" cargo test -p ml-dqn --features cuda --lib -- --test-threads=1
                    run_tests "dqn-ml-lib" cargo test -p ml --features cuda --lib -- --test-threads=1 dqn
                  fi
                  if [ "$SCOPE" = "integration" ] || [ "$SCOPE" = "all" ]; then
                    # --test-threads=1 for ALL GPU integration tests: parallel execution
                    # corrupts the CUDA primary context (cuDevicePrimaryCtxRetain fails
                    # when multiple threads race on context init/teardown).
                    run_tests "dqn-smoke" cargo test -p ml --features cuda --test smoke_test_real_data -- --test-threads=1
                    # Run each pipeline test in its own cargo test process.
                    # CUDA Graph capture corrupts the async memory pool, making
                    # cuMemAllocAsync fail with CUDA_ERROR_INVALID_VALUE in
                    # subsequent DQNTrainer instances within the same process.
                    run_tests "dqn-pipeline-train" cargo test -p ml --features cuda --test dqn_training_pipeline_test test_dqn_trains_on_es_fut -- --test-threads=1 --exact
                    run_tests "dqn-pipeline-loss" cargo test -p ml --features cuda --test dqn_training_pipeline_test test_dqn_loss_decreases -- --test-threads=1 --exact
                    run_tests "dqn-pipeline-ckpt" cargo test -p ml --features cuda --test dqn_training_pipeline_test test_dqn_checkpoint_save_load -- --test-threads=1 --exact
                    run_tests "dqn-pipeline-qval" cargo test -p ml --features cuda --test dqn_training_pipeline_test test_dqn_q_value_predictions -- --test-threads=1 --exact
                    run_tests "dqn-pipeline-eps" cargo test -p ml --features cuda --test dqn_training_pipeline_test test_dqn_epsilon_greedy -- --test-threads=1 --exact
                    run_tests "dqn-smoke-train" cargo test -p ml --features cuda --test dqn_training_smoke_test -- --test-threads=1
                    run_tests "dqn-early-stop" cargo test -p ml --features cuda --test dqn_early_stopping_termination_test -- --test-threads=1
                    run_tests "dqn-collapse" cargo test -p ml --features cuda --test dqn_action_collapse_fix_test -- --test-threads=1
                  fi
                  ;;
                ppo)
                  if [ "$SCOPE" = "lib" ] || [ "$SCOPE" = "all" ]; then
                    run_tests "ppo-lib" cargo test -p ml --features cuda --lib -- --test-threads=1 ppo
                  fi
                  if [ "$SCOPE" = "integration" ] || [ "$SCOPE" = "all" ]; then
                    run_tests "ppo-barrier" cargo test -p ml --features cuda --test barrier_optimization_test -- --test-threads=1
                  fi
                  ;;
                *)
                  # Supervised models (TFT, Mamba2, TGGN, TLOB, Liquid, KAN, xLSTM, Diffusion)
                  # Map model names to Rust module names where they differ
                  LIB_FILTER="$MODEL"
                  [ "$MODEL" = "tggn" ] && LIB_FILTER="tgnn"
                  if [ "$SCOPE" = "lib" ] || [ "$SCOPE" = "all" ]; then
                    run_tests "${MODEL}-lib" cargo test -p ml --features cuda --lib -- --test-threads=1 "$LIB_FILTER"
                  fi
                  if [ "$SCOPE" = "integration" ] || [ "$SCOPE" = "all" ]; then
                    run_tests "${MODEL}-gpu" cargo test -p ml --features cuda --test supervised_gpu_smoke_test -- --test-threads=1 "test_${MODEL}_gpu_smoke"
                    # Also run model-specific integration tests if they exist
                    if cargo test -p ml --features cuda --test "${MODEL}_integration" --no-run 2>/dev/null; then
                      run_tests "${MODEL}-integ" cargo test -p ml --features cuda --test "${MODEL}_integration" -- --test-threads=1
                    fi
                  fi
                  ;;
              esac
            done

            # --- Summary ---
            echo ""
            echo "========================================"
            echo "  GPU TEST RESULTS"
            echo "========================================"
            printf "$RESULTS" | while IFS=: read -r name status; do
              [ -z "$name" ] && continue
              printf "  %-25s %s\n" "$name" "$status"
            done
            echo "========================================"
            echo "  Total failures: $FAILURES"
            echo "========================================"

            # Write results for notification step
            mkdir -p /tmp/outputs
            printf "$RESULTS" > /tmp/outputs/results
            echo "$FAILURES" > /tmp/outputs/failures

            [ "$FAILURES" -gt 0 ] && exit 1 || exit 0

    # ── perf-benchmark: DQN epoch/s on 3Q data (performance regression guard) ──
    # Runs after tests pass. Trains DQN on 3Q of ES.FUT data and reports epoch time.
    # Uses the same binary compiled by compile-and-test (shared cargo-target PVC).
    # Fails the pipeline if epoch time exceeds 500ms (regression threshold for H100).
    - name: perf-benchmark
      nodeSelector:
        k8s.scaleway.com/pool-name: "{{workflow.parameters.gpu-pool}}"
      tolerations:
        - key: nvidia.com/gpu
          operator: Exists
          effect: NoSchedule
        - key: node.cilium.io/agent-not-ready
          operator: Exists
          effect: NoSchedule
      outputs:
        parameters:
          - name: epoch-ms
            valueFrom:
              path: /tmp/outputs/epoch-ms
              default: "9999"
      container:
        image: gitlab-registry.foxhunt.svc.cluster.local:5000/root/foxhunt/ci-builder:latest
        command: ["/bin/bash", "-c"]
        env:
          - name: SQLX_OFFLINE
            value: "true"
          - name: CARGO_TARGET_DIR
            value: /cargo-target
          - name: CARGO_HOME
            value: /cargo-target/cargo-home
          - name: CUDA_VISIBLE_DEVICES
            value: "0"
          - name: LD_LIBRARY_PATH
            value: /usr/local/nvidia/lib64:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64
        volumeMounts:
          - name: cargo-target
            mountPath: /cargo-target
          - name: test-data
            mountPath: /data/test-data
            readOnly: true
        resources:
          requests:
            nvidia.com/gpu: "1"
            cpu: "4"
            memory: 16Gi
          limits:
            nvidia.com/gpu: "1"
            cpu: "8"
            memory: 32Gi
        args:
          - |
            set -e
            cd /cargo-target/src
            export PATH="${CARGO_HOME}/bin:${PATH}"
            export CUDA_COMPUTE_CAP={{workflow.parameters.cuda-compute-cap}}

            echo "========================================"
            echo "  PERF BENCHMARK: DQN epoch/s (3Q ES.FUT)"
            echo "========================================"

            # Run 5 epochs × 100 steps on 3Q data (--train-months 3 fits in 1 quarter)
            # Use --step-months 3 to avoid walk-forward window splits
            OUTPUT=$(cargo run --release --example train_baseline_rl -p ml -- \
              --model dqn \
              --data-dir /data/test-data/ohlcv \
              --mbp10-data-dir /data/test-data/mbp10 \
              --trades-data-dir /data/test-data/trades \
              --symbol ES.FUT \
              --epochs 5 \
              --max-steps-per-epoch 100 \
              --batch-size 1024 \
              --train-months 3 --val-months 1 --test-months 1 --step-months 3 \
              2>&1)

            # Extract epoch times (skip epoch 1 which includes init)
            EPOCH_TIMES=$(echo "$OUTPUT" | grep "phase breakdown" | grep -v "Epoch 1/" | \
              sed 's/.*total=\([0-9]*\)ms.*/\1/' | head -4)

            if [ -z "$EPOCH_TIMES" ]; then
              echo "ERROR: No phase breakdown output found"
              echo "$OUTPUT" | tail -20
              mkdir -p /tmp/outputs
              echo "9999" > /tmp/outputs/epoch-ms
              exit 1
            fi

            # Compute average epoch time (epochs 2-5)
            SUM=0
            COUNT=0
            for T in $EPOCH_TIMES; do
              SUM=$((SUM + T))
              COUNT=$((COUNT + 1))
            done
            AVG=$((SUM / COUNT))

            echo ""
            echo "  Epoch times (ms, excl. epoch 1): $EPOCH_TIMES"
            echo "  Average: ${AVG}ms"
            echo ""

            mkdir -p /tmp/outputs
            echo "$AVG" > /tmp/outputs/epoch-ms

            # Regression guard: fail if avg epoch > 500ms on H100
            THRESHOLD=500
            if [ "$AVG" -gt "$THRESHOLD" ]; then
              echo "PERF REGRESSION: ${AVG}ms > ${THRESHOLD}ms threshold"
              echo "========================================"
              exit 1
            fi

            echo "  PASS: ${AVG}ms <= ${THRESHOLD}ms threshold"
            echo "========================================"

    # ── notify-result: post test outcome to Mattermost (onExit) ──
    - name: notify-result
      nodeSelector:
        k8s.scaleway.com/pool-name: platform
      container:
        image: curlimages/curl:8.12.1
        command: ["/bin/sh", "-c"]
        env:
          - name: WEBHOOK_URL
            valueFrom:
              secretKeyRef:
                name: notification-webhook
                key: webhook-url
                optional: true
        resources:
          requests:
            cpu: 50m
            memory: 32Mi
          limits:
            cpu: 100m
            memory: 64Mi
        args:
          - |
            NOTIFY="{{workflow.parameters.notify}}"
            if [ "$NOTIFY" != "true" ]; then
              echo "Notifications disabled, skipping"
              exit 0
            fi

            STATUS="{{workflow.status}}"
            NAME="{{workflow.name}}"

            if [ -z "$WEBHOOK_URL" ] || echo "$WEBHOOK_URL" | grep -q "PLACEHOLDER"; then
              echo "No webhook configured, skipping notification"
              exit 0
            fi

            if [ "$STATUS" = "Succeeded" ]; then
              EMOJI=":white_check_mark:"
            else
              EMOJI=":x:"
            fi

            PAYLOAD="{\"username\":\"Argo CI\",\"text\":\"${EMOJI} **GPU Tests** ${NAME} — ${STATUS} ({{workflow.duration}}s)\"}"

            curl -sf -X POST -H 'Content-Type: application/json' \
              -d "$PAYLOAD" "$WEBHOOK_URL" || echo "WARN: webhook post failed"