flowforge-map 0.1.0__py3-none-any.whl

flowforge/scheduler/executor.py

@@ -0,0 +1,390 @@
+"""Concurrent list scheduler (D4, D5, D7).
+
+One dispatcher loop owns all scheduling decisions; each running step is an asyncio task.
+
+  - ready steps sit in a heap whose key depends on the policy (see _priority)
+  - a step starts only when a global slot and its type slot are free
+    (a step whose identical call is already cached/in flight starts without a slot —
+    it only waits, it doesn't work)
+  - inside the task: rate-limit token → node.run under asyncio.wait_for(timeout);
+    transient errors retry with full-jitter backoff (or the server's Retry-After),
+    and every attempt takes a fresh token
+  - on failure every descendant is marked skipped; independent branches keep going
+
+Policies (benchmark strategies S1–S5 from D6):
+  sequential     k = 1, topological order
+  levels         BFS layers with a barrier: nothing from layer i+1 starts until layer i is done
+  greedy         unlimited concurrency, no type limits
+  fifo           k slots, first-ready first-served
+  critical_path  k slots, highest bottom level first (HLFET) — FlowForge's scheduler
+"""
+
+from __future__ import annotations
+
+import asyncio
+import heapq
+import itertools
+import math
+import random
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Literal
+
+from flowforge import templating
+from flowforge.nodes.base import Node, NodeError, TransientNodeError
+from flowforge.scheduler import graph
+from flowforge.scheduler.cache import CallCache, cache_key
+from flowforge.scheduler.critical_path import bottom_levels, critical_path
+from flowforge.scheduler.durations import estimate_weights
+from flowforge.scheduler.rate_limit import TokenBucket
+from flowforge.schema import Step, Workflow
+
+Policy = Literal["sequential", "levels", "greedy", "fifo", "critical_path"]
+POLICIES: tuple[Policy, ...] = ("sequential", "levels", "greedy", "fifo", "critical_path")
+
+
+class StepState(str, Enum):
+    PENDING = "pending"
+    READY = "ready"
+    RUNNING = "running"
+    SUCCEEDED = "succeeded"
+    FAILED = "failed"
+    SKIPPED = "skipped"
+
+
+TERMINAL = {StepState.SUCCEEDED, StepState.FAILED, StepState.SKIPPED}
+
+
+@dataclass
+class StepResult:
+    step_id: str
+    state: StepState = StepState.PENDING
+    output: Any = None
+    error: str | None = None
+    attempts: int = 0
+    cache_hit: bool = False
+    started_at: float | None = None   # seconds since run start
+    finished_at: float | None = None
+    call_ms: float | None = None      # the successful node call alone: no token wait, no backoff
+    answered_by: str | None = None    # connector that answered (or was last tried); differs on fallback (D10)
+
+    @property
+    def duration_ms(self) -> float:
+        if self.started_at is None or self.finished_at is None:
+            return 0.0
+        return (self.finished_at - self.started_at) * 1000
+
+
+@dataclass
+class RunResult:
+    workflow_id: str
+    policy: Policy
+    status: Literal["succeeded", "failed"]
+    makespan_ms: float
+    steps: dict[str, StepResult] = field(default_factory=dict)
+    predicted_critical_path: list[str] = field(default_factory=list)
+    predicted_critical_path_ms: float = 0.0
+    actual_critical_path: list[str] = field(default_factory=list)
+    actual_critical_path_ms: float = 0.0
+    api_calls: int = 0
+    cache_hits: int = 0
+
+
+Event = dict[str, Any]
+
+
+def backoff_s(attempt: int, error: Exception, rng: random.Random, base_s: float, cap_s: float) -> float:
+    """Full jitter: uniform(0, min(cap, base·2^attempt)); a server Retry-After wins."""
+    if isinstance(error, TransientNodeError) and error.retry_after_s is not None:
+        return error.retry_after_s
+    return rng.uniform(0, min(cap_s, base_s * 2**attempt))
+
+
+def node_for(step: Step, nodes: dict[str, Node]) -> Node:
+    """The step's connector node if it names one, else its type's default (D10)."""
+    return nodes[step.connector or step.type]
+
+
+def check_nodes(workflow: Workflow, nodes: dict[str, Node]) -> None:
+    """Reject a run up front if any step's connector or fallback can't be resolved (D10)."""
+    errors: list[str] = []
+    for step in workflow.steps:
+        key = step.connector or step.type
+        node = nodes.get(key)
+        if node is None:
+            what = f"connector '{key}'" if step.connector else f"step type '{key}'"
+            errors.append(f"step '{step.id}': no node registered for {what}")
+            continue
+        if node.type != step.type:
+            errors.append(f"step '{step.id}' has type '{step.type}' but connector '{key}' has type '{node.type}'")
+        if node.fallback is not None:
+            backup = nodes.get(node.fallback)
+            if backup is None:
+                errors.append(f"connector '{key}' names fallback '{node.fallback}', which is not registered")
+            elif backup.type != node.type:
+                errors.append(f"connector '{key}' falls back to '{node.fallback}' of a different type")
+    if errors:
+        raise ValueError("; ".join(errors))
+
+
+async def run_workflow(
+    workflow: Workflow,
+    nodes: dict[str, Node],
+    *,
+    policy: Policy = "critical_path",
+    use_cache: bool = True,
+    max_concurrency: int | None = None,
+    rate_limits: dict[str, TokenBucket] | None = None,
+    storage: Any | None = None,
+    history: dict[str, float] | None = None,
+    on_event: Callable[[Event], None] | None = None,
+    retry_base_s: float = 1.0,
+    retry_cap_s: float = 30.0,
+    seed: int | None = None,
+) -> RunResult:
+    """Execute `workflow`. `nodes` maps step type → the type's default Node, and connector
+    id → that connector's Node (D10); a step uses its `connector` if set, else its type.
+    `rate_limits` maps a node's rate_limit_key → shared TokenBucket. `storage` (optional)
+    enables the persistent cache, supplies duration history and records this run's
+    measured durations.
+    """
+    if policy not in POLICIES:
+        raise ValueError(f"unknown policy {policy!r}")
+    check_nodes(workflow, nodes)
+
+    dag = graph.build_dag(workflow)
+    topo = graph.topological_order(dag)
+    topo_index = {s: i for i, s in enumerate(topo)}
+    level = {s: i for i, layer in enumerate(graph.levels(dag)) for s in layer}
+
+    if history is None:
+        history = storage.duration_history(workflow.id) if storage is not None else {}
+    weights = estimate_weights(workflow, history)
+    bl = bottom_levels(dag, weights)
+    predicted = critical_path(dag, weights)
+
+    k = max_concurrency or workflow.max_concurrency
+    type_limits: dict[str, int] = dict(workflow.type_concurrency)
+    if policy == "sequential":
+        k = 1
+    elif policy == "greedy":
+        k, type_limits = math.inf, {}
+
+    steps = {s.id: s for s in workflow.steps}
+    results = {s.id: StepResult(s.id) for s in workflow.steps}
+    outputs: dict[str, Any] = {}
+    unmet = {s: len(dag.parents[s]) for s in dag.nodes}
+    cache = CallCache(storage)
+    rate_limits = rate_limits or {}
+    rng = random.Random(seed)
+    loop = asyncio.get_running_loop()
+    t0 = loop.time()
+    ready_seq = itertools.count()
+
+    def now() -> float:
+        return loop.time() - t0
+
+    def emit(event: Event) -> None:
+        if on_event is not None:
+            on_event({"t": round(now(), 4), **event})
+
+    def set_state(step_id: str, state: StepState, **extra: Any) -> None:
+        results[step_id].state = state
+        emit({"type": "step", "step_id": step_id, "state": state.value, **extra})
+
+    def _priority(step_id: str) -> tuple:
+        if policy == "critical_path":
+            return (-bl[step_id], topo_index[step_id])
+        if policy == "fifo":
+            return (next(ready_seq),)
+        if policy == "levels":
+            return (level[step_id], topo_index[step_id])
+        return (topo_index[step_id],)
+
+    ready: list[tuple[tuple, str]] = []
+
+    def make_ready(step_id: str) -> None:
+        set_state(step_id, StepState.READY)
+        heapq.heappush(ready, (_priority(step_id), step_id))
+
+    # --- one step's execution (runs as its own task) ---------------------------
+
+    def caching_for(step: Step, params: dict[str, Any]) -> tuple[bool, bool]:
+        """(dedupe within run, persist across runs) per D3."""
+        if not use_cache or step.cache is False:
+            return False, False
+        persistent = step.cache is True or node_for(step, nodes).cacheable_across_runs(params)
+        return True, persistent and storage is not None
+
+    async def attempts_on(node: Node, step: Step, params: dict[str, Any]) -> Any:
+        bucket = rate_limits.get(node.rate_limit_key) if node.rate_limit_key else None
+        for attempt in range(step.retries + 1):
+            if bucket is not None:
+                await bucket.acquire()
+            results[step.id].attempts += 1
+            results[step.id].answered_by = node.connector_id  # the last connector tried, so failures are attributed too
+            call_start = loop.time()
+            try:
+                output = await asyncio.wait_for(node.run(params), step.effective_timeout_s)
+                results[step.id].call_ms = (loop.time() - call_start) * 1000
+                return output
+            except (TransientNodeError, TimeoutError) as exc:
+                if attempt == step.retries:
+                    raise
+                delay = backoff_s(attempt, exc, rng, retry_base_s, retry_cap_s)
+                emit({"type": "retry", "step_id": step.id, "attempt": attempt + 1,
+                      "error": str(exc) or type(exc).__name__, "delay_s": round(delay, 3)})
+                await asyncio.sleep(delay)
+        raise AssertionError("unreachable")
+
+    async def call_with_retries(step: Step, params: dict[str, Any]) -> Any:
+        node = node_for(step, nodes)
+        try:
+            return await attempts_on(node, step, params)
+        except (TransientNodeError, TimeoutError) as exc:
+            if node.fallback is None:
+                raise
+            # one hop, only after transient exhaustion; permanent errors never get here (D10)
+            emit({"type": "fallback", "step_id": step.id, "from": node.connector_id, "to": node.fallback,
+                  "error": str(exc) or type(exc).__name__})
+            return await attempts_on(nodes[node.fallback], step, params)
+
+    async def execute(step: Step, params: dict[str, Any]) -> Any:
+        dedupe, persistent = caching_for(step, params)
+        if not dedupe:
+            return await call_with_retries(step, params)
+        value, hit = await cache.get_or_run(
+            cache_key(node_for(step, nodes).cache_namespace, params), lambda: call_with_retries(step, params),
+            persistent,
+        )
+        results[step.id].cache_hit = hit
+        return value
+
+    # --- dispatcher --------------------------------------------------------------
+
+    running: dict[asyncio.Task[Any], tuple[str, bool]] = {}  # task → (step id, holds a slot)
+    type_running: dict[str, int] = {}
+    slots_used = 0
+
+    def current_level() -> int:
+        return min((level[s] for s, r in results.items() if r.state not in TERMINAL), default=0)
+
+    def start(step_id: str, params: dict[str, Any], holds_slot: bool) -> None:
+        nonlocal slots_used
+        step = steps[step_id]
+        if holds_slot:
+            slots_used += 1
+            type_running[step.type] = type_running.get(step.type, 0) + 1
+        results[step_id].started_at = now()
+        set_state(step_id, StepState.RUNNING)
+        running[asyncio.create_task(execute(step, params))] = (step_id, holds_slot)
+
+    def dispatch() -> None:
+        """Start every ready step that can start now, best priority first."""
+        blocked: list[tuple[tuple, str]] = []
+        gate = current_level() if policy == "levels" else None
+        while ready:
+            item = heapq.heappop(ready)
+            step_id = item[1]
+            step = steps[step_id]
+            if gate is not None and level[step_id] > gate:
+                blocked.append(item)
+                continue
+            try:
+                params = templating.resolve(step.params, outputs)
+            except templating.TemplateError as exc:
+                fail(step_id, f"template error: {exc}")
+                continue
+            dedupe, _ = caching_for(step, params)
+            if dedupe and cache.has(cache_key(node_for(step, nodes).cache_namespace, params)):
+                start(step_id, params, holds_slot=False)
+                continue
+            type_full = type_running.get(step.type, 0) >= type_limits.get(step.type, math.inf)
+            if slots_used >= k or type_full:
+                blocked.append(item)
+                continue
+            start(step_id, params, holds_slot=True)
+        for item in blocked:
+            heapq.heappush(ready, item)
+
+    def fail(step_id: str, error: str) -> None:
+        results[step_id].error = error
+        results[step_id].finished_at = results[step_id].finished_at or now()
+        set_state(step_id, StepState.FAILED, error=error)
+        for d in sorted(graph.descendants(dag, step_id), key=topo_index.__getitem__):
+            if results[d].state not in TERMINAL:
+                set_state(d, StepState.SKIPPED, reason=f"upstream '{step_id}' failed")
+
+    def succeed(step_id: str, output: Any) -> None:
+        outputs[step_id] = output
+        results[step_id].output = output
+        # the output rides on the event so a live view can show it before the run ends (D16)
+        set_state(step_id, StepState.SUCCEEDED, cache_hit=results[step_id].cache_hit,
+                  duration_ms=round(results[step_id].duration_ms, 1), output=output)
+        for child in dag.children[step_id]:
+            unmet[child] -= 1
+            if unmet[child] == 0 and results[child].state == StepState.PENDING:
+                make_ready(child)
+
+    emit({"type": "run", "state": "started", "workflow_id": workflow.id, "policy": policy,
+          "predicted_critical_path": predicted.path, "predicted_critical_path_ms": predicted.length_ms,
+          "estimates": weights})
+    for s in topo:
+        if unmet[s] == 0:
+            make_ready(s)
+
+    try:
+        dispatch()
+        while running:
+            done, _ = await asyncio.wait(running, return_when=asyncio.FIRST_COMPLETED)
+            for task in done:
+                step_id, held = running.pop(task)
+                if held:
+                    slots_used -= 1
+                    type_running[steps[step_id].type] -= 1
+                results[step_id].finished_at = now()
+                exc = task.exception()
+                if exc is None:
+                    succeed(step_id, task.result())
+                elif isinstance(exc, (NodeError, TimeoutError)):
+                    fail(step_id, str(exc) or type(exc).__name__)
+                else:  # a bug in a node — still contained to this step
+                    fail(step_id, f"{type(exc).__name__}: {exc}")
+            dispatch()
+    finally:
+        for task in running:
+            task.cancel()
+
+    # --- results -------------------------------------------------------------------
+
+    stuck = [s for s, r in results.items() if r.state not in TERMINAL]
+    assert not stuck, f"scheduler bug: steps never finished: {stuck}"
+
+    measured = {s: r.duration_ms for s, r in results.items()}
+    actual = critical_path(dag, measured)
+    status: Literal["succeeded", "failed"] = (
+        "succeeded" if all(r.state == StepState.SUCCEEDED for r in results.values()) else "failed"
+    )
+    result = RunResult(
+        workflow_id=workflow.id,
+        policy=policy,
+        status=status,
+        makespan_ms=now() * 1000,
+        steps=results,
+        predicted_critical_path=predicted.path,
+        predicted_critical_path_ms=predicted.length_ms,
+        actual_critical_path=actual.path,
+        actual_critical_path_ms=actual.length_ms,
+        api_calls=sum(r.attempts for r in results.values()),
+        cache_hits=cache.hits,
+    )
+    if storage is not None:
+        storage.record_durations(
+            workflow.id,
+            {s: r.duration_ms for s, r in results.items()
+             if r.state == StepState.SUCCEEDED and not r.cache_hit},
+        )
+    emit({"type": "run", "state": status, "makespan_ms": round(result.makespan_ms, 1),
+          "actual_critical_path": actual.path})
+    return result

flowforge/scheduler/graph.py

@@ -0,0 +1,108 @@
+"""DAG construction, topological sort, cycle detection (D7)."""
+
+from __future__ import annotations
+
+import heapq
+from collections import deque
+from dataclasses import dataclass, field
+
+from flowforge.schema import Workflow
+
+
+class CycleError(ValueError):
+    def __init__(self, cycle: list[str]):
+        # cycle is closed: ["a", "b", "c", "a"]
+        super().__init__("cycle: " + " → ".join(cycle))
+        self.cycle = cycle
+
+
+@dataclass
+class DAG:
+    nodes: list[str]
+    parents: dict[str, list[str]] = field(default_factory=dict)   # step -> its depends_on
+    children: dict[str, list[str]] = field(default_factory=dict)  # step -> steps depending on it
+
+
+def build_dag(workflow: Workflow) -> DAG:
+    """Adjacency lists from `depends_on`. Raise CycleError if the graph is not acyclic."""
+    nodes = [s.id for s in workflow.steps]
+    parents = {s.id: list(dict.fromkeys(s.depends_on)) for s in workflow.steps}  # dedupe, keep order
+    children: dict[str, list[str]] = {n: [] for n in nodes}
+    for child in nodes:
+        for parent in parents[child]:
+            children[parent].append(child)
+
+    dag = DAG(nodes=nodes, parents=parents, children=children)
+    topological_order(dag)  # raises CycleError
+    return dag
+
+
+def topological_order(dag: DAG) -> list[str]:
+    """Kahn's algorithm. Ties are broken by original step order (a min-heap on the
+    step's index), so the output is deterministic: O((V + E) log V).
+
+    If fewer than V nodes come out, find and raise the actual cycle.
+    """
+    index = {n: i for i, n in enumerate(dag.nodes)}
+    indegree = {n: len(dag.parents[n]) for n in dag.nodes}
+    ready = [index[n] for n in dag.nodes if indegree[n] == 0]
+    heapq.heapify(ready)
+
+    order: list[str] = []
+    while ready:
+        node = dag.nodes[heapq.heappop(ready)]
+        order.append(node)
+        for child in dag.children[node]:
+            indegree[child] -= 1
+            if indegree[child] == 0:
+                heapq.heappush(ready, index[child])
+
+    if len(order) < len(dag.nodes):
+        leftover = [n for n in dag.nodes if indegree[n] > 0]
+        raise CycleError(_find_cycle(dag, set(leftover), leftover[0]))
+    return order
+
+
+def _find_cycle(dag: DAG, leftover: set[str], start: str) -> list[str]:
+    """Every node Kahn's left behind still has an unprocessed parent, which is itself
+    left over. So walking parent links from any leftover node stays inside the leftovers
+    and must eventually revisit a node — that repeat closes a cycle.
+    """
+    position: dict[str, int] = {}
+    walk: list[str] = []
+    node = start
+    while node not in position:
+        position[node] = len(walk)
+        walk.append(node)
+        node = next(p for p in dag.parents[node] if p in leftover)
+    # walk[position[node]:] follows child → parent; reverse it to read in execution order
+    cycle = walk[position[node]:][::-1]
+    return cycle + [cycle[0]]
+
+
+def levels(dag: DAG) -> list[list[str]]:
+    """BFS layers: level 0 = no parents, level i = longest parent chain of length i.
+
+    Used by benchmark strategy S2 (level-by-level with barriers).
+    Within a level, steps keep their original order.
+    """
+    level: dict[str, int] = {}
+    for node in topological_order(dag):
+        level[node] = 1 + max((level[p] for p in dag.parents[node]), default=-1)
+
+    layers: list[list[str]] = [[] for _ in range(max(level.values(), default=-1) + 1)]
+    for node in dag.nodes:
+        layers[level[node]].append(node)
+    return layers
+
+
+def descendants(dag: DAG, step_id: str) -> set[str]:
+    """Every step reachable from step_id — these get `skipped` when it fails."""
+    seen: set[str] = set()
+    queue = deque(dag.children[step_id])
+    while queue:
+        node = queue.popleft()
+        if node not in seen:
+            seen.add(node)
+            queue.extend(dag.children[node])
+    return seen

flowforge/scheduler/rate_limit.py

@@ -0,0 +1,36 @@
+"""Async token bucket (D5, D8)."""
+
+from __future__ import annotations
+
+import asyncio
+
+
+class TokenBucket:
+    """`rate_per_min` tokens refill continuously; bursts up to `capacity`.
+
+    acquire() waits with asyncio.sleep (never busy-loops) until a token is available.
+    Waiters queue on a lock while sleeping, so tokens are handed out first-come first-served.
+    """
+
+    def __init__(self, rate_per_min: float, capacity: int | None = None):
+        if rate_per_min <= 0:
+            raise ValueError("rate_per_min must be positive")
+        self.rate_per_s = rate_per_min / 60
+        self.capacity = capacity if capacity is not None else max(1, int(rate_per_min / 6))
+        self.tokens = float(self.capacity)
+        self._last: float | None = None
+        self._lock = asyncio.Lock()
+
+    def _refill(self, now: float) -> None:
+        if self._last is not None:
+            self.tokens = min(self.capacity, self.tokens + (now - self._last) * self.rate_per_s)
+        self._last = now
+
+    async def acquire(self) -> None:
+        loop = asyncio.get_running_loop()
+        async with self._lock:
+            self._refill(loop.time())
+            while self.tokens < 1:
+                await asyncio.sleep((1 - self.tokens) / self.rate_per_s)
+                self._refill(loop.time())
+            self.tokens -= 1

flowforge/schema.py

@@ -0,0 +1,81 @@
+"""Workflow JSON models and static validation (DECISIONS.md D2, D7).
+
+Validation here is everything that can be checked per-step or pairwise.
+Cycle detection needs the whole graph and lives in scheduler/graph.py.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+from flowforge.connectors.models import CONNECTOR_ID_PATTERN
+from flowforge.templating import referenced_steps
+
+StepType = Literal["llm", "mcp", "http", "local", "mock", "check"]   # check: a health check (D24)
+
+DEFAULT_TIMEOUT_S: dict[str, float] = {"llm": 60, "mcp": 30, "http": 30, "local": 60, "mock": 30, "check": 60}
+
+
+class Step(BaseModel):
+    model_config = ConfigDict(extra="forbid")  # catch typos like "depend_on"
+
+    id: str = Field(pattern=r"^[A-Za-z_][A-Za-z0-9_]*$")
+    type: StepType
+    connector: str | None = Field(default=None, pattern=CONNECTOR_ID_PATTERN)  # None = type default (D10)
+    depends_on: list[str] = Field(default_factory=list)
+    params: dict[str, Any] = Field(default_factory=dict)
+    estimated_ms: float | None = Field(default=None, gt=0)
+    timeout_s: float | None = Field(default=None, gt=0)
+    retries: int = Field(default=2, ge=0, le=10)
+    cache: bool | None = None  # None = use the per-type default (D3)
+    description: str = ""
+
+    @property
+    def effective_timeout_s(self) -> float:
+        return self.timeout_s if self.timeout_s is not None else DEFAULT_TIMEOUT_S[self.type]
+
+
+class Workflow(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    id: str
+    name: str = ""
+    max_concurrency: int = Field(default=4, ge=1)
+    type_concurrency: dict[StepType, int] = Field(default_factory=dict)
+    steps: list[Step] = Field(min_length=1)
+
+    @model_validator(mode="after")
+    def _check_references(self) -> Workflow:
+        errors: list[str] = []
+        ids = [s.id for s in self.steps]
+        known = set(ids)
+
+        seen: set[str] = set()
+        for sid in ids:
+            if sid in seen:
+                errors.append(f"duplicate step id '{sid}'")
+            seen.add(sid)
+
+        for step in self.steps:
+            deps = set(step.depends_on)
+            if step.id in deps:
+                errors.append(f"step '{step.id}' depends on itself")
+            for dep in deps - known:
+                errors.append(f"step '{step.id}' depends on unknown step '{dep}'")
+            for ref in referenced_steps(step.params) - deps:
+                errors.append(
+                    f"step '{step.id}' references '{{{{steps.{ref}...}}}}' "
+                    f"but '{ref}' is not in its depends_on"
+                )
+
+        if errors:
+            raise ValueError("; ".join(errors))
+        return self
+
+    def step(self, step_id: str) -> Step:
+        for s in self.steps:
+            if s.id == step_id:
+                return s
+        raise KeyError(step_id)