318 lines
10 KiB
Python
318 lines
10 KiB
Python
# Copyright 2026 Umbra Tools (umbratools.com)
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import time
|
|
import traceback
|
|
from dataclasses import dataclass, field
|
|
from enum import StrEnum
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
from pathlib import Path
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class PipelineState(StrEnum):
|
|
IDLE = "idle"
|
|
QUEUED = "queued"
|
|
RUNNING = "running"
|
|
PAUSED = "paused"
|
|
COMPLETED = "completed"
|
|
FAILED = "failed"
|
|
CANCELLED = "cancelled"
|
|
|
|
|
|
class PipelineStep(StrEnum):
|
|
UPLOAD = "upload"
|
|
SCAN = "scan"
|
|
DISCOVER = "discover"
|
|
VOICES = "voices"
|
|
SEGMENT = "segment"
|
|
ATTRIBUTE = "attribute"
|
|
GENERATE = "generate"
|
|
EXPORT = "export"
|
|
|
|
|
|
STEP_ORDER = [
|
|
PipelineStep.UPLOAD,
|
|
PipelineStep.SCAN,
|
|
PipelineStep.DISCOVER,
|
|
PipelineStep.VOICES,
|
|
PipelineStep.SEGMENT,
|
|
PipelineStep.ATTRIBUTE,
|
|
PipelineStep.GENERATE,
|
|
PipelineStep.EXPORT,
|
|
]
|
|
|
|
|
|
@dataclass
|
|
class PipelineProgress:
|
|
step: PipelineStep = PipelineStep.UPLOAD
|
|
state: PipelineState = PipelineState.IDLE
|
|
detail: str = ""
|
|
progress: int = 0
|
|
error: str = ""
|
|
started_at: float | None = None
|
|
updated_at: float | None = None
|
|
completed_at: float | None = None
|
|
chapters_completed: list[int] = field(default_factory=list)
|
|
total_chapters: int = 0
|
|
current_chapter: int = 0
|
|
completed_utterances: int = 0
|
|
total_utterances: int = 0
|
|
warnings: list[str] = field(default_factory=list)
|
|
errors: list[str] = field(default_factory=list)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"step": self.step.value,
|
|
"state": self.state.value,
|
|
"detail": self.detail,
|
|
"progress": self.progress,
|
|
"error": self.error,
|
|
"started_at": self.started_at,
|
|
"updated_at": self.updated_at,
|
|
"completed_at": self.completed_at,
|
|
"chapters_completed": self.chapters_completed,
|
|
"total_chapters": self.total_chapters,
|
|
"current_chapter": self.current_chapter,
|
|
"completed_utterances": self.completed_utterances,
|
|
"total_utterances": self.total_utterances,
|
|
"warnings": self.warnings,
|
|
"errors": self.errors,
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: dict[str, Any]) -> PipelineProgress:
|
|
raw_state = data.get("state") or data.get("status") or "idle"
|
|
if raw_state == "complete":
|
|
raw_state = "completed"
|
|
if raw_state == "error":
|
|
raw_state = "failed"
|
|
if raw_state in ("starting", "queued", "generating", "segmenting", "attributing", "saving"):
|
|
raw_state = "running"
|
|
step = data.get("step", "upload")
|
|
if step == "segmenting":
|
|
step = "segment"
|
|
return cls(
|
|
step=PipelineStep(step),
|
|
state=PipelineState(raw_state),
|
|
detail=data.get("detail", ""),
|
|
progress=data.get("progress", 0),
|
|
error=data.get("error", ""),
|
|
started_at=data.get("started_at"),
|
|
updated_at=data.get("updated_at"),
|
|
completed_at=data.get("completed_at"),
|
|
chapters_completed=data.get("chapters_completed", []),
|
|
total_chapters=data.get("total_chapters", 0),
|
|
current_chapter=data.get("current_chapter", 0),
|
|
completed_utterances=data.get("completed_utterances", 0),
|
|
total_utterances=data.get("total_utterances", 0),
|
|
warnings=data.get("warnings", []),
|
|
errors=data.get("errors", []),
|
|
)
|
|
|
|
|
|
class PipelineManager:
|
|
def __init__(self, data_dir: Path) -> None:
|
|
self.data_dir = data_dir
|
|
self._status_path = data_dir / "pipeline_status.json"
|
|
self._checkpoint_path = data_dir / "pipeline_checkpoint.json"
|
|
self._progress = PipelineProgress()
|
|
self._task: asyncio.Task | None = None
|
|
self._cancelled = False
|
|
self._paused = False
|
|
self._callbacks: list[Callable[[PipelineProgress], None]] = []
|
|
|
|
@property
|
|
def progress(self) -> PipelineProgress:
|
|
return self._progress
|
|
|
|
@property
|
|
def is_running(self) -> bool:
|
|
return self._progress.state == PipelineState.RUNNING
|
|
|
|
@property
|
|
def is_paused(self) -> bool:
|
|
return self._progress.state == PipelineState.PAUSED
|
|
|
|
def register_callback(self, callback: Callable[[PipelineProgress], None]) -> None:
|
|
self._callbacks.append(callback)
|
|
|
|
def _notify(self) -> None:
|
|
for cb in self._callbacks:
|
|
try:
|
|
cb(self._progress)
|
|
except Exception as e:
|
|
logger.warning("Callback error: %s", e)
|
|
|
|
def _write_status(self) -> None:
|
|
self._progress.updated_at = time.time()
|
|
with open(self._status_path, "w") as f:
|
|
json.dump(self._progress.to_dict(), f, indent=2)
|
|
self._notify()
|
|
|
|
def _save_checkpoint(self, data: dict[str, Any]) -> None:
|
|
checkpoint = {
|
|
"step": self._progress.step.value,
|
|
"state": self._progress.state.value,
|
|
"chapters_completed": self._progress.chapters_completed,
|
|
"total_chapters": self._progress.total_chapters,
|
|
"data": data,
|
|
"timestamp": time.time(),
|
|
}
|
|
with open(self._checkpoint_path, "w") as f:
|
|
json.dump(checkpoint, f, indent=2)
|
|
|
|
def load_checkpoint(self) -> dict[str, Any] | None:
|
|
if not self._checkpoint_path.exists():
|
|
return None
|
|
try:
|
|
with open(self._checkpoint_path) as f:
|
|
return json.load(f)
|
|
except Exception:
|
|
return None
|
|
|
|
def clear_checkpoint(self) -> None:
|
|
if self._checkpoint_path.exists():
|
|
self._checkpoint_path.unlink()
|
|
|
|
def update(
|
|
self,
|
|
step: PipelineStep | None = None,
|
|
state: PipelineState | None = None,
|
|
detail: str | None = None,
|
|
progress: int | None = None,
|
|
error: str | None = None,
|
|
warnings: list[str] | None = None,
|
|
errors: list[str] | None = None,
|
|
) -> None:
|
|
if step is not None:
|
|
self._progress.step = step
|
|
if state is not None:
|
|
self._progress.state = state
|
|
if detail is not None:
|
|
self._progress.detail = detail
|
|
if progress is not None:
|
|
self._progress.progress = progress
|
|
if error is not None:
|
|
self._progress.error = error
|
|
if warnings is not None:
|
|
self._progress.warnings = warnings
|
|
if errors is not None:
|
|
self._progress.errors = errors
|
|
self._write_status()
|
|
|
|
def start(
|
|
self,
|
|
pipeline_fn: Callable[..., Any],
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
if self.is_running:
|
|
logger.warning("Pipeline already running")
|
|
return
|
|
|
|
self._cancelled = False
|
|
self._paused = False
|
|
self._progress = PipelineProgress(
|
|
state=PipelineState.RUNNING,
|
|
started_at=time.time(),
|
|
)
|
|
self._write_status()
|
|
|
|
loop = asyncio.get_event_loop()
|
|
self._task = loop.create_task(self._run_wrapper(pipeline_fn, *args, **kwargs))
|
|
|
|
async def _run_wrapper(
|
|
self,
|
|
pipeline_fn: Callable[..., Any],
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
try:
|
|
await pipeline_fn(self, *args, **kwargs)
|
|
if not self._cancelled and not self._paused:
|
|
self._progress.state = PipelineState.COMPLETED
|
|
self._progress.completed_at = time.time()
|
|
self._write_status()
|
|
self.clear_checkpoint()
|
|
except asyncio.CancelledError:
|
|
self._progress.state = PipelineState.CANCELLED
|
|
self._write_status()
|
|
except Exception as e:
|
|
self._progress.state = PipelineState.FAILED
|
|
self._progress.error = f"{e}\n{traceback.format_exc()}"
|
|
self._write_status()
|
|
logger.error("Pipeline failed: %s", e)
|
|
|
|
def pause(self) -> None:
|
|
if not self.is_running:
|
|
return
|
|
self._paused = True
|
|
self._progress.state = PipelineState.PAUSED
|
|
self._write_status()
|
|
|
|
def resume(self, pipeline_fn: Callable[..., Any], *args: Any, **kwargs: Any) -> None:
|
|
if not self.is_paused:
|
|
return
|
|
self._paused = False
|
|
self._cancelled = False
|
|
self._progress.state = PipelineState.RUNNING
|
|
self._write_status()
|
|
|
|
loop = asyncio.get_event_loop()
|
|
self._task = loop.create_task(self._run_wrapper(pipeline_fn, *args, **kwargs))
|
|
|
|
def cancel(self) -> None:
|
|
if not self.is_running:
|
|
return
|
|
self._cancelled = True
|
|
if self._task and not self._task.done():
|
|
self._task.cancel()
|
|
|
|
def get_step_index(self, step: PipelineStep) -> int:
|
|
return STEP_ORDER.index(step)
|
|
|
|
def should_skip_step(self, step: PipelineStep) -> bool:
|
|
checkpoint = self.load_checkpoint()
|
|
if not checkpoint:
|
|
return False
|
|
checkpoint_step = PipelineStep(checkpoint.get("step", "upload"))
|
|
checkpoint_idx = self.get_step_index(step)
|
|
current_idx = self.get_step_index(checkpoint_step)
|
|
return checkpoint_idx < current_idx
|
|
|
|
def mark_chapter_complete(self, chapter_id: int) -> None:
|
|
if chapter_id not in self._progress.chapters_completed:
|
|
self._progress.chapters_completed.append(chapter_id)
|
|
self._save_checkpoint({
|
|
"chapters_completed": self._progress.chapters_completed,
|
|
})
|
|
|
|
def update_chapter_progress(
|
|
self, chapter_id: int, total: int, current: int
|
|
) -> None:
|
|
self._progress.total_chapters = total
|
|
self._progress.current_chapter = current
|
|
if total > 0:
|
|
self._progress.progress = int((current / total) * 100)
|
|
self._write_status()
|
|
|
|
def load_status(self) -> PipelineProgress:
|
|
if self._status_path.exists():
|
|
try:
|
|
with open(self._status_path) as f:
|
|
data = json.load(f)
|
|
self._progress = PipelineProgress.from_dict(data)
|
|
except Exception:
|
|
pass
|
|
return self._progress
|