strepsiades
207 lines
7.0 KiB
Python
207 lines
7.0 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from collections import deque
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any, Deque, Dict, List, Optional, Tuple
|
|
|
|
import sync_state_machine.domain # noqa: F401
|
|
from sync_state_machine.pipeline import (
|
|
FullSyncPipeline,
|
|
create_pipeline_from_config,
|
|
)
|
|
from sync_state_machine.config import (
|
|
PipelineRunConfig,
|
|
JsonlDataSourceConfig,
|
|
ApiDataSourceConfig,
|
|
apply_config_overrides,
|
|
resolve_log_file_path,
|
|
)
|
|
from sync_state_machine.pipeline.run_logger import clear_pipeline_logger_handlers, init_pipeline_logger
|
|
|
|
from .config_models import PipelineUiConfig, default_ui_config
|
|
|
|
|
|
@dataclass
|
|
class RunnerStatus:
|
|
running: bool = False
|
|
started_at: Optional[str] = None
|
|
ended_at: Optional[str] = None
|
|
last_error: str = ""
|
|
last_stats: Optional[Dict[str, Any]] = None
|
|
current_log_file: str = ""
|
|
|
|
|
|
class _UiLogHandler(logging.Handler):
|
|
def __init__(self, callback):
|
|
super().__init__(level=logging.INFO)
|
|
self._callback = callback
|
|
|
|
def emit(self, record: logging.LogRecord) -> None:
|
|
try:
|
|
self._callback(self.format(record) + "\n")
|
|
except Exception:
|
|
return
|
|
|
|
|
|
class PipelineRunnerService:
|
|
def __init__(self) -> None:
|
|
self._status = RunnerStatus(last_stats={})
|
|
self._config = default_ui_config()
|
|
self._task: Optional[asyncio.Task] = None
|
|
self._log_chunks: Deque[str] = deque(maxlen=5000)
|
|
self._lock = asyncio.Lock()
|
|
|
|
def get_default_config(self) -> PipelineUiConfig:
|
|
return default_ui_config()
|
|
|
|
def get_config(self) -> PipelineUiConfig:
|
|
return self._config
|
|
|
|
def set_config(self, cfg: PipelineUiConfig) -> PipelineUiConfig:
|
|
self._config = cfg
|
|
return self._config
|
|
|
|
def get_status(self) -> RunnerStatus:
|
|
return self._status
|
|
|
|
def get_logs(self, since: int = 0) -> Tuple[List[str], int]:
|
|
chunks = list(self._log_chunks)
|
|
if since < 0:
|
|
since = 0
|
|
if since >= len(chunks):
|
|
return [], len(chunks)
|
|
return chunks[since:], len(chunks)
|
|
|
|
async def start(self, cfg: Optional[PipelineUiConfig] = None) -> RunnerStatus:
|
|
async with self._lock:
|
|
if self._task and not self._task.done():
|
|
raise RuntimeError("Pipeline is already running")
|
|
|
|
if cfg is not None:
|
|
self._config = cfg
|
|
|
|
self._log_chunks.clear()
|
|
self._status = RunnerStatus(
|
|
running=True,
|
|
started_at=datetime.now().isoformat(timespec="seconds"),
|
|
ended_at=None,
|
|
last_error="",
|
|
last_stats={},
|
|
current_log_file="",
|
|
)
|
|
self._task = asyncio.create_task(self._run(self._config))
|
|
return self._status
|
|
|
|
async def stop(self) -> RunnerStatus:
|
|
async with self._lock:
|
|
task = self._task
|
|
if not task or task.done():
|
|
return self._status
|
|
task.cancel()
|
|
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
return self._status
|
|
|
|
def _append_log(self, data: str) -> int:
|
|
self._log_chunks.append(data)
|
|
return len(data)
|
|
|
|
def _append_log_callback(self, data: str) -> None:
|
|
self._append_log(data)
|
|
|
|
def write(self, data: str) -> int:
|
|
return self._append_log(data)
|
|
|
|
def flush(self) -> None:
|
|
return None
|
|
|
|
async def _run(self, cfg: PipelineUiConfig) -> None:
|
|
pipeline = None
|
|
ui_handler = None
|
|
app_logger = None
|
|
|
|
log_cfg = cfg.logging.model_copy(deep=True)
|
|
log_path = resolve_log_file_path(log_cfg)
|
|
self._status.current_log_file = str(log_path or "")
|
|
|
|
try:
|
|
app_logger = init_pipeline_logger(
|
|
log_file_path=log_path,
|
|
mirror_console=True,
|
|
replace_handlers=True,
|
|
)
|
|
ui_handler = _UiLogHandler(self._append_log_callback)
|
|
ui_handler.setFormatter(logging.Formatter("%(message)s"))
|
|
app_logger.addHandler(ui_handler)
|
|
|
|
if cfg.mode == "jsonl_to_api":
|
|
pipeline = await self._create_jsonl_to_api_pipeline(cfg)
|
|
else:
|
|
pipeline = await self._create_jsonl_to_jsonl_pipeline(cfg)
|
|
|
|
stats = await pipeline.run()
|
|
self._status.last_stats = stats
|
|
self._status.last_error = ""
|
|
except asyncio.CancelledError:
|
|
self._status.last_error = "cancelled"
|
|
raise
|
|
except Exception as exc:
|
|
self._status.last_error = f"{type(exc).__name__}: {exc}"
|
|
finally:
|
|
try:
|
|
if pipeline is not None:
|
|
await pipeline.close()
|
|
finally:
|
|
if app_logger is not None and ui_handler is not None:
|
|
app_logger.removeHandler(ui_handler)
|
|
try:
|
|
ui_handler.close()
|
|
except Exception:
|
|
pass
|
|
clear_pipeline_logger_handlers(app_logger)
|
|
self._status.running = False
|
|
self._status.ended_at = datetime.now().isoformat(timespec="seconds")
|
|
|
|
async def _create_jsonl_to_api_pipeline(self, cfg: PipelineUiConfig) -> FullSyncPipeline:
|
|
local_ds = cfg.local_datasource.model_copy(deep=True)
|
|
remote_ds = cfg.remote_datasource.model_copy(deep=True)
|
|
local_ds = JsonlDataSourceConfig.model_validate(local_ds.model_dump())
|
|
remote_ds = ApiDataSourceConfig.model_validate(remote_ds.model_dump())
|
|
|
|
run_cfg = PipelineRunConfig(
|
|
node_types=list(cfg.node_types),
|
|
local_datasource=local_ds,
|
|
remote_datasource=remote_ds,
|
|
strategies=[],
|
|
persist=cfg.persist.model_copy(deep=True),
|
|
logging=cfg.logging.model_copy(deep=True),
|
|
)
|
|
if cfg.strategies:
|
|
run_cfg = apply_config_overrides(run_cfg, {"strategies": dict(cfg.strategies)})
|
|
return await create_pipeline_from_config(run_cfg)
|
|
|
|
async def _create_jsonl_to_jsonl_pipeline(self, cfg: PipelineUiConfig) -> FullSyncPipeline:
|
|
local_ds = cfg.local_datasource.model_copy(deep=True)
|
|
remote_ds = cfg.remote_datasource.model_copy(deep=True)
|
|
local_ds = JsonlDataSourceConfig.model_validate(local_ds.model_dump())
|
|
remote_ds = JsonlDataSourceConfig.model_validate(remote_ds.model_dump())
|
|
|
|
run_cfg = PipelineRunConfig(
|
|
node_types=list(cfg.node_types),
|
|
local_datasource=local_ds,
|
|
remote_datasource=remote_ds,
|
|
strategies=[],
|
|
persist=cfg.persist.model_copy(deep=True),
|
|
logging=cfg.logging.model_copy(deep=True),
|
|
)
|
|
if cfg.strategies:
|
|
run_cfg = apply_config_overrides(run_cfg, {"strategies": dict(cfg.strategies)})
|
|
return await create_pipeline_from_config(run_cfg)
|