aethergraph 0.1.0a1__py3-none-any.whl

aethergraph/core/runtime/execution_context.py

@@ -0,0 +1,220 @@
+# aethergraph/core/execution/context.py
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from datetime import datetime
+import importlib
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from aethergraph.core.graph.task_node import TaskNodeRuntime
+
+from aethergraph.services.clock.clock import SystemClock
+from aethergraph.services.logger.std import StdLoggerService
+from aethergraph.services.resume.router import ResumeRouter
+
+from ..graph.graph_refs import GRAPH_INPUTS_NODE_ID, RESERVED_INJECTABLES
+from .bound_memory import BoundMemoryAdapter
+from .node_context import NodeContext
+from .node_services import NodeServices
+
+
+@dataclass
+class ExecutionContext:
+    run_id: str
+    graph_id: str | None
+    graph_inputs: dict[str, Any]
+    outputs_by_node: dict[str, dict[str, Any]]
+    services: NodeServices
+    logger_factory: StdLoggerService
+    clock: SystemClock
+    resume_payload: dict[str, Any] | None = None
+    should_run_fn: Callable[[], bool] | None = None
+    resume_router: ResumeRouter | None = None  # ResumeRouter
+
+    # Back-compat shim
+    bound_memory: BoundMemoryAdapter | None = None
+
+    def create_node_context(self, node: TaskNodeRuntime) -> NodeContext:
+        return NodeContext(
+            run_id=self.run_id,
+            graph_id=self.graph_id or "",
+            node_id=node.node_id,
+            services=self.services,
+            resume_payload=self.resume_payload,
+            # back-compat for old ctx.mem()
+            bound_memory=self.bound_memory,
+        )
+
+    # def as_node_context(self, ad) -> "NodeContext":
+    #     """ Create a NodeContext representing this execution context itself as a node.
+    #         Useful for ad-hoc contexts that don't have real nodes.
+    #     """
+    #     return NodeContext(
+    #         run_id=self.run_id,
+    #         graph_id=self.graph_id or "",
+    #         node_id="ad_",
+    #         services=self.services,
+    #         resume_payload=self.resume_payload,
+    #         # back-compat for old ctx.mem()
+    #         bound_memory=self.bound_memory,
+    #     )
+
+    # ----- helpers used by step forward() -----
+    def now(self) -> datetime:
+        return self.clock.now()
+
+    def resolve(self, logic_ref: str):
+        """Resolve a logic reference to a callable.
+        NOTE: This is not used anymore; prefer get_logic().
+        """
+        # fallback dotted import
+        mod, _, attr = logic_ref.rpartition(".")
+        return getattr(importlib.import_module(mod), attr)
+
+    def get_logic(self, logic_ref):
+        """Resolve a logic reference to a callable.
+        If a registry is available and the ref looks like a registry key, use it.
+        Otherwise, if a dotted path, import it.
+        Otherwise, return as-is (assumed callable).
+        Args:
+            logic_ref: A callable, dotted path string, or registry key string.
+        Returns:
+            The resolved callable.
+        """
+        if isinstance(logic_ref, str) and logic_ref.startswith("registry:") and self.registry:
+            # registry key
+            return self.registry.get_logic_ref(logic_ref)
+        if isinstance(logic_ref, str):
+            # dotted path fallback
+            mod, _, attr = logic_ref.rpartition(".")
+            return getattr(importlib.import_module(mod), attr)
+        return logic_ref
+
+    async def resolve_inputs(self, node) -> dict[str, Any]:
+        """
+        Materialize a node's input mapping by resolving:
+        - {"_type":"arg","key":K} → graph input value (or optional default)
+        - {"_type":"ref","from":NODE_ID,"key":OUT} → upstream node's output value
+        - {"_type":"context","key":K,"default":D} → memory value (or D if missing)
+        Works recursively over dicts/lists/tuples.
+
+        The function works as follows:
+         - If the value is a dict with "_type" of "arg", it looks up the graph input.
+         - If the value is a dict with "_type" of "ref", it looks up
+        the specified node's output.
+         - If the value is a dict without special keys, it recursively resolves
+           each key-value pair.
+         - If the value is a list or tuple, it recursively resolves each element.
+         - Otherwise, it returns the value as-is (assumed to be a constant).
+
+        Args:
+            node: The TaskNodeRuntime whose inputs to resolve.
+        Returns:
+            The fully resolved inputs dict for the node.
+        Raises:
+            KeyError: If a referenced graph input or node output is missing.
+        """
+        raw = getattr(node, "inputs", {}) or {}
+        # Grab optional defaults from the graph spec if available
+        opt_defaults: dict[str, Any] = {}
+        parent_graph = getattr(node, "_parent_graph", None)
+        if parent_graph and getattr(parent_graph, "spec", None):
+            # _io_inputs_optional is a dict[str, Any]
+            opt_defaults = getattr(parent_graph.spec, "inputs_optional", {}) or {}
+
+        # Allow a fallback to graph.state.node_outputs if scheduler hasn't copied yet
+        fallback_outputs = {}
+        if parent_graph and getattr(parent_graph, "state", None):
+            fallback_outputs = getattr(parent_graph.state, "node_outputs", {}) or {}
+
+        def _err_path(msg: str, path: str):
+            raise KeyError(
+                f"{msg} (node={getattr(node, 'node_id', getattr(node, 'id', '?'))}, path={path})"
+            )
+
+        def _resolve_arg(marker: dict[str, Any], path: str):
+            k = marker.get("key")
+            if k is None:
+                _err_path("Bad arg marker (missing 'key')", path)
+            if k in self.graph_inputs:
+                return self.graph_inputs[k]
+            if k in opt_defaults:
+                return opt_defaults[k]
+            # Helpful error: show known keys
+            known = list(self.graph_inputs.keys())
+            _err_path(f"Graph input '{k}' not provided (known inputs: {known})", path)
+
+        def _resolve_ref(marker: dict[str, Any], path: str):
+            src = marker.get("from")
+            out_key = marker.get("key")
+            if src is None or out_key is None:
+                _err_path("Bad ref marker (need 'from' and 'key')", path)
+
+            # Tolerate someone emitting a ref to the inputs sentinel
+            if src == GRAPH_INPUTS_NODE_ID:
+                # Interpret as an 'arg' reference to graph inputs
+                return _resolve_arg({"_type": "arg", "key": out_key}, path + ".__graph_inputs__")
+
+            # Primary source: env.outputs_by_node (scheduler publishes here)
+            if src in self.outputs_by_node:
+                outs = self.outputs_by_node[src] or {}
+                if out_key in outs:
+                    return outs[out_key]
+                _err_path(
+                    f"Upstream node '{src}' has no output key '{out_key}'. "
+                    f"Available: {list(outs.keys())}",
+                    path,
+                )
+
+            # Fallback: graph state (useful during tests or if scheduler filled it there)
+            if src in fallback_outputs:
+                outs = fallback_outputs[src] or {}
+                if out_key in outs:
+                    return outs[out_key]
+                _err_path(
+                    f"(fallback) Upstream node '{src}' has no output key '{out_key}'. "
+                    f"Available: {list(outs.keys())}",
+                    path,
+                )
+
+            _err_path(f"Upstream node '{src}' outputs not available yet", path)
+
+        def _resolve_any(val: Any, path: str):
+            # Handle dict markers
+            if isinstance(val, dict):
+                t = val.get("_type")
+                if t == "arg":
+                    return _resolve_arg(val, path)
+                if t == "ref":
+                    return _resolve_ref(val, path)
+                if t == "context":
+                    return self.memory.read(val["key"], val.get("default"))
+                # regular dict: recurse keys
+                return {k: _resolve_any(v, f"{path}.{k}") for k, v in val.items()}
+
+            # Handle list/tuple
+            if isinstance(val, list):
+                return [_resolve_any(v, f"{path}[{i}]") for i, v in enumerate(val)]
+            if isinstance(val, tuple):
+                return tuple(_resolve_any(v, f"{path}[{i}]") for i, v in enumerate(val))
+
+            # Pass-through literal
+            return val
+
+        # Make sure we don't mutate node.inputs
+        # materialized = _resolve_any(copy.deepcopy(raw), path="inputs")
+        materialized = _resolve_any(raw, path="inputs")
+
+        # Strip framework-reserved injectables from *user* inputs.
+        # We always inject these later from the execution context.
+        if isinstance(materialized, dict):
+            for k in list(materialized.keys()):
+                if k in RESERVED_INJECTABLES:
+                    materialized.pop(k, None)
+
+            # If someone put arguments under "kwargs", keep them;
+            # build_call_kwargs will flatten and then drop "kwargs".
+            # (No change needed here beyond not touching it.)
+        return materialized

aethergraph/core/runtime/graph_runner.py

@@ -0,0 +1,349 @@
+from __future__ import annotations
+
+import asyncio
+import threading
+from typing import Any
+import uuid
+
+from aethergraph.contracts.errors.errors import GraphHasPendingWaits
+from aethergraph.contracts.services.state_stores import GraphSnapshot
+from aethergraph.core.runtime.recovery import hash_spec, recover_graph_run
+from aethergraph.services.container.default_container import build_default_container
+from aethergraph.services.state_stores.graph_observer import PersistenceObserver
+from aethergraph.services.state_stores.resume_policy import (
+    assert_snapshot_json_only,
+)
+from aethergraph.services.state_stores.utils import snapshot_from_graph
+
+from ..execution.forward_scheduler import ForwardScheduler
+from ..execution.retry_policy import RetryPolicy
+from ..graph.graph_fn import GraphFunction
+from ..graph.graph_refs import resolve_any as _resolve_any
+from ..runtime.runtime_env import RuntimeEnv
+from ..runtime.runtime_services import ensure_services_installed
+from .run_registration import RunRegistrationGuard
+
+
+# ---------- env helpers ----------
+def _get_container():
+    # install once if not installed by sidecar/server
+    return ensure_services_installed(build_default_container)
+
+
+async def _attach_persistence(graph, env, spec, snapshot_every=1) -> PersistenceObserver:
+    """
+    Wire the centralized state_store to the graph via PersistenceObserver.
+    Returns the observer instance so caller can optionally force a final snapshot.
+    """
+    store = getattr(env.container, "state_store", None) or getattr(env, "state_store", None)
+    if not store:
+        # Safe no-op: resumability won't work but run still executes.
+        return None
+
+    obs = PersistenceObserver(
+        store=store,
+        artifact_store=getattr(env.container, "artifacts", None),
+        spec_hash=hash_spec(spec),
+        snapshot_every=snapshot_every,
+    )
+    graph.add_observer(obs)
+    return obs
+
+
+async def _build_env(
+    owner, inputs: dict[str, Any], **rt_overrides
+) -> tuple[RuntimeEnv, RetryPolicy, int]:
+    container = _get_container()
+    # apply optional overrides onto the container instance
+    for k, v in rt_overrides.items():
+        if v is not None and hasattr(container, k):
+            setattr(container, k, v)
+
+    run_id = rt_overrides.get("run_id") or f"run-{uuid.uuid4().hex[:8]}"
+    env = RuntimeEnv(
+        run_id=run_id,
+        graph_inputs=inputs,
+        outputs_by_node={},
+        container=container,
+    )
+    retry = rt_overrides.get("retry") or RetryPolicy()
+    max_conc = rt_overrides.get("max_concurrency", getattr(owner, "max_concurrency", 4))
+    return env, retry, max_conc
+
+
+# ---------- materialization ----------
+def _materialize_task_graph(target) -> Any:
+    """
+    Accept:
+      - TaskGraph instance (has io_signature attr)
+      - graph builder object with .build()
+      - a callable builder that returns a TaskGraph when invoked with no args
+    """
+    # already a TaskGraph
+    if hasattr(target, "io_signature"):
+        return target
+
+    # builder pattern with .build()
+    if hasattr(target, "build") and callable(target.build):
+        g = target.build()
+        if hasattr(g, "io_signature"):
+            return g
+
+    # callable builder that returns a TaskGraph
+    if callable(target):
+        g = target()
+        if hasattr(g, "io_signature"):
+            return g
+
+    raise TypeError(
+        "run_async: target must be a TaskGraph instance, a TaskGraph builder, "
+        "or a callable returning a TaskGraph."
+    )
+
+
+def _resolve_graph_outputs(
+    graph,
+    inputs: dict[str, Any],
+    env: RuntimeEnv,
+):
+    bindings = graph.io_signature().get("outputs", {}).get("bindings", {})
+
+    def _res(b):
+        return _resolve_any(b, graph_inputs=inputs, outputs_by_node=env.outputs_by_node)
+
+    try:
+        result = {k: _res(v) for k, v in bindings.items()}
+    except KeyError as e:
+        waiting = [
+            nid
+            for nid, n in graph.state.nodes.items()
+            if getattr(n, "status", "").startswith("WAITING_")
+        ]
+        continuations = []
+        if env.continuation_store and hasattr(env.continuation_store, "get"):
+            for nid in waiting:
+                cont = env.continuation_store.get(run_id=env.run_id, node_id=nid)
+                if cont:
+                    continuations.append(
+                        {
+                            "node_id": nid,
+                            "kind": cont.kind,
+                            "token": cont.token,
+                            "channel": cont.channel,
+                            "deadline": getattr(cont.deadline, "isoformat", lambda: None)(),
+                        }
+                    )
+        raise GraphHasPendingWaits(
+            "Graph quiesced with pending waits; outputs are not yet resolvable.",
+            waiting_nodes=waiting,
+            continuations=continuations,
+        ) from e
+
+    return next(iter(result.values())) if len(result) == 1 else result
+
+
+def _resolve_graph_outputs_or_waits(graph, inputs, env, *, raise_on_waits: bool = True):
+    try:
+        return _resolve_graph_outputs(graph, inputs, env)
+    except GraphHasPendingWaits as e:
+        if raise_on_waits:
+            raise
+        return {
+            "status": "waiting",
+            "waiting_nodes": e.waiting_nodes,
+            "continuations": e.continuations,
+        }
+
+
+def _seed_outputs_from_snapshot(env, snap: GraphSnapshot):
+    env.outputs_by_node = env.outputs_by_node or {}
+    nodes = snap.state.get("nodes", {})
+    for nid, ns in nodes.items():
+        outs = (ns or {}).get("outputs") or {}
+        if outs:
+            env.outputs_by_node[nid] = outs
+
+
+def _is_graph_complete(snap: GraphSnapshot) -> bool:
+    nodes = snap.state.get("nodes", {})
+    if not nodes:
+        return False
+
+    # completed if every node is DONE/SKIPPED or has outputs matching spec
+    def doneish(st):
+        s = (st or {}).get("status", "")
+        return s in ("DONE", "SKIPPED")
+
+    return all(doneish(ns) for ns in nodes.values())
+
+
+async def load_latest_snapshot_json(store, run_id: str) -> dict[str, Any] | None:
+    """
+    Returns the raw JSON dict of the latest snapshot (or None).
+    """
+    snap = await store.load_latest_snapshot(run_id)
+    if not snap:
+        return None
+    # JsonGraphStateStore serializes GraphSnapshot via snap.__dict__
+    # load_latest_snapshot already returns a GraphSnapshot(**jsondict).
+    # Convert back to plain JSON-ish dict:
+    return {
+        "run_id": snap.run_id,
+        "graph_id": snap.graph_id,
+        "rev": snap.rev,
+        "spec_hash": snap.spec_hash,
+        "state": snap.state,
+    }
+
+
+# ---------- public API ----------
+async def run_async(target, inputs: dict[str, Any] | None = None, **rt_overrides):
+    """
+    Generic async runner for TaskGraph or GraphFunction.
+    - GraphFunction → delegates to gf.run(env=..., **inputs)
+    - TaskGraph/builder → schedules and resolves graph-level outputs
+    """
+    inputs = inputs or {}
+    # GraphFunction path
+    if isinstance(target, GraphFunction):
+        env, retry, max_conc = await _build_env(target, inputs, **rt_overrides)
+        return await target.run(env=env, max_concurrency=max_conc, **inputs)
+
+    # TaskGraph path
+    graph = _materialize_task_graph(target)
+    env, retry, max_conc = await _build_env(graph, inputs, **rt_overrides)
+
+    # Extract spec for run/recovery ...
+    spec = getattr(graph, "spec", None) or getattr(graph, "get_spec", lambda: None)()
+    if spec is None:
+        spec = graph.spec
+
+    store = getattr(env.container, "state_store", None)
+    snap = None
+    assert store is None or hasattr(
+        store, "load_latest_snapshot"
+    ), "state_store must implement lo  ad_latest_snapshot(run_id)"
+
+    if store:
+        # 1) Attempt cold-resume (build a graph with hydrated state)
+        graph = await recover_graph_run(spec=spec, run_id=env.run_id, store=store)
+
+        # 2) Load raw JSON snapshot and ENFORCE strict policy
+        snap_json = await load_latest_snapshot_json(store, env.run_id)
+        if snap_json:
+            # keep for short-circuit + seeding
+            snap = await store.load_latest_snapshot(env.run_id)
+            # Short-circuit if already complete
+            if snap:
+                _seed_outputs_from_snapshot(env, snap)
+                if _is_graph_complete(snap):
+                    return _resolve_graph_outputs(graph, inputs, env)
+
+            # strict policy: block resume if any non-JSON / __aether_ref__ is present
+            assert_snapshot_json_only(env.run_id, snap_json, mode="reuse_only")
+    else:
+        graph = _materialize_task_graph(target)
+
+    # Bind/validate inputs
+    graph._validate_and_bind_inputs(inputs)
+
+    # Attach persistence observer + run (unchanged) ...
+    obs = await _attach_persistence(graph, env, spec, snapshot_every=1)
+
+    # get logger from env's container
+    from ..runtime.runtime_services import current_logger_factory
+
+    logger = current_logger_factory().for_scheduler()
+
+    sched = ForwardScheduler(
+        graph,
+        env,
+        retry_policy=retry,
+        max_concurrency=max_conc,
+        skip_dep_on_failure=True,
+        stop_on_first_error=True,
+        logger=logger,
+    )
+
+    # Register for resumes and run
+    with RunRegistrationGuard(run_id=env.run_id, scheduler=sched, container=env.container):
+        try:
+            await sched.run()
+        except asyncio.CancelledError:
+            raise
+        finally:
+            # FINAL SNAPSHOT on normal or cancelled exit (if store exists)
+            if store and obs:
+                artifacts = getattr(env.container, "artifacts", None)
+                snap = await snapshot_from_graph(
+                    run_id=graph.state.run_id or env.run_id,
+                    graph_id=graph.graph_id,
+                    rev=graph.state.rev,
+                    spec_hash=hash_spec(spec),
+                    state_obj=graph.state,
+                    artifacts=artifacts,
+                    allow_externalize=False,  # FIXME: artifact writer async loop error; set False to *avoid* writing artifacts during snapshot
+                    include_wait_spec=True,
+                )
+                await store.save_snapshot(snap)
+
+    # Resolve graph-level outputs (will raise GraphHasPendingWaits if waits)
+    return _resolve_graph_outputs_or_waits(graph, inputs, env, raise_on_waits=True)
+
+
+async def run_or_resume_async(
+    target, inputs: dict[str, Any], *, run_id: str | None = None, **rt_overrides
+):
+    """
+    If state exists for run_id → cold resume, else fresh run.
+    Exactly the same signature as run_async plus optional run_id.
+    """
+    if run_id is not None:
+        rt_overrides = dict(rt_overrides or {}, run_id=run_id)
+    return await run_async(target, inputs, **rt_overrides)
+
+
+# sync adapter (optional, safe in notebooks/servers)
+class _LoopThread:
+    def __init__(self):
+        self._ev = threading.Event()
+        self._thread = threading.Thread(target=self._worker, daemon=True)
+        self._loop = None
+        self._thread.start()
+        self._ev.wait()
+
+    def _worker(self):
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        self._loop = loop
+        self._ev.set()
+        loop.run_forever()
+
+    def submit_old(self, coro):
+        # this will block terminal until coro is done
+        fut = asyncio.run_coroutine_threadsafe(coro, self._loop)
+        return fut.result()
+
+    def submit(self, coro):
+        # this will allow KeyboardInterrupt to propagate -> still not perfect. Use async main if possible.
+        fut = asyncio.run_coroutine_threadsafe(coro, self._loop)
+        try:
+            return fut.result()
+        except KeyboardInterrupt:
+            # cancel the task in the loop thread and wait for cleanup
+            fut.cancel()
+
+            def _cancel_all():
+                for t in asyncio.all_tasks(self._loop):
+                    t.cancel()
+
+            self._loop.call_soon_threadsafe(_cancel_all)
+            raise
+
+
+_LOOP = _LoopThread()
+
+
+def run(target, inputs: dict[str, Any] | None = None, **rt_overrides):
+    inputs = inputs or {}
+    return _LOOP.submit(run_async(target, inputs, **rt_overrides))

aethergraph/core/runtime/lifecycle.py

@@ -0,0 +1,26 @@
+# lifecycle.py
+import asyncio
+
+from aethergraph.core.runtime.runtime_services import current_services
+
+
+async def start_all_services() -> None:
+    svc = current_services()
+    tasks = []
+    for _, inst in getattr(svc, "ext_services", {}).items():
+        start = getattr(inst, "start", None)
+        if asyncio.iscoroutinefunction(start):
+            tasks.append(start())
+    if tasks:
+        await asyncio.gather(*tasks)
+
+
+async def close_all_services() -> None:
+    svc = current_services()
+    tasks = []
+    for _, inst in getattr(svc, "ext_services", {}).items():
+        close = getattr(inst, "close", None)
+        if asyncio.iscoroutinefunction(close):
+            tasks.append(close())
+    if tasks:
+        await asyncio.gather(*tasks)