mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
bf2cb0e1df
Signed-off-by: CorgiBoyG <111257566+CorgiBoyG@users.noreply.github.com> Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
415 lines
17 KiB
Python
415 lines
17 KiB
Python
import asyncio
|
|
from datetime import UTC, datetime, timedelta
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
from test_terminal_store_contract import store as store
|
|
|
|
from deerflow.config.run_ownership_config import RunOwnershipConfig
|
|
from deerflow.runtime import RunManager, RunStatus
|
|
|
|
|
|
@pytest.fixture
|
|
def anyio_backend():
|
|
return "asyncio"
|
|
|
|
|
|
def manager_for(store):
|
|
return RunManager(store=store, worker_id="owner", run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True, lease_seconds=5, grace_seconds=0))
|
|
|
|
|
|
@pytest.mark.anyio
|
|
@pytest.mark.parametrize("strategy", ["interrupt", "rollback"])
|
|
async def test_same_owner_atomic_replacement_does_not_fence_its_old_run(store, strategy):
|
|
manager = manager_for(store)
|
|
first = await manager.create_or_reject("thread", user_id="review")
|
|
await manager.try_start(first.run_id)
|
|
second = await manager.create_or_reject("thread", multitask_strategy=strategy, user_id="review")
|
|
row = await store.get(first.run_id, user_id="review")
|
|
assert row["owner_worker_id"] == "owner"
|
|
assert row["status"] == "interrupted"
|
|
assert datetime.fromisoformat(row["lease_expires_at"]) > datetime.now(UTC)
|
|
assert second.status == RunStatus.pending
|
|
if strategy == "rollback":
|
|
await manager.set_status(first.run_id, RunStatus.error, error="Rolled back by user")
|
|
row = await store.get(first.run_id, user_id="review")
|
|
assert row["status"] == "error"
|
|
assert row["error"] == "Rolled back by user"
|
|
assert not first.ownership_lost, "successful same-owner atomic replacement was treated as lost ownership"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_fast_renewal_exception_retains_deadline_fencing():
|
|
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
|
|
|
store = MemoryRunStore()
|
|
manager = manager_for(store)
|
|
record = await manager.create_or_reject("thread")
|
|
await manager.try_start(record.run_id)
|
|
record.task = asyncio.create_task(asyncio.Event().wait())
|
|
expiry = datetime.now(UTC) + timedelta(milliseconds=100)
|
|
record.lease_expires_at = expiry.isoformat()
|
|
store._runs[record.run_id]["lease_expires_at"] = expiry.isoformat()
|
|
store.renew_lease = AsyncMock(side_effect=RuntimeError("unavailable"))
|
|
expired = asyncio.Event()
|
|
loop = asyncio.get_running_loop()
|
|
handle = loop.call_later(0.15, expired.set)
|
|
await manager.start_heartbeat()
|
|
try:
|
|
await manager._renew_leases()
|
|
await expired.wait()
|
|
assert datetime.now(UTC) > expiry
|
|
assert await store.claim_for_takeover(record.run_id, grace_seconds=0, error="peer recovery")
|
|
assert record.ownership_lost, "the failed lease check exited and left no deadline decision before the next heartbeat tick"
|
|
finally:
|
|
handle.cancel()
|
|
await manager.stop_heartbeat(timeout=1)
|
|
record.task.cancel()
|
|
await asyncio.gather(record.task, return_exceptions=True)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_shutdown_does_not_join_uncooperative_orphan_io_past_budget():
|
|
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
|
|
|
store = MemoryRunStore()
|
|
manager = manager_for(store)
|
|
entered, release, cancelled = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
|
|
|
async def resistant(*args, **kwargs):
|
|
entered.set()
|
|
while not release.is_set():
|
|
try:
|
|
await release.wait()
|
|
except asyncio.CancelledError:
|
|
cancelled.set()
|
|
return []
|
|
|
|
store.list_inflight_with_expired_lease = resistant
|
|
manager._schedule_orphan_reconciliation()
|
|
recovery = manager._orphan_recovery_task
|
|
await asyncio.wait_for(entered.wait(), 1)
|
|
shutdown = asyncio.create_task(manager.shutdown(timeout=0.01))
|
|
try:
|
|
await asyncio.wait_for(cancelled.wait(), 1)
|
|
done, _ = await asyncio.wait({shutdown}, timeout=0.05)
|
|
assert shutdown in done, "shutdown joined cancellation-resistant orphan I/O without a bound"
|
|
assert manager._orphan_recovery_task is recovery
|
|
assert not recovery.done()
|
|
manager._schedule_orphan_reconciliation()
|
|
assert manager._orphan_recovery_task is recovery
|
|
finally:
|
|
release.set()
|
|
await asyncio.wait_for(shutdown, 1)
|
|
await asyncio.wait_for(recovery, 1)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_cancelled_admission_compensation_cannot_write_after_expiry(store):
|
|
manager = manager_for(store)
|
|
record = await manager.create_or_reject("thread", user_id="review")
|
|
expired = (datetime.now(UTC) - timedelta(seconds=1)).isoformat()
|
|
await store.put(record.run_id, thread_id="thread", status="pending", user_id="review", owner_worker_id="owner", lease_expires_at=expired)
|
|
record.lease_expires_at = expired
|
|
await manager._close_cancelled_admission(record)
|
|
row = await store.get(record.run_id, user_id="review")
|
|
assert row["status"] == "pending", "unowned compensation terminalized a row after the owned write had lost its lease"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_atomic_attempt_is_registered_until_its_actual_claim_ack():
|
|
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
|
|
|
store = MemoryRunStore()
|
|
manager = manager_for(store)
|
|
old = await manager.create_or_reject("thread")
|
|
await manager.try_start(old.run_id)
|
|
committed, release = asyncio.Event(), asyncio.Event()
|
|
original = store.create_thread_operation_atomic
|
|
|
|
async def hold_ack(**kwargs):
|
|
assert manager._terminal_attempts.get(old.run_id), "attempt must precede Store I/O"
|
|
result = await original(**kwargs)
|
|
committed.set()
|
|
await release.wait()
|
|
return result
|
|
|
|
store.create_thread_operation_atomic = hold_ack
|
|
admission = asyncio.create_task(manager.create_or_reject("thread", multitask_strategy="interrupt"))
|
|
try:
|
|
await asyncio.wait_for(committed.wait(), 1)
|
|
assert manager._terminal_attempts.get(old.run_id)
|
|
assert not manager._has_terminal_ack(old)
|
|
release.set()
|
|
await asyncio.wait_for(admission, 1)
|
|
assert manager._has_terminal_ack(old)
|
|
assert old.error == "Cancelled by newer run"
|
|
assert not old.ownership_lost
|
|
assert not manager._terminal_attempts
|
|
finally:
|
|
release.set()
|
|
await asyncio.gather(admission, return_exceptions=True)
|
|
await manager.stop_heartbeat(timeout=0.1)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_unclaimed_local_row_cannot_borrow_atomic_admission_ack():
|
|
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
|
|
|
store = MemoryRunStore()
|
|
manager = manager_for(store)
|
|
old = await manager.create_or_reject("thread")
|
|
await manager.try_start(old.run_id)
|
|
original = store.create_thread_operation_atomic
|
|
|
|
async def omit_claim(**kwargs):
|
|
new, _ = await original(**kwargs)
|
|
return new, []
|
|
|
|
store.create_thread_operation_atomic = omit_claim
|
|
await manager.create_or_reject("thread", multitask_strategy="rollback")
|
|
assert old.status != RunStatus.interrupted
|
|
assert old.terminal_commit_proof is None
|
|
assert not manager._terminal_attempts
|
|
await manager.stop_heartbeat(timeout=0.1)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_atomic_exception_settles_attempt_without_cancelling_old_worker():
|
|
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
|
|
|
store = MemoryRunStore()
|
|
manager = manager_for(store)
|
|
old = await manager.create_or_reject("thread")
|
|
|
|
async def reject(**kwargs):
|
|
assert manager._terminal_attempts.get(old.run_id)
|
|
raise RuntimeError("admission rejected")
|
|
|
|
store.create_thread_operation_atomic = reject
|
|
with pytest.raises(RuntimeError, match="admission rejected"):
|
|
await manager.create_or_reject("thread", multitask_strategy="rollback")
|
|
assert not manager._terminal_attempts
|
|
assert old.terminal_commit_proof is None
|
|
assert not old.abort_event.is_set()
|
|
assert old.status == RunStatus.pending
|
|
await manager.stop_heartbeat(timeout=0.1)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_successful_renewal_invalidates_the_earlier_deadline_guard():
|
|
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
|
|
|
store = MemoryRunStore()
|
|
manager = manager_for(store)
|
|
record = await manager.create_or_reject("thread")
|
|
initial = (datetime.now(UTC) + timedelta(milliseconds=100)).isoformat()
|
|
record.lease_expires_at = initial
|
|
store._runs[record.run_id]["lease_expires_at"] = initial
|
|
original = store.renew_lease
|
|
store.renew_lease = AsyncMock(side_effect=RuntimeError("unavailable"))
|
|
await manager._renew_leases()
|
|
store.renew_lease = original
|
|
await manager._renew_leases()
|
|
elapsed = asyncio.Event()
|
|
asyncio.get_running_loop().call_later(0.15, elapsed.set)
|
|
await elapsed.wait()
|
|
assert not record.ownership_lost
|
|
assert record.lease_expires_at != initial
|
|
await manager.stop_heartbeat(timeout=0.1)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_expiry_fences_even_while_atomic_ack_holds_the_manager_lock():
|
|
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
|
|
|
store = MemoryRunStore()
|
|
manager = manager_for(store)
|
|
old = await manager.create_or_reject("thread")
|
|
original_atomic, original_renew = store.create_thread_operation_atomic, store.renew_lease
|
|
expiry = (datetime.now(UTC) + timedelta(milliseconds=100)).isoformat()
|
|
old.lease_expires_at = expiry
|
|
store._runs[old.run_id]["lease_expires_at"] = expiry
|
|
renewing, allow_renew, committed, release_ack = (asyncio.Event() for _ in range(4))
|
|
|
|
async def held_renew(*args, **kwargs):
|
|
renewing.set()
|
|
await allow_renew.wait()
|
|
return await original_renew(*args, **kwargs)
|
|
|
|
async def held_atomic(**kwargs):
|
|
result = await original_atomic(**kwargs)
|
|
committed.set()
|
|
await release_ack.wait()
|
|
return result
|
|
|
|
store.renew_lease, store.create_thread_operation_atomic = held_renew, held_atomic
|
|
renewal = asyncio.create_task(manager._renew_leases())
|
|
admission = None
|
|
try:
|
|
await asyncio.wait_for(renewing.wait(), 1)
|
|
admission = asyncio.create_task(manager.create_or_reject("thread", multitask_strategy="rollback"))
|
|
await asyncio.wait_for(committed.wait(), 1)
|
|
allow_renew.set()
|
|
done, _ = await asyncio.wait({renewal}, timeout=0.3)
|
|
assert renewal in done, "expiry fencing waited behind the atomic Store await under the manager lock"
|
|
assert old.ownership_lost
|
|
release_ack.set()
|
|
await asyncio.wait_for(admission, 1)
|
|
assert old.ownership_lost
|
|
assert not manager._has_terminal_ack(old)
|
|
assert not manager._terminal_attempts
|
|
finally:
|
|
allow_renew.set()
|
|
release_ack.set()
|
|
await asyncio.gather(renewal, *([admission] if admission else []), return_exceptions=True)
|
|
await manager.stop_heartbeat(timeout=0.1)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
@pytest.mark.parametrize("method", ["renew_lease", "update_lease"])
|
|
@pytest.mark.parametrize(
|
|
"status,owner,expired,allowed",
|
|
[
|
|
("interrupted", "owner", False, True),
|
|
("interrupted", "peer", False, False),
|
|
("interrupted", "owner", True, False),
|
|
("error", "owner", False, False),
|
|
("success", "owner", False, False),
|
|
],
|
|
)
|
|
async def test_interrupted_renewal_requires_opt_in_and_live_owner(store, method, status, owner, expired, allowed):
|
|
expiry = (datetime.now(UTC) + timedelta(seconds=-1 if expired else 30)).isoformat()
|
|
await store.put("run", thread_id="thread", status=status, user_id=None, owner_worker_id="owner", lease_expires_at=expiry)
|
|
kwargs = dict(owner_worker_id=owner, lease_expires_at=(datetime.now(UTC) + timedelta(seconds=60)).isoformat())
|
|
result = await getattr(store, method)("run", **kwargs)
|
|
assert not (result.renewed if method == "renew_lease" else result)
|
|
result = await getattr(store, method)("run", allow_interrupted=True, **kwargs)
|
|
assert (result.renewed if method == "renew_lease" else result) is allowed
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_atomic_claim_ack_is_a_snapshot_of_the_committed_outcome(store):
|
|
manager = manager_for(store)
|
|
old = await manager.create_or_reject("thread", user_id="review")
|
|
_, claims = await store.create_thread_operation_atomic(
|
|
run_id="replacement",
|
|
thread_id="thread",
|
|
owner_worker_id="owner",
|
|
lease_expires_at=old.lease_expires_at,
|
|
multitask_strategy="rollback",
|
|
user_id="review",
|
|
)
|
|
assert len(claims) == 1
|
|
assert claims[0]["run_id"] == old.run_id
|
|
assert claims[0]["owner_worker_id"] == "owner"
|
|
assert claims[0]["status"] == "interrupted"
|
|
assert claims[0]["error"] == "Cancelled by newer run"
|
|
assert await store.update_status(old.run_id, "error", error="Rolled back by user", owner_worker_id="owner")
|
|
assert claims[0]["status"] == "interrupted"
|
|
assert claims[0]["error"] == "Cancelled by newer run"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_atomic_replacement_rejects_live_peer_without_ack(store):
|
|
from deerflow.runtime.runs.manager import ConflictError
|
|
|
|
manager = manager_for(store)
|
|
old = await manager.create_or_reject("thread", user_id="review")
|
|
peer = RunManager(store=store, worker_id="peer", run_ownership_config=RunOwnershipConfig(heartbeat_enabled=True, lease_seconds=5, grace_seconds=0))
|
|
with pytest.raises(ConflictError):
|
|
await peer.create_or_reject("thread", multitask_strategy="rollback", user_id="review")
|
|
assert (await store.get(old.run_id, user_id="review"))["status"] == "pending"
|
|
assert old.terminal_commit_proof is None
|
|
assert not old.abort_event.is_set()
|
|
assert not peer._terminal_attempts
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_pending_caller_cancellation_closes_admitted_replacement():
|
|
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
|
|
|
store = MemoryRunStore()
|
|
manager = RunManager(store=store)
|
|
old = await manager.create_or_reject("thread")
|
|
original = store.create_thread_operation_atomic
|
|
|
|
async def cancel_before_return(**kwargs):
|
|
result = await original(**kwargs)
|
|
asyncio.current_task().cancel()
|
|
return result
|
|
|
|
store.create_thread_operation_atomic = cancel_before_return
|
|
task = asyncio.create_task(manager.create_or_reject("thread", multitask_strategy="interrupt"))
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
replacement = next(r for r in manager._runs.values() if r is not old)
|
|
assert (await store.get(replacement.run_id))["status"] == "interrupted"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_delayed_atomic_interrupted_ack_cannot_restart_acknowledged_error():
|
|
from deerflow.runtime.runs.manager import TerminalCommitProof
|
|
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
|
|
|
store = MemoryRunStore()
|
|
manager = manager_for(store)
|
|
record = await manager.create_or_reject("thread")
|
|
committed, release_ack = asyncio.Event(), asyncio.Event()
|
|
original = store.create_thread_operation_atomic
|
|
|
|
async def held_atomic(**kwargs):
|
|
result = await original(**kwargs)
|
|
committed.set()
|
|
await release_ack.wait()
|
|
return result
|
|
|
|
store.create_thread_operation_atomic = held_atomic
|
|
admission = asyncio.create_task(manager.create_or_reject("thread", multitask_strategy="rollback"))
|
|
try:
|
|
await asyncio.wait_for(committed.wait(), 1)
|
|
# A previously registered terminal writer can finish while admission
|
|
# awaits its older interrupted response under the manager lock.
|
|
outcome = TerminalCommitProof(RunStatus.error, "final error", None, None)
|
|
assert await manager._terminal_operation(record, outcome, lambda: store.update_status(record.run_id, "error", error="final error", owner_worker_id="owner"))
|
|
assert manager._has_final_terminal_ack(record)
|
|
release_ack.set()
|
|
await asyncio.wait_for(admission, 1)
|
|
assert record.status == RunStatus.error
|
|
assert record.error == "final error"
|
|
assert not record.abort_event.is_set()
|
|
assert not record.rollback_finalization_pending
|
|
assert not record.terminal_commit_pending
|
|
assert not manager._terminal_attempts
|
|
finally:
|
|
release_ack.set()
|
|
await asyncio.gather(admission, return_exceptions=True)
|
|
await manager.stop_heartbeat(timeout=0.1)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
@pytest.mark.parametrize("action", ["interrupt", "rollback"])
|
|
async def test_direct_cancel_keeps_terminal_confirmation_behind_rollback(store, action):
|
|
from deerflow.runtime.runs.manager import CancelOutcome
|
|
|
|
manager = manager_for(store)
|
|
record = await manager.create_or_reject("thread", user_id="review")
|
|
await manager.try_start(record.run_id)
|
|
record.task = asyncio.create_task(asyncio.Event().wait())
|
|
try:
|
|
assert await manager.cancel(record.run_id, action=action) == CancelOutcome.cancelled
|
|
assert (await store.get(record.run_id, user_id="review"))["status"] == "interrupted"
|
|
assert record.terminal_commit_proof.status == RunStatus.interrupted
|
|
assert record.rollback_finalization_pending is (action == "rollback")
|
|
assert record.terminal_status_persisted is (action == "interrupt")
|
|
|
|
if action == "rollback":
|
|
await manager.set_status(record.run_id, RunStatus.error, error="Rolled back by user")
|
|
assert record.terminal_status_persisted
|
|
assert not record.rollback_finalization_pending
|
|
assert record.terminal_commit_proof.status == RunStatus.error
|
|
finally:
|
|
record.task.cancel()
|
|
await asyncio.gather(record.task, return_exceptions=True)
|
|
await manager.stop_heartbeat(timeout=0.1)
|