# Argo Training Pipeline Redesign Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Replace 7 overlapping Argo workflow templates with a single unified `train` template that caches binaries by commit SHA on PVC and skips fxcache regeneration when valid.

**Architecture:** One workflow template (`train-template.yaml`) with a DAG: `ensure-binary` (cache-or-compile) and `gpu-warmup` run in parallel, then `ensure-fxcache` (cache-or-precompute), then `hyperopt` → `train-best` → `evaluate` → `upload-results`. One script (`argo-train.sh`) submits it. Five dead templates deleted.

**Tech Stack:** Argo Workflows, Kubernetes, bash, PVC storage

---

## File Structure

| Action | File | Responsibility |
|--------|------|----------------|
| Create | `infra/k8s/argo/train-template.yaml` | Unified training workflow template |
| Rewrite | `scripts/argo-train.sh` | Simplified CLI that submits the unified template |
| Delete | `infra/k8s/argo/compile-and-train-template.yaml` | Replaced by train-template |
| Delete | `infra/k8s/argo/train-dqn-template.yaml` | Unused duplicate |
| Delete | `infra/k8s/argo/train-baseline-rl-template.yaml` | Unused duplicate |
| Delete | `infra/k8s/argo/train-supervised-template.yaml` | Unused duplicate |
| Delete | `infra/k8s/argo/training-workflow-template.yaml` | Unused duplicate |
| Delete | `infra/k8s/argo/precompute-features-template.yaml` | Absorbed into ensure-fxcache step |
| Delete | `scripts/argo-precompute.sh` | Absorbed into ensure-fxcache step |
| Modify | `infra/k8s/argo/kustomization.yaml` | Remove deleted templates, add new one |

---

### Task 1: Create the unified `train-template.yaml`

**Files:**
- Create: `infra/k8s/argo/train-template.yaml`

- [ ] **Step 1: Write the template header, parameters, volumes, and DAG**

```yaml
# Unified training pipeline: ensure-binary → ensure-fxcache → train
#
# Smart caching:
#   - Binaries cached on PVC at /data/bin/{sha}/ — skips compile if present
#   - fxcache cached on feature-cache PVC — skips precompute if present
#   - GPU warmup runs in parallel with ensure-binary
#
# Usage:
#   ./scripts/argo-train.sh dqn
#   ./scripts/argo-train.sh dqn --sha abc1234 --epochs 100
---
apiVersion: argoproj.io/v1alpha1
kind: WorkflowTemplate
metadata:
  name: train
  namespace: foxhunt
  labels:
    app.kubernetes.io/name: train
    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: train
  securityContext:
    fsGroup: 0
  ttlStrategy:
    secondsAfterCompletion: 3600
  activeDeadlineSeconds: 28800

  arguments:
    parameters:
      - name: commit-sha
        value: HEAD
      - name: git-branch
        value: main
      - name: cuda-compute-cap
        value: "90"
      - name: model
        value: dqn
      - name: gpu-pool
        value: ci-training-h100
      - name: hyperopt-trials
        value: "20"
      - name: hyperopt-epochs
        value: "8"
      - name: train-epochs
        value: "50"
      - name: symbol
        value: ES.FUT
      - name: initial-capital
        value: "35000"
      - name: tx-cost-bps
        value: "0.1"
      - name: tick-size
        value: "0.25"
      - name: spread-ticks
        value: "1.0"
      - name: min-hold-bars
        value: "5"

  volumes:
    - name: git-ssh-key
      secret:
        secretName: argo-git-ssh-key
        defaultMode: 256
    - name: training-data
      persistentVolumeClaim:
        claimName: training-data-pvc
    - name: cargo-target-cuda
      persistentVolumeClaim:
        claimName: cargo-target-cuda
    - name: feature-cache
      persistentVolumeClaim:
        claimName: feature-cache-pvc

  volumeClaimTemplates:
    - metadata:
        name: workspace
      spec:
        accessModes: ["ReadWriteOnce"]
        storageClassName: scw-bssd
        resources:
          requests:
            storage: 5Gi

  templates:
    - name: pipeline
      dag:
        tasks:
          - name: ensure-binary
            template: ensure-binary
          - name: gpu-warmup
            template: gpu-warmup
          - name: ensure-fxcache
            template: ensure-fxcache
            dependencies: [ensure-binary]
          - name: hyperopt
            template: hyperopt
            dependencies: [ensure-binary, ensure-fxcache, gpu-warmup]
            when: "{{workflow.parameters.hyperopt-trials}} != 0"
          - name: train-best
            template: train-best
            dependencies: [ensure-binary, ensure-fxcache, gpu-warmup, hyperopt]
          - name: evaluate
            template: evaluate
            dependencies: [train-best]
          - name: upload-results
            template: upload-results
            dependencies: [evaluate]
```

- [ ] **Step 2: Write the `ensure-binary` template**

This is the key innovation — checks PVC cache first, compiles only on miss. Append to the same file:

```yaml
    # ── ensure-binary: cache-or-compile training binaries ──
    # Checks /data/bin/{sha}/ on PVC. If binaries exist, exits immediately.
    # If not, clones repo, builds, copies to /data/bin/{sha}/, cleans old SHAs.
    - name: ensure-binary
      nodeSelector:
        k8s.scaleway.com/pool-name: ci-compile-cpu
      tolerations:
        - key: node.cilium.io/agent-not-ready
          operator: Exists
          effect: NoSchedule
      container:
        image: gitlab-registry.foxhunt.svc.cluster.local:5000/root/foxhunt/ci-builder:latest
        imagePullPolicy: Always
        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: GITLAB_PAT
            valueFrom:
              secretKeyRef:
                name: gitlab-pat
                key: token
        resources:
          requests:
            cpu: "14"
            memory: 32Gi
          limits:
            cpu: "30"
            memory: 64Gi
        volumeMounts:
          - name: git-ssh-key
            mountPath: /etc/git-ssh
            readOnly: true
          - name: cargo-target-cuda
            mountPath: /cargo-target
          - name: training-data
            mountPath: /data
        args:
          - |
            set -e
            BRANCH="{{workflow.parameters.git-branch}}"
            REPO="ssh://git@gitlab-gitlab-shell.foxhunt.svc.cluster.local:2222/root/foxhunt.git"
            BUILD="/cargo-target/src"

            # Resolve commit SHA
            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

            git config --global --add safe.directory "$BUILD"

            SHA="{{workflow.parameters.commit-sha}}"
            if [ "$SHA" = "HEAD" ]; then
              # Need to resolve HEAD to actual SHA
              if [ -d "$BUILD/.git" ]; then
                cd "$BUILD" && git fetch origin
                SHA=$(git rev-parse "origin/$BRANCH")
                cd /
              else
                git clone --filter=blob:none --depth=1 "$REPO" /tmp/sha-resolve
                SHA=$(git -C /tmp/sha-resolve rev-parse HEAD)
                rm -rf /tmp/sha-resolve
              fi
            fi
            SHORT_SHA=$(echo "$SHA" | cut -c1-9)
            echo "Resolved SHA: $SHORT_SHA"

            # ── Cache check ──
            BIN_DIR="/data/bin/$SHORT_SHA"
            MODEL="{{workflow.parameters.model}}"
            case "$MODEL" in
              dqn|ppo) BINARIES="hyperopt_baseline_rl train_baseline_rl evaluate_baseline precompute_features" ;;
              *)       BINARIES="hyperopt_baseline_supervised train_baseline_supervised evaluate_supervised precompute_features" ;;
            esac

            ALL_CACHED=true
            for bin in $BINARIES; do
              if [ ! -x "$BIN_DIR/$bin" ]; then
                ALL_CACHED=false
                break
              fi
            done

            if $ALL_CACHED; then
              echo "=== Binary cache HIT for $SHORT_SHA ==="
              ls -lh "$BIN_DIR/"
              echo "$SHORT_SHA" > /tmp/sha
              exit 0
            fi

            echo "=== Binary cache MISS for $SHORT_SHA — compiling ==="

            # Checkout
            if [ -d "$BUILD/.git" ]; then
              cd "$BUILD"
              git fetch origin
              CURRENT=$(git rev-parse HEAD 2>/dev/null || echo "none")
              if [ "$CURRENT" != "$SHA" ]; then
                git checkout --force "$SHA"
                git clean -fd
              fi
            else
              git clone --filter=blob:none "$REPO" "$BUILD"
              cd "$BUILD"
              git checkout "$SHA"
            fi

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

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

            ML_EXAMPLE_ARGS=""
            for ex in $BINARIES; do
              ML_EXAMPLE_ARGS="$ML_EXAMPLE_ARGS --example $ex"
            done

            echo "=== Building: $BINARIES ==="
            cargo build --release -p ml --features ml/cuda $ML_EXAMPLE_ARGS

            mkdir -p "$BIN_DIR"
            for bin in $BINARIES; do
              cp "$CARGO_TARGET_DIR/release/examples/$bin" "$BIN_DIR/"
            done
            strip "$BIN_DIR/"*

            echo "=== Binaries cached at $BIN_DIR ==="
            ls -lh "$BIN_DIR/"

            # Cleanup: keep only last 5 SHA dirs
            ls -dt /data/bin/*/ 2>/dev/null | tail -n +6 | xargs rm -rf 2>/dev/null || true

            echo "$SHORT_SHA" > /tmp/sha
      outputs:
        parameters:
          - name: sha
            valueFrom:
              path: /tmp/sha
```

- [ ] **Step 3: Write the `ensure-fxcache` template**

```yaml
    # ── ensure-fxcache: run precompute if cache miss ──
    # The precompute binary computes its own cache key and skips if output exists.
    - name: ensure-fxcache
      nodeSelector:
        k8s.scaleway.com/pool-name: ci-compile-cpu
      tolerations:
        - key: node.cilium.io/agent-not-ready
          operator: Exists
          effect: NoSchedule
      container:
        image: ubuntu:24.04
        command: ["/bin/bash", "-c"]
        env:
          - name: RUST_LOG
            value: info
        resources:
          requests:
            cpu: "4"
            memory: 16Gi
          limits:
            cpu: "28"
            memory: 56Gi
        volumeMounts:
          - name: training-data
            mountPath: /data
            readOnly: true
          - name: feature-cache
            mountPath: /feature-cache
        args:
          - |
            set -e
            SHA="{{tasks.ensure-binary.outputs.parameters.sha}}"
            BINARY="/data/bin/$SHA/precompute_features"

            if [ ! -x "$BINARY" ]; then
              echo "ERROR: $BINARY not found"
              exit 1
            fi

            echo "=== Feature precompute (skip-if-exists) ==="
            $BINARY \
              --data-dir /data/futures-baseline \
              --mbp10-data-dir /data/futures-baseline-mbp10 \
              --trades-data-dir /data/futures-baseline-trades \
              --output-dir /feature-cache \
              --symbol {{workflow.parameters.symbol}} \
              --yes
```

- [ ] **Step 4: Write the `gpu-warmup` template**

```yaml
    # ── gpu-warmup: trigger GPU node autoscale in parallel ──
    - 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: busybox:1.37
        command: ["/bin/sh", "-c"]
        args:
          - |
            echo "GPU warmup: node scheduled, autoscale triggered"
        resources:
          requests:
            nvidia.com/gpu: "1"
            cpu: 100m
            memory: 64Mi
          limits:
            nvidia.com/gpu: "1"
            cpu: 200m
            memory: 128Mi
```

- [ ] **Step 5: Write `hyperopt`, `train-best`, `evaluate`, `upload-results`, `notify-result` templates**

These are copied from `compile-and-train-template.yaml` with one change: binaries are read from `/data/bin/{sha}/` instead of `/workspace/bin/`. Replace all `export PATH="/workspace/bin:$PATH"` with:

```bash
SHA="{{tasks.ensure-binary.outputs.parameters.sha}}"
export PATH="/data/bin/$SHA:$PATH"
```

And replace `feature-cache-dir` parameter references with the hardcoded `/feature-cache` path.

The `hyperopt` template (lines 411-503 of old file), `train-best` (505-587), `evaluate` (589-654), `upload-results` (656-704), and `notify-result` (706-747) are otherwise identical — copy them with the PATH fix above.

- [ ] **Step 6: Verify the template is valid YAML**

Run: `python3 -c "import yaml; yaml.safe_load(open('infra/k8s/argo/train-template.yaml'))" && echo "Valid YAML"`
Expected: `Valid YAML`

- [ ] **Step 7: Commit**

```bash
git add infra/k8s/argo/train-template.yaml
git commit -m "feat: unified train workflow template — cache-or-compile + cache-or-precompute"
```

---

### Task 2: Rewrite `scripts/argo-train.sh`

**Files:**
- Modify: `scripts/argo-train.sh`

- [ ] **Step 1: Replace the entire script**

```bash
#!/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=""

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)
  --watch              Follow workflow logs
  -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 ;;
    --watch)     WATCH=true; shift ;;
    -h|--help)   usage ;;
    *)           echo "Unknown option: $1"; usage ;;
  esac
done

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"

[[ -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"

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

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 ""
eval "$CMD"
```

- [ ] **Step 2: Verify script is executable**

Run: `bash -n scripts/argo-train.sh && echo "Syntax OK"`
Expected: `Syntax OK`

- [ ] **Step 3: Commit**

```bash
git add scripts/argo-train.sh
git commit -m "feat: rewrite argo-train.sh — single template, --sha for commit pinning"
```

---

### Task 3: Delete dead templates and scripts

**Files:**
- Delete: `infra/k8s/argo/compile-and-train-template.yaml`
- Delete: `infra/k8s/argo/train-dqn-template.yaml`
- Delete: `infra/k8s/argo/train-baseline-rl-template.yaml`
- Delete: `infra/k8s/argo/train-supervised-template.yaml`
- Delete: `infra/k8s/argo/training-workflow-template.yaml`
- Delete: `infra/k8s/argo/precompute-features-template.yaml`
- Delete: `infra/k8s/argo/train-ppo-template.yaml`
- Delete: `scripts/argo-precompute.sh`
- Modify: `infra/k8s/argo/kustomization.yaml`

- [ ] **Step 1: Delete the template files**

```bash
rm infra/k8s/argo/compile-and-train-template.yaml
rm infra/k8s/argo/train-dqn-template.yaml
rm infra/k8s/argo/train-baseline-rl-template.yaml
rm infra/k8s/argo/train-supervised-template.yaml
rm infra/k8s/argo/training-workflow-template.yaml
rm infra/k8s/argo/precompute-features-template.yaml
rm infra/k8s/argo/train-ppo-template.yaml
rm scripts/argo-precompute.sh
```

- [ ] **Step 2: Update kustomization.yaml**

Read `infra/k8s/argo/kustomization.yaml`, remove references to deleted templates, add `train-template.yaml`.

- [ ] **Step 3: Commit**

```bash
git add -A infra/k8s/argo/ scripts/
git commit -m "cleanup: delete 7 dead training templates + precompute script, update kustomization"
```

---

### Task 4: Deploy to cluster and delete old templates

**Files:**
- No file changes — cluster operations only

- [ ] **Step 1: Apply the new template to the cluster**

```bash
kubectl apply -f infra/k8s/argo/train-template.yaml
```

Expected: `workflowtemplate.argoproj.io/train created`

- [ ] **Step 2: Delete old templates from the cluster**

```bash
kubectl delete wftmpl -n foxhunt compile-and-train train-dqn train-baseline-rl train-supervised training-pipeline precompute-features train-ppo 2>/dev/null || true
```

- [ ] **Step 3: Verify only the new template exists (plus the non-training ones)**

```bash
kubectl get wftmpl -n foxhunt
```

Expected: `train` plus `compile-and-deploy`, `databento-download`, `gpu-test-pipeline`, `build-ci-image`, `ci-pipeline`

- [ ] **Step 4: Commit (no file changes — just verify)**

No commit needed. Cluster state is updated.

---

### Task 5: Test the new workflow — baseline DQN on H100

**Files:**
- No file changes — integration test

- [ ] **Step 1: Run a baseline DQN training (no hyperopt, 5 epochs)**

```bash
./scripts/argo-train.sh dqn --baseline --epochs 5 --watch
```

Expected: 
- `ensure-binary`: cache HIT (binary was compiled by the terminated workflow earlier) OR cache MISS → compile ~10 min
- `ensure-fxcache`: cache HIT (fxcache was generated earlier) OR run precompute ~4 min
- `gpu-warmup`: completes in ~3 min (node autoscale)
- `train-best`: runs 5 epochs on H100 with TF32 tensor cores
- Watch for epoch timing — should be ~10s/epoch on H100 (was 77s with old CUDA core path)

- [ ] **Step 2: Verify epoch timing in logs**

```bash
argo logs <workflow-name> -n foxhunt | grep 'duration_secs\|duration='
```

Expected: epoch times significantly faster than 77s (the old bf16/CUDA core baseline).

- [ ] **Step 3: Push all changes**

```bash
git push origin main
```
