mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-23 05:35:30 +00:00
Revert background worker cancellation changes
This commit is contained in:
酱牛肉
parent
7241e1053e
commit
4a2cf7d19c
2 changed files with 19 additions and 244 deletions
|
|
@ -5,7 +5,6 @@ import contextlib
|
|||
import random
|
||||
import threading
|
||||
import time
|
||||
from concurrent.futures import Future
|
||||
|
||||
from .base_job import BaseJob
|
||||
from ..component_registry import R
|
||||
|
|
@ -24,10 +23,7 @@ class BackgroundJob(BaseJob):
|
|||
supervisor cannot trigger a restart.
|
||||
|
||||
On close, the stop_event is set and the task is given up to
|
||||
``close_timeout`` seconds to exit gracefully; after that cancellation is
|
||||
requested on its event loop and close waits for the actual worker to exit.
|
||||
Cancellation of close is propagated only after worker cleanup finishes.
|
||||
Job bodies must cooperate with cancellation, including in thread-pool mode.
|
||||
``close_timeout`` seconds to exit gracefully; after that it is cancelled.
|
||||
If a single run survives at least ``attempt_reset_after`` seconds before
|
||||
crashing, the backoff attempt counter resets — so a long-stable job that
|
||||
eventually crashes restarts quickly rather than at the capped delay.
|
||||
|
|
@ -54,89 +50,41 @@ class BackgroundJob(BaseJob):
|
|||
self.attempt_reset_after: float = attempt_reset_after
|
||||
self.use_thread_pool: bool = use_thread_pool
|
||||
self._stop_event: asyncio.Event | threading.Event | None = None
|
||||
self._task: asyncio.Future | None = None
|
||||
self._executor_future: Future | None = None
|
||||
self._worker_lock = threading.Lock()
|
||||
self._worker_task: asyncio.Task | None = None
|
||||
self._task: asyncio.Task | None = None
|
||||
|
||||
async def _start(self) -> None:
|
||||
await super()._start()
|
||||
if self.use_thread_pool and self.app_context.thread_pool:
|
||||
self._stop_event = threading.Event()
|
||||
self._executor_future = self.app_context.thread_pool.submit(self._run_in_thread)
|
||||
self._task = asyncio.wrap_future(self._executor_future)
|
||||
self._task = asyncio.get_event_loop().run_in_executor(
|
||||
self.app_context.thread_pool,
|
||||
lambda: asyncio.run(self._run_with_supervisor()),
|
||||
)
|
||||
else:
|
||||
self._stop_event = asyncio.Event()
|
||||
self._task = asyncio.create_task(self._run_with_supervisor())
|
||||
|
||||
def _run_in_thread(self) -> None:
|
||||
"""Keep executor completion tied to actual event-loop shutdown."""
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
asyncio.run(self._run_thread_task())
|
||||
|
||||
async def _run_thread_task(self) -> None:
|
||||
with self._worker_lock:
|
||||
self._worker_task = asyncio.current_task()
|
||||
try:
|
||||
# close() may have set stop before the worker published its task.
|
||||
if not self._stop_event.is_set():
|
||||
await self._run_with_supervisor()
|
||||
finally:
|
||||
# Keep cancellation from targeting a loop that has already closed.
|
||||
with self._worker_lock:
|
||||
self._worker_task = None
|
||||
|
||||
def _cancel_task(self) -> None:
|
||||
if self._executor_future is None:
|
||||
self._task.cancel()
|
||||
return
|
||||
# Cancel queued work without waiting for another long-running job to free
|
||||
# a worker. If already running, cancel the coroutine on its own loop.
|
||||
self._executor_future.cancel()
|
||||
with self._worker_lock:
|
||||
if self._worker_task is not None:
|
||||
self._worker_task.get_loop().call_soon_threadsafe(self._worker_task.cancel)
|
||||
|
||||
async def _close(self) -> None:
|
||||
if self._stop_event is not None:
|
||||
self._stop_event.set()
|
||||
try:
|
||||
await self._shutdown_task()
|
||||
finally:
|
||||
await super()._close()
|
||||
await self._shutdown_task()
|
||||
await super()._close()
|
||||
|
||||
async def _shutdown_task(self) -> None:
|
||||
"""Wait for graceful exit, then request cancellation and await completion."""
|
||||
"""Wait close_timeout for graceful exit, then force-cancel."""
|
||||
if self._task is None:
|
||||
return
|
||||
|
||||
async def wait_for_exit() -> None:
|
||||
try:
|
||||
try:
|
||||
# A timeout must not cancel the executor's completion future.
|
||||
await asyncio.wait_for(asyncio.shield(self._task), timeout=self.close_timeout)
|
||||
except asyncio.TimeoutError:
|
||||
self._cancel_task()
|
||||
await self._task
|
||||
except asyncio.CancelledError:
|
||||
pass # The worker or queued executor future was cancelled.
|
||||
except Exception:
|
||||
self.logger.exception(f"Background task '{self.name}' raised during close")
|
||||
|
||||
# Shield the whole shutdown so repeated caller cancellations cannot
|
||||
# interrupt worker cleanup or cancel the executor's completion future.
|
||||
shutdown = asyncio.create_task(wait_for_exit())
|
||||
cancelled = None
|
||||
while not shutdown.done():
|
||||
try:
|
||||
await asyncio.shield(shutdown)
|
||||
except asyncio.CancelledError as exc:
|
||||
cancelled = exc
|
||||
shutdown.result()
|
||||
try:
|
||||
# shield prevents wait_for's cancellation from propagating to the task itself,
|
||||
# so a timeout here truly times out instead of cancelling silently.
|
||||
await asyncio.wait_for(asyncio.shield(self._task), timeout=self.close_timeout)
|
||||
except asyncio.TimeoutError:
|
||||
self._task.cancel()
|
||||
with contextlib.suppress(BaseException):
|
||||
await self._task
|
||||
except Exception:
|
||||
self.logger.exception(f"Background task '{self.name}' raised during close")
|
||||
self._task = None
|
||||
self._executor_future = None
|
||||
if cancelled is not None:
|
||||
raise cancelled
|
||||
|
||||
def _backoff_delay(self, attempt: int) -> float:
|
||||
"""Exponential backoff with ±50% jitter, capped at backoff_cap."""
|
||||
|
|
|
|||
|
|
@ -3,8 +3,6 @@
|
|||
# pylint: disable=protected-access,missing-function-docstring,missing-class-docstring,no-self-argument,unused-argument
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
|
|
@ -339,177 +337,6 @@ def test_shutdown_task_cancels_on_timeout():
|
|||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_thread_pool_shutdown_waits_for_worker_cleanup():
|
||||
async def run():
|
||||
started = threading.Event()
|
||||
events = []
|
||||
|
||||
class WaitingJob(BackgroundJob):
|
||||
async def __call__(self, **kwargs):
|
||||
started.set()
|
||||
try:
|
||||
# Bounded fallback ensures a regression cannot hang pytest.
|
||||
await asyncio.sleep(2)
|
||||
except asyncio.CancelledError:
|
||||
events.append("cancelled")
|
||||
raise
|
||||
finally:
|
||||
await asyncio.sleep(0.02)
|
||||
events.append("cleaned")
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||||
job = WaitingJob(app_context=SimpleNamespace(thread_pool=pool), use_thread_pool=True, close_timeout=0.01)
|
||||
await job.start()
|
||||
assert await asyncio.to_thread(started.wait, 1)
|
||||
executor_future = job._executor_future
|
||||
await asyncio.wait_for(job.close(), 1)
|
||||
assert events == ["cancelled", "cleaned"]
|
||||
assert executor_future.done()
|
||||
assert job._task is None
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("queued", [True, False], ids=["queued", "starting"])
|
||||
def test_thread_pool_shutdown_before_coroutine_starts(queued):
|
||||
async def run():
|
||||
entered = threading.Event()
|
||||
release = threading.Event()
|
||||
|
||||
class DelayedJob(BackgroundJob):
|
||||
def _run_in_thread(self):
|
||||
entered.set()
|
||||
release.wait(2)
|
||||
super()._run_in_thread()
|
||||
|
||||
def _cancel_task(self):
|
||||
super()._cancel_task()
|
||||
release.set()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||||
if queued:
|
||||
pool.submit(release.wait, 2)
|
||||
job = DelayedJob(app_context=SimpleNamespace(thread_pool=pool), use_thread_pool=True, close_timeout=0.01)
|
||||
body = AsyncMock()
|
||||
job._run_with_supervisor = body
|
||||
try:
|
||||
await job.start()
|
||||
if not queued:
|
||||
assert await asyncio.to_thread(entered.wait, 1)
|
||||
executor_future = job._executor_future
|
||||
await asyncio.wait_for(job.close(), 1)
|
||||
assert executor_future.done()
|
||||
assert executor_future.cancelled() == queued
|
||||
body.assert_not_awaited()
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_thread_pool", [False, True], ids=["async", "thread"])
|
||||
@pytest.mark.parametrize("phase", ["graceful", "forced"])
|
||||
def test_close_finishes_worker_cleanup_before_propagating_repeated_cancellation(use_thread_pool, phase):
|
||||
async def run():
|
||||
started, cleaning, release, finished = (threading.Event() for _ in range(4))
|
||||
|
||||
class WaitingJob(BackgroundJob):
|
||||
async def __call__(self, **kwargs):
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.sleep(1)
|
||||
finally:
|
||||
cleaning.set()
|
||||
deadline = asyncio.get_running_loop().time() + 1
|
||||
while not release.is_set() and asyncio.get_running_loop().time() < deadline:
|
||||
await asyncio.sleep(0.001)
|
||||
finished.set()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||||
job = WaitingJob(
|
||||
app_context=SimpleNamespace(thread_pool=pool),
|
||||
use_thread_pool=use_thread_pool,
|
||||
close_timeout=0.1 if phase == "graceful" else 0.01,
|
||||
)
|
||||
await job.start()
|
||||
worker = job._task
|
||||
executor_future = job._executor_future
|
||||
closing = None
|
||||
try:
|
||||
assert await asyncio.to_thread(started.wait, 1)
|
||||
closing = asyncio.create_task(job.close())
|
||||
if phase == "forced":
|
||||
assert await asyncio.to_thread(cleaning.wait, 1)
|
||||
else:
|
||||
await asyncio.sleep(0)
|
||||
for _ in range(3):
|
||||
closing.cancel()
|
||||
await asyncio.sleep(0)
|
||||
assert not closing.done()
|
||||
assert not worker.done()
|
||||
assert job.is_started
|
||||
assert await asyncio.to_thread(cleaning.wait, 1)
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.wait_for(closing, 1)
|
||||
assert finished.is_set()
|
||||
assert worker.done()
|
||||
assert executor_future is None or executor_future.done()
|
||||
assert job._task is None
|
||||
assert job._executor_future is None
|
||||
assert not job.is_started
|
||||
finally:
|
||||
release.set()
|
||||
if closing is not None:
|
||||
await asyncio.gather(closing, return_exceptions=True)
|
||||
await asyncio.gather(worker, return_exceptions=True)
|
||||
await job.close()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_thread_pool", [False, True], ids=["async", "thread"])
|
||||
def test_close_logs_worker_cleanup_failure_after_timeout(use_thread_pool):
|
||||
async def run():
|
||||
started = threading.Event()
|
||||
errors = []
|
||||
|
||||
class FailingCleanupJob(BackgroundJob):
|
||||
async def __call__(self, **kwargs):
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.sleep(1)
|
||||
finally:
|
||||
raise RuntimeError("worker cleanup failed")
|
||||
|
||||
def capture_error(_message):
|
||||
import sys
|
||||
|
||||
errors.append(sys.exception())
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||||
job = FailingCleanupJob(
|
||||
app_context=SimpleNamespace(thread_pool=pool),
|
||||
use_thread_pool=use_thread_pool,
|
||||
supervisor=False,
|
||||
close_timeout=0.01,
|
||||
)
|
||||
job.logger = MagicMock()
|
||||
job.logger.exception.side_effect = capture_error
|
||||
await job.start()
|
||||
try:
|
||||
assert await asyncio.to_thread(started.wait, 1)
|
||||
finally:
|
||||
await job.close()
|
||||
assert len(errors) == 1
|
||||
assert isinstance(errors[0], RuntimeError)
|
||||
assert str(errors[0]) == "worker cleanup failed"
|
||||
assert job._task is None
|
||||
assert job._executor_future is None
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_cron_uses_configured_timezone(monkeypatch):
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue