diff --git a/openviking/session/train/domain.py b/openviking/session/train/domain.py index 6e66158bf6..4c8b545425 100644 --- a/openviking/session/train/domain.py +++ b/openviking/session/train/domain.py @@ -64,12 +64,17 @@ class PolicySet: request_context: Any | None = field(default=None, repr=False, compare=False) @asynccontextmanager - async def lock(self): + async def lock(self, *, timeout_secs: float = 300.0): """Acquire a tree lock for the whole policy root directory. - Policy updates serialize on this lock so concurrent realtime/batch - training jobs plan and apply against a freshly reloaded policy set. - ``timeout=None`` means wait indefinitely until the lock is available. + Policy updates serialize on this lock. Keep the critical section + short and local (reload + apply): long-running work such as remote + LLM planning must run before acquiring the lock, otherwise concurrent + writers of the same root are blocked until they hit the acquisition + timeout. + + ``timeout_secs`` bounds how long acquisition waits for the current + holder before failing; it defaults to 300 seconds. """ if self.viking_fs is None: @@ -81,7 +86,9 @@ async def lock(self): raise RuntimeError("PolicySet.viking_fs must provide _uri_to_path for locking") path = uri_to_path(self.root_uri, ctx=self.request_context) - lease = await self.viking_fs._async_agfs.pathlock_acquire_tree(path, timeout_secs=300.0) + lease = await self.viking_fs._async_agfs.pathlock_acquire_tree( + path, timeout_secs=timeout_secs + ) try: yield lease finally: diff --git a/openviking/session/train/engine.py b/openviking/session/train/engine.py index 21d8c24d43..0c6af098df 100644 --- a/openviking/session/train/engine.py +++ b/openviking/session/train/engine.py @@ -23,6 +23,11 @@ SemanticGradient, ) +# Bounded number of optimistic re-plans when the policy set changed while the +# optimizer was planning. Keeps the (potentially minutes-long) LLM replanning +# from livelocking against sustained concurrent writers. +_PLAN_APPLY_MAX_REPLANS = 2 + @dataclass(slots=True) class PolicyTrainingEngine: @@ -84,17 +89,53 @@ async def plan_and_apply( policy_set: ExperienceSet, ctx: PipelineContext, ) -> tuple[PolicyUpdatePlan, PolicyApplyResult]: - async with policy_set.lock() as transaction_handle: - latest_policy_set = await policy_set.reload() + """Plan against a lock-free snapshot, then apply in a short critical section. + + ``policy_optimizer.plan`` may block on a remote LLM for minutes. Holding + the policy-root tree lock across that call starves every concurrent + writer of the same root until the lock acquisition timeout fails their + commits, so planning runs outside the lock. The lock is taken only to + re-check that the snapshot the plan was built against is still the + latest one before applying. When another writer committed meanwhile, + the plan is rebuilt against the fresh snapshot (bounded by + ``_PLAN_APPLY_MAX_REPLANS``); once the bound is exhausted the last plan + is applied under the lock with a conflict marker, and the updater's + per-item base-content guard reports stale targets as errors instead of + silently overwriting concurrent changes. + """ + + replans = 0 + while True: + planned_snapshot = await policy_set.reload() plan = await self.policy_optimizer.plan( gradients, - latest_policy_set, + planned_snapshot, ctx.optimization_context, ) - apply_result = await self.policy_updater.apply( - plan, - latest_policy_set, - ctx.apply_context or latest_policy_set.request_context, - transaction_handle=transaction_handle, - ) - return plan, apply_result + async with policy_set.lock() as transaction_handle: + latest_policy_set = await policy_set.reload() + snapshot_current = _policy_sets_match(latest_policy_set, planned_snapshot) + if snapshot_current or replans >= _PLAN_APPLY_MAX_REPLANS: + if not snapshot_current: + plan.metadata = { + **dict(plan.metadata or {}), + "optimistic_apply_after_conflict": True, + "replan_attempts": replans + 1, + } + apply_result = await self.policy_updater.apply( + plan, + latest_policy_set, + ctx.apply_context or latest_policy_set.request_context, + transaction_handle=transaction_handle, + ) + return plan, apply_result + # The policy set moved on while planning: the lock was released on + # exiting the context manager, so re-plan outside it against the + # fresh snapshot. + replans += 1 + + +def _policy_sets_match(current: ExperienceSet, planned: ExperienceSet) -> bool: + """Whether ``planned`` is still the latest committed view of the policy root.""" + + return current.root_uri == planned.root_uri and list(current.policies) == list(planned.policies) diff --git a/tests/session/train/test_fakes.py b/tests/session/train/test_fakes.py index cbcdd09cd6..2c624b3db0 100644 --- a/tests/session/train/test_fakes.py +++ b/tests/session/train/test_fakes.py @@ -3,7 +3,9 @@ from __future__ import annotations +import asyncio import threading +import time from types import SimpleNamespace from typing import Any @@ -217,6 +219,77 @@ def fake_request_context(account_id: str = "default", user_id: str = "u"): ) +class LockingAGFS: + """Async AGFS fake implementing tree path locks with real mutual exclusion. + + ``PolicySet.lock`` acquires tree locks through + ``viking_fs._async_agfs.pathlock_acquire_tree``; this fake enforces actual + exclusion between concurrent acquirers of the same path, honors the + ``timeout_secs`` deadline, and records every acquire/release plus the hold + interval of each lease so tests can assert which code ran while the lock + was held. + + ``lock_budget_secs`` caps how long any single acquisition may wait for a + holder regardless of the requested timeout; exceeding it raises + ``TimeoutError``. It exists to keep tests fast when simulating the + production 300s lock timeout. + """ + + def __init__(self, *, lock_budget_secs: float | None = None): + self.lock_budget_secs = lock_budget_secs + self._locks: dict[str, asyncio.Lock] = {} + self.acquire_calls: list[dict[str, Any]] = [] + self.release_calls: list[dict[str, Any]] = [] + self.holds: list[dict[str, float]] = [] + + def _lock_for(self, path: str) -> asyncio.Lock: + if path not in self._locks: + self._locks[path] = asyncio.Lock() + return self._locks[path] + + def is_locked(self, path: str) -> bool: + lock = self._locks.get(path) + return lock is not None and lock.locked() + + async def pathlock_acquire_tree( + self, + path: str, + timeout_secs: float = 0.0, + owner_lease_ref: Any = None, + **kwargs: Any, + ) -> dict[str, Any]: + del owner_lease_ref, kwargs + self.acquire_calls.append({"path": path, "timeout_secs": timeout_secs}) + deadline: float | None = None + if timeout_secs and timeout_secs > 0: + deadline = timeout_secs + if self.lock_budget_secs is not None: + deadline = min(deadline, self.lock_budget_secs) if deadline else self.lock_budget_secs + lock = self._lock_for(path) + try: + await asyncio.wait_for(lock.acquire(), deadline) + except asyncio.TimeoutError: + raise TimeoutError( + f"tree lock acquisition timed out after {timeout_secs}s: {path}" + ) from None + lease = {"path": path, "acquired_at": time.monotonic()} + return lease + + async def pathlock_release(self, lease: dict[str, Any], **kwargs: Any) -> None: + del kwargs + release = dict(lease) + release["released_at"] = time.monotonic() + self.release_calls.append(release) + self.holds.append( + { + "path": release["path"], + "start": release["acquired_at"], + "end": release["released_at"], + } + ) + self._lock_for(release["path"]).release() + + def _norm_path(path: str) -> str: if not path: return "/" diff --git a/tests/session/train/test_train_framework.py b/tests/session/train/test_train_framework.py index a45964b084..d1eb24d337 100644 --- a/tests/session/train/test_train_framework.py +++ b/tests/session/train/test_train_framework.py @@ -6,12 +6,13 @@ import asyncio import copy import json +import time from dataclasses import dataclass, field from pathlib import Path from typing import Any import pytest -from test_fakes import fake_request_context +from test_fakes import LockingAGFS, fake_request_context from openviking.message import Message, TextPart from openviking.session.memory.dataclass import StoredLink @@ -62,6 +63,7 @@ class DummyVikingFS: def __init__(self): self.reloads = 0 self.version = 1 + self._async_agfs = LockingAGFS() async def ls(self, uri: str, output: str = "original", ctx=None): del output, ctx @@ -338,7 +340,9 @@ async def test_default_policy_optimization_pipeline_runs_one_batch(): assert DummyUpdater.last_instance is not None assert len(DummyUpdater.last_instance.transaction_handles) == 1 assert DummyUpdater.last_instance.transaction_handles[0] is not None - assert initial_policy_set.viking_fs.reloads == 1 + # plan_and_apply reloads twice: once lock-free to plan against a snapshot, + # once under the tree lock to validate that snapshot before applying. + assert initial_policy_set.viking_fs.reloads == 2 assert len(result.epochs) == 1 assert result.epochs[0].epoch == 0 assert result.epochs[0].policy_snapshot_ids == ["snapshot-1"] @@ -1933,3 +1937,251 @@ def test_tau2_service_filter_parses_task_indices(): _task_indices_from_filters({"task_indices": 2}) with pytest.raises(ValueError, match="task index must be >= 0"): _task_indices_from_filters({"task_indices": [-1]}) + + +# --------------------------------------------------------------------------- +# plan_and_apply locking regression tests (issue #6095) +# --------------------------------------------------------------------------- + + +class SlowPlannerOptimizer: + """Optimizer whose plan() takes ``delay`` seconds and can bump the store. + + Records the wall-clock interval of every plan() call so tests can assert + the policy-root tree lock was not held across planning. + """ + + def __init__(self, viking_fs: DummyVikingFS, *, delay: float = 0.0, bump_version: bool = False): + self.viking_fs = viking_fs + self.delay = delay + self.bump_version = bump_version + self.calls = 0 + self.intervals: list[tuple[float, float]] = [] + + async def plan( + self, + gradients: list[DummyGradient], + policy_set: ExperienceSet, + context: Any, + ) -> PolicyUpdatePlan: + del policy_set, context + self.calls += 1 + start = time.monotonic() + try: + if self.delay: + await asyncio.sleep(self.delay) + if self.bump_version: + self.viking_fs.version += 1 + return PolicyUpdatePlan(metadata={"gradient_count": len(gradients)}) + finally: + self.intervals.append((start, time.monotonic())) + + +class RecordingUpdater: + """PolicyUpdater double recording every apply() call.""" + + def __init__(self): + self.applied: list[dict[str, Any]] = [] + + async def apply( + self, + plan: PolicyUpdatePlan, + policy_set: ExperienceSet, + context: Any, + *, + transaction_handle: Any = None, + ) -> PolicyApplyResult: + del context + self.applied.append( + { + "plan": plan, + "policy_set": policy_set, + "transaction_handle": transaction_handle, + "at": time.monotonic(), + } + ) + return PolicyApplyResult(updated_policy_set=policy_set, written_uris=[]) + + +def _plan_apply_engine(optimizer: Any, updater: Any) -> Any: + from openviking.session.train.engine import PolicyTrainingEngine + + return PolicyTrainingEngine( + rollout_analyzer=object(), + gradient_estimator=object(), + policy_optimizer=optimizer, + policy_updater=updater, + ) + + +def _lock_regression_gradient() -> DummyGradient: + return DummyGradient( + target_name="booking_duplicate_handling", + target_uri="viking://user/u/memories/experiences/booking_duplicate_handling.md", + base_version=1, + rationale="trajectory succeeded", + links=[], + confidence=0.9, + ) + + +def _lock_held_across( + interval: tuple[float, float], + holds: list[dict[str, float]], +) -> bool: + """Whether any single lock hold spans the whole planning interval. + + That is the pre-fix (#6095) signature: the tree lock was acquired before + planning and released after apply, so the hold contained the plan. Short + holds by *other* writers overlapping the plan are expected after the fix. + """ + + return any(hold["start"] <= interval[0] and interval[1] <= hold["end"] for hold in holds) + + +def _overlaps( + interval: tuple[float, float], + holds: list[dict[str, float]], +) -> bool: + return any(hold["start"] < interval[1] and interval[0] < hold["end"] for hold in holds) + + +@pytest.mark.asyncio +async def test_plan_and_apply_does_not_hold_policy_lock_across_slow_planning(): + """A slow planner must not block another writer on the policy tree lock. + + Regression for issue #6095: planning used to run inside the policy-root + tree lock, so a minutes-long LLM plan made every concurrent session commit + wait for the 300s lock acquisition timeout and fail terminally. The fake + lock budget (1s) stands in for that timeout: with the lock held across the + 2s plan, the second writer times out and this test fails. + """ + + fs = DummyVikingFS() + fs._async_agfs = LockingAGFS(lock_budget_secs=1.0) + policy_set = _policy_set(viking_fs=fs) + ctx = PipelineContext() + + slow_optimizer = SlowPlannerOptimizer(fs, delay=2.0) + fast_optimizer = SlowPlannerOptimizer(fs) + slow_updater, fast_updater = RecordingUpdater(), RecordingUpdater() + + async def run(engine): + return await engine.plan_and_apply( + gradients=[_lock_regression_gradient()], + policy_set=policy_set, + ctx=ctx, + ) + + slow_task = run(_plan_apply_engine(slow_optimizer, slow_updater)) + fast_task = run(_plan_apply_engine(fast_optimizer, fast_updater)) + await asyncio.wait_for(asyncio.gather(slow_task, fast_task), timeout=15.0) + + # Both writers applied exactly once. + assert len(slow_updater.applied) == 1 + assert len(fast_updater.applied) == 1 + # The fast writer planned and applied entirely inside the slow plan window. + slow_plan_interval = slow_optimizer.intervals[0] + assert fast_optimizer.intervals[0][1] <= slow_plan_interval[1] + assert fast_updater.applied[0]["at"] < slow_plan_interval[1] + # No acquisition ever hit the 1s lock budget (no TimeoutError bubbled up), + # and no single lock hold spans a whole planning window: short holds by + # the other writer inside the slow plan window are expected after the fix. + assert not _lock_held_across(slow_plan_interval, fs._async_agfs.holds) + assert not _lock_held_across(fast_optimizer.intervals[0], fs._async_agfs.holds) + + +@pytest.mark.asyncio +async def test_plan_and_apply_replans_bounded_when_snapshot_changes_during_planning(): + """A snapshot moved during planning triggers bounded re-plans, then a + guarded apply of the last plan against the freshest snapshot.""" + + from openviking.session.train.engine import _PLAN_APPLY_MAX_REPLANS + + fs = DummyVikingFS() + policy_set = _policy_set(viking_fs=fs) + optimizer = SlowPlannerOptimizer(fs, bump_version=True) + updater = RecordingUpdater() + engine = _plan_apply_engine(optimizer, updater) + + plan, apply_result = await asyncio.wait_for( + engine.plan_and_apply( + gradients=[_lock_regression_gradient()], + policy_set=policy_set, + ctx=PipelineContext(), + ), + timeout=15.0, + ) + + # One plan per attempt: initial + bounded replans. + assert optimizer.calls == _PLAN_APPLY_MAX_REPLANS + 1 + assert len(updater.applied) == 1 + assert apply_result.updated_policy_set is updater.applied[0]["policy_set"] + # The exhausted optimistic apply is marked so callers can audit it. + assert plan.metadata["optimistic_apply_after_conflict"] is True + assert plan.metadata["replan_attempts"] == _PLAN_APPLY_MAX_REPLANS + 1 + # apply received the freshest snapshot (every plan bumped the version), + # never the stale one the plan was built from. + applied_versions = [p.version for p in updater.applied[0]["policy_set"].policies] + assert applied_versions == [fs.version] + # The tree lock was never held while the optimizer was planning. + for interval in optimizer.intervals: + assert not _overlaps(interval, fs._async_agfs.holds) + + +@pytest.mark.asyncio +async def test_plan_and_apply_fast_path_plans_once_and_releases_lock(): + fs = DummyVikingFS() + policy_set = _policy_set(viking_fs=fs) + optimizer = SlowPlannerOptimizer(fs) + updater = RecordingUpdater() + engine = _plan_apply_engine(optimizer, updater) + + plan, _ = await asyncio.wait_for( + engine.plan_and_apply( + gradients=[_lock_regression_gradient()], + policy_set=policy_set, + ctx=PipelineContext(), + ), + timeout=15.0, + ) + + assert optimizer.calls == 1 + assert "optimistic_apply_after_conflict" not in plan.metadata + assert len(fs._async_agfs.acquire_calls) == 1 + assert len(fs._async_agfs.release_calls) == 1 + assert not fs._async_agfs.is_locked(fs._async_agfs.acquire_calls[0]["path"]) + + +@pytest.mark.asyncio +async def test_policy_set_lock_passes_timeout_and_releases_on_error(): + fs = DummyVikingFS() + agfs = fs._async_agfs + policy_set = _policy_set(viking_fs=fs) + path = None + + async with policy_set.lock(timeout_secs=5.0) as lease: + path = agfs.acquire_calls[-1]["path"] + assert agfs.acquire_calls[-1]["timeout_secs"] == 5.0 + assert agfs.is_locked(path) + assert lease is not None + # A second acquirer of the same tree fails once its timeout expires. + with pytest.raises(TimeoutError): + async with asyncio.timeout(0.2): + async with policy_set.lock(timeout_secs=0.2): + pass + assert not agfs.is_locked(path) + # The failed acquirer must not have released the first holder's lease. + assert len(agfs.release_calls) == 1 + + +@pytest.mark.asyncio +async def test_policy_set_lock_releases_when_body_raises(): + fs = DummyVikingFS() + policy_set = _policy_set(viking_fs=fs) + + with pytest.raises(RuntimeError, match="boom"): + async with policy_set.lock(): + raise RuntimeError("boom") + + assert not fs._async_agfs.is_locked(fs._async_agfs.acquire_calls[0]["path"])