feat(mesh): elastic worker pools and the event-loop fixes that let them scale (RFC 0002, PR 2)

jarviscore/__init__.py

@@ -60,6 +60,7 @@ async def run(self):
 from jarviscore.core.agent import Agent
 from jarviscore.core.profile import Profile
 from jarviscore.core.mesh import Mesh, MeshMode
+from jarviscore.core.pool import Pool
 
 # Execution profiles
 from jarviscore.profiles.autoagent import AutoAgent
@@ -137,6 +138,7 @@ async def run(self):
     "Profile",
     "Mesh",
     "MeshMode",
+    "Pool",
 
     # Profiles
     "AutoAgent",

jarviscore/core/pool.py

@@ -0,0 +1,94 @@
+"""Elastic worker pools: more instances of one agent while its work queues up.
+
+A pooled agent is one logical agent. Every member shares the agent's id, so a step
+one member started and checkpointed can be resumed by any other, and the step-claim
+lease stays the only thing that decides who runs a step. Members are separate
+instances because an agent's kernel, subagent cache and sandbox serve one step at
+a time.
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+from typing import Any, Callable, List
+
+from jarviscore.execution.governor import BULK, INTERACTIVE
+
+
+@dataclass(frozen=True)
+class Pool:
+    """How far one agent may scale out.
+
+    min: members kept warm, including the agent registered with ``mesh.add``.
+    max: most members running steps at once in this process.
+    idle_seconds: an idle member above ``min`` is retired after this long.
+    lane: ``"bulk"`` lets the LLM governor hold capacity back for interactive work.
+    """
+
+    min: int = 1
+    max: int = 8
+    idle_seconds: float = 60.0
+    lane: str = INTERACTIVE
+
+    def __post_init__(self) -> None:
+        if self.min < 1:
+            raise ValueError("Pool.min must be at least 1: the registered agent is a member")
+        if self.max < self.min:
+            raise ValueError("Pool.max must be at least Pool.min")
+        if self.idle_seconds < 0:
+            raise ValueError("Pool.idle_seconds cannot be negative")
+        if self.lane not in {INTERACTIVE, BULK}:
+            raise ValueError(f"Pool.lane must be {INTERACTIVE!r} or {BULK!r}")
+
+
+@dataclass
+class PoolState:
+    """A pool's members in this process."""
+
+    spec: Pool
+    factory: Callable[[], Any]
+    base: Any
+    members: List[Any] = field(default_factory=list)
+    idle: List[Any] = field(default_factory=list)
+    last_active: dict = field(default_factory=dict)
+    steps_started: int = 0
+
+    def __post_init__(self) -> None:
+        self.members = [self.base]
+        self.idle = [self.base]
+        self.last_active = {id(self.base): time.monotonic()}
+
+    @property
+    def busy(self) -> int:
+        return len(self.members) - len(self.idle)
+
+    def release(self, member: Any) -> None:
+        self.last_active[id(member)] = time.monotonic()
+        if member in self.members and member not in self.idle:
+            self.idle.append(member)
+
+    def retirable(self) -> List[Any]:
+        """Idle members beyond ``min`` that have waited out ``idle_seconds``."""
+        now = time.monotonic()
+        spare = len(self.members) - self.spec.min
+        retiring = []
+        for member in list(self.idle):
+            if spare <= 0:
+                break
+            if member is self.base:
+                continue
+            if now - self.last_active.get(id(member), now) >= self.spec.idle_seconds:
+                retiring.append(member)
+                spare -= 1
+        return retiring
+
+    def snapshot(self) -> dict:
+        return {
+            "members": len(self.members),
+            "busy": self.busy,
+            "min": self.spec.min,
+            "max": self.spec.max,
+            "lane": self.spec.lane,
+            "steps_started": self.steps_started,
+        }

jarviscore/execution/__init__.py

@@ -109,6 +109,7 @@
     FunctionRegistry,
     FunctionStatus,
     create_function_registry,
+    shared_function_registry,
     # Backward-compatible aliases
     CodeRegistry,
     create_code_registry,
@@ -192,6 +193,7 @@
     'FunctionRegistry',
     'FunctionStatus',
     'create_function_registry',
+    'shared_function_registry',
     # Backward-compatible aliases
     'CodeRegistry',
     'create_code_registry',

jarviscore/execution/code_registry.py

@@ -35,6 +35,7 @@
 import logging
 import os
 import re
+import threading
 from datetime import datetime
 from enum import Enum
 from pathlib import Path
@@ -1562,6 +1563,24 @@ def create_function_registry(
     return registry
 
 
+_SHARED_REGISTRIES: Dict[str, "FunctionRegistry"] = {}
+_SHARED_REGISTRIES_LOCK = threading.Lock()
+
+
+def shared_function_registry(storage_path: str = "./logs/function_registry") -> FunctionRegistry:
+    """The process's one registry for this directory, loaded and seeded once.
+
+    Agents writing the same directory see one index instead of diverging copies,
+    and a pool member joins in microseconds instead of reloading the catalogue.
+    """
+    key = str(Path(storage_path).resolve())
+    with _SHARED_REGISTRIES_LOCK:
+        registry = _SHARED_REGISTRIES.get(key)
+        if registry is None:
+            registry = _SHARED_REGISTRIES[key] = create_function_registry(storage_path)
+        return registry
+
+
 def _seed_shipped_atoms(registry: FunctionRegistry) -> None:
     """Load the packaged atoms so a fresh registry is not an empty one.
 
@@ -1594,7 +1613,7 @@ def _seed_shipped_atoms(registry: FunctionRegistry) -> None:
     )
     if failed:
         logger.warning("%d atom(s) in the catalogue could not be loaded", len(failed))
-    if not registered:
+    if not registry.function_metadata:
         # A silent no-op here is what made an empty registry look like normal
         # operation for as long as it did.
         logger.warning(

jarviscore/execution/governor.py

@@ -14,6 +14,7 @@
 import asyncio
 import logging
 import math
+import threading
 import time
 from contextlib import contextmanager
 from contextvars import ContextVar
@@ -179,8 +180,22 @@ def settle(self, second: int, reserved: int, actual: int) -> None:
         self._local.adjust(second, delta)
 
 
+class _Waiter:
+    __slots__ = ("lane", "loop", "future")
+
+    def __init__(self, lane: str, loop: asyncio.AbstractEventLoop) -> None:
+        self.lane = lane
+        self.loop = loop
+        self.future = loop.create_future()
+
+
 class AdaptiveConcurrency:
-    """Calls in flight for one deployment, found by slow start and AIMD."""
+    """Calls in flight for one deployment, found by slow start and AIMD.
+
+    One process can call the LLM from several event loops (sync work in a thread
+    runs its own), so the count is guarded by a thread lock and each waiter is
+    woken on its own loop.
+    """
 
     def __init__(self, initial: int, ceiling: int, bulk_reserve: float) -> None:
         self.initial = max(1, int(initial))
@@ -191,18 +206,8 @@ def __init__(self, initial: int, ceiling: int, bulk_reserve: float) -> None:
         self.in_flight = 0
         self.waiting = 0
         self._last_decrease = 0.0
-        self._condition: Optional[asyncio.Condition] = None
-        self._loop = None
-
-    def _wait_condition(self) -> asyncio.Condition:
-        loop = asyncio.get_running_loop()
-        if self._condition is None or self._loop is not loop:
-            # A new event loop (a restarted service, a test) starts with no calls in flight.
-            self._condition = asyncio.Condition()
-            self._loop = loop
-            self.in_flight = 0
-            self.waiting = 0
-        return self._condition
+        self._lock = threading.Lock()
+        self._waiters: list[_Waiter] = []
 
     def _capacity(self, lane: str) -> int:
         limit = max(1, int(self.limit))
@@ -211,24 +216,59 @@ def _capacity(self, lane: str) -> int:
         return limit
 
     async def acquire(self, lane: str) -> None:
-        condition = self._wait_condition()
-        async with condition:
+        with self._lock:
+            queued = any(waiter.lane == lane for waiter in self._waiters)
+            if not queued and self.in_flight < self._capacity(lane):
+                self.in_flight += 1
+                return
+            waiter = _Waiter(lane, asyncio.get_running_loop())
+            self._waiters.append(waiter)
             self.waiting += 1
-            try:
-                await condition.wait_for(lambda: self.in_flight < self._capacity(lane))
-            finally:
-                self.waiting -= 1
-            self.in_flight += 1
+        try:
+            await waiter.future
+        except BaseException:
+            with self._lock:
+                if waiter in self._waiters:
+                    self._waiters.remove(waiter)
+                    self.waiting -= 1
+                    raise
+            # The slot was handed over as this wait was cancelled.
+            if waiter.future.done() and not waiter.future.cancelled():
+                self._give_back()
+            raise
 
     async def release(self, outcome: str) -> None:
-        condition = self._wait_condition()
-        async with condition:
+        with self._lock:
             self.in_flight = max(0, self.in_flight - 1)
             if outcome == "ok":
                 self._increase()
             elif outcome == "congested":
                 self._decrease()
-            condition.notify_all()
+            self._wake()
+
+    def _give_back(self) -> None:
+        with self._lock:
+            self.in_flight = max(0, self.in_flight - 1)
+            self._wake()
+
+    def _wake(self) -> None:
+        """Hand free slots to waiters in arrival order; call with the lock held."""
+        for waiter in list(self._waiters):
+            if self.in_flight >= self._capacity(waiter.lane):
+                continue
+            self._waiters.remove(waiter)
+            self.waiting -= 1
+            self.in_flight += 1
+            try:
+                waiter.loop.call_soon_threadsafe(self._deliver, waiter)
+            except RuntimeError:  # its loop has closed
+                self.in_flight -= 1
+
+    def _deliver(self, waiter: _Waiter) -> None:
+        if waiter.future.done():  # cancelled before the slot arrived
+            self._give_back()
+        else:
+            waiter.future.set_result(None)
 
     def _increase(self) -> None:
         if self.limit < self.threshold:

jarviscore/profiles/autoagent.py

@@ -333,7 +333,7 @@ async def setup(self):
             create_coder_sandbox,
             create_autonomous_repair,
             create_result_handler,
-            create_function_registry
+            shared_function_registry
         )
 
         # 1. Initialize LLM (auto-detects providers)
@@ -397,7 +397,7 @@ async def setup(self):
         # 7. Initialize function registry (graduated, reusable generated functions)
         registry_dir = f"{log_dir}/function_registry"
         self._logger.info(f"Initializing function registry (dir: {registry_dir})...")
-        self.code_registry = create_function_registry(registry_dir)
+        self.code_registry = await asyncio.to_thread(shared_function_registry, registry_dir)
         # Confined code may read the registry but only registration may change it.
         self.sandbox.protected_paths = (Path(log_dir).resolve(),)
 

jarviscore/storage/redis_store.py

@@ -200,6 +200,7 @@ def save_step_output(self, workflow_id: str, step_id: str,
         }
         self._redis.hset(key, mapping={k: v for k, v in data.items() if v is not None})
         self._redis.expire(key, self._ttl_seconds)
+        self._index_step_output(workflow_id, step_id)
         return True
 
     def get_step_output(self, workflow_id: str,
@@ -222,18 +223,26 @@ def get_step_output(self, workflow_id: str,
                 pass
         return result
 
+    def _index_step_output(self, workflow_id: str, step_id: str, pipe=None) -> None:
+        target = pipe if pipe is not None else self._redis
+        index = f"step_output_ids:{workflow_id}"
+        target.sadd(index, step_id)
+        target.expire(index, self._ttl_seconds)
+
     def list_step_output_ids(self, workflow_id: str) -> List[str]:
         """Return all step IDs that have saved outputs for this workflow.
 
-        Uses SCAN (non-blocking) to find keys matching
-        step_output:{workflow_id}:*  and strips the prefix to return
-        just the step_id portion.
+        Reads the workflow's output index and graph instead of scanning the
+        keyspace: a SCAN walks every key in the database, so its cost grew with
+        every workflow ever run and it ran on each step dispatch.
         """
-        prefix = f"step_output:{workflow_id}:"
-        return [
-            k[len(prefix):]
-            for k in self._redis.scan_iter(match=f"{prefix}*")
-        ]
+        candidates = set(self._redis.smembers(f"step_output_ids:{workflow_id}") or ())
+        candidates.update(self._redis.hkeys(f"workflow_graph:{workflow_id}") or ())
+        ordered = sorted(candidates)
+        pipe = self._redis.pipeline()
+        for step_id in ordered:
+            pipe.exists(f"step_output:{workflow_id}:{step_id}")
+        return [step_id for step_id, found in zip(ordered, pipe.execute()) if found]
 
     # ------------------------------------------------------------------
     # Shared Context / Truth
@@ -2251,6 +2260,7 @@ def finish_claimed_step(
                 pipe.multi()
                 pipe.hset(output_key, mapping={k: v for k, v in mapping.items() if v is not None})
                 pipe.expire(output_key, self._ttl_seconds)
+                self._index_step_output(workflow_id, step_id, pipe)
                 pipe.hset(graph_key, mapping={step_id: json.dumps(data)})
                 if projection:
                     pipe.set(projection_key, json.dumps(projection, default=str))
@@ -2350,6 +2360,29 @@ def get_all_step_ids(self, workflow_id: str) -> List[str]:
         """Return all step IDs stored in the workflow DAG hash."""
         return list(self._redis.hkeys(f"workflow_graph:{workflow_id}"))
 
+    def get_active_workflow_graphs(self) -> Dict[str, Dict[str, Any]]:
+        """Every active, uncancelled workflow's step definitions, in one round trip."""
+        workflow_ids = sorted(self.get_active_workflows())
+        pipe = self._redis.pipeline(transaction=False)
+        for workflow_id in workflow_ids:
+            pipe.exists(f"workflow_cancelled:{workflow_id}")
+            pipe.hgetall(f"workflow_graph:{workflow_id}")
+        replies = pipe.execute()
+        graphs: Dict[str, Dict[str, Any]] = {}
+        for index, workflow_id in enumerate(workflow_ids):
+            cancelled, raw_steps = replies[2 * index], replies[2 * index + 1]
+            if cancelled:
+                self.unregister_active_workflow(workflow_id)
+                continue
+            steps = {}
+            for step_id, raw in (raw_steps or {}).items():
+                try:
+                    steps[step_id] = json.loads(raw)
+                except (json.JSONDecodeError, TypeError):
+                    continue
+            graphs[workflow_id] = steps
+        return graphs
+
     # ------------------------------------------------------------------
     # Workflow State (crash recovery)
     # ------------------------------------------------------------------