aethergraph 0.1.0a1__py3-none-any.whl

aethergraph/services/registry/unified_registry.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable, Mapping
+import re
+import threading
+from typing import Any
+
+try:
+    # Prefer packaging for correct PEP 440 / pre-release ordering
+    from packaging.version import Version
+
+    _has_packaging = True
+except Exception:
+    _has_packaging = False
+
+from .key_parsing import parse_ref
+from .registry_key import NS, Key
+
+# allow storing either the object, or a factory that returns the object on first use
+RegistryObject = Any
+RegistryFactory = Callable[[], Any]
+RegistryValue = RegistryObject | RegistryFactory
+
+
+class UnifiedRegistry:
+    """
+    Runtime-only registry: (nspace, name, version) -> object (or lazy factory).
+    Maintains a 'latest' pointer per (nspace, name).
+
+    Thread-safe for concurrent get/register operations.
+    """
+
+    def __init__(self, *, allow_overwrite: bool = True):
+        self._store: dict[tuple[str, str], dict[str, RegistryValue]] = {}
+        self._latest: dict[tuple[str, str], str] = {}
+        self._aliases: dict[tuple[str, str], dict[str, str]] = {}  # (ns,name) -> alias -> version
+        self._lock = threading.RLock()
+        self._allow_overwrite = allow_overwrite
+
+    # ---------- registration ----------
+
+    def register(self, *, nspace: str, name: str, version: str, obj: RegistryValue) -> None:
+        if nspace not in NS:
+            raise ValueError(f"Unknown namespace: {nspace}")
+        key = (nspace, name)
+        with self._lock:
+            versions = self._store.setdefault(key, {})
+            if (version in versions) and not self._allow_overwrite:
+                raise ValueError(
+                    f"{nspace}:{name}@{version} already registered and overwrite disabled"
+                )
+            versions[version] = obj
+            self._latest[key] = self._pick_latest(versions.keys())
+
+    def register_latest(
+        self, *, nspace: str, name: str, obj: RegistryValue, version: str = "0.0.0"
+    ) -> None:
+        # Explicit version anyway; also marks latest via _pick_latest
+        self.register(nspace=nspace, name=name, version=version, obj=obj)
+
+    def alias(self, *, nspace: str, name: str, tag: str, to_version: str) -> None:
+        """Define tag aliases like 'stable', 'canary' mapping to a concrete version."""
+        key = (nspace, name)
+        with self._lock:
+            if key not in self._store or to_version not in self._store[key]:
+                raise KeyError(f"Cannot alias to missing version: {nspace}:{name}@{to_version}")
+            m = self._aliases.setdefault(key, {})
+            m[tag] = to_version
+
+    # ---------- resolve ----------
+
+    def get(self, ref: str | Key) -> Any:
+        key = parse_ref(ref) if isinstance(ref, str) else ref
+        k = (key.nspace, key.name)
+        with self._lock:
+            versions = self._store.get(k)
+            if not versions:
+                raise KeyError(f"Not found: {key.canonical()}")
+
+            # resolve version: explicit → alias → latest
+            ver = key.version
+            ver = self._aliases.get(k, {}).get(ver, ver) if ver else self._latest.get(k)
+
+            if ver not in versions:
+                raise KeyError(f"Version not found: {key.nspace}:{key.name}@{ver}")
+
+            val = versions[ver]
+            # Materialize factories lazily (and cache)
+            if callable(val):
+                obj = val()
+                versions[ver] = obj
+                return obj
+            return val
+
+    # ---------- listing / admin ----------
+
+    def list(self, nspace: str | None = None) -> dict[str, str]:
+        """Return { 'ns:name': '<latest_version>' } optionally filtered."""
+        out: dict[str, str] = {}
+        with self._lock:
+            for (ns, name), _ in self._store.items():
+                if nspace and ns != nspace:
+                    continue
+                out[f"{ns}:{name}"] = self._latest.get((ns, name), "unknown")
+        return out
+
+    def list_versions(self, *, nspace: str, name: str) -> Iterable[str]:
+        k = (nspace, name)
+        with self._lock:
+            return tuple(sorted(self._store.get(k, {}).keys(), key=self._semver_sort_key))
+
+    def get_aliases(self, *, nspace: str, name: str) -> Mapping[str, str]:
+        with self._lock:
+            return dict(self._aliases.get((nspace, name), {}))
+
+    def unregister(self, *, nspace: str, name: str, version: str | None = None) -> None:
+        with self._lock:
+            k = (nspace, name)
+            if k not in self._store:
+                return
+            if version is None:
+                # remove all versions and aliases
+                self._store.pop(k, None)
+                self._latest.pop(k, None)
+                self._aliases.pop(k, None)
+                return
+            vers = self._store[k]
+            vers.pop(version, None)
+            # drop aliases pointing to this version
+            if k in self._aliases:
+                for tag, v in list(self._aliases[k].items()):
+                    if v == version:
+                        self._aliases[k].pop(tag, None)
+            # recompute latest
+            if vers:
+                self._latest[k] = self._pick_latest(vers.keys())
+            else:
+                self._store.pop(k, None)
+                self._latest.pop(k, None)
+                self._aliases.pop(k, None)
+
+    def clear(self) -> None:
+        with self._lock:
+            self._store.clear()
+            self._latest.clear()
+            self._aliases.clear()
+
+    # ---------- typed getters ----------
+
+    def get_tool(self, name: str, version: str | None = None) -> Any:
+        return self.get(Key(nspace="tool", name=name, version=version))
+
+    def get_graph(self, name: str, version: str | None = None) -> Any:
+        return self.get(Key(nspace="graph", name=name, version=version))
+
+    def get_graphfn(self, name: str, version: str | None = None) -> Any:
+        return self.get(Key(nspace="graphfn", name=name, version=version))
+
+    def get_agent(self, name: str, version: str | None = None) -> Any:
+        return self.get(Key(nspace="agent", name=name, version=version))
+
+    # ---------- helpers ----------
+
+    @staticmethod
+    def _semver_sort_key(v: str):
+        if _has_packaging:
+            try:
+                return Version(v)
+            except Exception:
+                # Fall back to naive
+                pass
+        # naive: split on dots and dashes, integers first
+        parts = []
+        for token in re.split(r"[.\-+]", v):
+            try:
+                parts.append((0, int(token)))
+            except ValueError:
+                parts.append((1, token))
+        return tuple(parts)
+
+    def _pick_latest(self, versions: Iterable[str]) -> str:
+        vs = list(versions)
+        if not vs:
+            return "0.0.0"
+        return sorted(vs, key=self._semver_sort_key)[-1]

aethergraph/services/resume/multi_scheduler_resume_bus.py

@@ -0,0 +1,65 @@
+import asyncio
+import hmac
+from logging import getLogger
+
+from aethergraph.contracts.services.continuations import AsyncContinuationStore
+from aethergraph.contracts.services.resume import ResumeBus
+from aethergraph.services.schedulers.registry import SchedulerRegistry
+
+log = getLogger(__name__)
+
+
+class MultiSchedulerResumeBus(ResumeBus):
+    def __init__(
+        self,
+        *,
+        registry: SchedulerRegistry,
+        store: AsyncContinuationStore,
+        delete_after_resume: bool = True,
+        logger=None,
+    ):
+        self.registry = registry
+        self.store = store
+        self.delete_after_resume = delete_after_resume
+        self.logger = logger or log
+
+    async def enqueue_resume(self, *, run_id: str, node_id: str, token: str, payload: dict) -> None:
+        cont = await self.store.get(run_id, node_id)
+        if not cont or not hmac.compare_digest(cont.token, token):
+            self.logger.warning(
+                "[multi-resume-bus] invalid continuation/token for %s/%s", run_id, node_id
+            )
+            return
+
+        sched = self.registry.get(run_id)
+        if not sched:
+            self.logger.warning("[multi-resume-bus] no scheduler for run_id=%s", run_id)
+            return
+
+        loop = getattr(sched, "loop", None)
+        if loop is None:
+            self.logger.error(
+                "[multi-resume-bus] scheduler.loop is not set yet for run_id=%s", run_id
+            )
+            return
+
+        # Always post to the scheduler's loop
+        fut = asyncio.run_coroutine_threadsafe(
+            sched.on_resume_event(run_id, node_id, payload), loop
+        )
+        try:
+            await asyncio.wrap_future(fut)
+        except Exception as e:
+            self.logger.error("[multi-resume-bus] dispatch failed: %s", e, exc_info=True)
+            return
+
+        if self.delete_after_resume:
+            try:
+                await self.store.delete(run_id, node_id)
+            except Exception as e:
+                self.logger.warning(
+                    f"[multi-resume-bus] failed to delete continuation for {run_id}/{node_id}: {e}"
+                )
+
+        sched.post_resume_event_threadsafe(run_id, node_id, payload)
+        return

aethergraph/services/resume/router.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import hmac
+from logging import getLogger
+from typing import Any
+
+from jsonschema import ValidationError, validate
+
+from aethergraph.contracts.services.continuations import AsyncContinuationStore
+from aethergraph.contracts.services.resume import ResumeBus
+
+log = getLogger(__name__)
+
+
+class ResumeRouter:
+    def __init__(
+        self, *, store: AsyncContinuationStore, runner: ResumeBus, logger=None, wait_registry=None
+    ):
+        self.store = store
+        self.runner = runner
+        self.logger = logger or log
+        self.waits = wait_registry
+
+    async def resume(self, run_id: str, node_id: str, token: str, payload: dict[str, Any]) -> None:
+        cont = await self.store.get(run_id, node_id)
+        if not cont:
+            self.logger.error("No continuation for %s/%s", run_id, node_id)
+            raise PermissionError("Invalid continuation or token")
+
+        if not hmac.compare_digest(token, cont.token):
+            self.logger.error("Invalid token for %s/%s", run_id, node_id)
+            raise PermissionError("Invalid continuation or token")
+
+        # Merge continuation payload (setup-time) with incoming resume payload (adapter-time)
+        base_payload = getattr(cont, "payload", None) or {}
+        full_payload: dict[str, Any] = {**base_payload, **(payload or {})}
+
+        # Cooperative fast path
+        if self.waits and token in getattr(self.waits, "_futs", {}):
+            try:
+                self.waits.resolve(token, full_payload)
+                self.logger.info("Resolved cooperative wait for %s/%s", run_id, node_id)
+                try:
+                    await self.store.delete(run_id, node_id)
+                except Exception as e:
+                    self.logger.warning(
+                        f"Failed to delete continuation after cooperative resolution: {e}"
+                    )
+                return
+            except Exception as e:
+                self.logger.error(
+                    "Error resolving cooperative wait for %s/%s: %s",
+                    run_id,
+                    node_id,
+                    e,
+                    exc_info=True,
+                )
+                raise
+
+        # Schema validate
+        if cont.resume_schema:
+            try:
+                validate(
+                    instance=payload, schema=cont.resume_schema
+                )  # validate incoming payload only
+            except ValidationError as e:
+                self.logger.error("Resume payload validation error: %s", e.message)
+                raise ValueError(f"Invalid resume payload: {e.message}") from e
+
+        # Hand off to scheduler bus
+        await self.runner.enqueue_resume(
+            run_id=run_id, node_id=node_id, token=token, payload=full_payload
+        )

aethergraph/services/schedulers/registry.py

@@ -0,0 +1,41 @@
+from contextlib import asynccontextmanager
+import threading
+
+from aethergraph.core.execution.global_scheduler import GlobalForwardScheduler
+
+
+class SchedulerRegistry:
+    def __init__(self):
+        self._by_run: dict[str, GlobalForwardScheduler] = {}
+        self._lock = threading.RLock()
+
+    def register(self, run_id: str, scheduler: GlobalForwardScheduler) -> None:
+        with self._lock:
+            self._by_run[run_id] = scheduler
+
+    def unregister(self, run_id: str) -> None:
+        with self._lock:
+            self._by_run.pop(run_id, None)
+
+    def get(self, run_id: str) -> GlobalForwardScheduler | None:
+        with self._lock:
+            return self._by_run.get(run_id)
+
+    def list_run_ids(self) -> dict[str, GlobalForwardScheduler]:
+        with self._lock:
+            return dict(self._by_run)
+
+
+@asynccontextmanager
+async def registered_scheduler(registry: SchedulerRegistry, run_id: str, scheduler):
+    registry.register(run_id, scheduler)
+    try:
+        yield
+    finally:
+        registry.unregister(run_id)
+
+
+"""# Example usage:
+async with registered_scheduler(SCHEDULERS, run_id, scheduler):
+    await scheduler.run()
+"""

aethergraph/services/secrets/base.py

@@ -0,0 +1,7 @@
+from typing import Protocol
+
+
+class Secrets(Protocol):
+    async def get(self, name: str) -> str | None:
+        """Retrieve the secret value by its name. Returns None if not found."""
+        ...

aethergraph/services/secrets/env.py

@@ -0,0 +1,8 @@
+import os
+
+from .base import Secrets
+
+
+class EnvSecrets(Secrets):
+    def get(self, name: str) -> str | None:
+        return os.getenv(name)

aethergraph/services/state_stores/externalize.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+import hashlib
+import io
+import json
+from typing import Any
+
+
+def _sha256_bytes(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
+def _maybe_numpy_to_bytes(obj: Any) -> bytes | None:
+    try:
+        import numpy as np
+
+        if isinstance(obj, np.ndarray):
+            buf = io.BytesIO()
+            # .npy
+            import numpy as _np
+
+            _np.save(buf, obj, allow_pickle=False)
+            return buf.getvalue()
+    except Exception:
+        pass
+    return None
+
+
+def _maybe_torch_to_bytes(obj: Any) -> bytes | None:
+    try:
+        import torch
+
+        if torch.is_tensor(obj):
+            buf = io.BytesIO()
+            torch.save(obj, buf)  # binary, portable within torch
+            return buf.getvalue()
+    except Exception:
+        pass
+    return None
+
+
+def _maybe_json_bytes(obj: Any) -> bytes | None:
+    # Only if JSON-serializable (pure)
+    try:
+        payload = json.dumps(obj, ensure_ascii=False).encode("utf-8")
+        return payload
+    except Exception:
+        return None
+
+
+def _pickle_fallback(obj: Any) -> bytes | None:
+    try:
+        import pickle
+
+        return pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
+    except Exception:
+        return None
+
+
+async def externalize_to_artifact(
+    obj: Any,
+    *,
+    run_id: str,
+    graph_id: str,
+    node_id: str,
+    tool_name: str | None,
+    tool_version: str | None,
+    artifacts,  # AsyncArtifactStore
+) -> dict[str, Any]:
+    """
+    Try to persist obj into artifact store; return a standard __aether_ref__ dict.
+    Priority: bytes | numpy | torch | json | pickle
+    """
+    binary: bytes | None = None
+    mime = "application/octet-stream"
+    planned_ext = ".bin"
+
+    # Already bytes?
+    if isinstance(obj, bytes | bytearray):
+        binary = bytes(obj)
+
+    if binary is None:
+        binary = _maybe_numpy_to_bytes(obj)
+        if binary is not None:
+            mime = "application/x-npy"
+            planned_ext = ".npy"
+
+    if binary is None:
+        binary = _maybe_torch_to_bytes(obj)
+        if binary is not None:
+            mime = "application/x-pytorch"
+            planned_ext = ".pt"
+
+    if binary is None:
+        json_bytes = _maybe_json_bytes(obj)
+        if json_bytes is not None:
+            # Use .json for nicer preview; still count as binary save
+            binary = json_bytes
+            mime = "application/json"
+            planned_ext = ".json"
+
+    if binary is None:
+        binary = _pickle_fallback(obj)
+        if binary is not None:
+            mime = "application/x-pickle"
+            planned_ext = ".pkl"
+
+    if binary is None:
+        # Give up: write a tiny JSON marker
+        a = await artifacts.save_json(
+            {"note": "unexternalizable-object", "repr": repr(obj)[:200]}, suggested_uri=None
+        )
+        return {
+            "__aether_ref__": a.uri,
+            "mime": "application/json",
+            "sha256": a.sha256,
+            "kind": a.kind,
+        }
+
+    sha = _sha256_bytes(binary)
+    # Use staged writer for atomicity FIXME: this write causes async loop errors, disable externalize for now
+    async with await artifacts.open_writer(
+        kind="blob",
+        run_id=run_id,
+        graph_id=graph_id,
+        node_id=node_id,
+        tool_name=tool_name or "externalize",
+        tool_version=tool_version or "0.1.0",
+        planned_ext=planned_ext,
+        pin=True,
+    ) as w:
+        await w.write(binary)  # FS wrapper should support .write
+        a = await w.commit(mime=mime, sha256=sha)
+
+    return {"__aether_ref__": a.uri, "mime": mime, "sha256": sha, "kind": a.kind}

aethergraph/services/state_stores/graph_observer.py

@@ -0,0 +1,131 @@
+# aethergraph/persist/observer.py
+from __future__ import annotations
+
+import time
+
+from aethergraph.contracts.services.artifacts import AsyncArtifactStore
+from aethergraph.contracts.services.state_stores import GraphStateStore, StateEvent
+from aethergraph.core.graph.task_node import NodeStatus
+from aethergraph.services.state_stores.serialize import _jsonish_outputs_with_refs
+from aethergraph.services.state_stores.utils import snapshot_from_graph
+
+
+class PersistenceObserver:
+    def __init__(
+        self,
+        *,
+        store: GraphStateStore,
+        artifact_store: AsyncArtifactStore,
+        spec_hash: str,
+        snapshot_every: int = 50,
+        min_interval_s: float = 5.0,
+    ):
+        self.store = store
+        self.artifact_store = artifact_store
+        self.spec_hash = spec_hash
+        self.snapshot_every = snapshot_every
+        self.min_interval_s = min_interval_s
+        self._event_count = 0
+        self._last_snap_ts = 0.0
+
+    async def on_node_status_change(self, runtime_node):
+        g = runtime_node._parent_graph
+        ev = StateEvent(
+            run_id=g.state.run_id or "unknown",
+            graph_id=g.graph_id,
+            rev=g.state.rev,
+            ts=time.time(),
+            kind="STATUS",
+            payload={
+                "node_id": runtime_node.node_id,
+                "status": runtime_node.state.status.name
+                if isinstance(runtime_node.state.status, NodeStatus)
+                else str(runtime_node.state.status),
+            },
+        )
+        await self.store.append_event(ev)
+        await self._maybe_snapshot(g)
+
+    async def on_node_output_change(self, runtime_node):
+        g = runtime_node._parent_graph
+        # make outputs JSON-safe for events (no externalization)
+        safe_outputs = await _jsonish_outputs_with_refs(
+            outputs=getattr(runtime_node.state, "outputs", None),
+            run_id=g.state.run_id or "unknown",
+            graph_id=g.graph_id,
+            node_id=runtime_node.node_id,
+            tool_name=getattr(runtime_node.state, "tool_name", None)
+            or getattr(getattr(runtime_node, "spec", None), "tool_name", None),
+            tool_version=getattr(runtime_node.state, "tool_version", None)
+            or getattr(getattr(runtime_node, "spec", None), "tool_version", None),
+            artifacts=None,  # ← keep events self-contained
+            allow_externalize=False,  # ← do not write artifacts from events
+        )
+
+        ev = StateEvent(
+            run_id=g.state.run_id or "unknown",
+            graph_id=g.graph_id,
+            rev=g.state.rev,
+            ts=time.time(),
+            kind="OUTPUT",
+            payload={
+                "node_id": runtime_node.node_id,
+                "outputs": safe_outputs or {},  # ✅ JSON-safe
+            },
+        )
+        await self.store.append_event(ev)
+        await self._maybe_snapshot(g)
+
+    async def on_inputs_bound(self, graph):
+        # also sanitize inputs for events (in case user passed non-JSON)
+        safe_inputs = await _jsonish_outputs_with_refs(
+            outputs=getattr(graph.state, "_bound_inputs", None),
+            run_id=graph.state.run_id or "unknown",
+            graph_id=graph.graph_id,
+            node_id="__graph_inputs__",
+            tool_name=None,
+            tool_version=None,
+            artifacts=None,
+            allow_externalize=False,
+        )
+        ev = StateEvent(
+            run_id=graph.state.run_id or "unknown",
+            graph_id=graph.graph_id,
+            rev=graph.state.rev,
+            ts=time.time(),
+            kind="INPUTS_BOUND",
+            payload={"inputs": safe_inputs or {}},  # JSON-safe
+        )
+        await self.store.append_event(ev)
+        await self._maybe_snapshot(graph)
+
+    async def on_patch_applied(self, graph, patch):
+        ev = StateEvent(
+            run_id=graph.state.run_id or "unknown",
+            graph_id=graph.graph_id,
+            rev=graph.state.rev,
+            ts=time.time(),
+            kind="PATCH",
+            payload={"patch": patch.__dict__},
+        )
+        await self.store.append_event(ev)
+        await self._maybe_snapshot(graph)
+
+    async def _maybe_snapshot(self, graph):
+        self._event_count += 1
+        now = time.time()
+        if (self._event_count % self.snapshot_every == 0) and (
+            now - self._last_snap_ts >= self.min_interval_s
+        ):
+            snap = await snapshot_from_graph(
+                run_id=graph.state.run_id or "unknown",
+                graph_id=graph.graph_id,
+                rev=graph.state.rev,
+                spec_hash=self.spec_hash,
+                state_obj=graph.state,
+                artifacts=self.artifact_store,
+                allow_externalize=False,  # keep snapshots JSON-only (opaque refs)
+                include_wait_spec=True,
+            )
+            await self.store.save_snapshot(snap)
+            self._last_snap_ts = now