luke_scribe/src/luke_scribe/jobqueue/worker.py

"""Batch 워커 — 큐에서 Job을 소비해 전사 파이프라인 실행.

- no-fork/장수명: 프로세스 수명 동안 EngineOwner가 모델을 보유.
- 리스(하트비트) 갱신 + progress throttle(≥1s 또는 의미 있는 delta).
- 세그먼트 경계 취소 + 종료 경로 임시파일 정리.
- OOM 강등: ``attempted_profiles``를 Job에 영속화 (재시도 2회 캡).
- 스타트업 reconciler: stale processing 복구.
"""

from __future__ import annotations

import threading
import time
from dataclasses import dataclass

from ..config import Settings
from ..results.store import ResultStore
from .broker import JobBroker, make_broker
from .cancel import CancellationToken
from .jobs import Job, JobStatus

LEASE_TTL_SEC = 60.0
PROGRESS_INTERVAL_SEC = 1.0


@dataclass
class WorkerCallbacks:
    on_job_done: callable | None = None  # (job, result) — 테스트 훅


class Worker:
    def __init__(
        self,
        *,
        settings: Settings | None = None,
        broker: JobBroker | None = None,
        store: ResultStore | None = None,
        owner=None,  # EngineOwner
        worker_id: str = "w1",
        callbacks: WorkerCallbacks | None = None,
        ingestor=None,  # AudioIngestor (테스트 주입용)
    ) -> None:
        self.settings = settings or Settings()
        self.broker = broker or make_broker(self.settings)
        self.store = store or ResultStore(self.settings.results_root)
        from ..engine.owner import EngineOwner

        self.owner = owner or EngineOwner.get(self.settings)
        self.worker_id = worker_id
        self.callbacks = callbacks or WorkerCallbacks()
        self.ingestor = ingestor
        self._stop = threading.Event()
        self._active_job: Job | None = None

    def start_reconcile(self) -> list[str]:
        return self.broker.reconcile_stale(LEASE_TTL_SEC)

    def run_forever(self) -> None:
        while not self._stop.is_set():
            if self._claim_and_process() is None:
                time.sleep(0.2)

    def drain(self, max_jobs: int = 100) -> int:
        """큐가 빌 때까지 처리 후 종료 (테스트/일회성 실행용). 처리한 job 수 반환."""
        processed = 0
        while processed < max_jobs:
            if self._claim_and_process() is None:
                break
            processed += 1
        return processed

    def _claim_and_process(self) -> Job | None:
        """job 1개 claim/처리. 처리할 게 없으면 None."""
        job = self.broker.claim_next(self.worker_id, LEASE_TTL_SEC)
        if job is None:
            return None
        self._active_job = job
        try:
            self._process(job)
        finally:
            self._active_job = None
        return job

    def _process(self, job: Job) -> None:
        from ..engine.base import TranscriptionOptions
        from ..pipeline.batch import BatchPipeline

        token = CancellationToken()
        token.cancel() if job.cancel_requested else None

        pipeline = BatchPipeline(
            settings=self.settings,
            store=self.store,
            owner=self.owner,
            token=token,
            ingestor=self.ingestor,
        )

        def progress_cb(processed_sec: float, total_sec: float) -> None:
            self._emit_progress(job, processed_sec, total_sec)

        try:
            # API의 TranscribeOptions에는 엔진과 무관한 키가 포함될 수 있다
            # (formats/timestamps/glossary_id/post_correction/diarize 등) —
            # 엔진 계약 필드만 골라 전달한다.
            options = TranscriptionOptions(
                **{
                    k: v
                    for k, v in (job.options or {}).items()
                    if k in TranscriptionOptions.__slots__
                }
            )
            result = pipeline.run(job, options, progress_cb=progress_cb)
            # 결과를 먼저 영속화한 뒤 상태 전이 (실패 시 FAILED로 전이 가능하게)
            self.store.write_result(job.id, result)
            current = self._transition(job, JobStatus.COMPLETED)
            if self.settings.delete_source and job.source_path:
                self.store.delete_derived(job.id)
            self.store.write_source_metadata(current.id, current.to_meta())
            if self.callbacks.on_job_done:
                self.callbacks.on_job_done(current, result)
        except Exception as exc:
            from ..errors import CancelledError

            if isinstance(exc, CancelledError) or token.is_cancelled():
                current = self._transition(job, JobStatus.CANCELLED)
                current.cancelled_at = time.time()
            else:
                current = self._transition(job, JobStatus.FAILED)
                current.error_code = getattr(exc, "code", "transcription_failed")
                current.error_message = str(exc)
                current.failed_at = time.time()
            # 오류/완료 필드가 포함된 meta를 재저장 (Redis: transition이 별도 인스턴스 반환)
            self.store.write_source_metadata(current.id, current.to_meta())

    def _transition(self, job: Job, to: JobStatus) -> Job:
        """브로커 전이. Redis는 새 인스턴스를 반환하므로 그 인스턴스를 돌려받는다."""
        try:
            return self.broker.transition(job, to)
        except Exception:
            # 전이가 이미 terminal 상태로 막혔으면 원본 상태 그대로
            return job

    def _emit_progress(self, job: Job, processed_sec: float, total_sec: float) -> None:
        job.progress = min(1.0, processed_sec / total_sec) if total_sec else None
        job.processed_sec = processed_sec
        job.total_sec = total_sec
        job.refresh_lease(LEASE_TTL_SEC)
        self.broker.save_meta(job)

    def stop(self) -> None:
        self._stop.set()