508 lines
22 KiB
Python
508 lines
22 KiB
Python
"""
|
|
worker.py - Background job pipeline for YouTube SIPOC Tool
|
|
Handles pause/stop/resume, per-artifact saving, duplicate skipping,
|
|
random delays, retry logic, and resume-on-failure.
|
|
"""
|
|
import time
|
|
import random
|
|
import logging
|
|
import threading
|
|
from typing import Dict, Callable, Optional
|
|
|
|
import database as db
|
|
import extractors
|
|
import config
|
|
|
|
logger = logging.getLogger("worker")
|
|
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
# Job Control Registry
|
|
# Stores threading.Event objects per job_id for pause/stop/skip signals.
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
|
|
_job_controls: Dict[str, Dict] = {}
|
|
_registry_lock = threading.Lock()
|
|
|
|
|
|
def _get_controls(job_id: str) -> Dict:
|
|
with _registry_lock:
|
|
if job_id not in _job_controls:
|
|
_job_controls[job_id] = {
|
|
"pause": threading.Event(), # set = paused
|
|
"stop": threading.Event(), # set = stop requested
|
|
"skip": threading.Event(), # set = skip current video
|
|
}
|
|
return _job_controls[job_id]
|
|
|
|
|
|
def _cleanup_controls(job_id: str):
|
|
with _registry_lock:
|
|
_job_controls.pop(job_id, None)
|
|
|
|
|
|
def pause_job(job_id: str):
|
|
_get_controls(job_id)["pause"].set()
|
|
db.update_job_status(job_id, "paused")
|
|
db.add_log(job_id, "Job paused by user.", level="INFO")
|
|
|
|
|
|
def resume_job(job_id: str):
|
|
_get_controls(job_id)["pause"].clear()
|
|
db.update_job_status(job_id, "running")
|
|
db.add_log(job_id, "Job resumed by user.", level="INFO")
|
|
|
|
|
|
def stop_job(job_id: str):
|
|
_get_controls(job_id)["stop"].set()
|
|
_get_controls(job_id)["pause"].clear() # Unpause so thread can exit
|
|
db.update_job_status(job_id, "stopped")
|
|
db.add_log(job_id, "Job stopped by user.", level="INFO")
|
|
|
|
|
|
def skip_video(job_id: str):
|
|
"""Signal the worker to skip the current video and move to the next."""
|
|
_get_controls(job_id)["skip"].set()
|
|
db.add_log(job_id, "Skip requested for current video.", level="INFO")
|
|
|
|
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
# SSE Event Emitter
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
|
|
class EventEmitter:
|
|
"""Thread-safe event queue for Server-Sent Events."""
|
|
|
|
def __init__(self):
|
|
self._listeners: Dict[str, list] = {}
|
|
self._lock = threading.Lock()
|
|
|
|
def subscribe(self, job_id: str) -> "EventQueue":
|
|
queue = EventQueue()
|
|
with self._lock:
|
|
if job_id not in self._listeners:
|
|
self._listeners[job_id] = []
|
|
self._listeners[job_id].append(queue)
|
|
return queue
|
|
|
|
def unsubscribe(self, job_id: str, queue: "EventQueue"):
|
|
with self._lock:
|
|
if job_id in self._listeners:
|
|
try:
|
|
self._listeners[job_id].remove(queue)
|
|
except ValueError:
|
|
pass
|
|
|
|
def emit(self, job_id: str, event_type: str, data: dict):
|
|
with self._lock:
|
|
listeners = list(self._listeners.get(job_id, []))
|
|
for q in listeners:
|
|
q.put(event_type, data)
|
|
|
|
|
|
class EventQueue:
|
|
"""Per-subscriber event queue with thread-safe get."""
|
|
|
|
def __init__(self):
|
|
import queue
|
|
self._q = queue.Queue(maxsize=500)
|
|
|
|
def put(self, event_type: str, data: dict):
|
|
try:
|
|
self._q.put_nowait({"type": event_type, "data": data})
|
|
except Exception:
|
|
pass # Drop events if queue is full
|
|
|
|
def get(self, timeout: float = 30.0):
|
|
try:
|
|
return self._q.get(timeout=timeout)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
# Global emitter instance
|
|
emitter = EventEmitter()
|
|
|
|
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
# Worker Thread Entry Point
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
|
|
def run_job(job_id: str):
|
|
"""
|
|
Main worker function. Runs in a background thread.
|
|
Resolves inputs, iterates videos, processes each artifact,
|
|
emits SSE events, handles pause/stop/skip/retry.
|
|
"""
|
|
controls = _get_controls(job_id)
|
|
job = db.get_job(job_id)
|
|
if not job:
|
|
logger.error(f"Job {job_id} not found in database.")
|
|
return
|
|
|
|
db.update_job_status(job_id, "running")
|
|
_emit(job_id, "job_status", {"job_id": job_id, "status": "running"})
|
|
db.add_log(job_id, "Job started.", level="INFO")
|
|
|
|
try:
|
|
# ── Step 1: Resolve input to video list ──────────────────────────────
|
|
_emit(job_id, "progress", {"message": "Resolving input URLs…"})
|
|
db.add_log(job_id, f"Resolving input: type={job['input_type']}", level="INFO")
|
|
|
|
try:
|
|
videos = extractors.resolve_input(job["input_type"], job["input_value"])
|
|
except Exception as e:
|
|
db.add_log(job_id, f"Failed to resolve input: {e}", level="ERROR")
|
|
db.update_job_status(job_id, "failed")
|
|
_emit(job_id, "job_status", {"job_id": job_id, "status": "failed", "error": str(e)})
|
|
return
|
|
|
|
if not videos:
|
|
db.add_log(job_id, "No videos found for input.", level="WARN")
|
|
db.update_job_status(job_id, "done")
|
|
_emit(job_id, "job_status", {"job_id": job_id, "status": "done"})
|
|
return
|
|
|
|
# Deduplicate by video_id (preserve order)
|
|
seen = set()
|
|
unique_videos = []
|
|
for v in videos:
|
|
if v["video_id"] not in seen:
|
|
seen.add(v["video_id"])
|
|
unique_videos.append(v)
|
|
videos = unique_videos
|
|
|
|
# ── Step 2: Register videos in DB ────────────────────────────────────
|
|
for i, v in enumerate(videos):
|
|
db.upsert_video(job_id, v["video_id"], v["url"], v.get("title", ""), i)
|
|
db.update_job_progress(job_id, len(videos), 0)
|
|
_emit(job_id, "video_list", {
|
|
"job_id": job_id,
|
|
"total": len(videos),
|
|
"videos": [{"video_id": v["video_id"], "title": v.get("title", ""), "position": i}
|
|
for i, v in enumerate(videos)]
|
|
})
|
|
db.add_log(job_id, f"Found {len(videos)} unique videos.", level="INFO")
|
|
|
|
# ── Step 3: Process each video ────────────────────────────────────────
|
|
selections = job["selections"]
|
|
done_count = 0
|
|
|
|
for idx, video in enumerate(videos):
|
|
# Check stop
|
|
if controls["stop"].is_set():
|
|
db.add_log(job_id, "Job stopped before processing video.", level="INFO")
|
|
break
|
|
|
|
# Handle pause (blocking)
|
|
_wait_if_paused(controls)
|
|
if controls["stop"].is_set():
|
|
break
|
|
|
|
# Reset skip flag for this video
|
|
controls["skip"].clear()
|
|
|
|
video_id = video["video_id"]
|
|
db.add_log(job_id, f"Processing video {idx+1}/{len(videos)}: {video_id}", level="INFO")
|
|
_emit(job_id, "video_start", {"job_id": job_id, "video_id": video_id, "position": idx})
|
|
|
|
# ── Step 3a: Fetch metadata (once per video) ──────────────────
|
|
meta = {}
|
|
try:
|
|
_emit(job_id, "artifact_start", {
|
|
"job_id": job_id, "video_id": video_id, "artifact_type": "metadata"
|
|
})
|
|
meta = extractors.fetch_metadata(video_id)
|
|
title = meta.get("title", "")
|
|
if title:
|
|
db.update_video_title(job_id, video_id, title)
|
|
_emit(job_id, "video_title", {"job_id": job_id, "video_id": video_id, "title": title})
|
|
db.add_log(job_id, f"Metadata fetched for {video_id}: {title or '(no title)'}", level="INFO")
|
|
except Exception as e:
|
|
db.add_log(job_id, f"Failed to fetch metadata for {video_id}: {e}", level="ERROR")
|
|
_emit(job_id, "video_error", {"job_id": job_id, "video_id": video_id, "error": str(e)})
|
|
db.update_video_status(job_id, video_id, "failed")
|
|
continue
|
|
|
|
# ── Step 3b: Initialize artifact records ──────────────────────
|
|
for artifact_type in selections:
|
|
db.upsert_artifact(job_id, video_id, artifact_type)
|
|
|
|
# ── Step 3c: Process each selected artifact ───────────────────
|
|
video_failed = False
|
|
for artifact_type in selections:
|
|
# Check stop or skip
|
|
if controls["stop"].is_set():
|
|
break
|
|
if controls["skip"].is_set():
|
|
db.add_log(job_id, f"Skipping remaining artifacts for {video_id} (skip requested).", level="INFO")
|
|
break
|
|
|
|
# Handle pause
|
|
_wait_if_paused(controls)
|
|
if controls["stop"].is_set():
|
|
break
|
|
|
|
# Skip if already done
|
|
if db.artifact_is_done(job_id, video_id, artifact_type):
|
|
db.add_log(job_id, f"[SKIP-DUP] {artifact_type} for {video_id} already done.", level="INFO")
|
|
_emit(job_id, "artifact_skip", {
|
|
"job_id": job_id, "video_id": video_id, "artifact_type": artifact_type,
|
|
"reason": "already_done"
|
|
})
|
|
continue
|
|
|
|
# Process with retry
|
|
success = _process_artifact_with_retry(
|
|
job_id=job_id,
|
|
video_id=video_id,
|
|
artifact_type=artifact_type,
|
|
meta=meta,
|
|
job=job,
|
|
controls=controls,
|
|
)
|
|
if not success:
|
|
video_failed = True
|
|
|
|
# Random delay between artifacts (bot avoidance)
|
|
if not controls["stop"].is_set() and not controls["skip"].is_set():
|
|
delay = random.uniform(job["delay_min"], job["delay_max"])
|
|
db.add_log(job_id, f"Waiting {delay:.1f}s before next artifact…", level="DEBUG")
|
|
_interruptible_sleep(delay, controls)
|
|
|
|
# ── Step 3d: Mark video done/failed ──────────────────────────
|
|
if not controls["stop"].is_set():
|
|
status = "failed" if video_failed else "done"
|
|
db.update_video_status(job_id, video_id, status)
|
|
done_count += 1
|
|
db.update_job_progress(job_id, len(videos), done_count)
|
|
_emit(job_id, "video_done", {
|
|
"job_id": job_id,
|
|
"video_id": video_id,
|
|
"status": status,
|
|
"done": done_count,
|
|
"total": len(videos),
|
|
})
|
|
|
|
# Random delay between videos
|
|
if idx < len(videos) - 1 and not controls["stop"].is_set():
|
|
delay = random.uniform(job["delay_min"], job["delay_max"])
|
|
db.add_log(job_id, f"Waiting {delay:.1f}s before next video…", level="DEBUG")
|
|
_interruptible_sleep(delay, controls)
|
|
|
|
# ── Step 4: Finalize job ──────────────────────────────────────────────
|
|
if controls["stop"].is_set():
|
|
final_status = "stopped"
|
|
else:
|
|
stats = db.get_job_stats(job_id)
|
|
final_status = "done" if stats["artifact_failed"] == 0 else "done_with_errors"
|
|
|
|
db.update_job_status(job_id, final_status)
|
|
stats = db.get_job_stats(job_id)
|
|
db.add_log(job_id, f"Job finished with status: {final_status}. Stats: {stats}", level="INFO")
|
|
_emit(job_id, "job_status", {"job_id": job_id, "status": final_status, "stats": stats})
|
|
|
|
except Exception as e:
|
|
logger.exception(f"Unexpected error in job {job_id}: {e}")
|
|
db.add_log(job_id, f"Unexpected error: {e}", level="ERROR")
|
|
db.update_job_status(job_id, "failed")
|
|
_emit(job_id, "job_status", {"job_id": job_id, "status": "failed", "error": str(e)})
|
|
finally:
|
|
_cleanup_controls(job_id)
|
|
|
|
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
# Artifact Processing with Retry
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
|
|
def _process_artifact_with_retry(
|
|
job_id: str,
|
|
video_id: str,
|
|
artifact_type: str,
|
|
meta: Dict,
|
|
job: Dict,
|
|
controls: Dict,
|
|
) -> bool:
|
|
"""
|
|
Attempt to process one artifact, retrying up to MAX_RETRIES times.
|
|
Returns True on success, False on final failure.
|
|
"""
|
|
max_retries = config.MAX_RETRIES
|
|
output_folder = job["output_folder"]
|
|
resolution = job.get("resolution", "best")
|
|
api_key = job.get("openrouter_model", "") # model field
|
|
# We store the API key in the job separately; retrieve from job dict
|
|
# (api_key is set in the job via the POST body - see api_server.py)
|
|
openrouter_api_key = job.get("_openrouter_api_key", "") # injected by api_server
|
|
openrouter_model = job.get("openrouter_model", config.OPENROUTER_MODEL)
|
|
summary_prompt = job.get("summary_prompt", config.SUMMARY_SYSTEM_PROMPT)
|
|
|
|
for attempt in range(1, max_retries + 1):
|
|
if controls["stop"].is_set() or controls["skip"].is_set():
|
|
return False
|
|
|
|
db.mark_artifact_started(job_id, video_id, artifact_type)
|
|
_emit(job_id, "artifact_start", {
|
|
"job_id": job_id,
|
|
"video_id": video_id,
|
|
"artifact_type": artifact_type,
|
|
"attempt": attempt,
|
|
})
|
|
db.add_log(job_id, f"[{artifact_type}] Starting (attempt {attempt}/{max_retries})", level="INFO",
|
|
video_id=video_id, artifact_type=artifact_type)
|
|
|
|
try:
|
|
file_path = _run_extractor(
|
|
artifact_type=artifact_type,
|
|
video_id=video_id,
|
|
meta=meta,
|
|
output_folder=output_folder,
|
|
resolution=resolution,
|
|
openrouter_api_key=openrouter_api_key,
|
|
openrouter_model=openrouter_model,
|
|
summary_prompt=summary_prompt,
|
|
)
|
|
db.update_artifact_status(job_id, video_id, artifact_type, "done", file_path=file_path)
|
|
db.add_log(job_id, f"[{artifact_type}] Done → {file_path}", level="INFO",
|
|
video_id=video_id, artifact_type=artifact_type)
|
|
_emit(job_id, "artifact_done", {
|
|
"job_id": job_id,
|
|
"video_id": video_id,
|
|
"artifact_type": artifact_type,
|
|
"file_path": file_path,
|
|
})
|
|
return True
|
|
|
|
except Exception as e:
|
|
err_msg = str(e)
|
|
logger.warning(f"[{job_id}][{video_id}][{artifact_type}] Attempt {attempt} failed: {e}")
|
|
db.add_log(job_id, f"[{artifact_type}] Attempt {attempt} failed: {e}", level="WARN",
|
|
video_id=video_id, artifact_type=artifact_type)
|
|
|
|
if attempt < max_retries:
|
|
# Exponential backoff between retries
|
|
backoff = min(2 ** attempt * 3, 30) + random.uniform(0, 2)
|
|
db.add_log(job_id, f"[{artifact_type}] Retry in {backoff:.1f}s…", level="INFO",
|
|
video_id=video_id, artifact_type=artifact_type)
|
|
_interruptible_sleep(backoff, controls)
|
|
else:
|
|
db.update_artifact_status(
|
|
job_id, video_id, artifact_type, "failed", error_msg=err_msg
|
|
)
|
|
db.add_log(job_id, f"[{artifact_type}] FAILED after {max_retries} attempts: {e}", level="ERROR",
|
|
video_id=video_id, artifact_type=artifact_type)
|
|
_emit(job_id, "artifact_failed", {
|
|
"job_id": job_id,
|
|
"video_id": video_id,
|
|
"artifact_type": artifact_type,
|
|
"error": err_msg,
|
|
"attempts": max_retries,
|
|
})
|
|
return False
|
|
|
|
return False
|
|
|
|
|
|
def _run_extractor(
|
|
artifact_type: str,
|
|
video_id: str,
|
|
meta: Dict,
|
|
output_folder: str,
|
|
resolution: str,
|
|
openrouter_api_key: str,
|
|
openrouter_model: str,
|
|
summary_prompt: str,
|
|
) -> str:
|
|
"""Dispatch to the correct extractor function."""
|
|
if artifact_type == "title":
|
|
file_path, _ = extractors.extract_title(video_id, meta, output_folder)
|
|
return file_path
|
|
elif artifact_type == "stats":
|
|
return extractors.extract_stats(video_id, meta, output_folder)
|
|
elif artifact_type == "transcript":
|
|
return extractors.extract_transcript(video_id, meta, output_folder)
|
|
elif artifact_type == "summary":
|
|
return extractors.extract_summary(
|
|
video_id, meta, output_folder,
|
|
api_key=openrouter_api_key,
|
|
model=openrouter_model,
|
|
system_prompt=summary_prompt,
|
|
)
|
|
elif artifact_type == "thumbnail":
|
|
return extractors.extract_thumbnail(video_id, meta, output_folder)
|
|
elif artifact_type == "mp3":
|
|
return extractors.extract_mp3(video_id, meta, output_folder)
|
|
elif artifact_type == "mp4":
|
|
return extractors.extract_mp4(video_id, meta, output_folder, resolution=resolution)
|
|
else:
|
|
raise ValueError(f"Unknown artifact type: {artifact_type}")
|
|
|
|
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
# Helpers
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
|
|
def _emit(job_id: str, event_type: str, data: dict):
|
|
"""Emit an SSE event."""
|
|
emitter.emit(job_id, event_type, data)
|
|
|
|
|
|
def _wait_if_paused(controls: Dict):
|
|
"""Block worker thread while paused, check every 0.5s."""
|
|
while controls["pause"].is_set() and not controls["stop"].is_set():
|
|
time.sleep(0.5)
|
|
|
|
|
|
def _interruptible_sleep(seconds: float, controls: Dict):
|
|
"""Sleep for `seconds` but wake early if stop/skip is set."""
|
|
end = time.monotonic() + seconds
|
|
while time.monotonic() < end:
|
|
if controls["stop"].is_set() or controls["skip"].is_set():
|
|
return
|
|
time.sleep(min(0.3, end - time.monotonic()))
|
|
|
|
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
# Job Launcher
|
|
# ─────────────────────────────────────────────────────────────────────────────
|
|
|
|
def launch_job(job_id: str, openrouter_api_key: str = ""):
|
|
"""
|
|
Launch job in a daemon background thread.
|
|
openrouter_api_key is injected at runtime (not stored in DB for security).
|
|
"""
|
|
# Inject API key into job dict via a side channel stored in controls
|
|
controls = _get_controls(job_id)
|
|
controls["_openrouter_api_key"] = openrouter_api_key
|
|
|
|
# Patch the run_job function to use the injected key
|
|
def _patched_run():
|
|
job = db.get_job(job_id)
|
|
if job:
|
|
job["_openrouter_api_key"] = openrouter_api_key
|
|
run_job_with_key(job_id, openrouter_api_key)
|
|
|
|
t = threading.Thread(target=_patched_run, daemon=True, name=f"job-{job_id[:8]}")
|
|
t.start()
|
|
return t
|
|
|
|
|
|
def run_job_with_key(job_id: str, openrouter_api_key: str):
|
|
"""Wrapper that injects the API key into job processing."""
|
|
# Monkey-patch the job dict to include the API key during this run
|
|
original_get_job = db.get_job
|
|
|
|
def patched_get_job(jid):
|
|
result = original_get_job(jid)
|
|
if result and jid == job_id:
|
|
result["_openrouter_api_key"] = openrouter_api_key
|
|
return result
|
|
|
|
# Temporarily patch db.get_job for this thread
|
|
import database
|
|
_orig = database.get_job
|
|
database.get_job = patched_get_job
|
|
try:
|
|
run_job(job_id)
|
|
finally:
|
|
database.get_job = _orig
|