# Deploy ACE-Step 1.5 (Apache 2.0 music generation foundation model)
# to irv-ml1.
#
# Builds the image locally from ace-step/ACE-Step via docker buildx
# git URL context. ~10-15 min cold build (CUDA 12.6 runtime + torch +
# transformers + spacy + audio deps). First /generate triggers the
# model download (~5-10 GB) into the bind-mounted HF cache + warmup.
#
# Pre-req (user runs once):
#   ssh -t irv-ml1 'sudo mkdir -p /worktank/ace-step/{checkpoints,outputs,logs,hf_cache} \
#       /opt/docker/compose/ace-step && \
#       sudo chown -R lkraven:lkraven /worktank/ace-step /opt/docker/compose/ace-step'
#
# Usage:
#   scripts/elway irv-ml1 --playbook playbooks/deploy-ace-step.yaml
#
# Idempotent — every step is creates-/when-gated; rerun is safe.

vars:
  compose_dir: /opt/docker/compose/ace-step
  worktank_root: /worktank/ace-step
  host_port: "8210"

steps:
  # ── sanity checks (dirs were created by user-side sudo prep) ────────

  - name: Verify /worktank/ace-step exists and is writable
    shell: test -w {{ worktank_root }}
    changed_when: "false"

  - name: Verify compose dir exists and is writable
    shell: test -w {{ compose_dir }}
    changed_when: "false"

  # ── deploy compose + env ────────────────────────────────────────────

  - name: Upload compose.yaml
    upload:
      src: stacks/ace-step/compose.yaml
      dest: "{{ compose_dir }}/compose.yaml"
      mode: "0644"

  - name: Upload Dockerfile (patched for cu126 torch resolution)
    upload:
      src: stacks/ace-step/Dockerfile
      dest: "{{ compose_dir }}/Dockerfile"
      mode: "0644"

  - name: Upload patched infer-api.py (fixes upstream 18-arg tuple bug)
    upload:
      src: stacks/ace-step/infer-api.py
      dest: "{{ compose_dir }}/infer-api.py"
      mode: "0644"

  - name: Seed .env from template (only if absent)
    upload:
      src: stacks/ace-step/.env.example
      dest: "{{ compose_dir }}/.env"
      mode: "0644"
    when: "[ ! -f {{ compose_dir }}/.env ]"

  # ── build + bring up ────────────────────────────────────────────────

  - name: docker compose build (~10-15 min first time; cached after)
    shell: |
      set -o pipefail
      cd {{ compose_dir }} && docker compose build 2>&1 \
        | grep -vE '^#[0-9]+ |^ => |^=> |Collecting|Downloading|Requirement|Using cached|Installing collected|Successfully (installed|built)|━'

  - name: docker compose up -d
    shell: cd {{ compose_dir }} && docker compose up -d

  - name: Wait for /health to respond (allow ~10 min for first model download + warmup)
    shell: |
      for i in $(seq 1 120); do
        curl -sf -o /dev/null --max-time 3 http://localhost:{{ host_port }}/health && exit 0
        sleep 5
      done
      exit 1
    changed_when: "false"

verify:
  - name: /health returns 200
    shell: curl -sf -o /dev/null http://localhost:{{ host_port }}/health
    changed_when: "false"

  - name: Container is running
    shell: docker inspect ace-step --format '{{.State.Status}}' | grep -q running
    changed_when: "false"