feat: initial Tauri desktop shell for BatchTube v1.0
This commit is contained in:
Agent Zero
507
backend/worker.py
Normal file
507
backend/worker.py
Normal file
@@ -0,0 +1,507 @@
|
||||
"""
|
||||
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
|
||||
Reference in New Issue
Block a user