# Vertex AI Custom Job — CXR-VLM EVALUATION on L4
#
# Submit with:
#   gcloud ai custom-jobs create \
#     --region=us-central1 \
#     --display-name=cxr-vlm-eval-run1 \
#     --config=scripts/vertex_eval_job.yaml
#
# Watch logs (streaming):
#   gcloud ai custom-jobs stream-logs <JOB_ID> --region=us-central1
#
# This reuses the SAME Docker image you trained with — no rebuild needed.
# Quota: uses CustomModelTrainingNvidiaL4Gpus (NOT global GPU quota), so it
# works even when Workbench / Compute Engine GPU creation is blocked.
#
# Output:
#   - Predictions + metrics_summary.json land under
#     /workspace/results/<RUN_ID>/ inside the container
#   - With UPLOAD_RESULTS_TO_HF=1 they're also uploaded to
#     hieu3636/cxr-vlm-runs/<RUN_ID>/results/
# (After the job ends the container is destroyed, so HF upload is what you'll
#  actually retrieve afterward.)
#
# Typical wall-clock on L4 for MIMIC-CXR_resized (10k test samples, all tasks):
#   ~1.5-2.5 hours total (3-5x faster than T4).

workerPoolSpecs:
- machineSpec:
    machineType: g2-standard-8           # 8 vCPU, 32GB RAM, 1×L4 24GB
    acceleratorType: NVIDIA_L4
    acceleratorCount: 1
  replicaCount: 1
  diskSpec:
    bootDiskType: pd-ssd
    bootDiskSizeGb: 200                  # enough for MIMIC-CXR_resized (~40GB extracted) + Vicuna cache
  containerSpec:
    imageUri: us-central1-docker.pkg.dev/cxr-vlm-thesis/cxr-vlm/cxr-vlm-env:cu128
    # Same bootstrap pattern as the training job: pull code from HF Hub,
    # then exec the eval entrypoint.
    command:
    - bash
    - -c
    - |
      set -e
      echo "[bootstrap] downloading code from HF Hub …"
      python -c "
      from huggingface_hub import snapshot_download
      import os
      snapshot_download('hieu3636/cxr-vlm-code',
                        repo_type='model',
                        token=os.environ['HF_TOKEN'],
                        local_dir='/workspace/code')
      "
      cd /workspace/code
      echo "[bootstrap] exec scripts/gcp_eval_entrypoint.py …"
      exec python scripts/gcp_eval_entrypoint.py
    env:
    # ── Required ────────────────────────────────────────────────────────────
    - name: HF_TOKEN
      value: _tSxYqkUSvesmWQOtoyrNpsHmfZlNFVUDRb
    - name: DATASET_NAME
      value: MIMIC-CXR_resized           # 'IU-Xray' | 'MIMIC-CXR' | 'MIMIC-CXR_resized'
    - name: RUN_ID
      value: MIMIC-CXR_resized_run_3     # which run on HF_RUNS_REPO to evaluate
    # ── Optional ────────────────────────────────────────────────────────────
    - name: HF_RUNS_REPO
      value: hieu3636/cxr-vlm-runs
    - name: CKPT_PICK
      value: best                        # 'best' | 'last'
    # ── Which tasks to (re-)evaluate — per-task toggles like the notebook ──
    # Re-running a subset MERGES into the existing metrics_summary.json, so the
    # other tasks' numbers are preserved. Below: re-eval ONLY findings +
    # impression (VQA off). Flip EVAL_VQA to "1" to include VQA.
    - name: EVAL_FINDINGS
      value: "1"
    - name: EVAL_IMPRESSION
      value: "1"
    - name: EVAL_VQA
      value: "0"
    # - name: EVAL_REPORT
    #   value: "0"                       # only for merged report_mode runs
    # Optional explicit override (e.g. "all" / "vqa"); blank → derive from EVAL_*.
    # - name: TASK
    #   value: all
    - name: SPLIT
      value: test
    # PNU prompt-condition source. 'oracle' = GT chex_* labels from JSON;
    # 'predicted' = run the Stage-0 CheXpert classifier per image (realistic).
    - name: PNU_SOURCE
      value: predicted                   # 'oracle' | 'predicted'
    # Classifier checkpoint on HF_RUNS_REPO (only used when PNU_SOURCE=predicted).
    - name: CHEXPERT_CKPT_PATH
      value: chexpert_classifier/chexpert_mimic_resized.pt
    - name: MAX_NEW_TOKENS
      value: "300"
    - name: UPLOAD_RESULTS_TO_HF
      value: "1"                         # '0' to disable HF upload
    # Leave REPORT_MODE / IMAGE_MODE blank → use the run's saved snapshot.
    # Set explicitly only if you need to override (rare).
    # - name: REPORT_MODE
    #   value: split_cascade
    # - name: IMAGE_MODE
    #   value: all_views_split
    # Leave BATCH_SIZE blank → auto from GPU profile (L4 -> 4).
    # - name: BATCH_SIZE
    #   value: "8"
    # LLM-as-judge for VQA (extra cost; needs OPENAI_API_KEY).
    # - name: LLM_JUDGE
    #   value: "1"
    # - name: OPENAI_API_KEY
    #   value: sk-...
    # - name: LLM_JUDGE_MODEL
    #   value: gpt-4o-mini

scheduling:
  # STANDARD on-demand. Eval is short (~2h) so spot interruption isn't worth
  # the risk — evaluate.py doesn't checkpoint mid-task. If preempted you'd
  # restart from the beginning.
  strategy: STANDARD
  # 6h ceiling. MIMIC-CXR_resized on L4 finishes in ~2h; raise this if
  # you're evaluating the full MIMIC-CXR with VQA.
  timeout: 21600s