#cloud-config
#
# Benchmark VM provisioning for Unstructured-IO/core-product
#
# Document processing pipeline -- Python 3.12, uv-based multi-package workspace.
# Private repo: requires SSH agent forwarding for clone (not done in cloud-init).
#
# Two-phase setup:
#   Phase 1 (cloud-init): packages, hyperfine, uv, bench scripts
#   Phase 2 (manual):     ssh -A, clone, make install, baseline benchmarks
#
# Usage:
#   az vm create ... --custom-data infra/cloud-init.yaml
#   bash infra/vm-manage.sh ssh       # connects with -A for agent forwarding
#   bash ~/setup.sh                    # clone + install + verify
#
# VM: Azure Standard_E4s_v5 (4 vCPU, 32 GB RAM, memory-optimized)
#   Matches production pod profile (1 CPU request, 32 GB RAM limit).
#   Use taskset -c 0 to pin benchmarks to 1 core for production-realistic results.
#   Non-burstable ensures consistent CPU -- no thermal throttling or turbo variability.

package_update: true
packages:
  - git
  - build-essential
  - curl
  - wget
  - jq
  - tesseract-ocr
  - libtesseract-dev
  - libleptonica-dev
  - poppler-utils
  - libmagic1
  - libgl1
  - libglib2.0-0
  - pandoc
  - libreoffice-core

write_files:
  # --- Benchmark: unit tests (fast, no ML models) ---
  - path: /home/azureuser/bench/bench_tests.sh
    owner: azureuser:azureuser
    permissions: "0755"
    defer: true
    content: |
      #!/usr/bin/env bash
      set -euo pipefail

      cd ~/core-product
      PYTHON=.venv/bin/python

      echo "=== core-product unit tests (fast, no ML models) ==="
      echo ""

      $PYTHON -m pytest unstructured_prop/tests/unit -n auto -m "not slow" -q 2>&1 | tail -5
      $PYTHON -m pytest unstructured_inference_prop/tests/unit -n auto -m "not slow" -q 2>&1 | tail -5

  # --- Benchmark: A/B branch comparison ---
  - path: /home/azureuser/bench/bench_compare.sh
    owner: azureuser:azureuser
    permissions: "0755"
    defer: true
    content: |
      #!/usr/bin/env bash
      set -euo pipefail
      BRANCH="${1:?Usage: bench_compare.sh <branch-or-commit>}"
      TS=$(date +%Y%m%d-%H%M%S)
      OUTDIR="$HOME/results/${BRANCH//\//-}-${TS}"
      mkdir -p "$OUTDIR"

      cd ~/core-product
      git fetch origin
      git checkout "$BRANCH"

      # Rebuild after switching branches
      export PATH="$HOME/.local/bin:$PATH"
      export VIRTUAL_ENV=$PWD/.venv
      make install

      PYTHON=.venv/bin/python

      echo "=== Benchmarking branch: $BRANCH ==="

      # Unit tests (fast, no ML models)
      $PYTHON -m pytest unstructured_prop/tests/unit -n auto -m "not slow" -q 2>&1 | tee "$OUTDIR/test_output.txt"

      echo ""
      echo "Results saved to $OUTDIR/"
      ls -la "$OUTDIR/"

  # --- Benchmark: side-by-side two branches ---
  - path: /home/azureuser/bench/bench_ab.sh
    owner: azureuser:azureuser
    permissions: "0755"
    defer: true
    content: |
      #!/usr/bin/env bash
      set -euo pipefail
      BASE="${1:?Usage: bench_ab.sh <base-branch> <opt-branch>}"
      OPT="${2:?Usage: bench_ab.sh <base-branch> <opt-branch>}"

      echo "=== A/B comparison: $BASE vs $OPT ==="
      bash ~/bench/bench_compare.sh "$BASE"
      bash ~/bench/bench_compare.sh "$OPT"

      echo ""
      echo "Compare results in ~/results/"
      ls ~/results/

  # --- Post-provision setup (run manually after ssh -A) ---
  - path: /home/azureuser/setup.sh
    owner: azureuser:azureuser
    permissions: "0755"
    defer: true
    content: |
      #!/usr/bin/env bash
      set -euo pipefail
      export PATH="$HOME/.local/bin:$PATH"

      echo "=== Cloning core-product ==="
      git clone git@github.com:Unstructured-IO/core-product.git ~/core-product
      cd ~/core-product

      echo "=== Creating shared venv ==="
      uv venv --python 3.12
      export VIRTUAL_ENV=$PWD/.venv

      echo "=== Installing dependencies ==="
      # Requires UV_INDEX_UNSTRUCTURED_USERNAME and UV_INDEX_UNSTRUCTURED_PASSWORD
      # to be set for private PyPI access
      make install

      echo "=== Creating results directory ==="
      mkdir -p ~/results

      echo "=== Verifying installation ==="
      .venv/bin/python -c 'from unstructured.partition.pdf import partition_pdf; print("OK")'

      echo "=== Done ==="

runcmd:
  - wget -q https://github.com/sharkdp/hyperfine/releases/download/v1.19.0/hyperfine_1.19.0_amd64.deb -O /tmp/hyperfine.deb
  - dpkg -i /tmp/hyperfine.deb
  # Install uv (phase 1 -- no git clone, that requires SSH agent forwarding)
  - su - azureuser -c 'curl -LsSf https://astral.sh/uv/install.sh | sh'
