Jonathan
b1757b631c
The Floor bar used a file-count poller that jumps 0→~99% for small/fast albums (streamrip reports only 0/1). But yt-dlp and slskd already emit real byte-level progress via the on_progress callback the pipeline was discarding (lambda _pct: None). - Thread a throttled on_progress → Job.downloadProgress (write on a >=1% move or every 0.5s; clamped 0..1; write failures swallowed). Called synchronously from adapter.download on the pipeline thread, so it reuses the main conn safely. - The file-count poller stays as the FALLBACK: it skips writing while a `reports` flag is set (an adapter is driving real progress), and the flag is cleared per attempt so a fall-through to streamrip still gets the estimate. - Smooth bars for YouTube + Soulseek now; Qobuz/streamrip keeps the file-count estimate until its per-track callback is hooked later. worker 216 tests / 7-skip (on_progress write/throttle/clamp covered; existing pipeline tests now exercise the fake's on_progress end-to-end). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
331 lines
14 KiB
Python
331 lines
14 KiB
Python
import os
|
|
import shutil
|
|
import threading
|
|
import time
|
|
from dataclasses import replace
|
|
from typing import Sequence
|
|
|
|
import psycopg
|
|
|
|
from lyra_worker.adapters.base import SourceAdapter
|
|
from lyra_worker.confidence import score_confidence
|
|
from lyra_worker.library import album_dir, import_album, staging_dir
|
|
from lyra_worker.quality import quality_class
|
|
from lyra_worker.ranker import rank_candidates
|
|
from lyra_worker.types import Candidate, MBTarget
|
|
|
|
_AUDIO_EXT = {".flac", ".mp3", ".m4a", ".opus", ".ogg", ".aac", ".wav"}
|
|
|
|
|
|
def _count_staged_audio(staging: str) -> int:
|
|
"""Count audio files anywhere under the staging dir (streamrip nests them in a subfolder)."""
|
|
n = 0
|
|
for _root, _dirs, files in os.walk(staging):
|
|
for f in files:
|
|
if os.path.splitext(f)[1].lower() in _AUDIO_EXT:
|
|
n += 1
|
|
return n
|
|
|
|
|
|
def _poll_download_progress(job_id: str, staging: str, expected: int, stop: "threading.Event",
|
|
reports: "threading.Event") -> None:
|
|
"""Until `stop`, periodically write Job.downloadProgress = files-in-staging / expected
|
|
(capped 0.99) — the FALLBACK for adapters (streamrip/Qobuz) that report no byte-level
|
|
progress. Skips writing while `reports` is set, i.e. the in-flight adapter is driving
|
|
progress itself via on_progress (yt-dlp/slskd). Uses its own short-lived connection
|
|
(psycopg conns aren't shareable across threads). A missing DSN or expected<=0 just no-ops."""
|
|
dsn = os.environ.get("DATABASE_URL")
|
|
if not dsn or expected <= 0:
|
|
return
|
|
conn = psycopg.connect(dsn)
|
|
try:
|
|
while not stop.is_set():
|
|
if not reports.is_set(): # the adapter isn't reporting real progress → estimate
|
|
frac = min(_count_staged_audio(staging) / expected, 0.99)
|
|
with conn.cursor() as cur:
|
|
cur.execute('UPDATE "Job" SET "downloadProgress" = %s WHERE id = %s', (frac, job_id))
|
|
conn.commit()
|
|
stop.wait(1.5)
|
|
except Exception as e: # a progress poller must never affect the job
|
|
print(f"pipeline: download progress poller error: {e}", flush=True)
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def _set_state(conn: psycopg.Connection, job_id: str, state: str, stage: str) -> None:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
'UPDATE "Job" SET state = %s, "currentStage" = %s, "updatedAt" = now() WHERE id = %s',
|
|
(state, stage, job_id),
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def _set_download_progress(conn: psycopg.Connection, job_id: str, frac: float) -> None:
|
|
with conn.cursor() as cur:
|
|
cur.execute('UPDATE "Job" SET "downloadProgress" = %s WHERE id = %s', (frac, job_id))
|
|
conn.commit()
|
|
|
|
|
|
def _make_on_progress(conn: psycopg.Connection, job_id: str, reports: "threading.Event"):
|
|
"""Build the on_progress(pct) callback threaded into adapter.download. It marks the
|
|
adapter as reporting real byte-level progress (so the file-count poller yields) and
|
|
throttle-writes Job.downloadProgress. Called synchronously from adapter.download on the
|
|
pipeline's own thread, so it safely reuses `conn`."""
|
|
state = {"frac": -1.0, "t": 0.0}
|
|
|
|
def on_progress(pct: float) -> None:
|
|
reports.set()
|
|
frac = 0.0 if pct < 0 else 1.0 if pct > 1 else float(pct)
|
|
now = time.monotonic()
|
|
# throttle DB writes: on a >=1% move or every 0.5s, not once per received byte
|
|
if frac - state["frac"] >= 0.01 or (now - state["t"]) >= 0.5:
|
|
state["frac"] = frac
|
|
state["t"] = now
|
|
try:
|
|
_set_download_progress(conn, job_id, frac)
|
|
except Exception as e: # a progress write must never fail the download
|
|
print(f"pipeline: on_progress write failed: {e}", flush=True)
|
|
|
|
return on_progress
|
|
|
|
|
|
def _request_id(conn: psycopg.Connection, job_id: str) -> str:
|
|
with conn.cursor() as cur:
|
|
cur.execute('SELECT "requestId" FROM "Job" WHERE id = %s', (job_id,))
|
|
return cur.fetchone()[0]
|
|
|
|
|
|
def _fail(conn: psycopg.Connection, job_id: str, reason: str) -> None:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
'UPDATE "Job" SET state = \'needs_attention\', error = %s, "updatedAt" = now() WHERE id = %s',
|
|
(reason, job_id),
|
|
)
|
|
cur.execute(
|
|
'UPDATE "Request" SET status = \'needs_attention\' WHERE id = (SELECT "requestId" FROM "Job" WHERE id = %s)',
|
|
(job_id,),
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def _load_target(conn: psycopg.Connection, job_id: str) -> MBTarget:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
'SELECT artist, album FROM "Request" WHERE id = (SELECT "requestId" FROM "Job" WHERE id = %s)',
|
|
(job_id,),
|
|
)
|
|
artist, album = cur.fetchone()
|
|
return MBTarget(artist=artist, album=album)
|
|
|
|
|
|
def _already_in_library(conn: psycopg.Connection, target: MBTarget) -> bool:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
'SELECT 1 FROM "LibraryItem" WHERE artist = %s AND album = %s LIMIT 1',
|
|
(target.artist, target.album),
|
|
)
|
|
return cur.fetchone() is not None
|
|
|
|
|
|
def _is_upgrade_job(conn: psycopg.Connection, job_id: str) -> bool:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
'SELECT "monitoredReleaseId" IS NOT NULL FROM "Request" '
|
|
'WHERE id = (SELECT "requestId" FROM "Job" WHERE id = %s)',
|
|
(job_id,),
|
|
)
|
|
row = cur.fetchone()
|
|
return bool(row and row[0])
|
|
|
|
|
|
def _already_in_library_at_cutoff(conn: psycopg.Connection, target: MBTarget, cutoff: int) -> bool:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
'SELECT 1 FROM "LibraryItem" WHERE artist = %s AND album = %s '
|
|
'AND "qualityClass" >= %s LIMIT 1',
|
|
(target.artist, target.album, cutoff),
|
|
)
|
|
return cur.fetchone() is not None
|
|
|
|
|
|
def _library_quality(conn: psycopg.Connection, target: MBTarget) -> int | None:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
'SELECT "qualityClass" FROM "LibraryItem" WHERE artist = %s AND album = %s',
|
|
(target.artist, target.album),
|
|
)
|
|
row = cur.fetchone()
|
|
return row[0] if row else None
|
|
|
|
|
|
def _persist_candidates(conn: psycopg.Connection, job_id: str, candidates: list[Candidate]) -> None:
|
|
with conn.cursor() as cur:
|
|
for c in candidates:
|
|
cur.execute(
|
|
'INSERT INTO "Candidate" (id, "jobId", source, format, "qualityClass", '
|
|
'"trackCount", confidence, "sourceRef", chosen, "createdAt") '
|
|
"VALUES (gen_random_uuid()::text, %s, %s, %s, %s, %s, %s, %s, false, now())",
|
|
(job_id, c.source, c.quality.fmt, quality_class(c.quality),
|
|
c.track_count, c.confidence, c.source_ref),
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def _mark_chosen(conn: psycopg.Connection, job_id: str, source_ref: str) -> None:
|
|
"""Mark exactly one candidate (the one now being downloaded) as chosen, clearing any
|
|
prior choice. Called at each download attempt so `chosen` reflects the source in flight."""
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
'UPDATE "Candidate" SET chosen = ("sourceRef" = %s) WHERE "jobId" = %s',
|
|
(source_ref, job_id),
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def _import(conn: psycopg.Connection, job_id: str, target: MBTarget,
|
|
winner: Candidate, path: str) -> None:
|
|
request_id = _request_id(conn, job_id)
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
'INSERT INTO "LibraryItem" (id, "requestId", artist, album, path, source, '
|
|
'format, "qualityClass", "importedAt") '
|
|
"VALUES (gen_random_uuid()::text, %s, %s, %s, %s, %s, %s, %s, now()) "
|
|
'ON CONFLICT (artist, album) DO UPDATE SET '
|
|
' "requestId" = EXCLUDED."requestId", path = EXCLUDED.path, source = EXCLUDED.source, '
|
|
' format = EXCLUDED.format, "qualityClass" = EXCLUDED."qualityClass", '
|
|
' "importedAt" = now() '
|
|
'WHERE EXCLUDED."qualityClass" > "LibraryItem"."qualityClass"',
|
|
(request_id, target.artist, target.album, path, winner.source,
|
|
winner.quality.fmt, quality_class(winner.quality)),
|
|
)
|
|
cur.execute(
|
|
"UPDATE \"Request\" SET status = 'completed' WHERE id = %s", (request_id,)
|
|
)
|
|
conn.commit()
|
|
|
|
|
|
def run_pipeline(
|
|
conn: psycopg.Connection,
|
|
job_id: str,
|
|
adapters: Sequence[SourceAdapter],
|
|
resolver=None,
|
|
tagger=None,
|
|
min_confidence: float = 0.7,
|
|
dest_root: str = "/music",
|
|
staging_root: str | None = None,
|
|
upgrade_cutoff: int | None = None,
|
|
) -> None:
|
|
"""Real staged acquisition using source-agnostic adapters. Fakes in this plan."""
|
|
names = [a.name for a in adapters]
|
|
if len(names) != len(set(names)):
|
|
raise ValueError(f"adapter names must be unique, got {names}")
|
|
|
|
# 1. intake
|
|
_set_state(conn, job_id, "matching", "intake")
|
|
target = _load_target(conn, job_id)
|
|
if resolver is not None:
|
|
resolved = resolver.resolve(target.artist, target.album)
|
|
if resolved is not None:
|
|
target = resolved
|
|
# dedupe: a monitor-driven upgrade job proceeds unless the existing copy already meets
|
|
# the cutoff; a plain request short-circuits on any existing copy (slice-1 behavior).
|
|
if upgrade_cutoff is not None and _is_upgrade_job(conn, job_id):
|
|
_dedupe_hit = _already_in_library_at_cutoff(conn, target, upgrade_cutoff)
|
|
else:
|
|
_dedupe_hit = _already_in_library(conn, target)
|
|
if _dedupe_hit:
|
|
_set_state(conn, job_id, "imported", "import")
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"UPDATE \"Request\" SET status = 'completed' WHERE id = %s",
|
|
(_request_id(conn, job_id),),
|
|
)
|
|
conn.commit()
|
|
return
|
|
|
|
# 2. match — search all adapters; score + persist EVERY candidate found
|
|
_set_state(conn, job_id, "matching", "match")
|
|
found: list[Candidate] = []
|
|
for adapter in adapters:
|
|
try:
|
|
results = adapter.search(target)
|
|
except Exception as e: # a down source contributes no candidates, never crashes the job
|
|
print(f"pipeline: adapter {adapter.name} search failed: {e}", flush=True)
|
|
continue
|
|
for c in results:
|
|
found.append(replace(c, source_tier=adapter.tier, confidence=score_confidence(target, c)))
|
|
_persist_candidates(conn, job_id, found)
|
|
|
|
# 3. rank
|
|
_set_state(conn, job_id, "matched", "rank")
|
|
ranked = rank_candidates(target, found, min_confidence)
|
|
if not ranked:
|
|
_fail(conn, job_id, "no candidate above confidence threshold")
|
|
return
|
|
|
|
# 4. download (fall-through) into an isolated per-job staging dir
|
|
_set_state(conn, job_id, "downloading", "download")
|
|
by_source = {a.name: a for a in adapters}
|
|
staging = staging_dir(staging_root or f"{dest_root}/.staging", job_id)
|
|
winner = None
|
|
result = None
|
|
_set_download_progress(conn, job_id, 0.0) # reset for this run
|
|
expected = target.track_count or (ranked[0].track_count if ranked else 0)
|
|
_stop = threading.Event()
|
|
_reports = threading.Event() # set while the in-flight adapter reports real byte progress
|
|
_poller = threading.Thread(
|
|
target=_poll_download_progress, args=(job_id, staging, expected, _stop, _reports), daemon=True
|
|
)
|
|
_poller.start()
|
|
try:
|
|
try:
|
|
for candidate in ranked:
|
|
adapter = by_source.get(candidate.source)
|
|
if adapter is None:
|
|
continue
|
|
_mark_chosen(conn, job_id, candidate.source_ref) # reflect the source now in flight
|
|
shutil.rmtree(staging, ignore_errors=True) # clean slate per attempt
|
|
_reports.clear() # this attempt hasn't reported yet → poller estimates until it does
|
|
result = adapter.download(candidate, staging, _make_on_progress(conn, job_id, _reports))
|
|
if result.ok:
|
|
winner = candidate
|
|
break
|
|
finally:
|
|
_stop.set()
|
|
_poller.join(timeout=3)
|
|
if winner is None or result is None or not result.ok:
|
|
_fail(conn, job_id, "all downloads failed")
|
|
return
|
|
_set_download_progress(conn, job_id, 1.0) # download done — UI shows 100% into Finishing
|
|
|
|
# 5. verify completeness on the staging copy, then promote + tag
|
|
_set_state(conn, job_id, "tagging", "tag")
|
|
expected = target.track_count if target.track_count is not None else winner.track_count
|
|
if result.track_count < expected:
|
|
_fail(conn, job_id, "incomplete download")
|
|
return
|
|
|
|
final = album_dir(dest_root, target)
|
|
existing_q = _library_quality(conn, target)
|
|
new_q = quality_class(winner.quality)
|
|
_set_state(conn, job_id, "imported", "import")
|
|
if existing_q is None or new_q > existing_q:
|
|
import_album(staging, final) # clean audio(+cover) move, atomically replacing any prior copy
|
|
if tagger is not None:
|
|
try:
|
|
tagger.tag_album(final, target)
|
|
except Exception as e: # a tagging failure must not discard a good download
|
|
print(f"pipeline: tagging failed for job {job_id}: {e}", flush=True)
|
|
_import(conn, job_id, target, winner, final)
|
|
else:
|
|
# an existing copy is already at >= this quality — keep it untouched, just complete the request
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"UPDATE \"Request\" SET status = 'completed' WHERE id = %s",
|
|
(_request_id(conn, job_id),),
|
|
)
|
|
conn.commit()
|
|
finally:
|
|
shutil.rmtree(staging, ignore_errors=True) # a partial/failed download never persists
|