Spaces:
Runtime error
Runtime error
Refactor VideoTemplateEditor component: remove AI editing features and related state management, streamline template loading, and update styles for improved layout. Enhance CSS for responsive design and visual consistency across various components, including dropdowns, sync stats, and quiz elements. Remove deprecated premium badge styles and adjust layout for chat components.
66de9c8 | """Worker: runs queued jobs from the Postgres jobs table plus scheduled email batches. | |
| Run: python -m app.worker | |
| - Claims jobs with FOR UPDATE SKIP LOCKED, runs up to WORKER_CONCURRENCY at once in threads. | |
| - Failed jobs retry with backoff up to max_attempts; jobs whose worker died (no heartbeat for | |
| JOB_STALE_SECONDS) are requeued. | |
| - Periodic tasks take a Postgres advisory lock, so with several workers only one runs each tick. | |
| """ | |
| import logging | |
| import os | |
| import signal | |
| import socket | |
| import threading | |
| import time | |
| import uuid | |
| from collections.abc import Callable | |
| from concurrent.futures import Future, ThreadPoolExecutor | |
| from contextlib import contextmanager | |
| from dataclasses import dataclass, field | |
| from datetime import UTC, datetime | |
| import httpx | |
| from sqlalchemy import select, text | |
| from sqlalchemy.orm import Session | |
| from app.chat_memory.organize import forget_orphans, organize_chat, rebuild_chat | |
| from app.config import settings | |
| from app.db import SessionLocal, engine | |
| from app.models import ( | |
| Artifact, | |
| Chat, | |
| Document, | |
| Job, | |
| Message, | |
| Source, | |
| Subscription, | |
| UpdateEmail, | |
| UpdateEmailSend, | |
| Upload, | |
| User, | |
| VoiceConsent, | |
| VoiceProfile, | |
| Workspace, | |
| WorkspaceMember, | |
| ) | |
| from app.pipeline import flashcards, generate, infographic, launchkit, media, quiz, report, slides | |
| from app.pipeline.generate import NothingToDo | |
| from app.pipeline.ingest import FeedNotFound, entry_from_upload, extract_article, ingest_source, store_entries | |
| from app.pipeline.memory import learn_from_message | |
| from app.services import launchpad, report_suggestions, tts, video_quota | |
| from app.services.email import email_service | |
| from app.services.email_verification import purge_old_codes | |
| from app.services.jobs import ( | |
| claim_next, | |
| create_job, | |
| pending_job, | |
| record_usage, | |
| retry_or_fail, | |
| stale_jobs, | |
| update_job, | |
| ) | |
| from app.services.notestack_voices import save_notestack_voice | |
| from app.services.plans import effective_plan, post_room | |
| from app.services.renderer import PermanentJobError, request_render | |
| from app.services.social.health import social_health | |
| from app.services.storage import keys, storage | |
| from app.services.uploads import sweep_orphan_uploads | |
| log = logging.getLogger("notestack.worker") | |
| class Done: | |
| result: dict = field(default_factory=dict) | |
| message: str = "Done" | |
| HANDED_OFF = object() # the job continues elsewhere (renderer) and reports back via callbacks | |
| # Handlers: (db, job) -> Done | HANDED_OFF. Raise to trigger a retry. | |
| def _after_import(db: Session, job: Job, workspace_id: uuid.UUID, changed: list[str]) -> None: | |
| """New or changed posts are mapped into the topic constellation and scored for resurfacing, and the | |
| first import builds the writing voice, so none of these need a button press.""" | |
| if not changed or not settings.llm_api_key: | |
| return | |
| create_job(db, workspace_id, "topics", {"document_ids": changed}, max_attempts=2) | |
| create_job(db, workspace_id, "ideas", {"document_ids": changed}, max_attempts=2) | |
| if not pending_job(db, workspace_id, "resurface_scan"): | |
| create_job(db, workspace_id, "resurface_scan", {}, max_attempts=2) | |
| vp = db.scalar(select(VoiceProfile).where(VoiceProfile.workspace_id == workspace_id)) | |
| if not (vp and vp.profile_json) and not pending_job(db, workspace_id, "voice_profile"): | |
| create_job(db, workspace_id, "voice_profile", {"document_ids": []}, max_attempts=2) | |
| def handle_ingest(db: Session, job: Job): | |
| source = db.get(Source, uuid.UUID(job.params["source_id"])) | |
| if not source: | |
| return Done({"error": "source deleted"}, "Source no longer exists") | |
| plan = effective_plan(db, db.get(Workspace, source.workspace_id)) | |
| # The plan limit is one total across every source, link and upload: this source gets what the others leave. | |
| room = post_room(db, source.workspace_id, plan, source.id) | |
| try: | |
| result = ingest_source(db, source, job, max_posts=room) | |
| except Exception as exc: | |
| db.rollback() | |
| source.sync_status = "error" | |
| source.sync_error = str(exc)[:500] | |
| db.commit() | |
| if isinstance(exc, httpx.HTTPStatusError) and exc.response.status_code in (401, 403, 404, 410): | |
| raise PermanentJobError(f"The feed returned {exc.response.status_code}.") from exc | |
| if isinstance(exc, FeedNotFound): | |
| raise PermanentJobError(str(exc)) from exc | |
| raise | |
| _after_import(db, job, source.workspace_id, result.get("changed", [])) | |
| if result.get("locked"): | |
| return Done(result, f"Found {result['found']} posts. Your latest {result['available']} are indexed; " | |
| f"upgrade to index the other {result['locked']}") | |
| return Done(result, f"All posts in orbit: {result['indexed']} new or updated") | |
| def handle_memory(db: Session, job: Job): | |
| """Decide whether the writer's chat message holds a lasting note, and update their memory.""" | |
| msg = db.get(Message, uuid.UUID(job.params["message_id"])) | |
| chat = db.get(Chat, msg.chat_id) if msg else None | |
| if not msg or msg.role != "user" or not chat or chat.workspace_id != job.workspace_id: | |
| return Done({"saved": []}, "Message no longer exists") | |
| result = learn_from_message(db, job, job.workspace_id, msg.content) | |
| return Done(result, "Notes updated" if result["saved"] else "Nothing to save") | |
| def handle_chat_memory(db: Session, job: Job): | |
| """File a chat's new rounds under topics (summary, keywords, exact quotes) in the notebook's chat memory.""" | |
| chat = db.get(Chat, uuid.UUID(job.params["chat_id"])) | |
| if not chat or chat.workspace_id != job.workspace_id: | |
| return Done({"rounds": 0}, "Chat no longer exists") | |
| result = organize_chat(db, job, chat) | |
| return Done(result, f"Filed {result['rounds']} rounds" if result["rounds"] else "Nothing new to file") | |
| def handle_chat_memory_rebuild(db: Session, job: Job): | |
| """Rebuild chat memory from the messages. Params: notebook_id, and optionally chat_id for just one chat.""" | |
| notebook_id = uuid.UUID(job.params["notebook_id"]) | |
| chats = db.scalars(select(Chat).where(Chat.notebook_id == notebook_id, Chat.workspace_id == job.workspace_id)) | |
| only = job.params.get("chat_id") | |
| done = [rebuild_chat(db, job, c) for c in chats if not only or str(c.id) == only] | |
| return Done({"chats": len(done), "rounds": sum(d["rounds"] for d in done)}, f"Rebuilt {len(done)} chats") | |
| def handle_chat_memory_cleanup(db: Session, job: Job): | |
| """Remove chat memory whose chat is gone. Params: notebook_id.""" | |
| removed = forget_orphans(db, job.workspace_id, uuid.UUID(job.params["notebook_id"])) | |
| return Done({"removed": removed}, f"Removed {removed} orphaned chats") | |
| def _require_room(db: Session, source: Source, url: str) -> None: | |
| """Re-reading a post that is already indexed is fine; a new one needs room under the plan's post limit.""" | |
| known = db.scalar(select(Document.id).where(Document.source_id == source.id, Document.url == url, | |
| Document.path.is_not(None))) | |
| plan = effective_plan(db, db.get(Workspace, source.workspace_id)) | |
| if not known and post_room(db, source.workspace_id, plan) <= 0: | |
| raise PermanentJobError(f"Your plan indexes up to {plan.indexed_posts} posts, and they are all used. " | |
| "Remove a post or upgrade to add more.") | |
| def handle_import_url(db: Session, job: Job): | |
| source = db.get(Source, uuid.UUID(job.params["source_id"])) | |
| update_job(db, job, progress=0.1, message="Fetching the article") | |
| try: | |
| entry = extract_article(job.params["url"]) | |
| except FeedNotFound as exc: | |
| raise PermanentJobError(str(exc)) from exc | |
| _require_room(db, source, entry.url) | |
| indexed, _, changed = store_entries(db, source, [entry]) | |
| _after_import(db, job, source.workspace_id, [str(i) for i in changed]) | |
| return Done({"indexed": indexed, "document_ids": [str(i) for i in changed]}, f"{entry.title[:80]} is in orbit") | |
| def handle_import_upload(db: Session, job: Job): | |
| source = db.get(Source, uuid.UUID(job.params["source_id"])) | |
| upload = db.get(Upload, uuid.UUID(job.params["upload_id"])) | |
| update_job(db, job, progress=0.1, message=f"Reading {upload.filename}") | |
| entry = entry_from_upload(upload.filename, upload.content_type, storage.get_bytes(upload.key)) | |
| if not entry.sections and not entry.html: | |
| raise PermanentJobError("We could not read any text in that file.") | |
| _require_room(db, source, entry.url) | |
| indexed, _, changed = store_entries(db, source, [entry]) | |
| _after_import(db, job, source.workspace_id, [str(i) for i in changed]) | |
| return Done({"indexed": indexed, "document_ids": [str(i) for i in changed]}, f"{entry.title[:80]} is in orbit") | |
| def handle_topics(db: Session, job: Job): | |
| return Done(generate.extract_topics(db, job, job.workspace_id, job.params.get("document_ids")), "Topic map ready") | |
| def handle_ideas(db: Session, job: Job): | |
| result = generate.extract_ideas(db, job, job.workspace_id, job.params.get("document_ids")) | |
| if result["extracted"]: | |
| report_suggestions.queue_warmup(db, job.workspace_id) # their notebooks' suggested report templates are written now | |
| return Done(result, f"Read {result['extracted']} posts for their ideas") | |
| def handle_report_templates(db: Session, job: Job): | |
| written = report_suggestions.warm_workspace(db, job.workspace_id) | |
| return Done({"written": written}, f"Prepared report templates for {written} notebooks") | |
| def handle_voice_profile(db: Session, job: Job): | |
| try: | |
| result = generate.build_voice_profile(db, job, job.workspace_id, job.params.get("document_ids") or []) | |
| except NothingToDo as exc: | |
| raise PermanentJobError(str(exc)) from exc | |
| return Done(result, "Voice profile ready") | |
| def handle_voice_clone(db: Session, job: Job): | |
| consent = db.get(VoiceConsent, uuid.UUID(job.params["consent_id"])) | |
| if not consent or consent.revoked_at: | |
| return Done({}, "Consent was revoked") | |
| sample_keys = job.params.get("sample_keys") or [consent.sample_key] | |
| update_job(db, job, progress=0.15, message=f"Sending {len(sample_keys)} recording(s) to ElevenLabs") | |
| samples = [] | |
| for key in sample_keys: | |
| head = storage.head(key) or {} | |
| samples.append((key.rsplit("/", 1)[-1], storage.get_bytes(key), head.get("ContentType") or "audio/webm")) | |
| owner = db.get(User, consent.user_id) | |
| try: | |
| voice_id = tts.clone_voice(f"{(owner.name if owner and owner.name else 'My')} (Notestack)", samples, | |
| remove_background_noise=bool(job.params.get("remove_background_noise", True))) | |
| except tts.TTSError as exc: | |
| raise PermanentJobError(str(exc)) from exc | |
| consent.elevenlabs_voice_id = voice_id | |
| db.commit() | |
| # Save the new voice to the workspace's voices (and keep it as the default host A), then record a short preview | |
| # so the writer can judge it straight away. | |
| update_job(db, job, progress=0.7, message="Recording a preview in your voice") | |
| vp = db.scalar(select(VoiceProfile).where(VoiceProfile.workspace_id == consent.workspace_id)) | |
| if not vp: | |
| vp = VoiceProfile(workspace_id=consent.workspace_id) | |
| db.add(vp) | |
| vp.host_voices = {**(vp.host_voices or {}), "host_a": voice_id} | |
| save_notestack_voice(db, consent.workspace_id, voice_id, "My voice") | |
| db.commit() | |
| preview_text = job.params.get("preview_text") or ( | |
| "Hi, this is my Notestack voice. From now on, my audio overviews and videos can sound like me.") | |
| preview_key = None | |
| try: | |
| clip = tts.synthesize(consent.workspace_id, preview_text[:600], voice_id) | |
| preview_key = storage.put_bytes(keys.artifact(consent.workspace_id, consent.id, "voice-preview", "mp3"), | |
| clip, "audio/mpeg") | |
| record_usage(db, workspace_id=consent.workspace_id, kind="tts", provider="elevenlabs", | |
| quantity=round(tts.mp3_seconds(clip), 1), unit="seconds", job=job) | |
| except tts.TTSError: | |
| log.warning("voice preview failed for %s", voice_id, exc_info=True) | |
| return Done({"voice_id": voice_id, "preview_key": preview_key}, "Your voice is ready and set as host A") | |
| def handle_resurface(db: Session, job: Job): | |
| return Done(generate.score_evergreen(db, job, job.workspace_id), "Evergreen scores updated") | |
| def _artifact_handler(fn: Callable[[Session, Job, Artifact], object], done_message: str): | |
| def handler(db: Session, job: Job): | |
| artifact = db.get(Artifact, uuid.UUID(job.params["artifact_id"])) | |
| if not artifact: | |
| return Done({}, "Artifact was deleted") | |
| artifact.status = "generating" | |
| db.commit() | |
| try: | |
| result = fn(db, job, artifact) | |
| except (NothingToDo, tts.TTSError) as exc: | |
| raise PermanentJobError(str(exc)) from exc | |
| if result is None: | |
| return HANDED_OFF # the renderer reports back through /api/internal/jobs/{id}/progress | |
| return Done(result if isinstance(result, dict) else {}, done_message) | |
| return handler | |
| def handle_render(db: Session, job: Job): | |
| """Direct render of caller supplied props (POST /api/artifacts/{id}/render).""" | |
| artifact = db.get(Artifact, uuid.UUID(job.params["artifact_id"])) | |
| update_job(db, job, progress=0.01, message="T-minus: preparing render") | |
| request_render(job, artifact, job.params["composition"], job.params.get("props") or {}) | |
| return HANDED_OFF | |
| def handle_publish_post(db: Session, job: Job): | |
| item = launchpad.publish_job(db, job.params["item_id"]) | |
| return Done({"status": item.status if item else "gone"}, "Published" if item and item.status == "posted" | |
| else f"Post {item.status}" if item else "Post deleted") | |
| HANDLERS: dict[str, Callable[[Session, Job], object]] = { | |
| "ingest": handle_ingest, | |
| "import_url": handle_import_url, | |
| "import_upload": handle_import_upload, | |
| "topics": handle_topics, | |
| "publish_post": handle_publish_post, | |
| "ideas": handle_ideas, | |
| "report_templates": handle_report_templates, | |
| "memory_update": handle_memory, | |
| "chat_memory": handle_chat_memory, | |
| "chat_memory_rebuild": handle_chat_memory_rebuild, | |
| "chat_memory_cleanup": handle_chat_memory_cleanup, | |
| "voice_profile": handle_voice_profile, | |
| "voice_clone": handle_voice_clone, | |
| "resurface_scan": handle_resurface, | |
| "summary": _artifact_handler(generate.summarize_notebook, "Summary ready"), | |
| "mind_map": _artifact_handler(generate.build_mind_map, "Mind Constellation ready"), | |
| "quiz": _artifact_handler(quiz.build_quiz, "Quiz ready"), | |
| "flashcards": _artifact_handler(flashcards.build_flashcards, "Flashcards ready"), | |
| "report": _artifact_handler(report.build_report, "Report ready"), | |
| "infographic": _artifact_handler(infographic.build_infographic, "Infographic ready"), | |
| "slide_deck": _artifact_handler(slides.build_slide_deck, "Slide deck ready"), | |
| "report_block": report.handle_report_block, | |
| "audio_overview": _artifact_handler(media.audio_overview, "Audio overview ready"), | |
| "video": _artifact_handler(media.make_video, "Audiogram ready"), # audiograms; videos come from blog2video | |
| "quote_card": _artifact_handler(media.make_quote_card, "Quote card ready"), | |
| "carousel": _artifact_handler(media.render_carousel, "Carousel ready"), | |
| "launch_kit": _artifact_handler(launchkit.build_launch_kit, "Launch Kit ready"), | |
| "render": handle_render, | |
| } | |
| def mark_artifact_failed(db: Session, job: Job, error: str) -> None: | |
| # A failed "add a visual" job leaves the report as it was; the job's own error tells the page what went wrong. | |
| if job.artifact_id and job.kind != "report_block": | |
| artifact = db.get(Artifact, job.artifact_id) | |
| if artifact: | |
| artifact.status = "failed" | |
| artifact.content_json = {**(artifact.content_json or {}), "error": error[:500]} | |
| db.commit() | |
| def run_job( | |
| job_id: uuid.UUID, | |
| session_factory: Callable[[], Session] = SessionLocal, | |
| handlers: dict | None = None, | |
| ) -> None: | |
| handlers = handlers or HANDLERS | |
| with session_factory() as db: | |
| job = db.get(Job, job_id) | |
| if not job: | |
| return | |
| handler = handlers.get(job.kind) | |
| if not handler: | |
| update_job(db, job, status="failed", error=f"No handler for job kind {job.kind!r}") | |
| return | |
| try: | |
| outcome = handler(db, job) | |
| except PermanentJobError as exc: | |
| db.rollback() | |
| job = db.get(Job, job_id) | |
| update_job(db, job, status="failed", error=str(exc), message=str(exc)[:500]) | |
| mark_artifact_failed(db, job, str(exc)) | |
| return | |
| except Exception as exc: | |
| # Roll back first: after a failed flush, touching `job` raises PendingRollbackError. | |
| db.rollback() | |
| job = db.get(Job, job_id) | |
| log.exception("job %s (%s) attempt %s failed", job_id, job.kind, job.attempts) | |
| if not retry_or_fail(db, job, f"{type(exc).__name__}: {exc}"): | |
| mark_artifact_failed(db, job, f"{type(exc).__name__}: {exc}") | |
| return | |
| if outcome is HANDED_OFF: | |
| return | |
| done = outcome if isinstance(outcome, Done) else Done() | |
| update_job(db, job, status="done", progress=1.0, message=done.message, result=done.result) | |
| def recover_stale(session_factory: Callable[[], Session] = SessionLocal) -> int: | |
| with session_factory() as db: | |
| jobs = stale_jobs(db) | |
| for job in jobs: | |
| log.warning("job %s stale (worker %s), requeueing", job.id, job.locked_by) | |
| retry_or_fail(db, job, "Worker stopped responding") | |
| db.commit() | |
| return len(jobs) | |
| # Periodic tasks | |
| def single_flight(name: str): | |
| """Postgres advisory lock on a dedicated connection; yields False if another worker holds it.""" | |
| if engine.dialect.name != "postgresql": | |
| yield True | |
| return | |
| with engine.connect() as conn: | |
| got = conn.execute(text("SELECT pg_try_advisory_lock(hashtext(:n))"), {"n": name}).scalar() | |
| try: | |
| yield bool(got) | |
| finally: | |
| if got: | |
| conn.execute(text("SELECT pg_advisory_unlock(hashtext(:n))"), {"n": name}) | |
| conn.commit() | |
| def topics_backfill(session_factory: Callable[[], Session] = SessionLocal) -> None: | |
| """Every indexed post gets topics: one whose tagging failed (the model timed out, rate limited, answered nothing) | |
| or that came before posts were tagged on import is asked for again here, one topics job per workspace that has | |
| some, unless one is already queued. A post that keeps failing stops being asked (generate.MAX_TAG_FAILURES).""" | |
| if not settings.llm_api_key: | |
| return | |
| with session_factory() as db: | |
| rows = db.execute(select(Document.workspace_id, Document.metadata_json).where(Document.path.is_not(None))) | |
| for workspace_id in {ws for ws, meta in rows if generate.needs_topics(meta)}: | |
| if not pending_job(db, workspace_id, "topics"): | |
| create_job(db, workspace_id, "topics", {}, max_attempts=2) | |
| def update_email_batch(session_factory: Callable[[], Session] = SessionLocal) -> None: | |
| """Sends one daily batch per scheduled campaign at its send hour (UTC).""" | |
| now = datetime.now(UTC) | |
| with session_factory() as db: | |
| campaigns = db.scalars(select(UpdateEmail).where(UpdateEmail.status.in_(["scheduled", "running"]))).all() | |
| for campaign in campaigns: | |
| hour = campaign.send_hour if campaign.send_hour >= 0 else settings.update_email_send_hour | |
| if now.hour != hour: | |
| continue | |
| if campaign.last_batch_at and campaign.last_batch_at.date() == now.date(): | |
| continue | |
| already = select(UpdateEmailSend.user_id).where(UpdateEmailSend.update_email_id == campaign.id) | |
| q = select(User).where( | |
| User.is_active.is_(True), User.email_unsubscribed.is_(False), User.id.not_in(already) | |
| ) | |
| if campaign.user_filter != "all": | |
| plans = {"paid": ["writer", "studio"]}.get(campaign.user_filter, [campaign.user_filter]) | |
| q = ( | |
| q.join(WorkspaceMember, WorkspaceMember.user_id == User.id) | |
| .join(Subscription, Subscription.workspace_id == WorkspaceMember.workspace_id) | |
| .where(Subscription.plan.in_(plans)) | |
| ) | |
| users = db.scalars(q.limit(campaign.batch_size)).unique().all() | |
| campaign.status = "running" | |
| for user in users: | |
| ok = email_service.send_blast_email( | |
| str(user.id), user.email, user.name, campaign.subject, campaign.body | |
| ) | |
| db.add(UpdateEmailSend(update_email_id=campaign.id, user_id=user.id, status="sent" if ok else "failed")) | |
| campaign.sent_count += int(ok) | |
| campaign.failed_count += int(not ok) | |
| campaign.total_users = campaign.sent_count + campaign.failed_count | |
| campaign.last_batch_at = now | |
| if len(users) < campaign.batch_size: | |
| campaign.status = "completed" | |
| db.commit() | |
| def daily_cleanup(session_factory: Callable[[], Session] = SessionLocal) -> None: | |
| with session_factory() as db: | |
| purge_old_codes(db) | |
| sweep_orphan_uploads(session_factory) # upload files whose record was deleted, and unfinished uploads | |
| class Periodic: | |
| name: str | |
| every_seconds: float | |
| fn: Callable[[], object] | |
| next_at: float = 0.0 | |
| # Run on the wall clock's multiples of every_seconds (1800: at :00 and :30 UTC), not every_seconds after the last | |
| # run. The first run is still at start-up, to catch up on anything already due. | |
| align: bool = False | |
| def schedule_next(self, now: float, wall: float | None = None) -> None: | |
| if not self.align: | |
| self.next_at = now + self.every_seconds | |
| return | |
| wall = time.time() if wall is None else wall | |
| # A couple of seconds past the boundary, so a post scheduled exactly on it is already due. | |
| self.next_at = now + (self.every_seconds - wall % self.every_seconds) + ALIGN_GRACE_SECONDS | |
| ALIGN_GRACE_SECONDS = 2 | |
| PERIODIC = [ | |
| Periodic("recover_stale", 60, recover_stale), | |
| # Posts are scheduled on half hours (in the user's zone), so they go out on the :00 / :30 run. Zones on a :45 | |
| # offset (Nepal) get slots between runs: those go out at the next run, up to 15 minutes late. | |
| Periodic("publish_due", 1800, launchpad.publish_due, align=True), | |
| Periodic("social_health", 600, social_health), | |
| Periodic("sync_engagement", 3600, launchpad.sync_engagement), | |
| Periodic("update_email_batch", 300, update_email_batch), | |
| Periodic("topics_backfill", 900, topics_backfill), | |
| Periodic("daily_cleanup", 24 * 3600, daily_cleanup), | |
| Periodic("video_period_reset", 3600, video_quota.reset_due_video_periods), | |
| Periodic("video_refund_sweep", 600, video_quota.sweep_failed_videos), | |
| Periodic("video_capacity_check", 3600, video_quota.check_capacity), | |
| ] | |
| def run_periodic(now: float) -> None: | |
| for task in PERIODIC: | |
| if now < task.next_at: | |
| continue | |
| task.schedule_next(now) | |
| try: | |
| with single_flight(task.name) as got: | |
| if got: | |
| task.fn() | |
| except Exception: | |
| log.exception("periodic task %s failed", task.name) | |
| # Main loop | |
| def _log_crash(job_id: uuid.UUID) -> Callable[[Future], None]: | |
| """run_job handles job errors itself; anything escaping it would otherwise vanish inside the Future and leave | |
| the job 'running' until recover_stale.""" | |
| def callback(future: Future) -> None: | |
| if exc := future.exception(): | |
| log.error("job %s crashed outside its error handling", job_id, exc_info=exc) | |
| return callback | |
| def run_loop(stop: threading.Event, worker_id: str | None = None) -> None: | |
| """Claim and run jobs until `stop` is set. Used by the CLI and, in development, by the API.""" | |
| worker_id = worker_id or f"{socket.gethostname()}:{os.getpid()}" | |
| pool = ThreadPoolExecutor(max_workers=settings.worker_concurrency, thread_name_prefix="job") | |
| running: set[Future] = set() | |
| log.info("worker %s up, concurrency=%s", worker_id, settings.worker_concurrency) | |
| while not stop.is_set(): | |
| running = {f for f in running if not f.done()} | |
| run_periodic(time.monotonic()) | |
| claimed = False | |
| while len(running) < settings.worker_concurrency and not stop.is_set(): | |
| try: | |
| with SessionLocal() as db: | |
| job_id = claim_next(db, worker_id) | |
| except Exception: | |
| log.exception("claim failed") | |
| break | |
| if not job_id: | |
| break | |
| future = pool.submit(run_job, job_id) | |
| future.add_done_callback(_log_crash(job_id)) | |
| running.add(future) | |
| claimed = True | |
| if not claimed: | |
| stop.wait(settings.worker_poll_seconds) | |
| log.info("worker %s stopping, waiting for %d running jobs", worker_id, len(running)) | |
| pool.shutdown(wait=True) | |
| def start_background(stop: threading.Event) -> threading.Thread: | |
| thread = threading.Thread(target=run_loop, args=(stop, f"api:{os.getpid()}"), name="worker", daemon=True) | |
| thread.start() | |
| return thread | |
| def main() -> None: | |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s") | |
| stop = threading.Event() | |
| for sig in (signal.SIGINT, signal.SIGTERM): | |
| signal.signal(sig, lambda *_: stop.set()) | |
| run_loop(stop) | |
| if __name__ == "__main__": | |
| main() | |