mrleo1nid
75ae2092ad
ci / api / api (push) Successful in 56s
ci / contract / contract (сгенерированный код совпадает с openapi.yaml) (push) Successful in 1m23s
ci / db / db (push) Successful in 1m8s
ci / recommender / recommender (push) Failing after 40s
ci / mobile / mobile (Metro-бандл Android) (push) Successful in 1m25s
ci / ts / ts (packages + web) (push) Successful in 7m20s
239 lines
9.6 KiB
Python
239 lines
9.6 KiB
Python
"""Воркер аудиоанализа: берёт задачи из analysis_jobs, считает признаки и пишет их в БД."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import subprocess
|
||
import time
|
||
from concurrent.futures import Future, ProcessPoolExecutor
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import numpy as np
|
||
import psycopg
|
||
|
||
from pipu_rec.analysis.features import ANALYZER_VER, SAMPLE_RATE, Features, analyze
|
||
|
||
log = logging.getLogger(__name__)
|
||
|
||
MAX_ATTEMPTS = 3
|
||
MAX_DURATION_S = 15 * 60
|
||
IDLE_SLEEP_S = 10.0
|
||
MISSING_SLEEP_S = 60.0 # файлы пропали целиком — ждём, пока смонтируют диск
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class Job:
|
||
track_id: str
|
||
path: str
|
||
|
||
|
||
def decode(path: str, ffmpeg: str = "ffmpeg") -> np.ndarray:
|
||
"""Моно float32 22 050 Гц; путь — отдельный аргумент, без shell."""
|
||
out = subprocess.run(
|
||
[
|
||
ffmpeg,
|
||
"-v",
|
||
"error",
|
||
"-t",
|
||
str(MAX_DURATION_S),
|
||
"-i",
|
||
path,
|
||
"-ac",
|
||
"1",
|
||
"-ar",
|
||
str(SAMPLE_RATE),
|
||
"-f",
|
||
"f32le",
|
||
"-",
|
||
],
|
||
capture_output=True,
|
||
check=False,
|
||
timeout=300,
|
||
)
|
||
if out.returncode != 0:
|
||
raise RuntimeError(out.stderr.decode(errors="replace").strip()[:500] or "ffmpeg failed")
|
||
samples = np.frombuffer(out.stdout, dtype=np.float32)
|
||
if samples.size < SAMPLE_RATE:
|
||
raise RuntimeError("audio is shorter than 1 s")
|
||
return samples
|
||
|
||
|
||
def analyze_file(path: str, ffmpeg: str) -> Features:
|
||
"""Выполняется в процессе пула."""
|
||
return analyze(decode(path, ffmpeg))
|
||
|
||
|
||
def vector_literal(v: np.ndarray) -> str:
|
||
return "[" + ",".join(f"{x:.6f}" for x in v.tolist()) + "]"
|
||
|
||
|
||
class Worker:
|
||
def __init__(
|
||
self, dsn: str, library_root: str, host_root: str, workers: int, ffmpeg: str = "ffmpeg"
|
||
) -> None:
|
||
self.dsn = dsn
|
||
self.library_root = Path(library_root)
|
||
self.host_root = Path(host_root)
|
||
self.workers = max(1, workers)
|
||
self.ffmpeg = ffmpeg
|
||
|
||
def resolve(self, source_id: str, source_type: str, local_path: str | None, rel_path: str) -> str:
|
||
"""Путь к файлу внутри контейнера анализатора — те же точки монтирования, что у api."""
|
||
if source_type == "local":
|
||
return str(self.host_root.joinpath(local_path or "", rel_path))
|
||
return str(self.library_root.joinpath(source_id, rel_path))
|
||
|
||
def claim(self, conn: psycopg.Connection[Any], limit: int) -> list[Job]:
|
||
"""Забирает задачи; SKIP LOCKED — несколько анализаторов не возьмут одну и ту же."""
|
||
rows = conn.execute(
|
||
"""
|
||
UPDATE analysis_jobs j SET status = 'running', attempts = attempts + 1, updated_at = now()
|
||
WHERE j.track_id IN (
|
||
SELECT a.track_id FROM analysis_jobs a
|
||
JOIN tracks t ON t.id = a.track_id
|
||
WHERE a.status = 'queued'
|
||
AND EXISTS (
|
||
SELECT 1 FROM media_files f JOIN library_sources s ON s.id = f.source_id
|
||
WHERE f.track_id = t.id AND s.status = 'mounted'
|
||
AND f.missing_since IS NULL AND f.scan_error IS NULL
|
||
AND NOT library_path_excluded(f.rel_path, s.excluded_dirs)
|
||
)
|
||
ORDER BY a.updated_at
|
||
LIMIT %s
|
||
FOR UPDATE OF a SKIP LOCKED)
|
||
RETURNING j.track_id::text
|
||
""",
|
||
(limit,),
|
||
).fetchall()
|
||
conn.commit()
|
||
if not rows:
|
||
return []
|
||
files = conn.execute(
|
||
"""
|
||
SELECT t.id::text, f.source_id::text, f.source_type, f.local_path, f.rel_path
|
||
FROM tracks t
|
||
CROSS JOIN LATERAL (
|
||
SELECT s.id AS source_id, s.type::text AS source_type, s.local_path, mf.rel_path
|
||
FROM media_files mf JOIN library_sources s ON s.id = mf.source_id
|
||
WHERE mf.track_id = t.id AND s.status = 'mounted'
|
||
AND mf.missing_since IS NULL AND mf.scan_error IS NULL
|
||
AND NOT library_path_excluded(mf.rel_path, s.excluded_dirs)
|
||
ORDER BY (mf.id = t.primary_file_id) DESC NULLS LAST, mf.id
|
||
LIMIT 1
|
||
) f
|
||
WHERE t.id = ANY(%s::uuid[])
|
||
""",
|
||
([r[0] for r in rows],),
|
||
).fetchall()
|
||
return [Job(track_id=r[0], path=self.resolve(r[1], r[2], r[3], r[4])) for r in files]
|
||
|
||
def save(self, conn: psycopg.Connection[Any], track_id: str, f: Features) -> None:
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO track_features (track_id, analyzer_ver, bpm, bpm_confidence, key, scale,
|
||
loudness_lufs, energy, danceability, valence, acousticness, intro_ms, analyzed_at)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, now())
|
||
ON CONFLICT (track_id) DO UPDATE SET analyzer_ver = EXCLUDED.analyzer_ver, bpm = EXCLUDED.bpm,
|
||
bpm_confidence = EXCLUDED.bpm_confidence, key = EXCLUDED.key, scale = EXCLUDED.scale,
|
||
loudness_lufs = EXCLUDED.loudness_lufs, energy = EXCLUDED.energy,
|
||
danceability = EXCLUDED.danceability, valence = EXCLUDED.valence,
|
||
acousticness = EXCLUDED.acousticness, intro_ms = EXCLUDED.intro_ms, analyzed_at = now()
|
||
""",
|
||
(
|
||
track_id,
|
||
ANALYZER_VER,
|
||
f.bpm,
|
||
f.bpm_confidence,
|
||
f.key,
|
||
f.scale,
|
||
f.loudness_lufs,
|
||
f.energy,
|
||
f.danceability,
|
||
f.valence,
|
||
f.acousticness,
|
||
f.intro_ms,
|
||
),
|
||
)
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO track_embeddings (track_id, audio, updated_at) VALUES (%s, %s::vector, now())
|
||
ON CONFLICT (track_id) DO UPDATE SET audio = EXCLUDED.audio, updated_at = now()
|
||
""",
|
||
(track_id, vector_literal(f.embedding)),
|
||
)
|
||
conn.execute(
|
||
"UPDATE analysis_jobs SET status = 'done', error = NULL, updated_at = now() WHERE track_id = %s",
|
||
(track_id,),
|
||
)
|
||
conn.commit()
|
||
|
||
def fail(self, conn: psycopg.Connection[Any], track_id: str, err: str) -> None:
|
||
conn.execute(
|
||
"""
|
||
UPDATE analysis_jobs SET status = CASE WHEN attempts >= %s THEN 'failed' ELSE 'queued' END,
|
||
error = %s, updated_at = now()
|
||
WHERE track_id = %s
|
||
""",
|
||
(MAX_ATTEMPTS, err[:500], track_id),
|
||
)
|
||
conn.commit()
|
||
|
||
def postpone(self, conn: psycopg.Connection[Any], track_ids: list[str]) -> None:
|
||
"""Файла нет на диске (не смонтирован?) — обратно в очередь, попытка не засчитывается."""
|
||
conn.execute(
|
||
"""
|
||
UPDATE analysis_jobs SET status = 'queued', attempts = greatest(attempts - 1, 0),
|
||
error = 'file not found — is the disk mounted?', updated_at = now()
|
||
WHERE track_id = ANY(%s::uuid[])
|
||
""",
|
||
(track_ids,),
|
||
)
|
||
conn.commit()
|
||
|
||
def requeue_stale(self, conn: psycopg.Connection[Any]) -> None:
|
||
"""Задачи, брошенные упавшим воркером, и признаки старой версии анализатора — снова в очередь."""
|
||
conn.execute(
|
||
"UPDATE analysis_jobs SET status = 'queued' "
|
||
"WHERE status = 'running' AND updated_at < now() - interval '30 minutes'"
|
||
)
|
||
conn.execute(
|
||
"""
|
||
UPDATE analysis_jobs a SET status = 'queued', attempts = 0, updated_at = now()
|
||
FROM track_features f WHERE f.track_id = a.track_id AND f.analyzer_ver <> %s AND a.status = 'done'
|
||
""",
|
||
(ANALYZER_VER,),
|
||
)
|
||
conn.commit()
|
||
|
||
def run_forever(self) -> None:
|
||
log.info("analyzer started: workers=%d version=%s", self.workers, ANALYZER_VER)
|
||
with psycopg.connect(self.dsn) as conn, ProcessPoolExecutor(self.workers) as pool:
|
||
self.requeue_stale(conn)
|
||
while True:
|
||
jobs = self.claim(conn, self.workers * 2)
|
||
missing = [j for j in jobs if not Path(j.path).exists()]
|
||
if missing:
|
||
self.postpone(conn, [j.track_id for j in missing])
|
||
jobs = [j for j in jobs if j not in missing]
|
||
if not jobs:
|
||
log.warning(
|
||
"analysis: %d files not found, waiting — is the disk mounted?", len(missing)
|
||
)
|
||
time.sleep(MISSING_SLEEP_S)
|
||
continue
|
||
if not jobs:
|
||
time.sleep(IDLE_SLEEP_S)
|
||
continue
|
||
futures: list[tuple[Job, Future[Features]]] = [
|
||
(j, pool.submit(analyze_file, j.path, self.ffmpeg)) for j in jobs
|
||
]
|
||
for job, fut in futures:
|
||
try:
|
||
self.save(conn, job.track_id, fut.result(timeout=600))
|
||
except Exception as exc:
|
||
log.warning("analysis failed: track=%s err=%s", job.track_id, exc)
|
||
conn.rollback()
|
||
self.fail(conn, job.track_id, str(exc))
|