auto-workflow 0.1.0__py3-none-any.whl

auto_workflow/dag.py

@@ -0,0 +1,80 @@
+"""Internal DAG representation and cycle detection."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+
+from .exceptions import CycleDetectedError
+
+
+@dataclass(slots=True)
+class Node:
+    name: str
+    upstream: set[str] = field(default_factory=set)
+    downstream: set[str] = field(default_factory=set)
+
+
+class DAG:
+    def __init__(self) -> None:
+        self.nodes: dict[str, Node] = {}
+
+    def add_node(self, name: str) -> None:
+        if name not in self.nodes:
+            self.nodes[name] = Node(name=name)
+
+    def add_edge(self, upstream: str, downstream: str) -> None:
+        self.add_node(upstream)
+        self.add_node(downstream)
+        self.nodes[upstream].downstream.add(downstream)
+        self.nodes[downstream].upstream.add(upstream)
+
+    def topological_sort(self) -> list[str]:
+        # Kahn's algorithm with deterministic ordering
+        in_degree = {n: len(node.upstream) for n, node in self.nodes.items()}
+        ready = sorted([n for n, d in in_degree.items() if d == 0])
+        order: list[str] = []
+        while ready:
+            # pop the smallest name for stable results
+            current = ready.pop(0)
+            order.append(current)
+            for child in sorted(self.nodes[current].downstream):
+                in_degree[child] -= 1
+                if in_degree[child] == 0:
+                    # keep ready sorted
+                    from bisect import insort
+
+                    insort(ready, child)
+        if len(order) != len(self.nodes):
+            # Return nodes that still have in-degree > 0 deterministically
+            remaining = sorted([n for n, d in in_degree.items() if d > 0])
+            raise CycleDetectedError(remaining)
+        return order
+
+    def subgraph(self, names: Iterable[str]) -> DAG:
+        sg = DAG()
+        for name in names:
+            if name not in self.nodes:
+                continue
+            sg.add_node(name)
+            for d in self.nodes[name].downstream:
+                if d in names:
+                    sg.add_edge(name, d)
+        return sg
+
+    # Export utilities
+    def to_dot(self) -> str:
+        lines = ["digraph G {"]
+        for name, node in self.nodes.items():
+            if not node.downstream:
+                lines.append(f'  "{name}";')
+            for d in node.downstream:
+                lines.append(f'  "{name}" -> "{d}";')
+        lines.append("}")
+        return "\n".join(lines)
+
+    def to_dict(self) -> dict:
+        return {
+            name: {"upstream": sorted(n.upstream), "downstream": sorted(n.downstream)}
+            for name, n in self.nodes.items()
+        }

auto_workflow/events.py

@@ -0,0 +1,22 @@
+"""Lightweight event bus."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+_subscribers: dict[str, list[Callable[[dict[str, Any]], None]]] = {}
+logger = logging.getLogger("auto_workflow.events")
+
+
+def subscribe(event: str, callback: Callable[[dict[str, Any]], None]) -> None:
+    _subscribers.setdefault(event, []).append(callback)
+
+
+def emit(event: str, payload: dict[str, Any]) -> None:
+    for cb in _subscribers.get(event, []):
+        try:
+            cb(payload)
+        except Exception as e:  # pragma: no cover
+            logger.debug("event subscriber error", exc_info=e)

auto_workflow/exceptions.py

@@ -0,0 +1,42 @@
+"""Domain-specific exceptions for auto_workflow."""
+
+from __future__ import annotations
+
+
+class AutoWorkflowError(Exception):
+    """Base exception."""
+
+
+class CycleDetectedError(AutoWorkflowError):
+    def __init__(self, cycle: list[str]):
+        super().__init__(f"Cycle detected in DAG: {' -> '.join(cycle)}")
+        self.cycle = cycle
+
+
+class TaskExecutionError(AutoWorkflowError):
+    def __init__(self, task_name: str, original: BaseException):
+        super().__init__(f"Task '{task_name}' failed: {original!r}")
+        self.task_name = task_name
+        self.original = original
+
+
+class TimeoutError(TaskExecutionError):
+    pass
+
+
+class RetryExhaustedError(TaskExecutionError):
+    pass
+
+
+class InvalidGraphError(AutoWorkflowError):
+    pass
+
+
+class AggregateTaskError(AutoWorkflowError):
+    """Raised when multiple tasks fail under AGGREGATE failure policy."""
+
+    def __init__(self, errors: list[TaskExecutionError]):
+        self.errors = errors
+        summary = "; ".join(f"{e.task_name}: {e.original!r}" for e in errors[:5])
+        more = "" if len(errors) <= 5 else f" (+{len(errors) - 5} more)"
+        super().__init__(f"Multiple task failures: {summary}{more}")

auto_workflow/execution.py

@@ -0,0 +1,45 @@
+"""Execution helpers for process-based task offload."""
+
+from __future__ import annotations
+
+import atexit
+from concurrent.futures import ProcessPoolExecutor
+from contextlib import suppress
+from typing import Any
+
+import cloudpickle
+
+from .config import load_config
+
+_SHARED_POOL: ProcessPoolExecutor | None = None
+
+
+def _shutdown_pool() -> None:
+    global _SHARED_POOL
+    if _SHARED_POOL is not None:
+        with suppress(Exception):
+            _SHARED_POOL.shutdown(wait=True, cancel_futures=True)
+        _SHARED_POOL = None
+
+
+def get_process_pool() -> ProcessPoolExecutor:
+    global _SHARED_POOL
+    if _SHARED_POOL is None:
+        cfg = load_config()
+        max_workers = cfg.get("process_pool_max_workers")
+        if isinstance(max_workers, str):
+            try:
+                max_workers = int(max_workers) if max_workers.isdigit() else None
+            except Exception:
+                max_workers = None
+        if not isinstance(max_workers, int) or max_workers <= 0:
+            max_workers = None
+        _SHARED_POOL = ProcessPoolExecutor(max_workers=max_workers)
+        # Ensure clean shutdown on interpreter exit
+        atexit.register(_shutdown_pool)
+    return _SHARED_POOL
+
+
+def run_pickled(payload: bytes) -> Any:  # executed in worker process
+    fn, args, kwargs = cloudpickle.loads(payload)
+    return fn(*args, **kwargs)

auto_workflow/fanout.py

@@ -0,0 +1,59 @@
+"""Fan-out helper utilities."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
+from .build import TaskInvocation, current_build_context
+
+
+class DynamicFanOut(list):  # placeholder container recognized by scheduler
+    def __init__(
+        self,
+        task_def,
+        source_invocation: TaskInvocation | DynamicFanOut,
+        max_concurrency: int | None,
+        ctx,
+    ):
+        super().__init__()
+        self._task_def = task_def
+        self._source = source_invocation
+        self._max_conc = max_concurrency
+        self._expanded = False
+        self._ctx = ctx
+
+    def expand(self, values: Iterable[Any]):
+        if self._expanded:
+            return
+        for v in values:
+            self.append(
+                self._ctx.register(self._task_def.name, self._task_def.fn, (v,), {}, self._task_def)
+            )
+        self._expanded = True
+
+
+def fan_out(task_def, iterable: Iterable[Any], *, max_concurrency: int | None = None) -> list[Any]:
+    """Create multiple task invocations from an iterable.
+
+    max_concurrency is reserved for future scheduling throttling; currently unused.
+    """
+    ctx = current_build_context()
+    out = []
+    if ctx is None:
+        # immediate execution path
+        return [task_def(item) for item in iterable]
+    from .fanout import DynamicFanOut  # self-import safe
+
+    if isinstance(iterable, TaskInvocation):  # dynamic runtime fan-out
+        df = DynamicFanOut(task_def, iterable, max_concurrency, ctx)
+        ctx.dynamic_fanouts.append(df)
+        return df
+    if isinstance(iterable, DynamicFanOut):  # nested dynamic
+        # nested placeholder - track as root only if original source TaskInvocation (first-level)
+        df = DynamicFanOut(task_def, iterable, max_concurrency, ctx)
+        ctx.dynamic_fanouts.append(df)
+        return df
+    for item in iterable:
+        out.append(task_def(item))
+    return out

auto_workflow/flow.py

@@ -0,0 +1,165 @@
+"""Flow abstraction and decorator."""
+
+from __future__ import annotations
+
+import asyncio
+import uuid
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+from .build import BuildContext, collect_invocations, replace_invocations
+from .config import load_config
+from .context import RunContext, set_context
+from .events import emit
+from .scheduler import FailurePolicy, execute_dag
+from .tracing import get_tracer
+
+
+@dataclass(slots=True)
+class Flow:
+    name: str
+    build_fn: Callable[..., Any]
+
+    async def _run_internal(
+        self,
+        *args: Any,
+        params: dict[str, Any] | None = None,
+        failure_policy: str = FailurePolicy.FAIL_FAST,
+        max_concurrency: int | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        cfg = load_config()
+        if max_concurrency is None:
+            mc = cfg.get("max_dynamic_tasks")
+            if isinstance(mc, str):
+                if mc.isdigit():
+                    try:
+                        mc = int(mc)
+                    except ValueError:  # pragma: no cover
+                        mc = None
+                else:
+                    mc = None  # non-numeric string -> ignore
+            max_concurrency = mc if isinstance(mc, int) and mc > 0 else None
+        ctx = RunContext(run_id=str(uuid.uuid4()), flow_name=self.name, params=params or {})
+        set_context(ctx)
+        tracer = get_tracer()
+        emit("flow_started", {"flow": self.name, "run_id": ctx.run_id})
+        async with tracer.span(f"flow:{self.name}", run_id=ctx.run_id):
+            with BuildContext() as bctx:
+                structure = self.build_fn(*args, **kwargs)
+                dynamic_roots = list(bctx.dynamic_fanouts)
+            invocations = collect_invocations(structure)
+        if not invocations:
+            # trivial, return original structure (no tasks used)
+            emit("flow_completed", {"flow": self.name, "run_id": ctx.run_id, "tasks": 0})
+            return structure
+        # Execute via scheduler
+        results = await execute_dag(
+            invocations,
+            failure_policy=failure_policy,
+            max_concurrency=max_concurrency,
+            dynamic_roots=dynamic_roots,
+        )  # type: ignore
+        emit("flow_completed", {"flow": self.name, "run_id": ctx.run_id, "tasks": len(invocations)})
+        return replace_invocations(structure, results)
+
+    def run(
+        self,
+        *args: Any,
+        params: dict[str, Any] | None = None,
+        failure_policy: str = FailurePolicy.FAIL_FAST,
+        max_concurrency: int | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        return asyncio.run(
+            self._run_internal(
+                *args,
+                params=params,
+                failure_policy=failure_policy,
+                max_concurrency=max_concurrency,
+                **kwargs,
+            )
+        )
+
+    def describe(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
+        """Return a JSON-serializable representation of the DAG that would be built.
+        Does not execute any tasks.
+        """
+        with BuildContext():
+            structure = self.build_fn(*args, **kwargs)
+        invocations = collect_invocations(structure)
+        # Attach a synthetic upstream hint for nodes consuming dynamic placeholders
+        from .fanout import DynamicFanOut
+
+        for inv in invocations:
+            for arg in list(inv.args) + list(inv.kwargs.values()):
+                if isinstance(arg, DynamicFanOut):
+                    # depend on source invocation so scheduling waits for expansion
+                    # children edges injected later
+                    src = arg._source
+                    if hasattr(src, "name"):
+                        inv.upstream.add(src.name)
+        nodes = []
+        for inv in invocations:
+            nodes.append(
+                {
+                    "id": inv.name,
+                    "task": inv.task_name,
+                    "upstream": sorted(inv.upstream),
+                    "persist": getattr(inv.definition, "persist", False),
+                    "run_in": getattr(inv.definition, "run_in", "async"),
+                    "retries": getattr(inv.definition, "retries", 0),
+                }
+            )
+        return {
+            "flow": self.name,
+            "nodes": nodes,
+            "count": len(nodes),
+        }
+
+    def export_dot(self, *args: Any, **kwargs: Any) -> str:
+        from .dag import DAG
+
+        with BuildContext():
+            structure = self.build_fn(*args, **kwargs)
+        invocations = collect_invocations(structure)
+        from .fanout import DynamicFanOut
+
+        # ensure upstream dependency on dynamic source(s)
+        for inv in invocations:
+            for arg in list(inv.args) + list(inv.kwargs.values()):
+                if isinstance(arg, DynamicFanOut):
+                    src = arg._source
+                    if hasattr(src, "name"):
+                        inv.upstream.add(src.name)
+        dag = DAG()
+        for inv in invocations:
+            dag.add_node(inv.name)
+            for up in inv.upstream:
+                dag.add_edge(up, inv.name)
+        return dag.to_dot()
+
+    def export_graph(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
+        from .dag import DAG
+
+        with BuildContext():
+            structure = self.build_fn(*args, **kwargs)
+        invocations = collect_invocations(structure)
+        dag = DAG()
+        for inv in invocations:
+            dag.add_node(inv.name)
+            for up in inv.upstream:
+                dag.add_edge(up, inv.name)
+        return dag.to_dict()
+
+
+def flow(
+    _fn: Callable[..., Any] | None = None, *, name: str | None = None
+) -> Callable[[Callable[..., Any]], Flow]:
+    def wrap(fn: Callable[..., Any]) -> Flow:
+        return Flow(name=name or fn.__name__, build_fn=fn)
+
+    if _fn is not None:
+        return wrap(_fn)
+    return wrap

auto_workflow/lifecycle.py

@@ -0,0 +1,16 @@
+"""Lifecycle helpers for graceful shutdown of runtime components."""
+
+from __future__ import annotations
+
+from contextlib import suppress
+
+from .execution import _shutdown_pool  # internal but intentional for lifecycle
+
+
+def shutdown() -> None:
+    """Gracefully shut down background resources (process pool, etc.).
+
+    Safe to call multiple times.
+    """
+    with suppress(Exception):
+        _shutdown_pool()

auto_workflow/logging_middleware.py

@@ -0,0 +1,191 @@
+"""Structured logging for flows and tasks.
+
+Provides:
+- structured_logging_middleware: logs task completion (ok/err) with duration.
+- register_structured_logging(): registers middleware and event subscribers to log
+    flow start/completion and task start events.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+from contextlib import suppress
+from typing import Any
+
+from .context import get_context
+from .events import subscribe
+
+logger = logging.getLogger("auto_workflow.tasks")
+
+
+def _now_iso() -> str:
+    import datetime as _dt
+
+    return _dt.datetime.now(tz=_dt.UTC).isoformat()
+
+
+async def structured_logging_middleware(nxt, task_def, args, kwargs):
+    ctx = None
+    with suppress(Exception):  # outside flow safe
+        ctx = get_context()
+    start = time.time()
+    meta: dict[str, Any] = {
+        "task": task_def.name,
+        "run_id": getattr(ctx, "run_id", None),
+        "flow": getattr(ctx, "flow_name", None),
+        "ts": _now_iso(),
+    }
+    try:
+        result = await nxt()
+        duration = (time.time() - start) * 1000.0
+        meta["duration_ms"] = duration
+        logger.info(json.dumps({"event": "task_ok", **meta}))
+        return result
+    except Exception as e:
+        duration = (time.time() - start) * 1000.0
+        meta["duration_ms"] = duration
+        meta["error"] = repr(e)
+        logger.error(json.dumps({"event": "task_err", **meta}))
+        raise
+
+
+# Event subscribers for flow/task lifecycle
+def _on_flow_started(payload: dict[str, Any]) -> None:
+    logger.info(
+        json.dumps(
+            {
+                "event": "flow_started",
+                "flow": payload.get("flow"),
+                "run_id": payload.get("run_id"),
+                "ts": _now_iso(),
+            }
+        )
+    )
+
+
+def _on_flow_completed(payload: dict[str, Any]) -> None:
+    logger.info(
+        json.dumps(
+            {
+                "event": "flow_completed",
+                "flow": payload.get("flow"),
+                "run_id": payload.get("run_id"),
+                "tasks": payload.get("tasks"),
+                "ts": _now_iso(),
+            }
+        )
+    )
+
+
+def _on_task_started(payload: dict[str, Any]) -> None:
+    logger.info(
+        json.dumps(
+            {
+                "event": "task_started",
+                "task": payload.get("task"),
+                "node": payload.get("node"),
+                "ts": _now_iso(),
+            }
+        )
+    )
+
+
+_registered = False
+
+
+def register_structured_logging() -> None:
+    """Register structured logging for flow and task lifecycle.
+
+    - Subscribes to flow_started, flow_completed, and task_started events.
+    - Registers the task middleware to log completion with duration.
+    """
+    global _registered
+    if _registered:
+        return
+    subscribe("flow_started", _on_flow_started)
+    subscribe("flow_completed", _on_flow_completed)
+    subscribe("task_started", _on_task_started)
+    # Defer import to avoid cycles
+    from .middleware import register as _register
+
+    _register(structured_logging_middleware)
+    _registered = True
+
+
+def enable_default_logging(level: str = "INFO") -> None:
+    """Attach a simple stdout handler to the auto_workflow logger if none present.
+
+    This makes structured JSON lines visible when running scripts.
+    """
+    if not logger.handlers:
+        h = logging.StreamHandler()
+        h.setFormatter(logging.Formatter("%(message)s"))
+        # mark so we can replace with pretty handler if requested
+        h._aw_default = True
+        logger.addHandler(h)
+    try:
+        logger.setLevel(getattr(logging, level.upper(), logging.INFO))
+    except Exception:
+        logger.setLevel(logging.INFO)
+    # Avoid duplicate propagation to root
+    logger.propagate = False
+
+
+class StructuredPrettyFormatter(logging.Formatter):
+    def __init__(self, datefmt: str | None = None) -> None:
+        super().__init__(datefmt=datefmt)
+
+    def format(self, record: logging.LogRecord) -> str:
+        # Try to parse JSON message emitted by structured logger
+        try:
+            data = json.loads(record.getMessage())
+        except Exception:
+            return super().format(record)
+        ts = self.formatTime(record, self.datefmt)
+        lvl = record.levelname
+        parts = [ts, lvl]
+        ev = data.get("event")
+        if ev:
+            parts.append(ev)
+        # Common fields
+        kvs = []
+        for k in ("flow", "run_id", "task", "node"):
+            v = data.get(k)
+            if v is not None:
+                kvs.append(f"{k}={v}")
+        if "duration_ms" in data:
+            try:
+                kvs.append(f"duration={float(data['duration_ms']):.1f}ms")
+            except Exception:
+                kvs.append(f"duration={data['duration_ms']}ms")
+        if "error" in data:
+            kvs.append(f"error={data['error']}")
+        if kvs:
+            parts.append(" ".join(kvs))
+        return " | ".join(parts)
+
+
+def enable_pretty_logging(level: str = "INFO") -> None:
+    """Replace the default JSON-line handler with a human-friendly formatter.
+
+    Safe to call multiple times.
+    """
+    # Remove our default handler if present
+    new_handlers = []
+    for h in logger.handlers:
+        if getattr(h, "_aw_default", False):
+            continue
+        new_handlers.append(h)
+    logger.handlers = new_handlers
+    # Add pretty handler
+    h = logging.StreamHandler()
+    h.setFormatter(StructuredPrettyFormatter(datefmt="%Y-%m-%d %H:%M:%S%z"))
+    h._aw_pretty = True
+    logger.addHandler(h)
+    try:
+        logger.setLevel(getattr(logging, level.upper(), logging.INFO))
+    except Exception:
+        logger.setLevel(logging.INFO)
+    logger.propagate = False

auto_workflow/metrics_provider.py

@@ -0,0 +1,35 @@
+"""Pluggable metrics provider abstraction."""
+
+from __future__ import annotations
+
+from typing import Any, Protocol
+
+
+class MetricsProvider(Protocol):  # pragma: no cover - interface
+    def inc(self, name: str, value: float = 1.0, **labels: Any) -> None: ...
+    def observe(self, name: str, value: float, **labels: Any) -> None: ...
+
+
+class InMemoryMetrics(MetricsProvider):
+    def __init__(self) -> None:
+        self.counters: dict[str, float] = {}
+        self.histograms: dict[str, list[float]] = {}
+
+    def inc(self, name: str, value: float = 1.0, **labels: Any) -> None:
+        key = name
+        self.counters[key] = self.counters.get(key, 0.0) + value
+
+    def observe(self, name: str, value: float, **labels: Any) -> None:
+        self.histograms.setdefault(name, []).append(value)
+
+
+_provider: MetricsProvider = InMemoryMetrics()
+
+
+def set_metrics_provider(p: MetricsProvider) -> None:
+    global _provider
+    _provider = p
+
+
+def get_metrics_provider() -> MetricsProvider:
+    return _provider