vh
888ba6a714
Host service (runs as llmuser, owns /opt/fluxgym + GPU access) that runs
sd-scripts SDXL LoRA training on demand for arbo — the infra-ops half of the
in-arbo LoRA training Phase 1 ownership split (vh/arbo
docs/contracts/in-arbo-lora-training-phase1.contract.md §4.1/§2).
- Fixed-invocation only (INV-T7): bounded params -> one sd-scripts command
shape; every param range/allowlist/path-containment checked before spawn;
bad request = 422, never a silent downgrade. 14 unit tests green.
- Thin supervisor: never imports torch; subprocesses the fluxgym venv's
accelerate. 1-job-at-a-time (arbo lease is the serializer, 409 is backstop).
Durable job records + boot reconciliation (§4.3).
- API: POST /train, GET /train/{id}[/log], POST /train/{id}/cancel,
GET /gpu-status (per-device VRAM + tts_on_3090 co-OOM signal), GET /healthz.
- Wire-shape (§7 resolved with comfy-dev): shared /worktank/arbo/train handoff
(group arbotrain, setgid 2770); worker binds 0.0.0.0:8203, arbo reaches via
host.docker.internal:host-gateway (reachability proven on 172.20.0.1:8203);
device-aware TTS steering via /gpu-status.
Deployed to irv-ml1 via playbooks/deploy-lora-training-worker.yaml (elway,
idempotent); systemd unit active; /healthz + /gpu-status verified live.
93 lines
4.5 KiB
Python
93 lines
4.5 KiB
Python
"""Static configuration for the LoRA training worker.
|
|
|
|
Everything load-bearing is an explicit module constant here (explicit-over-implicit): the
|
|
paths the worker is allowed to touch, the fixed sd-scripts invocation surface, the tier
|
|
table, and the SDXL-LoRA hyperparameters. Nothing about a training run is free-form — the
|
|
worker only ever builds ONE command shape (INV-T7), and every knob that shapes it is
|
|
visible in this file.
|
|
|
|
The worker runs as `llmuser` on irv-ml1. It never imports torch / sd-scripts; it only
|
|
*subprocesses* the fluxgym venv's `accelerate launch`. So this process stays tiny and the
|
|
fluxgym venv stays pristine.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
# ---- Network ---------------------------------------------------------------------------
|
|
# 0.0.0.0 (not 127.0.0.1): arbo runs CONTAINERIZED on the traefik-net bridge and reaches
|
|
# the host worker via host.docker.internal:host-gateway — host loopback is unreachable from
|
|
# that bridge. irv-ml1 is WireGuard-only + ACL'd, so 0.0.0.0 exposure is bounded to the tunnel.
|
|
HOST = os.environ.get("LORA_WORKER_HOST", "0.0.0.0")
|
|
PORT = int(os.environ.get("LORA_WORKER_PORT", "8203"))
|
|
|
|
# ---- Fluxgym / sd-scripts (the ONLY thing the worker executes) -------------------------
|
|
FLUXGYM_ROOT = Path("/opt/fluxgym")
|
|
FLUXGYM_VENV = FLUXGYM_ROOT / ".venv"
|
|
ACCELERATE_BIN = FLUXGYM_VENV / "bin" / "accelerate"
|
|
SD_SCRIPTS_DIR = FLUXGYM_ROOT / "sd-scripts"
|
|
SDXL_TRAIN_SCRIPT = SD_SCRIPTS_DIR / "sdxl_train_network.py"
|
|
|
|
# ---- Allowed path roots (INV-T7 path safety) -------------------------------------------
|
|
# dataset_dir + output_dir MUST live under the shared handoff root (the group-shared,
|
|
# setgid /worktank/arbo/train that both arbo-container-uid and llmuser can rw). base_model
|
|
# must live under a known model root. Anything else → rejected before a process is spawned.
|
|
HANDOFF_ROOT = Path(os.environ.get("LORA_WORKER_HANDOFF_ROOT", "/worktank/arbo/train"))
|
|
ALLOWED_MODEL_ROOTS = tuple(
|
|
Path(p)
|
|
for p in os.environ.get(
|
|
"LORA_WORKER_MODEL_ROOTS",
|
|
"/opt/fluxgym/models:/worktank/comfyui:/worktank/models:/worktank/arbo",
|
|
).split(":")
|
|
if p
|
|
)
|
|
|
|
# ---- Worker state + logs (survives a worker restart for boot reconciliation, §4.3) -----
|
|
STATE_DIR = Path(os.environ.get("LORA_WORKER_STATE_DIR", "/opt/lora-training-worker/state"))
|
|
LOG_DIR = Path(os.environ.get("LORA_WORKER_LOG_DIR", "/opt/lora-training-worker/logs"))
|
|
JOBS_FILE = STATE_DIR / "jobs.json"
|
|
MAX_RETAINED_JOBS = 50 # keep terminal jobs queryable for arbo's status proxy
|
|
|
|
# ---- Tier table (§4.6) -----------------------------------------------------------------
|
|
# tier -> (max_train_steps, network_dim, resolution). Quality (1024) is A6000-only; the
|
|
# device-fit guard lives in invocation.build_command (quality on the 3090 → rejected).
|
|
TIERS = {
|
|
"fast": {"steps": 400, "dim": 16, "resolution": 768},
|
|
"balanced": {"steps": 1500, "dim": 32, "resolution": 768},
|
|
"quality": {"steps": 3000, "dim": 32, "resolution": 1024},
|
|
}
|
|
|
|
# ---- Device model (CUDA_DEVICE_ORDER=PCI_BUS_ID) ---------------------------------------
|
|
# Under PCI_BUS_ID (which the recipe pins), index 0 = RTX 3090, index 1 = RTX A6000.
|
|
# Quality (1024) will not fit the 3090's ~24GB LoRA envelope → allowed on the A6000 only.
|
|
DEVICE_3090 = 0
|
|
DEVICE_A6000 = 1
|
|
QUALITY_ONLY_DEVICES = (DEVICE_A6000,)
|
|
|
|
# ---- SDXL-LoRA hyperparameters (agent-discretion defaults; confirm vs Sindra) ----------
|
|
# The contract (§4.6) pins the flag SET + tier->(steps,dim,res) but NOT lr/scheduler/batch.
|
|
# These are standard, conservative SDXL-LoRA values, surfaced here so they're one-line
|
|
# auditable + tunable. Flagged to comfy-dev to cross-check against the proven Sindra runs.
|
|
SDXL_HPARAMS = {
|
|
"learning_rate": "1e-4",
|
|
"lr_scheduler": "cosine",
|
|
"lr_warmup_steps": "0",
|
|
"train_batch_size_lean": "1", # 3090
|
|
"train_batch_size_full": "2", # A6000
|
|
"network_alpha_ratio": 1.0, # alpha = dim * ratio
|
|
"min_snr_gamma": "5",
|
|
"noise_offset": "0.1",
|
|
"save_precision": "bf16",
|
|
"mixed_precision": "bf16",
|
|
"max_data_loader_n_workers": "2",
|
|
}
|
|
|
|
# ---- TTS-liveness heuristic (GET /gpu-status → tts_on_3090) -----------------------------
|
|
# arbo's scheduler steers a lean train off the 3090 when TTS is live there. We report the
|
|
# raw per-device VRAM + a boolean derived from (a) a compute-app on the 3090 whose cmdline
|
|
# matches a TTS marker, else (b) a used-MB floor fallback.
|
|
TTS_CMDLINE_MARKERS = ("chatterbox", "omnivoice", "tts_server", "chatterbox_fast")
|
|
TTS_3090_USED_MB_FLOOR = int(os.environ.get("LORA_WORKER_TTS_FLOOR_MB", "2000"))
|