356 lines
13 KiB
Python
356 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""Dispatcher daemon: picks queued jobs/stitches and runs them on available workers.
|
|
|
|
Env:
|
|
COSMA_QC_DB : SQLite path (default /var/lib/cosma-qc/jobs.db)
|
|
COSMA_QC_WORKERS : JSON list of workers [{host, ssh_alias, gpu, vram_mib,
|
|
frames_dir, lingbot_path}]
|
|
COSMA_QC_FPS : extraction fps (default 3)
|
|
COSMA_QC_IMG_H : image height (default 294)
|
|
COSMA_QC_IMG_W : image width (default 518)
|
|
|
|
Jobs lifecycle:
|
|
queued → extracting → running → done → [triggers per_auv stitch]
|
|
↘ error
|
|
Stitch lifecycle:
|
|
queued → running → done → [triggers cross_auv stitch if all per_auv done]
|
|
↘ error
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import shlex
|
|
import sqlite3
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from contextlib import closing
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
|
|
def _now_iso() -> str:
|
|
return datetime.now(timezone.utc).isoformat(timespec="seconds")
|
|
|
|
DB_PATH = Path(os.environ.get("COSMA_QC_DB", "/var/lib/cosma-qc/jobs.db"))
|
|
FPS = int(os.environ.get("COSMA_QC_FPS", "3"))
|
|
IMG_H = int(os.environ.get("COSMA_QC_IMG_H", "294"))
|
|
IMG_W = int(os.environ.get("COSMA_QC_IMG_W", "518"))
|
|
POLL_S = int(os.environ.get("COSMA_QC_POLL_S", "4"))
|
|
STITCH_SCRIPT = Path(__file__).parent / "stitch.py"
|
|
|
|
DEFAULT_WORKERS = [
|
|
{
|
|
"host": "192.168.0.87", "ssh_alias": "gpu", "gpu": "RTX 3060 12GB",
|
|
"vram_mib": 11913,
|
|
"frames_dir": "/home/floppyrj45/cosma-qc-frames",
|
|
"lingbot_path": "/home/floppyrj45/ai-video/lingbot-map",
|
|
"viser_port_base": 8100,
|
|
},
|
|
{
|
|
"host": "192.168.0.84", "ssh_alias": "cosma-vm", "gpu": "RTX 3090 24GB",
|
|
"vram_mib": 24576,
|
|
"frames_dir": "/home/floppyrj45/cosma-qc-frames",
|
|
"lingbot_path": "/home/floppyrj45/ai-video/lingbot-map",
|
|
"viser_port_base": 8100,
|
|
},
|
|
]
|
|
WORKERS = json.loads(os.environ.get("COSMA_QC_WORKERS", json.dumps(DEFAULT_WORKERS)))
|
|
|
|
|
|
def db() -> sqlite3.Connection:
|
|
conn = sqlite3.connect(DB_PATH, isolation_level=None)
|
|
conn.execute("PRAGMA journal_mode=WAL")
|
|
conn.row_factory = sqlite3.Row
|
|
return conn
|
|
|
|
|
|
def ssh(alias: str, cmd: str, timeout: int = 30) -> tuple[int, str, str]:
|
|
p = subprocess.run(
|
|
["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=5", alias, cmd],
|
|
capture_output=True, text=True, timeout=timeout,
|
|
)
|
|
return p.returncode, p.stdout, p.stderr
|
|
|
|
|
|
def worker_free_vram_mib(worker: dict) -> int:
|
|
rc, out, _ = ssh(worker["ssh_alias"], "nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits")
|
|
try:
|
|
return int(out.strip().splitlines()[0]) if rc == 0 else 0
|
|
except Exception:
|
|
return 0
|
|
|
|
|
|
def pick_worker(estimated_vram_mib: int) -> dict | None:
|
|
best = None
|
|
for w in WORKERS:
|
|
free = worker_free_vram_mib(w)
|
|
if free >= estimated_vram_mib and (best is None or free > best[0]):
|
|
best = (free, w)
|
|
return best[1] if best else None
|
|
|
|
|
|
def estimate_vram_mib(frame_count: int) -> int:
|
|
return int(3500 + 13 * frame_count)
|
|
|
|
|
|
def set_status(job_id: int, **fields):
|
|
keys = list(fields.keys())
|
|
vals = [fields[k] for k in keys]
|
|
q = "UPDATE jobs SET " + ", ".join(f"{k}=?" for k in keys) + " WHERE id=?"
|
|
with closing(db()) as conn:
|
|
conn.execute(q, (*vals, job_id))
|
|
|
|
|
|
def set_stitch_status(stitch_id: int, **fields):
|
|
keys = list(fields.keys())
|
|
vals = [fields[k] for k in keys]
|
|
q = "UPDATE stitches SET " + ", ".join(f"{k}=?" for k in keys) + " WHERE id=?"
|
|
with closing(db()) as conn:
|
|
conn.execute(q, (*vals, stitch_id))
|
|
|
|
|
|
def count_frames(worker: dict, frames_dir: str) -> int:
|
|
rc, out, _ = ssh(worker["ssh_alias"], f"ls {shlex.quote(frames_dir)} 2>/dev/null | wc -l")
|
|
try:
|
|
return int(out.strip()) if rc == 0 else 0
|
|
except Exception:
|
|
return 0
|
|
|
|
|
|
def do_extract(job: sqlite3.Row, worker: dict) -> str:
|
|
videos = json.loads(job["video_paths"])
|
|
frames_dir = f"{worker['frames_dir']}/job_{job['id']}"
|
|
ssh(worker["ssh_alias"], f"mkdir -p {shlex.quote(frames_dir)}")
|
|
idx = 0
|
|
for v in videos:
|
|
vf = f"fps={FPS},scale={IMG_W}:{IMG_H}"
|
|
pattern = f"{frames_dir}/frame_%06d.jpg"
|
|
cmd = (
|
|
f"ffmpeg -hide_banner -loglevel error -i {shlex.quote(v)} "
|
|
f"-vf {shlex.quote(vf)} -start_number {idx} -q:v 4 "
|
|
f"{shlex.quote(pattern)}"
|
|
)
|
|
rc, _, err = ssh(worker["ssh_alias"], cmd, timeout=3600)
|
|
if rc != 0:
|
|
raise RuntimeError(f"ffmpeg failed on {v}: {err[:200]}")
|
|
idx = count_frames(worker, frames_dir)
|
|
set_status(job["id"], frame_count=idx)
|
|
return frames_dir
|
|
|
|
|
|
def do_reconstruct(job: sqlite3.Row, worker: dict, frames_dir: str) -> tuple[str, str, str]:
|
|
port = worker["viser_port_base"] + job["id"]
|
|
log = f"/tmp/cosma-qc-job-{job['id']}.log"
|
|
ckpt = f"{worker['lingbot_path']}/checkpoints/lingbot-map/lingbot-map-long.pt"
|
|
ply_path = f"{frames_dir}/reconstruction.ply"
|
|
cmd = (
|
|
f"cd {shlex.quote(worker['lingbot_path'])} && source .venv/bin/activate && "
|
|
f"python3 demo.py --model_path {shlex.quote(ckpt)} "
|
|
f"--image_folder {shlex.quote(frames_dir)} --port {port} "
|
|
f"--use_sdpa --mode windowed --window_size 16 --overlap_size 2 --offload_to_cpu "
|
|
f"--save_ply {shlex.quote(ply_path)} "
|
|
f"> {log} 2>&1"
|
|
)
|
|
rc, _, err = ssh(worker["ssh_alias"], cmd, timeout=3 * 3600)
|
|
if rc != 0:
|
|
tail = ssh(worker["ssh_alias"], f"tail -30 {log}")[1]
|
|
raise RuntimeError(f"demo.py failed: {err[:200]}\n---\n{tail[:800]}")
|
|
viser_url = f"http://{worker['host']}:{port}"
|
|
return viser_url, log, ply_path
|
|
|
|
|
|
def _maybe_create_per_auv_stitch(job_id: int):
|
|
with closing(db()) as conn:
|
|
job = conn.execute("SELECT * FROM jobs WHERE id=?", (job_id,)).fetchone()
|
|
if not job:
|
|
return
|
|
acq_id, auv = job["acquisition_id"], job["auv"]
|
|
total = conn.execute(
|
|
"SELECT COUNT(*) FROM jobs WHERE acquisition_id=? AND auv=?", (acq_id, auv)
|
|
).fetchone()[0]
|
|
done = conn.execute(
|
|
"SELECT COUNT(*) FROM jobs WHERE acquisition_id=? AND auv=? AND status='done'", (acq_id, auv)
|
|
).fetchone()[0]
|
|
if total == 0 or done < total:
|
|
return
|
|
existing = conn.execute(
|
|
"SELECT id FROM stitches WHERE acquisition_id=? AND level='per_auv' AND auv=?", (acq_id, auv)
|
|
).fetchone()
|
|
if existing:
|
|
return
|
|
job_ids = [r["id"] for r in conn.execute(
|
|
"SELECT id FROM jobs WHERE acquisition_id=? AND auv=?", (acq_id, auv)
|
|
).fetchall()]
|
|
conn.execute(
|
|
"INSERT INTO stitches (acquisition_id, level, auv, input_job_ids) VALUES (?,?,?,?)",
|
|
(acq_id, "per_auv", auv, json.dumps(job_ids))
|
|
)
|
|
print(f" → Stitch per_auv créé pour {auv} acq#{acq_id}")
|
|
|
|
|
|
def _maybe_create_cross_auv_stitch(stitch_id: int):
|
|
with closing(db()) as conn:
|
|
st = conn.execute("SELECT * FROM stitches WHERE id=?", (stitch_id,)).fetchone()
|
|
if not st:
|
|
return
|
|
acq_id = st["acquisition_id"]
|
|
n_auvs = conn.execute(
|
|
"SELECT COUNT(DISTINCT auv) FROM jobs WHERE acquisition_id=?", (acq_id,)
|
|
).fetchone()[0]
|
|
if n_auvs < 2:
|
|
return
|
|
total_per_auv = conn.execute(
|
|
"SELECT COUNT(*) FROM stitches WHERE acquisition_id=? AND level='per_auv'", (acq_id,)
|
|
).fetchone()[0]
|
|
done_per_auv = conn.execute(
|
|
"SELECT COUNT(*) FROM stitches WHERE acquisition_id=? AND level='per_auv' AND status='done'", (acq_id,)
|
|
).fetchone()[0]
|
|
if total_per_auv == 0 or done_per_auv < n_auvs:
|
|
return
|
|
existing = conn.execute(
|
|
"SELECT id FROM stitches WHERE acquisition_id=? AND level='cross_auv'", (acq_id,)
|
|
).fetchone()
|
|
if existing:
|
|
return
|
|
stitch_ids = [r["id"] for r in conn.execute(
|
|
"SELECT id FROM stitches WHERE acquisition_id=? AND level='per_auv'", (acq_id,)
|
|
).fetchall()]
|
|
conn.execute(
|
|
"INSERT INTO stitches (acquisition_id, level, input_stitch_ids, input_job_ids) VALUES (?,?,?,?)",
|
|
(acq_id, "cross_auv", json.dumps(stitch_ids), "[]")
|
|
)
|
|
print(f" → Stitch cross_auv créé pour acq#{acq_id}")
|
|
|
|
|
|
def deploy_stitch_script(worker: dict):
|
|
subprocess.run(
|
|
["scp", str(STITCH_SCRIPT), f"{worker['ssh_alias']}:/tmp/cosma-stitch.py"],
|
|
capture_output=True, timeout=30
|
|
)
|
|
|
|
|
|
def run_one_stitch(stitch: sqlite3.Row):
|
|
stitch_id = stitch["id"]
|
|
worker = pick_worker(2000)
|
|
if not worker:
|
|
worker = WORKERS[0]
|
|
|
|
with closing(db()) as conn:
|
|
if stitch["level"] == "per_auv":
|
|
job_ids = json.loads(stitch["input_job_ids"] or "[]")
|
|
if job_ids:
|
|
rows = conn.execute(
|
|
f"SELECT ply_path FROM jobs WHERE id IN ({','.join('?'*len(job_ids))})",
|
|
job_ids
|
|
).fetchall()
|
|
else:
|
|
rows = []
|
|
ply_paths = [r["ply_path"] for r in rows if r["ply_path"]]
|
|
else:
|
|
stitch_ids = json.loads(stitch["input_stitch_ids"] or "[]")
|
|
if stitch_ids:
|
|
rows = conn.execute(
|
|
f"SELECT output_ply FROM stitches WHERE id IN ({','.join('?'*len(stitch_ids))})",
|
|
stitch_ids
|
|
).fetchall()
|
|
else:
|
|
rows = []
|
|
ply_paths = [r["output_ply"] for r in rows if r["output_ply"]]
|
|
|
|
if len(ply_paths) < 2:
|
|
set_stitch_status(stitch_id, status="error",
|
|
error=f"Pas assez de PLY disponibles ({len(ply_paths)})",
|
|
finished_at=_now_iso())
|
|
return
|
|
|
|
out_ply = f"{worker['frames_dir']}/stitch_{stitch_id}.ply"
|
|
deploy_stitch_script(worker)
|
|
|
|
cmd = (
|
|
f"source {shlex.quote(worker['lingbot_path'])}/.venv/bin/activate && "
|
|
f"python3 /tmp/cosma-stitch.py {shlex.quote(out_ply)} "
|
|
+ " ".join(shlex.quote(p) for p in ply_paths)
|
|
+ f" > /tmp/cosma-stitch-{stitch_id}.log 2>&1"
|
|
)
|
|
|
|
set_stitch_status(stitch_id, status="running", worker_host=worker["host"], started_at=_now_iso())
|
|
try:
|
|
rc, _, err = ssh(worker["ssh_alias"], cmd, timeout=4 * 3600)
|
|
except Exception as e:
|
|
set_stitch_status(stitch_id, status="error", error=str(e)[:500], finished_at=_now_iso())
|
|
return
|
|
|
|
if rc == 0:
|
|
set_stitch_status(stitch_id, status="done", output_ply=out_ply, finished_at=_now_iso())
|
|
_maybe_create_cross_auv_stitch(stitch_id)
|
|
else:
|
|
tail = ssh(worker["ssh_alias"], f"tail -20 /tmp/cosma-stitch-{stitch_id}.log")[1]
|
|
set_stitch_status(stitch_id, status="error",
|
|
error=f"{err[:200]}\n{tail[:600]}",
|
|
finished_at=_now_iso())
|
|
|
|
|
|
def run_one(job: sqlite3.Row) -> bool:
|
|
"""Returns True if a worker was picked and work started."""
|
|
job_id = job["id"]
|
|
estimated = estimate_vram_mib(job["frame_count"] or 400)
|
|
worker = pick_worker(estimated)
|
|
if not worker:
|
|
return False
|
|
set_status(job_id, status="extracting", worker_host=worker["host"], started_at=_now_iso())
|
|
try:
|
|
frames_dir = do_extract(job, worker)
|
|
frame_count = count_frames(worker, frames_dir)
|
|
set_status(job_id, frames_dir=frames_dir, frame_count=frame_count,
|
|
status="running", progress=0)
|
|
viser_url, log, ply_path = do_reconstruct(job, worker, frames_dir)
|
|
set_status(job_id, status="done", viser_url=viser_url, ply_path=ply_path,
|
|
progress=100, log_tail=log, finished_at=_now_iso())
|
|
_maybe_create_per_auv_stitch(job_id)
|
|
except Exception as e:
|
|
set_status(job_id, status="error", error=str(e)[:2000], finished_at=_now_iso())
|
|
return True
|
|
|
|
|
|
def pop_queued() -> sqlite3.Row | None:
|
|
with closing(db()) as conn:
|
|
return conn.execute(
|
|
"SELECT * FROM jobs WHERE status='queued' ORDER BY created_at LIMIT 1"
|
|
).fetchone()
|
|
|
|
|
|
def pop_queued_stitch() -> sqlite3.Row | None:
|
|
with closing(db()) as conn:
|
|
return conn.execute(
|
|
"SELECT * FROM stitches WHERE status='queued' ORDER BY created_at LIMIT 1"
|
|
).fetchone()
|
|
|
|
|
|
def main():
|
|
print(f"cosma-qc dispatcher · DB={DB_PATH} · workers={[w['host'] for w in WORKERS]}")
|
|
while True:
|
|
job = pop_queued()
|
|
if job:
|
|
print(f"→ job #{job['id']} ({job['auv']}/{job['gopro_serial']}/{job['segment_label']})")
|
|
if not run_one(job):
|
|
print(" ↳ pas de worker dispo, retry dans 30s")
|
|
time.sleep(30)
|
|
continue
|
|
stitch = pop_queued_stitch()
|
|
if stitch:
|
|
label = f"{stitch['level']} {stitch['auv'] or ''} acq#{stitch['acquisition_id']}"
|
|
print(f"→ stitch #{stitch['id']} ({label})")
|
|
run_one_stitch(stitch)
|
|
continue
|
|
time.sleep(POLL_S)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except KeyboardInterrupt:
|
|
sys.exit(0)
|