# Deploy the LoRA training worker to irv-ml1 (arbo in-arbo LoRA training Phase 1, §4.1).
#
# Two-part deploy (elway upload is single-file, so code lands via rsync first):
#   1. Stage the code (run from the eshpfi-management repo root, as infra-ops):
#        rsync -a --delete \
#          --exclude .venv --exclude __pycache__ --exclude state --exclude logs \
#          services/lora-training-worker/ \
#          infra-ops@10.100.79.3:/tmp/lora-training-worker-stage/
#   2. Run this playbook (privileged on-box install + health-gate):
#        scripts/elway irv-ml1 --playbook playbooks/deploy-lora-training-worker.yaml
#
# Idempotent: a second run shows mostly ok/skipped. Runs as infra-ops (NOPASSWD sudo on irv-ml1).

vars:
  stage_dir: /tmp/lora-training-worker-stage
  install_dir: /opt/lora-training-worker
  handoff_dir: /worktank/arbo/train
  loras_publish_dir: /storetank/arbo/models/loras/trained
  worker_user: llmuser
  arbo_user: lkraven          # arbo container runs as uid 1000 = host lkraven
  group: arbotrain
  port: "8203"

steps:
  - name: Create the shared handoff group
    shell: getent group {{ group }} >/dev/null || groupadd {{ group }}
    sudo: true
    changed_when: "false"   # groupadd-or-noop; report ok either way

  - name: Add the worker user (llmuser) to the handoff group
    shell: id -nG {{ worker_user }} | tr ' ' '\n' | grep -qx {{ group }} || usermod -aG {{ group }} {{ worker_user }}
    sudo: true
    when: "! id -nG {{ worker_user }} | tr ' ' '\\n' | grep -qx {{ group }}"

  - name: Add the arbo-container user (lkraven) to the handoff group
    shell: usermod -aG {{ group }} {{ arbo_user }}
    sudo: true
    when: "! id -nG {{ arbo_user }} | tr ' ' '\\n' | grep -qx {{ group }}"

  - name: Create the shared handoff dir (group-owned, setgid 2770)
    shell: mkdir -p {{ handoff_dir }}
    sudo: true
    creates: "{{ handoff_dir }}"

  - name: Set handoff dir group + setgid perms
    shell: chgrp {{ group }} {{ handoff_dir }} && chmod 2770 {{ handoff_dir }}
    sudo: true
    changed_when: "false"

  - name: Create the Phase-2 LoRA publish dir (ComfyUI loras/trained, group-writable)
    # 2775 (not 2770): world-readable + traversable so ComfyUI (uid 1025 comfytoo) can list +
    # load; group arbotrain + group-WRITE so the worker (llmuser) can publish into it. setgid
    # propagates the group to per-train subdirs (the Phase-1 group-write lesson).
    shell: mkdir -p {{ loras_publish_dir }} && chgrp {{ group }} {{ loras_publish_dir }} && chmod 2775 {{ loras_publish_dir }}
    sudo: true
    changed_when: "false"

  - name: Create the install dir owned by the worker user
    shell: mkdir -p {{ install_dir }} && chown {{ worker_user }}:{{ worker_user }} {{ install_dir }}
    sudo: true
    creates: "{{ install_dir }}"

  - name: Sync staged code into the install dir (worker-owned)
    shell: >
      rsync -a --delete
      --exclude .venv --exclude __pycache__ --exclude state --exclude logs
      {{ stage_dir }}/ {{ install_dir }}/
      && chown -R {{ worker_user }}:{{ worker_user }} {{ install_dir }}
    sudo: true

  - name: Ensure state + logs dirs exist (worker-writable)
    shell: mkdir -p {{ install_dir }}/state {{ install_dir }}/logs && chown {{ worker_user }}:{{ worker_user }} {{ install_dir }}/state {{ install_dir }}/logs
    sudo: true
    creates: "{{ install_dir }}/logs"

  - name: Build the worker venv + install deps (as llmuser; prefer uv, fall back to python3 -m venv)
    shell: >
      sudo -u {{ worker_user }} bash -lc '
      cd {{ install_dir }} &&
      if command -v uv >/dev/null 2>&1; then
        uv venv .venv && uv pip install --python .venv/bin/python . ;
      else
        python3 -m venv .venv && .venv/bin/pip install -q --upgrade pip && .venv/bin/pip install -q . ;
      fi'
    sudo: true
    creates: "{{ install_dir }}/.venv/bin/uvicorn"

  - name: Install the systemd unit
    upload:
      src: services/lora-training-worker/lora-training-worker.service
      dest: /etc/systemd/system/lora-training-worker.service
      mode: "0644"
    sudo: true

  - name: Reload systemd + enable the worker
    shell: systemctl daemon-reload && systemctl enable lora-training-worker.service
    sudo: true
    changed_when: "false"

  - name: Restart the worker to pick up the synced code
    shell: systemctl restart lora-training-worker.service
    sudo: true
    changed_when: "false"

  - name: Give the service a moment to bind
    shell: sleep 3
    changed_when: "false"

verify:
  - name: Worker health endpoint responds ok
    shell: curl -fsS http://127.0.0.1:{{ port }}/healthz
    changed_when: "false"

  - name: gpu-status reports both devices
    shell: curl -fsS http://127.0.0.1:{{ port }}/gpu-status | grep -q '"index"'
    changed_when: "false"

  - name: Service is enabled + active
    shell: systemctl is-active lora-training-worker.service
    sudo: true
    changed_when: "false"