Files
vh b617a8b674 feat(lora-worker): add optional train_id to POST /train (explicit publish-path namespace)
comfy-dev's explicit-over-implicit call: arbo now sends train_id, so the
worker no longer derives the loras/trained/{train_id}/ namespace from
output_dir.parent (which coupled it to arbo's handoff layout). train_id is
optional + path-safe-validated; when present it wins, else the path
derivation remains as the fallback. Wired through TrainRequest ->
validate_request -> published_relative_path -> _publish_lora. 18 tests green.
2026-07-07 01:49:24 -07:00
..

LoRA training worker

A small host service on irv-ml1 that runs sd-scripts SDXL LoRA training on demand for arbo. It is the infra-ops half of the "in-arbo LoRA training, Phase 1" split:

  • infra-ops (this service) owns the training worker — the ONLY thing that invokes accelerate launch. Runs as llmuser (the only user that can exec /opt/fluxgym/.venv and owns the GPU-training surface).
  • comfy-dev (arbo) owns dispatch, the GPU scheduler + serving-pause, the training-job model + durable history + status proxy, and the Studio UI.

Canonical contract: vh/arbodocs/contracts/in-arbo-lora-training-phase1.contract.md (§4.1 worker API, §2 ownership, §4.6 recipe, §7 open items). This service implements §4.1.

Design posture

  • Fixed invocation only (INV-T7). arbo sends bounded parameters; the worker builds exactly one command shape (sdxl_train_network.py with the §4.6 lean flag set). There is no path from a request to a free-form argument — see worker/invocation.py. Every param is range/allowlist/path-containment checked before a process is spawned; a bad request is a 422, never a silent downgrade.
  • Thin supervisor. The worker never imports torch/sd-scripts. It subprocesses the fluxgym venv's accelerate, so this process stays tiny and the fluxgym venv stays pristine.
  • One job at a time (§4.1). arbo's lease is the real serializer (INV-T2); the worker's 409 is a backstop.
  • Durable (§4.3 boot reconciliation). Job records persist to JSON; on restart a non-terminal job whose OS process is gone is marked failed so arbo never sees a phantom training after a worker blip.

Wire-shape (the §7 open items, resolved with comfy-dev)

  1. Mount / handoff. Shared host dir /worktank/arbo/train/{train_id}/, group arbotrain (setgid 2770), members lkraven (arbo container uid = 1000) + llmuser (worker, uid 1001) → both rw, new files inherit the group. arbo adds one bind to its compose: - /worktank/arbo/train:/worktank/arbo/train. arbo (container) writes dataset_dir + the kohya folder; the worker (llmuser) reads it and writes output_dir; arbo reads the .safetensors for download. Never via /data (arbo-only volume).
  2. Bind. Worker binds 0.0.0.0:8203. arbo (containerized, traefik-net bridge — host loopback unreachable) reaches it via host.docker.internal:8203 with extra_hosts: ["host.docker.internal:host-gateway"] in arbo's compose.
  3. TTS-liveness signal. GET /gpu-status reports per-device VRAM (index 0 = 3090, index 1 = A6000 under CUDA_DEVICE_ORDER=PCI_BUS_ID) + tts_on_3090. arbo polls it before assigning device_index and steers a lean train to the A6000 when TTS is live on the 3090.

API (§4.1)

Method + path Purpose
POST /train dispatch a train (409 if busy, 422 on invalid params) → {worker_job_id}
GET /train/{id} {status, step, total_steps, loss, eta_s, lora_path?, error?}
GET /train/{id}/log tail of the run log
POST /train/{id}/cancel best-effort kill → cancelled
GET /gpu-status per-device VRAM + tts_on_3090
GET /healthz liveness + active job id

POST /train body: {dataset_dir, base_model_path, output_dir, output_name, trigger, subject_class, repeats, tier, device_index, seed}. tier ∈ {fast, balanced, quality} → (steps, dim, res) = (400,16,768) / (1500,32,768) / (3000,32,1024). quality (1024) requires device_index=1 (A6000) — quality on the 3090 is a 422.

Recipe hyperparameters

The contract §4.6 pins the flag set + tier→(steps,dim,res) but not lr/scheduler/batch. Those live as auditable constants in worker/config.py::SDXL_HPARAMS (lr 1e-4, cosine, adamw8bit, min_snr_gamma 5, noise_offset 0.1, batch 1/2 lean/full). These are agent-discretion defaults pending a cross-check against the proven Sindra runs — adjust in one place if Sindra used different values.

Deploy

Privileged host setup (group, dirs, /opt install, venv, systemd) is delegated to infra-ops via the elway playbook (needs NOPASSWD sudo on irv-ml1, which infra-ops has):

# from the eshpfi-management repo root
scripts/elway irv-ml1 --playbook playbooks/deploy-lora-training-worker.yaml

The playbook is idempotent: creates the arbotrain group + /worktank/arbo/train (2770 setgid), rsyncs the code to /opt/lora-training-worker, builds the worker venv with uv, installs + enables the systemd unit, and health-gates on GET /healthz.

Test

cd services/lora-training-worker
uv run --with pytest python -m pytest -q     # invocation-builder (INV-T7) unit tests

Status

  • Worker code + INV-T7 invocation builder + tests (14 green).
  • Box deploy (elway playbook) + live /healthz + /gpu-status smoke.
  • First real end-to-end train once arbo dispatches a studio dataset (arbo backend built + dormant; waits on this worker).
  • Sindra hyperparameter cross-check.