From b6fb1909a69631f1b4262cbc8c0e9fd1d5242e7d Mon Sep 17 00:00:00 2001 From: akushonkamen Date: Fri, 16 Oct 2026 09:38:41 +0800 Subject: [PATCH] fix(train): plan policy updates outside the policy-root tree lock plan_and_apply held the policy-root tree lock across reload -> policy_optimizer.plan (remote LLM, minutes-long in production traces) -> apply, so concurrent session commits of the same root blocked until the 300s pathlock acquisition timeout and failed terminally (#6095). Plan against a lock-free snapshot instead, then take the lock only to re-validate that snapshot before applying. If another writer committed meanwhile, re-plan against the fresh snapshot (bounded at 2 replans); once the bound is exhausted, apply the last plan under the lock with a conflict marker and let the updater's per-item base-content guard report stale targets as errors instead of silently overwriting them. PolicySet.lock now takes a timeout_secs parameter (default 300s) and its docstring no longer promises a wait-indefinitely None mode that the implementation never had. Tests: add LockingAGFS (tree-lock fake with real mutual exclusion, deadline handling and hold-interval recording) plus regression tests asserting that (1) a 2s plan no longer blocks a concurrent writer under a 1s lock budget, (2) snapshot conflicts trigger bounded re-plans followed by a guarded apply against the freshest snapshot, (3) the uncontended fast path plans once and releases the lock, and (4) lock timeouts are passed through and leases released on error. Giving the DummyVikingFS fake the _async_agfs attribute that lock() requires also fixes the 14 train framework tests that already failed on unmodified main. Patch prepared with AI assistance and human review. --- openviking/session/train/domain.py | 17 +- openviking/session/train/engine.py | 61 ++++- tests/session/train/test_fakes.py | 73 ++++++ tests/session/train/test_train_framework.py | 256 +++++++++++++++++++- 4 files changed, 390 insertions(+), 17 deletions(-) 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"])