stitch-durable 0.1.0__py3-none-any.whl

stitch_durable/__init__.py

@@ -0,0 +1,65 @@
+from .alerts import Alert, Alerter
+from .critical import CriticalPolicy, CriticalRegistry
+from .errors import (
+    CriticalStepFailed,
+    HealthCheckFailed,
+    NoProviderAvailable,
+    Park,
+    PermanentError,
+    PipelineHalted,
+    ProviderDown,
+    ProviderError,
+    QuotaExhausted,
+    RateLimited,
+    StepFailed,
+    TransientError,
+)
+from .execution.context import ExecutionContext
+from .execution.reporting import Reporter, RunStatistics
+from .execution.step_decorator import step
+from .providers.breaker_tuning import BreakerTuning
+from .providers.provider_pool import ProviderPool
+from .records import AuditRecord
+from .retry_policy import DEFAULT_RETRY_POLICY, NO_RETRY, RetryPolicy
+from .storage.checkpoint_store import CheckpointStore
+from .tracing import TraceRenderer
+from .workflows.recovery import WorkflowRecovery
+from .workflows.registry import WorkflowRegistry
+from .workflows.workflow import Workflow, workflow
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Alert",
+    "Alerter",
+    "CriticalPolicy",
+    "CriticalRegistry",
+    "CriticalStepFailed",
+    "HealthCheckFailed",
+    "PipelineHalted",
+    "AuditRecord",
+    "BreakerTuning",
+    "CheckpointStore",
+    "DEFAULT_RETRY_POLICY",
+    "ExecutionContext",
+    "NO_RETRY",
+    "NoProviderAvailable",
+    "Park",
+    "PermanentError",
+    "ProviderDown",
+    "ProviderError",
+    "ProviderPool",
+    "QuotaExhausted",
+    "RateLimited",
+    "Reporter",
+    "RetryPolicy",
+    "RunStatistics",
+    "StepFailed",
+    "TraceRenderer",
+    "TransientError",
+    "Workflow",
+    "WorkflowRecovery",
+    "WorkflowRegistry",
+    "step",
+    "workflow",
+]

stitch_durable/alerts.py

@@ -0,0 +1,15 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Alert:
+    severity: str
+    title: str
+    run_id: str
+    step_name: str
+    detail: dict
+
+
+class Alerter:
+    def send(self, alert: Alert) -> None:
+        return

stitch_durable/critical.py

@@ -0,0 +1,46 @@
+import asyncio
+import inspect
+from dataclasses import dataclass
+from typing import Any, Callable
+
+from .errors import describe_error
+
+
+@dataclass(frozen=True)
+class CriticalPolicy:
+    health_check: Callable[[], Any] | None = None
+    health_check_timeout_seconds: float = 10.0
+
+    async def verify(self) -> str | None:
+        if self.health_check is None:
+            return None
+        try:
+            async with asyncio.timeout(self.health_check_timeout_seconds):
+                outcome = await self._call_health_check()
+        except Exception as error:
+            return describe_error(error)
+        return "the health check returned False" if outcome is False else None
+
+    async def _call_health_check(self) -> Any:
+        if inspect.iscoroutinefunction(self.health_check):
+            return await self.health_check()
+        return await asyncio.to_thread(self.health_check)
+
+
+class CriticalRegistry:
+    _default = None
+
+    def __init__(self):
+        self._policies: dict[str, CriticalPolicy] = {}
+
+    @classmethod
+    def default(cls) -> "CriticalRegistry":
+        if cls._default is None:
+            cls._default = cls()
+        return cls._default
+
+    def register(self, step_name: str, policy: CriticalPolicy) -> None:
+        self._policies[step_name] = policy
+
+    def find(self, step_name: str) -> CriticalPolicy | None:
+        return self._policies.get(step_name)

stitch_durable/errors.py

@@ -0,0 +1,76 @@
+class TransientError(Exception):
+    def __init__(self, message: str, retry_after: float | None = None):
+        super().__init__(message)
+        self.retry_after = retry_after
+
+
+class PermanentError(Exception):
+    pass
+
+
+class ProviderError(Exception):
+    def __init__(self, message: str, retry_after: float | None = None):
+        super().__init__(message)
+        self.retry_after = retry_after
+
+
+class QuotaExhausted(ProviderError):
+    pass
+
+
+class RateLimited(ProviderError):
+    pass
+
+
+class ProviderDown(ProviderError):
+    pass
+
+
+class NoProviderAvailable(Exception):
+    def __init__(self, message: str, retry_at: float | None = None):
+        super().__init__(message)
+        self.retry_at = retry_at
+
+
+class Park(Exception):
+    def __init__(self, until: float, reason: str = ""):
+        super().__init__(reason or "parked")
+        self.until = until
+        self.reason = reason
+
+
+class StepFailed(Exception):
+    def __init__(self, step_name: str, cause: BaseException, transient: bool):
+        super().__init__(f"{step_name}: {describe_error(cause)}")
+        self.step_name = step_name
+        self.cause = cause
+        self.transient = transient
+
+
+class PipelineHalted(Exception):
+    pass
+
+
+class CriticalStepFailed(PipelineHalted):
+    def __init__(self, step_name: str, item_key: str, cause: BaseException):
+        super().__init__(
+            f"critical step {step_name} failed for {item_key}: {describe_error(cause)}"
+        )
+        self.step_name = step_name
+        self.item_key = item_key
+        self.cause = cause
+
+
+class HealthCheckFailed(PipelineHalted):
+    def __init__(self, step_name: str, reason: str):
+        super().__init__(
+            f"the run is halted and the resource behind {step_name} is still unhealthy: {reason}"
+        )
+        self.step_name = step_name
+        self.reason = reason
+
+
+def describe_error(error: BaseException) -> str:
+    if isinstance(error, TimeoutError):
+        return "timed out"
+    return str(error) or type(error).__name__

stitch_durable/execution/batch_runner.py

@@ -0,0 +1,200 @@
+import asyncio
+import time
+from dataclasses import dataclass
+from typing import Any, AsyncIterable, AsyncIterator, Awaitable, Callable, Iterable
+
+from ..alerts import Alert, Alerter
+from ..critical import CriticalRegistry
+from ..errors import CriticalStepFailed, HealthCheckFailed
+from ..retry_policy import NO_RETRY, RetryPolicy
+from ..storage.checkpoint_store import CheckpointStore
+from .context import ExecutionContext
+from .item_processor import ItemProcessor
+from .reporting import Reporter, RunStatistics
+from .run_session import RunSession
+
+PARKED_WAKE_UP_MARGIN_SECONDS = 0.05
+WORKER_STOP_SIGNAL = object()
+
+
+@dataclass(frozen=True)
+class BatchOptions:
+    concurrency: int = 1
+    item_retry_policy: RetryPolicy = NO_RETRY
+    wait_for_parked: bool = False
+    maximum_parked_wait_seconds: float | None = None
+    record_audit: bool = True
+    parameters: dict | None = None
+    alerter: Alerter | None = None
+
+
+class BatchRunner:
+    def __init__(
+        self,
+        workflow_name: str,
+        batch_label: str,
+        handler: Callable[[Any], Awaitable[Any]],
+        key_function: Callable[[Any], str],
+        store: CheckpointStore,
+        reporter: Reporter | None = None,
+        options: BatchOptions = BatchOptions(),
+    ):
+        self.workflow_name = workflow_name
+        self.options = options
+        self.store = store
+        run_id = f"{workflow_name}:{batch_label}"
+        self.session = RunSession(
+            run_id, store, reporter or Reporter(), options.record_audit, options.alerter
+        )
+        self.processor = ItemProcessor(
+            self.session, handler, key_function, options.item_retry_policy
+        )
+
+    async def run(self, items: Iterable[Any] | AsyncIterable[Any]) -> RunStatistics:
+        await self._require_healthy_resource()
+        previous_status = self.store.runs.open_run(
+            self.session.run_id, self.workflow_name, self.options.parameters or {}
+        )
+        self.session.emit(
+            "run_started",
+            run_id=self.session.run_id,
+            concurrency=self.options.concurrency,
+            resume=self._summarize_previous_run(previous_status),
+        )
+        token = ExecutionContext.bind_run(self.session)
+        try:
+            await self._process_items(items)
+            if self.options.wait_for_parked:
+                await self._retry_parked_items()
+        except CriticalStepFailed as failure:
+            self._record_halt(failure)
+            raise
+        finally:
+            ExecutionContext.release_run(token)
+        return self._finish_run()
+
+    async def _require_healthy_resource(self) -> None:
+        halt = self.store.runs.get_halt(self.session.run_id)
+        if halt is None:
+            return
+        step_name = halt[0]
+        policy = CriticalRegistry.default().find(step_name)
+        problem = await policy.verify() if policy else None
+        if problem is not None:
+            self.session.emit("run_blocked", step=step_name, reason=problem)
+            raise HealthCheckFailed(step_name, problem)
+        self.session.emit("run_recovered", step=step_name)
+        title = f"{step_name} is healthy again, the run resumes"
+        self.session.send_alert(
+            Alert("recovered", title, self.session.run_id, step_name, {})
+        )
+
+    def _record_halt(self, failure: CriticalStepFailed) -> None:
+        self.store.runs.mark_halted(
+            self.session.run_id, failure.step_name, str(failure)
+        )
+        self.session.emit(
+            "run_halted",
+            step=failure.step_name,
+            item=failure.item_key,
+            error=str(failure),
+        )
+
+    async def _process_items(self, items: Iterable[Any] | AsyncIterable[Any]) -> None:
+        try:
+            await self._feed_workers(items)
+        except BaseExceptionGroup as group:
+            critical, _ = group.split(CriticalStepFailed)
+            if critical is None:
+                raise
+            raise critical.exceptions[0] from None
+
+    async def _feed_workers(self, items: Iterable[Any] | AsyncIterable[Any]) -> None:
+        concurrency = self.options.concurrency
+        queue: asyncio.Queue = asyncio.Queue(maxsize=concurrency)
+        async with asyncio.TaskGroup() as workers:
+            for _ in range(concurrency):
+                workers.create_task(self._work_through(queue))
+            async for item in self._as_async_iterator(items):
+                await queue.put(item)
+            for _ in range(concurrency):
+                await queue.put(WORKER_STOP_SIGNAL)
+
+    async def _work_through(self, queue: asyncio.Queue) -> None:
+        while True:
+            item = await queue.get()
+            if item is WORKER_STOP_SIGNAL:
+                return
+            await self.processor.process(item)
+
+    async def _retry_parked_items(self) -> None:
+        deadline = self._parked_wait_deadline()
+        while True:
+            parked_items = self.store.items.list_parked(self.session.run_id)
+            if not parked_items:
+                return
+            due_at = min(parked.retry_at for parked in parked_items)
+            delay_seconds = max(0.0, due_at - time.time())
+            if deadline is not None and time.time() + delay_seconds > deadline:
+                return
+            self.session.emit(
+                "cooldown_wait",
+                count=len(parked_items),
+                delay=delay_seconds,
+                until=due_at,
+            )
+            await asyncio.sleep(delay_seconds + PARKED_WAKE_UP_MARGIN_SECONDS)
+            await self._process_due_items(parked_items)
+
+    async def _process_due_items(self, parked_items: list) -> None:
+        cutoff = time.time() + PARKED_WAKE_UP_MARGIN_SECONDS
+        due_values = [
+            parked.input_value for parked in parked_items if parked.retry_at <= cutoff
+        ]
+        self.session.emit("cooldown_over", count=len(due_values))
+        await self._process_items(due_values)
+
+    def _parked_wait_deadline(self) -> float | None:
+        limit = self.options.maximum_parked_wait_seconds
+        if limit is None:
+            return None
+        return time.time() + limit
+
+    def _finish_run(self) -> RunStatistics:
+        statistics = self.session.statistics
+        statistics.parked_remaining = len(
+            self.store.items.list_parked(self.session.run_id)
+        )
+        if statistics.items_failed == 0 and statistics.parked_remaining == 0:
+            self.store.runs.finish_run(self.session.run_id)
+        statistics.elapsed_seconds = time.monotonic() - statistics.started_at
+        self.session.emit("run_finished", stats=statistics)
+        return statistics
+
+    def _summarize_previous_run(self, previous_status: str | None) -> dict | None:
+        if previous_status is None:
+            return None
+        run_id = self.session.run_id
+        statuses = [status for _, status in self.store.items.list_statuses(run_id)]
+        interrupted = [
+            (record.item_key, record.step_name)
+            for record in self.store.steps.list_for_run(run_id)
+            if record.status == "running"
+        ]
+        return {
+            "previous": previous_status,
+            "done": statuses.count("done"),
+            "failed": statuses.count("failed"),
+            "parked": statuses.count("parked"),
+            "interrupted": interrupted,
+        }
+
+    async def _as_async_iterator(
+        self, items: Iterable[Any] | AsyncIterable[Any]
+    ) -> AsyncIterator[Any]:
+        if isinstance(items, AsyncIterable):
+            async for item in items:
+                yield item
+            return
+        for item in items:
+            yield item

stitch_durable/execution/context.py

@@ -0,0 +1,98 @@
+import contextvars
+from dataclasses import dataclass
+from typing import Any
+
+from ..storage.checkpoint_store import CheckpointStore
+
+
+@dataclass(frozen=True)
+class StepIdentity:
+    item_key: str
+    step_id: str
+    sequence: int
+
+
+@dataclass(frozen=True)
+class StepContext:
+    run_id: str
+    item_key: str
+    step_name: str
+    attempt: int
+
+
+class ItemContext:
+    def __init__(self, item_key: str):
+        self.item_key = item_key
+        self.degraded = False
+        self._sequence = 0
+        self._call_counts: dict[str, int] = {}
+
+    def register_step_call(self, step_name: str) -> StepIdentity:
+        call_count = self._call_counts.get(step_name, 0) + 1
+        self._call_counts[step_name] = call_count
+        self._sequence += 1
+        step_id = step_name if call_count == 1 else f"{step_name}#{call_count}"
+        return StepIdentity(self.item_key, step_id, self._sequence)
+
+
+class ExecutionContext:
+    _run = contextvars.ContextVar("stitch_run", default=None)
+    _item = contextvars.ContextVar("stitch_item", default=None)
+    _step = contextvars.ContextVar("stitch_step", default=None)
+
+    @classmethod
+    def bind_run(cls, session) -> contextvars.Token:
+        return cls._run.set(session)
+
+    @classmethod
+    def release_run(cls, token: contextvars.Token) -> None:
+        cls._run.reset(token)
+
+    @classmethod
+    def bind_item(cls, item_context: ItemContext) -> contextvars.Token:
+        return cls._item.set(item_context)
+
+    @classmethod
+    def release_item(cls, token: contextvars.Token) -> None:
+        cls._item.reset(token)
+
+    @classmethod
+    def bind_step(cls, step_context: StepContext) -> contextvars.Token:
+        return cls._step.set(step_context)
+
+    @classmethod
+    def release_step(cls, token: contextvars.Token) -> None:
+        cls._step.reset(token)
+
+    @classmethod
+    def current_run(cls):
+        return cls._run.get()
+
+    @classmethod
+    def current_item(cls) -> ItemContext | None:
+        return cls._item.get()
+
+    @classmethod
+    def current_step(cls) -> StepContext | None:
+        return cls._step.get()
+
+    @classmethod
+    def require_store(cls) -> CheckpointStore:
+        session = cls._run.get()
+        if session is None:
+            raise RuntimeError("this call is only valid inside a running workflow")
+        return session.store
+
+    @classmethod
+    def report_event(cls, event: str, **data: Any) -> None:
+        session = cls._run.get()
+        if session is None:
+            return
+        item_context = cls._item.get()
+        step_context = cls._step.get()
+        if item_context is not None:
+            data.setdefault("item", item_context.item_key)
+        if step_context is not None:
+            data.setdefault("step", step_context.step_name)
+            data.setdefault("attempt", step_context.attempt)
+        session.emit(event, **data)

stitch_durable/execution/item_processor.py

@@ -0,0 +1,169 @@
+import asyncio
+import time
+from dataclasses import dataclass
+from typing import Any, Awaitable, Callable
+
+from ..errors import CriticalStepFailed, Park, StepFailed
+from ..records import ItemRecord
+from ..retry_policy import RetryPolicy
+from ..serialization import to_json_text
+from .context import ExecutionContext, ItemContext
+from .run_session import RunSession
+
+
+@dataclass
+class ItemAttempt:
+    result: Any = None
+    should_retry: bool = False
+
+
+class ItemProcessor:
+    def __init__(
+        self,
+        session: RunSession,
+        handler: Callable[[Any], Awaitable[Any]],
+        key_function: Callable[[Any], str],
+        item_retry_policy: RetryPolicy,
+    ):
+        self.session = session
+        self.handler = handler
+        self.key_function = key_function
+        self.item_retry_policy = item_retry_policy
+        self.store = session.store
+        self.statistics = session.statistics
+
+    async def process(self, value: Any) -> Any:
+        item_key = self.key_function(value)
+        existing = self.store.items.get(self.session.run_id, item_key)
+        if existing is not None and existing.status == "done":
+            return self._skip_finished_item(item_key, existing)
+        input_text = to_json_text(value)
+        previous_status = existing.status if existing else None
+        attempt = 0
+        while True:
+            attempt += 1
+            outcome = await self._attempt(
+                value, item_key, input_text, attempt, previous_status
+            )
+            if not outcome.should_retry:
+                return outcome.result
+
+    async def _attempt(
+        self,
+        value: Any,
+        item_key: str,
+        input_text: str | None,
+        attempt: int,
+        previous_status: str | None,
+    ) -> ItemAttempt:
+        item_context = ItemContext(item_key)
+        token = ExecutionContext.bind_item(item_context)
+        self._record_item_start(item_key, input_text, attempt, previous_status)
+        started_at = time.monotonic()
+        try:
+            result = await self.handler(value)
+        except Park as error:
+            return ItemAttempt(self._park_item(item_key, input_text, error))
+        except StepFailed as error:
+            return await self._handle_step_failure(item_key, error, attempt)
+        except CriticalStepFailed:
+            raise
+        except Exception as error:
+            return ItemAttempt(
+                self._fail_item(item_key, f"{type(error).__name__}: {error}")
+            )
+        finally:
+            ExecutionContext.release_item(token)
+        return ItemAttempt(
+            self._finish_item(
+                item_key, result, item_context, time.monotonic() - started_at
+            )
+        )
+
+    async def _handle_step_failure(
+        self, item_key: str, error: StepFailed, attempt: int
+    ) -> ItemAttempt:
+        if not self._should_retry_item(error, attempt):
+            return ItemAttempt(self._fail_item(item_key, str(error)))
+        await self._wait_before_item_retry(item_key, attempt, error)
+        return ItemAttempt(should_retry=True)
+
+    def _skip_finished_item(self, item_key: str, existing: ItemRecord) -> Any:
+        self.statistics.items_skipped += 1
+        self.session.emit("item_skipped", item=item_key, degraded=existing.degraded)
+        return existing.result
+
+    def _record_item_start(
+        self,
+        item_key: str,
+        input_text: str | None,
+        attempt: int,
+        previous_status: str | None,
+    ) -> None:
+        self.store.items.mark_running(self.session.run_id, item_key, input_text)
+        self.session.emit(
+            "item_started",
+            item=item_key,
+            attempt=attempt,
+            resumed=previous_status is not None,
+            was=previous_status,
+        )
+
+    def _should_retry_item(self, error: StepFailed, attempt: int) -> bool:
+        return error.transient and attempt < self.item_retry_policy.maximum_attempts
+
+    async def _wait_before_item_retry(
+        self, item_key: str, attempt: int, error: StepFailed
+    ) -> None:
+        delay_seconds = self.item_retry_policy.delay_before_next_attempt(
+            attempt, error.cause
+        )
+        self.statistics.item_retries += 1
+        self.session.emit(
+            "item_retry",
+            item=item_key,
+            attempt=attempt,
+            attempts=self.item_retry_policy.maximum_attempts,
+            delay=delay_seconds,
+            error=str(error),
+        )
+        await asyncio.sleep(delay_seconds)
+
+    def _park_item(self, item_key: str, input_text: str | None, error: Park) -> None:
+        if input_text is None:
+            return self._fail_item(
+                item_key, f"cannot park, input is not JSON: {error.reason}"
+            )
+        self.store.items.mark_parked(
+            self.session.run_id, item_key, error.until, error.reason
+        )
+        self.statistics.items_parked += 1
+        self.session.emit(
+            "item_parked", item=item_key, until=error.until, reason=error.reason
+        )
+
+    def _fail_item(self, item_key: str, error_text: str) -> None:
+        self.store.items.mark_failed(self.session.run_id, item_key, error_text)
+        self.statistics.items_failed += 1
+        self.session.emit("item_failed", item=item_key, error=error_text)
+
+    def _finish_item(
+        self,
+        item_key: str,
+        result: Any,
+        item_context: ItemContext,
+        elapsed_seconds: float,
+    ) -> Any:
+        self.store.items.mark_done(
+            self.session.run_id, item_key, to_json_text(result), item_context.degraded
+        )
+        self.statistics.items_done += 1
+        if item_context.degraded:
+            self.statistics.items_degraded += 1
+        self.session.emit(
+            "item_done",
+            item=item_key,
+            elapsed=elapsed_seconds,
+            degraded=item_context.degraded,
+        )
+        return result