mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-27 08:28:00 +00:00
fix(runtime): close cancelled replacement admission (#4472)
This commit is contained in:
DanielWalnut
GitHub
parent
1c7531242c
commit
bb9f67aaf1
@ -498,6 +498,7 @@ JSONL event stores when `GATEWAY_WORKERS > 1`.
|
||||
- When a persistent `RunStore` is configured, `get()` and `list_by_thread()` hydrate historical runs from the store. In-memory records win for the same `run_id` so task, abort, and stream-control state stays attached to active local runs.
|
||||
- Thread metadata status switches to `running` only after `RunManager.try_start()` succeeds. Pending-cancelled runs therefore skip the old `running` projection, while clients may observe the prior thread status during the short worker-startup window.
|
||||
- `cancel()` returns a :class:`~deerflow.runtime.CancelOutcome` enum: `cancelled` (local cancel), `taken_over` (non-owning worker claimed the run because the owner's lease expired — marks it as `error`), `lease_valid_elsewhere` (owner's lease is still alive — caller should return 409 + `Retry-After`), `not_active_locally` (heartbeat disabled, preserving the old 409 path), `not_cancellable` (terminal state), or `unknown` (not found in memory or store). `create_or_reject(..., multitask_strategy="interrupt"|"rollback")` persists interrupted status through `RunStore.update_status()`, matching normal `set_status()` transitions.
|
||||
- Interrupt/rollback admission registers the replacement before its best-effort persistence of locally interrupted predecessors. If the admitting caller is cancelled during that post-registration await, `RunManager` drains a shielded replacement cleanup before propagating `CancelledError`, including across repeated cancellation. The cleanup normally persists `interrupted`; if that best-effort transition fails, it retries the active-to-`interrupted` store transition strictly and verifies the result with the replacement's captured owner identity. A concurrent peer terminal transition wins and is synchronized back into the local record rather than being overwritten or deleted.
|
||||
- Store-only hydrated runs are readable history. In multi-worker mode with heartbeat enabled, cancel on a store-only run can take over (mark `error`) when the owner's lease has expired past the grace window; otherwise it fails with 409 + `Retry-After`. In single-worker mode (heartbeat off), store-only runs still return 409.
|
||||
- Startup/orphan reconciliation must claim stale active rows with `RunStore.claim_for_takeover()`, not a plain `update_status()`. The final claim re-checks `status` and lease expiry atomically, so a heartbeat renewal between the candidate scan and the recovery write keeps the run active.
|
||||
- Run admission and independent checkpoint writes are first-class thread operations. `runs.operation_kind` distinguishes user-visible `run` rows from internal `checkpoint_write` reservations, while every active kind shares the existing durable active-thread uniqueness constraint. New operation kinds must go through `RunStore.create_thread_operation_atomic()` and `RunManager.reserve_thread_operation()` rather than adding another lock or metadata marker. Live and lease-less reservations are non-interruptible; an expired leased reservation can be reclaimed immediately by interrupt/rollback admission without waiting for orphan reconciliation. Lease-less rows stay fail-closed because the store cannot distinguish a stale row from a live checkpoint writer in another heartbeat-disabled worker; a rare failed delete therefore requires startup reconciliation, and heartbeat-disabled multi-worker deployment remains unsupported. Reservation bodies are attached to their caller task so loss detected by lease renewal cancels the checkpoint writer before it can continue after takeover; the context manager translates that lease-loss cancellation to `ConflictError` after cleanup so Gateway mutation routes return a retryable 409 instead of dropping the HTTP request. The cleanup scope begins immediately after durable admission, including the await that attaches the caller task, so cancellation cannot strand a locally renewed pending reservation. A failed renewal is revalidated under the manager lock before cancellation; if the reservation completed and unregistered while the store update was in flight, its request task must not be cancelled after the checkpoint write. Reservations are excluded from run history/reporting and from run-only helpers such as `list_by_thread()` and `has_inflight()`, release uses the captured owner rather than ambient user context, and local cleanup still runs when the best-effort store delete fails. `RunStore.create_run_atomic()` remains a deprecated compatibility shim for external stores that only admit normal runs; new stores must implement `create_thread_operation_atomic()` to support internal operation kinds.
|
||||
|
||||
@ -1092,6 +1092,51 @@ class RunManager:
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
async def _close_cancelled_admission(self, record: RunRecord) -> None:
|
||||
"""Terminalize an unseen replacement and confirm its durable state."""
|
||||
await self.cancel(record.run_id)
|
||||
if self._store is None:
|
||||
return
|
||||
|
||||
stored = await self._call_store_with_retry(
|
||||
"verify cancelled admission",
|
||||
record.run_id,
|
||||
lambda: self._store.get(record.run_id, user_id=record.user_id),
|
||||
)
|
||||
active_statuses = (RunStatus.pending.value, RunStatus.running.value)
|
||||
if stored is not None and stored.get("status") in active_statuses:
|
||||
# `_persist_status` is deliberately best-effort. This compensation
|
||||
# path needs a strict second CAS attempt because the caller never
|
||||
# receives the record and no worker can attach after it returns.
|
||||
# A peer terminal transition wins the CAS and is preserved below.
|
||||
await self._call_store_with_retry(
|
||||
"terminalize cancelled admission",
|
||||
record.run_id,
|
||||
lambda: self._store.update_status(record.run_id, RunStatus.interrupted.value),
|
||||
)
|
||||
stored = await self._call_store_with_retry(
|
||||
"verify terminal cancelled admission",
|
||||
record.run_id,
|
||||
lambda: self._store.get(record.run_id, user_id=record.user_id),
|
||||
)
|
||||
if stored is not None and stored.get("status") in active_statuses:
|
||||
raise RuntimeError(f"Cancelled admission {record.run_id} remains active in the run store")
|
||||
|
||||
if stored is None:
|
||||
async with self._lock:
|
||||
if self._runs.get(record.run_id) is record:
|
||||
self._runs.pop(record.run_id, None)
|
||||
self._unindex_run_locked(record.run_id, record.thread_id)
|
||||
return
|
||||
|
||||
stored_status = RunStatus(stored.get("status") or RunStatus.pending.value)
|
||||
async with self._lock:
|
||||
if self._runs.get(record.run_id) is record:
|
||||
record.status = stored_status
|
||||
record.error = stored.get("error")
|
||||
record.stop_reason = stored.get("stop_reason")
|
||||
record.updated_at = _now_iso()
|
||||
|
||||
async def _admit_thread_operation(
|
||||
self,
|
||||
thread_id: str,
|
||||
@ -1260,9 +1305,29 @@ class RunManager:
|
||||
interrupted_records.append(r)
|
||||
|
||||
# Outside the lock: persist interrupted status for locally-cancelled
|
||||
# runs. Store-side claimed rows are already finalised.
|
||||
for interrupted_record in interrupted_records:
|
||||
await self._persist_status(interrupted_record, RunStatus.interrupted)
|
||||
# runs. Store-side claimed rows are already finalised. Cancellation at
|
||||
# this point happens after the replacement was admitted, so close that
|
||||
# new run before propagating cancellation to the caller.
|
||||
try:
|
||||
for interrupted_record in interrupted_records:
|
||||
await self._persist_status(interrupted_record, RunStatus.interrupted)
|
||||
except asyncio.CancelledError:
|
||||
cleanup = asyncio.create_task(self._close_cancelled_admission(record))
|
||||
cleanup.set_name(f"deerflow-close-cancelled-admission-{record.run_id}")
|
||||
while not cleanup.done():
|
||||
try:
|
||||
await asyncio.shield(cleanup)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception:
|
||||
break
|
||||
try:
|
||||
cleanup.result()
|
||||
except asyncio.CancelledError:
|
||||
logger.error("Cancelled admission cleanup task was itself cancelled for run %s", record.run_id)
|
||||
except Exception:
|
||||
logger.exception("Failed to close run %s after admission was cancelled", record.run_id)
|
||||
raise
|
||||
|
||||
logger.info("Run created: run_id=%s thread_id=%s", run_id, thread_id)
|
||||
return record
|
||||
|
||||
@ -993,6 +993,245 @@ async def test_create_or_reject_does_not_interrupt_old_run_when_new_run_store_wr
|
||||
assert stored_old["status"] == "running"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("strategy", ["interrupt", "rollback"])
|
||||
async def test_create_or_reject_cancellation_after_registration_interrupts_replacement(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
strategy: str,
|
||||
) -> None:
|
||||
"""Cancellation after admission must not leave the replacement active."""
|
||||
store = MemoryRunStore()
|
||||
manager = RunManager(store=store)
|
||||
old = await manager.create("thread-1")
|
||||
await manager.set_status(old.run_id, RunStatus.running)
|
||||
persist_started = asyncio.Event()
|
||||
release_persist = asyncio.Event()
|
||||
original_persist_status = manager._persist_status
|
||||
|
||||
async def blocking_persist_status(record: Any, status: RunStatus, **kwargs: Any) -> bool:
|
||||
persist_started.set()
|
||||
await asyncio.wait_for(release_persist.wait(), timeout=1)
|
||||
return await original_persist_status(record, status, **kwargs)
|
||||
|
||||
monkeypatch.setattr(manager, "_persist_status", blocking_persist_status)
|
||||
create_task = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy=strategy))
|
||||
await asyncio.wait_for(persist_started.wait(), timeout=1)
|
||||
create_task.cancel()
|
||||
release_persist.set()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
_ = await create_task
|
||||
|
||||
records = await manager.list_by_thread("thread-1")
|
||||
replacement = next(record for record in records if record.run_id != old.run_id)
|
||||
stored_replacement = await store.get(replacement.run_id)
|
||||
assert not await manager.has_inflight("thread-1")
|
||||
assert replacement.status == RunStatus.interrupted
|
||||
assert replacement.abort_event.is_set()
|
||||
assert stored_replacement is not None
|
||||
assert stored_replacement["status"] == RunStatus.interrupted.value
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_or_reject_repeated_cancellation_drains_replacement_cleanup(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Repeated cancellation must not abandon the durable cleanup task."""
|
||||
store = MemoryRunStore()
|
||||
manager = RunManager(store=store)
|
||||
old = await manager.create("thread-1")
|
||||
await manager.set_status(old.run_id, RunStatus.running)
|
||||
old_persist_started = asyncio.Event()
|
||||
release_old_persist = asyncio.Event()
|
||||
replacement_persist_started = asyncio.Event()
|
||||
release_replacement_persist = asyncio.Event()
|
||||
original_persist_status = manager._persist_status
|
||||
|
||||
async def staged_persist_status(record: Any, status: RunStatus, **kwargs: Any) -> bool:
|
||||
if record.run_id == old.run_id:
|
||||
old_persist_started.set()
|
||||
await asyncio.wait_for(release_old_persist.wait(), timeout=1)
|
||||
else:
|
||||
replacement_persist_started.set()
|
||||
await asyncio.wait_for(release_replacement_persist.wait(), timeout=1)
|
||||
return await original_persist_status(record, status, **kwargs)
|
||||
|
||||
monkeypatch.setattr(manager, "_persist_status", staged_persist_status)
|
||||
create_task = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy="interrupt"))
|
||||
await asyncio.wait_for(old_persist_started.wait(), timeout=1)
|
||||
replacement = next(record for record in manager._runs.values() if record.run_id != old.run_id)
|
||||
|
||||
create_task.cancel()
|
||||
release_old_persist.set()
|
||||
await asyncio.wait_for(replacement_persist_started.wait(), timeout=1)
|
||||
create_task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not create_task.done()
|
||||
|
||||
release_replacement_persist.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
_ = await asyncio.wait_for(create_task, timeout=1)
|
||||
|
||||
stored_replacement = await store.get(replacement.run_id)
|
||||
assert replacement.status == RunStatus.interrupted
|
||||
assert stored_replacement is not None
|
||||
assert stored_replacement["status"] == RunStatus.interrupted.value
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_or_reject_retries_replacement_when_cancel_status_cannot_persist(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A failed best-effort update must get a strict durable retry."""
|
||||
|
||||
class FailFirstReplacementInterruptStore(MemoryRunStore):
|
||||
failed = False
|
||||
|
||||
async def get(self, run_id: str, *, user_id: str | None = None) -> dict[str, Any] | None:
|
||||
raw = await super().get(run_id)
|
||||
if raw is not None and raw.get("status") == RunStatus.pending.value and raw.get("user_id") is not None and user_id != raw.get("user_id"):
|
||||
raise RuntimeError("replacement lookup was not owner-scoped")
|
||||
return await super().get(run_id, user_id=user_id)
|
||||
|
||||
async def update_status(self, run_id: str, status: str, **kwargs: Any) -> bool:
|
||||
row = await super().get(run_id)
|
||||
if not self.failed and status == RunStatus.interrupted.value and row is not None and row.get("status") == RunStatus.pending.value:
|
||||
self.failed = True
|
||||
raise RuntimeError("replacement status write failed")
|
||||
return await super().update_status(run_id, status, **kwargs)
|
||||
|
||||
store = FailFirstReplacementInterruptStore()
|
||||
manager = RunManager(store=store)
|
||||
old = await manager.create("thread-1", user_id="owner-1")
|
||||
await manager.set_status(old.run_id, RunStatus.running)
|
||||
old_persist_started = asyncio.Event()
|
||||
release_old_persist = asyncio.Event()
|
||||
original_persist_status = manager._persist_status
|
||||
|
||||
async def block_old_persist(record: Any, status: RunStatus, **kwargs: Any) -> bool:
|
||||
if record.run_id != old.run_id:
|
||||
return await original_persist_status(record, status, **kwargs)
|
||||
old_persist_started.set()
|
||||
await asyncio.wait_for(release_old_persist.wait(), timeout=1)
|
||||
return await original_persist_status(record, status, **kwargs)
|
||||
|
||||
monkeypatch.setattr(manager, "_persist_status", block_old_persist)
|
||||
create_task = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy="interrupt", user_id="owner-1"))
|
||||
await asyncio.wait_for(old_persist_started.wait(), timeout=1)
|
||||
replacement = next(record for record in manager._runs.values() if record.run_id != old.run_id)
|
||||
create_task.cancel()
|
||||
release_old_persist.set()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
_ = await asyncio.wait_for(create_task, timeout=1)
|
||||
|
||||
stored_replacement = await store.get(replacement.run_id, user_id="owner-1")
|
||||
assert stored_replacement is not None
|
||||
assert stored_replacement["status"] == RunStatus.interrupted.value
|
||||
assert replacement.status == RunStatus.interrupted
|
||||
assert not await manager.has_inflight("thread-1")
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_or_reject_cleanup_failure_preserves_caller_cancellation(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Cleanup IO failure must not replace the caller's CancelledError."""
|
||||
|
||||
class FailingReplacementCleanupStore(MemoryRunStore):
|
||||
replacement_update_failed = False
|
||||
|
||||
async def update_status(self, run_id: str, status: str, **kwargs: Any) -> bool:
|
||||
row = await super().get(run_id)
|
||||
if status == RunStatus.interrupted.value and row is not None and row.get("status") == RunStatus.pending.value:
|
||||
self.replacement_update_failed = True
|
||||
raise RuntimeError("replacement status unavailable")
|
||||
return await super().update_status(run_id, status, **kwargs)
|
||||
|
||||
async def get(self, run_id: str, *, user_id: str | None = None) -> dict[str, Any] | None:
|
||||
row = await super().get(run_id, user_id=user_id)
|
||||
if self.replacement_update_failed and row is not None and row.get("status") == RunStatus.pending.value:
|
||||
raise RuntimeError("replacement verification unavailable")
|
||||
return row
|
||||
|
||||
store = FailingReplacementCleanupStore()
|
||||
manager = RunManager(store=store)
|
||||
old = await manager.create("thread-1")
|
||||
await manager.set_status(old.run_id, RunStatus.running)
|
||||
old_persist_started = asyncio.Event()
|
||||
release_old_persist = asyncio.Event()
|
||||
original_persist_status = manager._persist_status
|
||||
|
||||
async def block_old_persist(record: Any, status: RunStatus, **kwargs: Any) -> bool:
|
||||
if record.run_id != old.run_id:
|
||||
return await original_persist_status(record, status, **kwargs)
|
||||
old_persist_started.set()
|
||||
await asyncio.wait_for(release_old_persist.wait(), timeout=1)
|
||||
return await original_persist_status(record, status, **kwargs)
|
||||
|
||||
monkeypatch.setattr(manager, "_persist_status", block_old_persist)
|
||||
create_task = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy="interrupt"))
|
||||
await asyncio.wait_for(old_persist_started.wait(), timeout=1)
|
||||
create_task.cancel()
|
||||
release_old_persist.set()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
_ = await asyncio.wait_for(create_task, timeout=1)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_or_reject_preserves_peer_terminal_status_during_cancel_retry(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A peer terminal transition must win the strict cancellation retry."""
|
||||
|
||||
class PeerWinsReplacementInterruptStore(MemoryRunStore):
|
||||
replacement_attempts = 0
|
||||
|
||||
async def update_status(self, run_id: str, status: str, **kwargs: Any) -> bool:
|
||||
row = await self.get(run_id)
|
||||
if status == RunStatus.interrupted.value and row is not None and row.get("status") == RunStatus.pending.value:
|
||||
self.replacement_attempts += 1
|
||||
if self.replacement_attempts == 1:
|
||||
raise RuntimeError("replacement status write failed")
|
||||
await super().update_status(run_id, RunStatus.error.value, error="peer takeover")
|
||||
return False
|
||||
return await super().update_status(run_id, status, **kwargs)
|
||||
|
||||
store = PeerWinsReplacementInterruptStore()
|
||||
manager = RunManager(store=store)
|
||||
old = await manager.create("thread-1")
|
||||
await manager.set_status(old.run_id, RunStatus.running)
|
||||
old_persist_started = asyncio.Event()
|
||||
release_old_persist = asyncio.Event()
|
||||
original_persist_status = manager._persist_status
|
||||
|
||||
async def block_old_persist(record: Any, status: RunStatus, **kwargs: Any) -> bool:
|
||||
if record.run_id != old.run_id:
|
||||
return await original_persist_status(record, status, **kwargs)
|
||||
old_persist_started.set()
|
||||
await asyncio.wait_for(release_old_persist.wait(), timeout=1)
|
||||
return await original_persist_status(record, status, **kwargs)
|
||||
|
||||
monkeypatch.setattr(manager, "_persist_status", block_old_persist)
|
||||
create_task = asyncio.create_task(manager.create_or_reject("thread-1", multitask_strategy="interrupt"))
|
||||
await asyncio.wait_for(old_persist_started.wait(), timeout=1)
|
||||
replacement = next(record for record in manager._runs.values() if record.run_id != old.run_id)
|
||||
create_task.cancel()
|
||||
release_old_persist.set()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
_ = await asyncio.wait_for(create_task, timeout=1)
|
||||
|
||||
stored_replacement = await store.get(replacement.run_id)
|
||||
assert stored_replacement is not None
|
||||
assert stored_replacement["status"] == RunStatus.error.value
|
||||
assert stored_replacement["error"] == "peer takeover"
|
||||
assert replacement.status == RunStatus.error
|
||||
assert replacement.error == "peer takeover"
|
||||
assert not await manager.has_inflight("thread-1")
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_create_or_reject_rollback_persists_interrupted_status_to_store():
|
||||
"""rollback strategy should persist interrupted status for old runs."""
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user