auto-workflow 0.1.0__py3-none-any.whl

auto_workflow/middleware.py

@@ -0,0 +1,59 @@
+"""Middleware chaining for task execution (extensible)."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+from .events import emit
+
+TaskCallable = Callable[[], Awaitable[Any]]
+Middleware = Callable[[TaskCallable, Any, tuple, dict], Awaitable[Any]]
+
+_registry: list[Middleware] = []
+
+
+def register(mw: Middleware) -> None:
+    _registry.append(mw)
+
+
+def clear() -> None:  # pragma: no cover
+    _registry.clear()
+
+
+async def _call_chain(
+    index: int, core: TaskCallable, task_def: Any, args: tuple, kwargs: dict
+) -> Any:
+    if index == len(_registry):
+        return await core()
+    mw = _registry[index]
+    executing_core = False
+
+    async def nxt():
+        nonlocal executing_core
+        executing_core = True
+        return await _call_chain(index + 1, core, task_def, args, kwargs)
+
+    try:
+        return await mw(lambda: nxt(), task_def, args, kwargs)
+    except Exception as e:  # noqa: BLE001
+        # If exception occurred during/after core execution, propagate (task error)
+        if executing_core:
+            raise
+        # Otherwise classify as middleware error and continue chain
+        emit(
+            "middleware_error",
+            {
+                "task": getattr(task_def, "name", "unknown"),
+                "middleware_index": index,
+                "error": repr(e),
+            },
+        )
+        return await _call_chain(index + 1, core, task_def, args, kwargs)
+
+
+def get_task_middleware_chain() -> Callable[[TaskCallable, Any, tuple, dict], Awaitable[Any]]:
+    async def runner(core: TaskCallable, task_def: Any, args: tuple, kwargs: dict) -> Any:
+        return await _call_chain(0, core, task_def, args, kwargs)
+
+    return runner

auto_workflow/scheduler.py

@@ -0,0 +1,362 @@
+"""Simple topological scheduler executing TaskInvocations respecting dependencies."""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from typing import Any
+
+from .artifacts import get_store
+from .build import TaskInvocation
+from .cache import get_result_cache
+from .dag import DAG
+from .events import emit
+from .exceptions import AggregateTaskError, TaskExecutionError
+from .metrics_provider import get_metrics_provider
+from .middleware import get_task_middleware_chain
+from .tracing import get_tracer
+
+# result cache implementation centralized in auto_workflow.cache
+
+
+class FailurePolicy:
+    FAIL_FAST = "fail_fast"
+    CONTINUE = "continue"
+    AGGREGATE = "aggregate"
+
+
+async def execute_dag(
+    invocations: list[TaskInvocation],
+    *,
+    failure_policy: str = FailurePolicy.FAIL_FAST,
+    max_concurrency: int | None = None,
+    cancel_event: asyncio.Event | None = None,
+    dynamic_roots: list[Any] | None = None,
+) -> dict[str, Any]:
+    dag = DAG()
+    for inv in invocations:
+        dag.add_node(inv.name)
+        for up in inv.upstream:
+            dag.add_edge(up, inv.name)
+    order = dag.topological_sort()
+    inv_map = {inv.name: inv for inv in invocations}
+    results: dict[str, Any] = {}
+    cache = get_result_cache()
+    # simple ready set processing
+    remaining_deps = {name: set(inv_map[name].upstream) for name in order}
+    ready = [n for n, deps in remaining_deps.items() if not deps]
+    # removed unused 'running' variable
+    pending_tasks: dict[str, asyncio.Task[Any]] = {}
+
+    sem = asyncio.Semaphore(max_concurrency or len(order))
+    task_errors: list[TaskExecutionError] = []
+
+    # For tasks with cache_ttl we may have multiple identical invocations ready simultaneously
+    # (e.g., synchronous functions now offloaded to threads). Use an in-flight map to deduplicate.
+    inflight: dict[str, asyncio.Future[Any]] = {}
+    inflight_lock = asyncio.Lock()
+    # Per dynamic fan-out concurrency control
+    placeholder_sems: dict[int, asyncio.Semaphore] = {}
+    child_to_placeholder: dict[str, int] = {}
+
+    async def schedule(name: str) -> None:
+        inv = inv_map[name]
+        # prepare args by replacing dependencies
+        resolved_args = _hydrate(inv.args, results)
+        resolved_kwargs = _hydrate(inv.kwargs, results)
+        # caching (per task definition attributes)
+        cache_ttl = getattr(inv.definition, "cache_ttl", None)
+        cache_key = inv.definition.cache_key(*resolved_args, **resolved_kwargs)
+        if cache_ttl is not None:
+            # Atomic check-and-register for in-flight de-dup
+            async with inflight_lock:
+                cached = cache.get(cache_key, cache_ttl)
+                if cached is not None:
+                    from contextlib import suppress
+
+                    with suppress(Exception):
+                        get_metrics_provider().inc("cache_hits")
+                    results[name] = cached
+                    return
+                existing = inflight.get(cache_key)
+                if existing is None:
+                    # register placeholder before any execution begins so followers can await
+                    inflight[cache_key] = asyncio.get_running_loop().create_future()
+                    existing = None
+            # If another identical task already running, await its future
+            if existing is not None:
+                from contextlib import suppress
+
+                with suppress(Exception):
+                    get_metrics_provider().inc("dedup_joins")
+                try:
+                    value = await existing
+                except Exception as e:  # propagate underlying failure
+                    raise e
+                results[name] = value
+                return
+        try:
+            emit("task_started", {"task": inv.task_name, "node": name})
+            start = time.time()
+            async with sem:
+                tracer = get_tracer()
+                async with tracer.span(f"task:{inv.task_name}", node=name):
+
+                    async def core_run():
+                        return await inv.definition.run(*resolved_args, **resolved_kwargs)
+
+                    # apply middleware chain
+                    value = await get_task_middleware_chain()(
+                        core_run, inv.definition, resolved_args, resolved_kwargs
+                    )
+            duration = time.time() - start
+            # artifact persistence if requested
+            if getattr(inv.definition, "persist", False):
+                store = get_store()
+                ref = store.put(value)
+                value = ref
+            results[name] = value
+            if cache_ttl is not None:
+                cache.set(cache_key, value)
+                from contextlib import suppress
+
+                with suppress(Exception):
+                    get_metrics_provider().inc("cache_sets")
+                # resolve and cleanup inflight placeholder atomically
+                async with inflight_lock:
+                    fut = inflight.get(cache_key)
+                    if fut and not fut.done():
+                        fut.set_result(value)
+                    inflight.pop(cache_key, None)
+            emit(
+                "task_succeeded",
+                {"task": inv.task_name, "node": name, "duration_ms": duration * 1000.0},
+            )
+            mp = get_metrics_provider()
+            mp.inc("tasks_succeeded")
+            mp.observe("task_duration_ms", duration * 1000.0)
+        except Exception as e:  # noqa: BLE001
+            te = TaskExecutionError(inv.task_name, e)
+            if cache_ttl is not None:
+                async with inflight_lock:
+                    fut = inflight.get(cache_key)
+                    if fut and not fut.done():
+                        fut.set_exception(e)
+                    inflight.pop(cache_key, None)
+            if failure_policy == FailurePolicy.FAIL_FAST:
+                emit("task_failed", {"task": inv.task_name, "node": name, "error": repr(e)})
+                # record failure result so downstream inspection doesn't KeyError before propagation
+                results[name] = te
+                mp = get_metrics_provider()
+                mp.inc("tasks_failed")
+                raise te from None
+            task_errors.append(te)
+            results[name] = te
+            emit("task_failed", {"task": inv.task_name, "node": name, "error": repr(e)})
+            mp = get_metrics_provider()
+            mp.inc("tasks_failed")
+
+    # Map consumer invocation name -> list of DynamicFanOut placeholders it references
+    from .fanout import DynamicFanOut
+
+    consumer_placeholders: dict[str, list[DynamicFanOut]] = {}
+    all_placeholders: list[DynamicFanOut] = []
+    if dynamic_roots:
+        for p in dynamic_roots:
+            if isinstance(p, DynamicFanOut):
+                all_placeholders.append(p)
+    for inv in invocations:
+        for arg in list(inv.args) + list(inv.kwargs.values()):
+            if isinstance(arg, DynamicFanOut):
+                consumer_placeholders.setdefault(inv.name, []).append(arg)
+                all_placeholders.append(arg)
+
+    while ready or pending_tasks:
+        if cancel_event and cancel_event.is_set():
+            # Cancel all running tasks
+            for t in pending_tasks.values():
+                t.cancel()
+            # Wait for cancellation to propagate
+            if pending_tasks:
+                await asyncio.gather(*pending_tasks.values(), return_exceptions=True)
+            break
+        # Determine which ready nodes are truly runnable (no unexpanded placeholders they depend on)
+        runnable: list[str] = []
+        for node in list(ready):
+            placeholders = consumer_placeholders.get(node, [])
+            gate = False
+            for p in placeholders:
+                # Need expansion complete and all child results materialized
+                if (not p._expanded) or any(child.name not in results for child in p):
+                    gate = True
+                    break
+            if gate:
+                continue
+            runnable.append(node)
+        # Enforce priority ordering (higher first) and deterministic tie-break by node name
+        runnable.sort(
+            key=lambda n: (
+                -getattr(inv_map[n].definition, "priority", 0),
+                n,
+            )
+        )
+        for node in runnable:
+            # Skip if already scheduled (defensive against accidental duplicates in ready)
+            if node in pending_tasks:
+                # Remove duplicate occurrence
+                ready.remove(node)
+                continue
+            ready.remove(node)
+            # If node belongs to a dynamic placeholder with a concurrency limit, wrap schedule
+            pid = child_to_placeholder.get(node)
+            if pid is not None and pid in placeholder_sems:
+                sem_p = placeholder_sems[pid]
+
+                async def schedule_with_limit(n=node, s=sem_p):
+                    async with s:
+                        await schedule(n)
+
+                pending_tasks[node] = asyncio.create_task(schedule_with_limit())
+            else:
+                pending_tasks[node] = asyncio.create_task(schedule(node))
+        if not pending_tasks:
+            break
+        done, _ = await asyncio.wait(pending_tasks.values(), return_when=asyncio.FIRST_COMPLETED)
+        finished_names = [n for n, t in list(pending_tasks.items()) if t in done]
+        for fname in finished_names:
+            task = pending_tasks.pop(fname)
+            exc = task.exception()
+            if exc:
+                raise exc
+            # After task success, check if any dynamic fan-outs depend on it
+            from .fanout import DynamicFanOut
+
+            # Evaluate expansion conditions for all placeholders (supports nesting)
+            for placeholder in list(all_placeholders):
+                if placeholder._expanded:
+                    continue
+                src = placeholder._source
+                ready_to_expand = False
+                if isinstance(src, TaskInvocation) and src.name == fname:
+                    ready_to_expand = True
+                else:
+                    # Nested: expand when all children have results
+                    if (
+                        isinstance(src, DynamicFanOut)
+                        and src._expanded
+                        and all(child.name in results for child in src)
+                    ):
+                        ready_to_expand = True
+                if ready_to_expand:
+                    # Derive nested source value: collect each child result into a list
+                    if isinstance(src, TaskInvocation):
+                        source_value = results[src.name]
+                    else:
+                        source_value = [results[c.name] for c in src]
+                    if not isinstance(source_value, (list, tuple, set)):
+                        raise TaskExecutionError(
+                            getattr(src, "task_name", "dynamic"),
+                            RuntimeError("Dynamic fan_out source must return an iterable"),
+                        )
+                    placeholder.expand(source_value)
+                    for child_inv in placeholder:
+                        # Track per-placeholder concurrency if specified
+                        if getattr(placeholder, "_max_conc", None):
+                            pid = id(placeholder)
+                            if pid not in placeholder_sems:
+                                placeholder_sems[pid] = asyncio.Semaphore(placeholder._max_conc)
+                            child_to_placeholder[child_inv.name] = pid
+                        if child_inv.name not in dag.nodes:
+                            dag.add_node(child_inv.name)
+                        # Edge from all upstream of src's children or src itself
+                        if isinstance(src, TaskInvocation):
+                            dag.add_edge(src.name, child_inv.name)
+                            # Source already finished, so dependency is satisfied immediately
+                            remaining_deps[child_inv.name] = set()
+                            ready.append(child_inv.name)
+                        else:
+                            # depend on all children of src
+                            deps = {c.name for c in src}
+                            for d in deps:
+                                dag.add_edge(d, child_inv.name)
+                            # All deps completed (readiness condition); mark none remaining
+                            remaining_deps[child_inv.name] = set()
+                            ready.append(child_inv.name)
+                        # ensure inv_map aware
+                        if child_inv.name not in inv_map:
+                            inv_map[child_inv.name] = child_inv
+                    # Update consumers
+                    for consumer in invocations:
+                        if consumer is src:
+                            continue
+                        replaced = False
+
+                        def _walk(o, target=placeholder):  # bind loop var
+                            nonlocal replaced
+                            if o is target:
+                                replaced = True
+                                return list(target)
+                            # Do not traverse into unrelated DynamicFanOut placeholders;
+                            # treat them as atomic until their own expansion pass.
+                            from .fanout import DynamicFanOut as DynamicFanOutPlaceholder
+
+                            if isinstance(o, DynamicFanOutPlaceholder):
+                                return o
+                            if isinstance(o, list):
+                                return [_walk(x, target) for x in o]
+                            if isinstance(o, tuple):
+                                return tuple(_walk(x, target) for x in o)
+                            if isinstance(o, dict):
+                                return {k: _walk(v, target) for k, v in o.items()}
+                            return o
+
+                        consumer.args = (
+                            _walk(list(consumer.args)) if consumer.args else consumer.args
+                        )
+                        consumer.kwargs = (
+                            _walk(consumer.kwargs) if consumer.kwargs else consumer.kwargs
+                        )
+                        if replaced:
+                            for child_inv in placeholder:
+                                dag.add_edge(child_inv.name, consumer.name)
+                                remaining_deps.setdefault(consumer.name, set()).add(child_inv.name)
+                    # Placeholders remain for nested detection; do not delete
+            # update downstream readiness
+            # Use set() to guard against duplicate edges causing repeated scheduling attempts
+            for child in set(dag.nodes[fname].downstream):
+                remaining_deps[child].discard(fname)
+                # Skip scheduling if any upstream failed and policy is not CONTINUE
+                upstream_failed = any(
+                    isinstance(results.get(up), TaskExecutionError)
+                    for up in dag.nodes[child].upstream
+                    if up in results
+                )
+                if upstream_failed and failure_policy != FailurePolicy.CONTINUE:
+                    results[child] = TaskExecutionError(child, RuntimeError("Upstream failed"))
+                    continue
+                if (
+                    not remaining_deps[child]
+                    and child not in ready
+                    and child not in pending_tasks
+                    and child not in results
+                ):
+                    ready.append(child)
+    if failure_policy == FailurePolicy.AGGREGATE and task_errors:
+        raise AggregateTaskError(task_errors)
+    return results
+
+
+def _hydrate(struct: Any, results: dict[str, Any]) -> Any:
+    from .build import TaskInvocation
+
+    if isinstance(struct, TaskInvocation):
+        return results[struct.name]
+    if isinstance(struct, list):
+        return [_hydrate(s, results) for s in struct]
+    if isinstance(struct, tuple):
+        return tuple(_hydrate(s, results) for s in struct)
+    if isinstance(struct, set):
+        return {_hydrate(s, results) for s in struct}
+    if isinstance(struct, dict):
+        return {k: _hydrate(v, results) for k, v in struct.items()}
+    return struct

auto_workflow/secrets.py

@@ -0,0 +1,45 @@
+"""Secrets provider abstraction."""
+
+from __future__ import annotations
+
+import os
+from typing import Protocol
+
+
+class SecretsProvider(Protocol):  # pragma: no cover - interface
+    def get(self, key: str) -> str | None: ...
+
+
+class EnvSecrets(SecretsProvider):
+    def get(self, key: str) -> str | None:
+        return os.environ.get(key)
+
+
+class StaticMappingSecrets(SecretsProvider):
+    def __init__(self, data: dict[str, str]):
+        self.data = data
+
+    def get(self, key: str) -> str | None:
+        return self.data.get(key)
+
+
+class DummyVaultSecrets(SecretsProvider):  # placeholder for future HashiCorp Vault integration
+    def __init__(self, prefix: str = ""):
+        self.prefix = prefix
+        # In a real implementation, store client/session
+
+    def get(self, key: str) -> str | None:  # pragma: no cover - placeholder
+        # Would query Vault; here just environment fallback with prefix
+        return os.environ.get(self.prefix + key)
+
+
+_provider: SecretsProvider = EnvSecrets()
+
+
+def set_secrets_provider(p: SecretsProvider) -> None:
+    global _provider
+    _provider = p
+
+
+def secret(key: str) -> str | None:
+    return _provider.get(key)

auto_workflow/task.py

@@ -0,0 +1,158 @@
+"""Task abstraction and decorator."""
+
+from __future__ import annotations
+
+import asyncio
+import builtins
+import functools
+import inspect
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+from .build import current_build_context
+from .events import emit
+from .exceptions import RetryExhaustedError, TimeoutError
+from .tracing import get_tracer
+from .types import TaskFn
+from .utils import default_cache_key, maybe_await
+
+
+@dataclass(slots=True)
+class TaskDefinition:
+    name: str
+    fn: TaskFn
+    original_fn: TaskFn | None = None
+    retries: int = 0
+    retry_backoff: float = 0.0
+    retry_jitter: float = 0.0
+    timeout: float | None = None
+    cache_ttl: int | None = None
+    cache_key_fn: Callable[..., str] = default_cache_key
+    tags: set[str] = field(default_factory=set)
+    run_in: str = "async"  # one of: async, thread, process
+    persist: bool = False  # store large result via artifact store (future)
+    priority: int = 0  # higher runs earlier when ready
+
+    async def run(self, *args: Any, **kwargs: Any) -> Any:
+        attempt = 0
+        while True:
+
+            async def _invoke():
+                tracer = get_tracer()
+                async with tracer.span(f"task:{self.name}"):
+                    if self.run_in == "async":
+                        return await maybe_await(self.fn(*args, **kwargs))
+                    if self.run_in == "thread":
+                        return await asyncio.to_thread(self.fn, *args, **kwargs)
+                    if self.run_in == "process":
+                        from .execution import get_process_pool
+
+                        loop = asyncio.get_running_loop()
+                        pool = get_process_pool()
+                        import cloudpickle
+
+                        from .execution import run_pickled
+
+                        fn_bytes = cloudpickle.dumps((self.fn, args, kwargs))
+                        return await loop.run_in_executor(
+                            pool, functools.partial(run_pickled, fn_bytes)
+                        )
+                    raise RuntimeError(f"Unknown run_in mode: {self.run_in}")
+
+            try:
+                if self.timeout:
+                    return await asyncio.wait_for(_invoke(), timeout=self.timeout)
+                return await _invoke()
+            except builtins.TimeoutError as te:
+                err: Exception = TimeoutError(self.name, te)
+            except Exception as e:  # noqa: BLE001
+                err = e
+            # error path
+            if attempt < self.retries:
+                attempt += 1
+                emit("task_retry", {"task": self.name, "attempt": attempt, "max": self.retries})
+                sleep_dur = self.retry_backoff * (2 ** (attempt - 1))
+                if self.retry_jitter:
+                    import random
+
+                    sleep_dur += random.uniform(0, self.retry_jitter)
+                await asyncio.sleep(sleep_dur)
+                continue
+            if isinstance(err, TimeoutError):
+                raise err
+            raise RetryExhaustedError(self.name, err) from err
+
+    def cache_key(self, *args: Any, **kwargs: Any) -> str:
+        return self.cache_key_fn(self.fn, args, kwargs)
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:  # building vs immediate execution
+        ctx = current_build_context()
+        if ctx is None:
+            # immediate execution outside a flow build (synchronous helper)
+            tracer = get_tracer()
+
+            async def _run():
+                async with tracer.span(f"task:{self.name}"):
+                    val = await self.run(*args, **kwargs)
+                    if self.persist:
+                        from .artifacts import get_store
+
+                        store = get_store()
+                        ref = store.put(val)
+                        return ref
+                    return val
+
+            return asyncio.run(_run())
+        return ctx.register(self.name, self.fn, args, kwargs, self)
+
+
+def task(
+    _fn: TaskFn | None = None,
+    *,
+    name: str | None = None,
+    retries: int = 0,
+    retry_backoff: float = 0.0,
+    retry_jitter: float = 0.0,
+    timeout: float | None = None,
+    cache_ttl: int | None = None,
+    cache_key_fn: Callable[..., str] = default_cache_key,
+    tags: set[str] | None = None,
+    run_in: str | None = None,
+    persist: bool = False,
+    priority: int = 0,
+) -> Callable[[TaskFn], TaskDefinition]:
+    """Decorator to define a task.
+
+    Auto executor selection:
+    - If run_in explicitly provided -> honored.
+    - If function is an async def -> defaults to "async".
+    - Else (synchronous callable) -> defaults to "thread" to avoid blocking the event loop.
+    """
+
+    def wrap(fn: TaskFn) -> TaskDefinition:
+        inferred_run_in = run_in
+        if inferred_run_in is None:
+            inferred_run_in = (
+                "async" if inspect.iscoroutinefunction(fn) else "thread"
+            )  # safe default for blocking sync functions
+        td_obj = TaskDefinition(
+            name=name or fn.__name__,
+            fn=fn,
+            original_fn=fn,
+            retries=retries,
+            retry_backoff=retry_backoff,
+            retry_jitter=retry_jitter,
+            timeout=timeout,
+            cache_ttl=cache_ttl,
+            cache_key_fn=cache_key_fn,
+            tags=tags or set(),
+            run_in=inferred_run_in,
+            persist=persist,
+            priority=priority,
+        )
+        return td_obj
+
+    if _fn is not None:
+        return wrap(_fn)
+    return wrap

auto_workflow/tracing.py

@@ -0,0 +1,29 @@
+"""Tracing scaffold (OpenTelemetry friendly)."""
+
+from __future__ import annotations
+
+import time
+from contextlib import asynccontextmanager
+from typing import Any
+
+
+class DummyTracer:
+    @asynccontextmanager
+    async def span(self, name: str, **attrs: Any):  # pragma: no cover simple scaffold
+        start = time.time()
+        try:
+            yield {"start": start, "name": name, **attrs}
+        finally:
+            _ = time.time() - start
+
+
+_tracer: DummyTracer = DummyTracer()
+
+
+def get_tracer() -> DummyTracer:
+    return _tracer
+
+
+def set_tracer(t: DummyTracer) -> None:
+    global _tracer
+    _tracer = t

auto_workflow/types.py

@@ -0,0 +1,27 @@
+"""Common type aliases and Protocols."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import Protocol, TypeVar, runtime_checkable
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+TaskFn = Callable[..., R] | Callable[..., Awaitable[R]]
+
+
+@runtime_checkable
+class SupportsHash(Protocol):
+    def __hash__(self) -> int: ...  # pragma: no cover
+
+
+CacheKey = str
+
+
+class CancelledSentinel:
+    def __repr__(self) -> str:  # pragma: no cover
+        return "<CANCELLED>"
+
+
+CANCELLED = CancelledSentinel()