From 41a9a8ed8178079a81e7cc98feb3625a5f8d8087 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 14 Aug 2026 16:45:44 +0800 Subject: [PATCH 1/2] revert(worker): rely on visibility-timeout recovery The bounded gevent cancellation raced with an active PyMuPDF subprocess, removed its workspace, and still missed the ECS stop deadline. Keep the independent Redis visibility watchdog as the accepted interruption fallback. --- .../worker/app/core/gevent_worker_shutdown.py | 180 -------- apps/worker/app/core/worker_bootstrap.py | 91 ++--- .../contract/test_worker_shutdown_contract.py | 384 ------------------ .../support/worker_shutdown_process_app.py | 92 ----- 4 files changed, 42 insertions(+), 705 deletions(-) delete mode 100644 apps/worker/app/core/gevent_worker_shutdown.py delete mode 100644 apps/worker/tests/support/worker_shutdown_process_app.py diff --git a/apps/worker/app/core/gevent_worker_shutdown.py b/apps/worker/app/core/gevent_worker_shutdown.py deleted file mode 100644 index 02567c912..000000000 --- a/apps/worker/app/core/gevent_worker_shutdown.py +++ /dev/null @@ -1,180 +0,0 @@ -"""Bound normal Celery warm shutdown for the gevent worker pool.""" - -from __future__ import annotations - -from collections.abc import Callable -from typing import cast - -import gevent -from celery.concurrency.gevent import TaskPool as GeventTaskPool -from celery.worker import WorkController, state -from celery.worker.request import Request -from gevent import Greenlet -from loguru import logger - -_TerminateJob = Callable[[GeventTaskPool, int, int | None], None] - - -class _ShutdownCancellationPool: - """Expose Celery's pool cancellation interface only for shutdown.""" - - def __init__( - self, - pool: GeventTaskPool, - terminate_job: _TerminateJob, - ) -> None: - self._pool: GeventTaskPool = pool - self._terminate_job: _TerminateJob = terminate_job - - def terminate_job(self, pid: int, signal: int | None = None) -> None: - self._terminate_job(self._pool, pid, signal) - - -def _ignore_reconnect_cancellation( - self: GeventTaskPool, - pid: int, - signal: int | None = None, -) -> None: - """Keep the existing broker-reconnect behavior for gevent tasks.""" - logger.warning( - "Ignoring gevent task cancellation outside bounded worker shutdown " - f"(pid={pid}); relying on RedisJobLock for redelivery deduplication" - ) - - -class GeventWorkerShutdownController: - """Cancel unacknowledged active tasks after a bounded warm shutdown. - - Celery's ``worker_soft_shutdown_timeout`` does not bound normal SIGTERM. - Celery applies that setting only after entering its cold/SIGQUIT path. That - path is unsafe for this worker because Celery calls its patched ``sleep`` - directly from the gevent signal callback, which raises - ``BlockingSwitchOutError``. This controller keeps normal warm SIGTERM and - schedules its own non-blocking timer instead. - - The worker also intentionally ignores Celery cancellation requests caused - by broker reconnects. The original gevent cancellation method is retained - privately and exposed only to this explicit shutdown path, so reconnects - cannot accidentally kill useful work while Redis redelivery is settling. - """ - - def __init__( - self, - worker: WorkController, - timeout_seconds: float, - ) -> None: - self._worker: WorkController = worker - self._timeout_seconds: float = timeout_seconds - self._shutdown_timer: Greenlet | None = None - self._has_scheduled_shutdown: bool = False - self._original_terminate_job: _TerminateJob = ( - self._patch_reconnect_cancellation() - ) - - def schedule(self) -> None: - """Schedule the one bounded cancellation pass for warm SIGTERM.""" - if self._has_scheduled_shutdown: - logger.info("Bounded worker shutdown is already scheduled") - return - - self._has_scheduled_shutdown = True - - # ``worker_shutting_down`` is emitted inside Celery's signal handler. - # Waiting or killing a greenlet there would try to switch out of the - # gevent hub callback and reproduce the staging BlockingSwitchOutError. - # ``spawn_later`` only arms a timer here; its callback runs in a normal - # greenlet where cooperative cancellation is safe. - self._shutdown_timer = gevent.spawn_later( - self._timeout_seconds, - self._cancel_unacknowledged_active_tasks, - ) - logger.warning( - "Scheduled bounded warm shutdown cancellation in " - f"{self._timeout_seconds:g} seconds" - ) - - def close(self) -> None: - """Disarm a pending timer after the worker finishes naturally.""" - shutdown_timer: Greenlet | None = self._shutdown_timer - if shutdown_timer is not None and not shutdown_timer.dead: - shutdown_timer.kill(block=False) - self._shutdown_timer = None - - @staticmethod - def _patch_reconnect_cancellation() -> _TerminateJob: - original_terminate_job: _TerminateJob = cast( - _TerminateJob, - getattr( - GeventTaskPool, - "_original_terminate_job", - GeventTaskPool.terminate_job, - ), - ) - - if not hasattr(GeventTaskPool, "_original_terminate_job"): - setattr( - GeventTaskPool, - "_original_terminate_job", - original_terminate_job, - ) - - # Celery uses ``pool.terminate_job`` both for broker reconnect recovery - # and for deliberate task cancellation. Replacing the class method - # only for a moment would race with reconnect handling, so the public - # pool behavior stays a no-op and shutdown uses the private adapter - # above to reach the saved original method. - GeventTaskPool.terminate_job = _ignore_reconnect_cancellation - logger.info( - "Patched gevent TaskPool.terminate_job for reconnect-safe recovery" - ) - return original_terminate_job - - @staticmethod - def _should_cancel(request: Request) -> bool: - if not request.task.acks_late: - return True - return not request.acknowledged - - def _cancel_unacknowledged_active_tasks(self) -> None: - requests_to_cancel: tuple[Request, ...] = tuple( - request - for request in state.active_requests - if self._should_cancel(request) - ) - if not requests_to_cancel: - logger.info( - "Bounded warm shutdown completed without active task cancellation" - ) - return - - pool: object = self._worker.pool - if not isinstance(pool, GeventTaskPool): - logger.error( - "Cannot cancel active tasks during bounded shutdown: " - f"expected gevent pool, got {type(pool).__name__}" - ) - return - - cancellation_pool = _ShutdownCancellationPool( - pool, - self._original_terminate_job, - ) - logger.warning( - "Bounded warm shutdown timeout expired; cancelling " - f"{len(requests_to_cancel)} unacknowledged active task(s)" - ) - for request in requests_to_cancel: - # Request.cancel performs Celery's normal task-ready bookkeeping, - # but the adapter invokes gevent's saved original cancellation. - # The broker connection then closes normally, allowing Kombu to - # restore each unacknowledged reservation for another worker. - try: - request.cancel(cancellation_pool) - except Exception as exc: - # The greenlet is terminated before Celery records its retry - # event. A result-backend outage must not abort this loop and - # leave later active tasks running until ECS force-kills us. - logger.warning( - "Cancelled task but could not record shutdown retry " - f"(task_id={request.id}): {exc}" - ) diff --git a/apps/worker/app/core/worker_bootstrap.py b/apps/worker/app/core/worker_bootstrap.py index 0947784f6..2b99effbd 100644 --- a/apps/worker/app/core/worker_bootstrap.py +++ b/apps/worker/app/core/worker_bootstrap.py @@ -5,16 +5,13 @@ import subprocess import sys -from celery.signals import worker_init, worker_shutdown, worker_shutting_down -from celery.worker import WorkController +from celery.signals import worker_init, worker_shutdown from loguru import logger -from app.core.gevent_worker_shutdown import GeventWorkerShutdownController from shared.core.celery_app import celery_app from shared.core.logging import setup_logging from shared.services.worker_health import start_worker_heartbeat, stop_worker_heartbeat -_worker_shutdown_controller: GeventWorkerShutdownController | None = None _CHILD_PROCESS_TERM_TIMEOUT_SECONDS: float = 5 _CHILD_PROCESS_KILL_TIMEOUT_SECONDS: float = 5 @@ -44,28 +41,54 @@ def _stop_child_process( @worker_init.connect -def init_worker( - sender: WorkController | None = None, - **kwargs: object, -) -> None: +def init_worker(**kwargs: object) -> None: """Initialize structured logging and sync Redis when worker process starts.""" - global _worker_shutdown_controller - setup_logging(service_name="knowhere-worker") start_worker_heartbeat() - if sender is None: - logger.error("Cannot configure bounded worker shutdown without worker sender") - else: - _worker_shutdown_controller = GeventWorkerShutdownController( - worker=sender, - timeout_seconds=float(celery_app.conf.worker_soft_shutdown_timeout), - ) + # Do not cancel gevent tasks from Celery's reconnect or shutdown lifecycle. + # Staging proved that forced greenlet cancellation races with parser child + # processes and still misses the ECS stop deadline. Interrupted reservations + # instead recover through the independent Redis visibility watchdog. + try: + from celery.concurrency.gevent import TaskPool as GeventTaskPool + + if not hasattr(GeventTaskPool, "_original_terminate_job"): + + def _ignore_gevent_cancellation( + self: GeventTaskPool, + pid: int, + signal: int | None = None, + ) -> None: + logger.warning( + f"gevent pool cannot kill greenlet (pid={pid}), " + "relying on visibility recovery and RedisJobLock for " + "redelivery deduplication" + ) + + setattr( + GeventTaskPool, + "_original_terminate_job", + getattr( + GeventTaskPool, + "terminate_job", + None, + ), + ) + GeventTaskPool.terminate_job = _ignore_gevent_cancellation + logger.info( + "Patched gevent TaskPool.terminate_job for visibility recovery" + ) + except Exception as exc: + logger.warning(f"Could not patch gevent TaskPool: {exc}") try: - from shared.services.redis.redis_sync_service import SyncRedisServiceFactory + from shared.services.redis.redis_sync_service import ( + SyncRedisService, + SyncRedisServiceFactory, + ) - service = SyncRedisServiceFactory.get_service() + service: SyncRedisService = SyncRedisServiceFactory.get_service() if service.ping(): logger.info("Worker sync Redis connection verified") else: @@ -74,39 +97,9 @@ def init_worker( logger.warning(f"Worker sync Redis init deferred: {exc}") -@worker_shutting_down.connect -def begin_bounded_worker_shutdown( - sender: str | None = None, - sig: str | None = None, - how: str | None = None, - **kwargs: object, -) -> None: - """Bound ECS SIGTERM without entering Celery's gevent-unsafe cold path.""" - if sig != "SIGTERM" or how != "Warm": - return - - shutdown_controller: GeventWorkerShutdownController | None = ( - _worker_shutdown_controller - ) - if shutdown_controller is None: - logger.error("Cannot schedule bounded worker shutdown before worker init") - return - - shutdown_controller.schedule() - - @worker_shutdown.connect def shutdown_worker(**kwargs: object) -> None: """Clean up shared resources on worker shutdown.""" - global _worker_shutdown_controller - - shutdown_controller: GeventWorkerShutdownController | None = ( - _worker_shutdown_controller - ) - _worker_shutdown_controller = None - if shutdown_controller is not None: - shutdown_controller.close() - try: stop_worker_heartbeat() logger.info("Worker heartbeat stopped") diff --git a/apps/worker/tests/contract/test_worker_shutdown_contract.py b/apps/worker/tests/contract/test_worker_shutdown_contract.py index 7c818991d..34e39827e 100644 --- a/apps/worker/tests/contract/test_worker_shutdown_contract.py +++ b/apps/worker/tests/contract/test_worker_shutdown_contract.py @@ -1,391 +1,16 @@ from __future__ import annotations import json -import os -import signal -import subprocess -import sys -import time -from collections.abc import Callable from pathlib import Path -from types import SimpleNamespace -from typing import cast import pytest -from celery import Celery from pydantic import ValidationError -_PROCESS_CONTRACT_QUEUE_NAME: str = "worker_shutdown_contract" -_PROCESS_CONTRACT_TASK_NAME: str = ( - "worker_shutdown_contract.run_blocking_task" -) -_PROCESS_CONTRACT_SOFT_TIMEOUT_SECONDS: float = 0.25 -_PROCESS_CONTRACT_EXIT_DEADLINE_SECONDS: float = 5.0 - - -def _wait_for_started_tasks( - marker_path: Path, - worker_process: subprocess.Popen[str], - task_count: int, - timeout_seconds: float, -) -> bool: - deadline: float = time.monotonic() + timeout_seconds - while time.monotonic() < deadline: - if marker_path.exists() and len( - marker_path.read_text(encoding="utf-8").splitlines() - ) >= task_count: - return True - if worker_process.poll() is not None: - return False - time.sleep(0.05) - return marker_path.exists() and len( - marker_path.read_text(encoding="utf-8").splitlines() - ) >= task_count - - -def _stop_process(worker_process: subprocess.Popen[str]) -> str: - if worker_process.poll() is None: - worker_process.kill() - worker_process.wait(timeout=5) - - if worker_process.stdout is None: - return "" - return worker_process.stdout.read() - - -def _start_shutdown_contract_worker( - repository_root: Path, - broker_directory: Path, - started_marker_path: Path, - heartbeat_path: Path, - concurrency: int, - shutdown_timeout_seconds: float, - fail_result_backend: bool, -) -> subprocess.Popen[str]: - process_environment: dict[str, str] = os.environ.copy() - python_paths: tuple[str, ...] = ( - str(repository_root / "apps/worker"), - str(repository_root / "apps/worker/tests"), - str(repository_root / "packages/shared-python"), - ) - existing_python_path: str | None = process_environment.get("PYTHONPATH") - process_environment["PYTHONPATH"] = os.pathsep.join( - (*python_paths, existing_python_path) - if existing_python_path - else python_paths - ) - process_environment["WORKER_HEARTBEAT_FILE"] = str(heartbeat_path) - process_environment["WORKER_SHUTDOWN_CONTRACT_BROKER_DIRECTORY"] = str( - broker_directory - ) - process_environment["WORKER_SHUTDOWN_CONTRACT_STARTED_MARKER"] = str( - started_marker_path - ) - process_environment["WORKER_SHUTDOWN_CONTRACT_TIMEOUT_SECONDS"] = str( - shutdown_timeout_seconds - ) - if fail_result_backend: - process_environment[ - "WORKER_SHUTDOWN_CONTRACT_FAIL_RESULT_BACKEND" - ] = "1" - else: - process_environment.pop( - "WORKER_SHUTDOWN_CONTRACT_FAIL_RESULT_BACKEND", - None, - ) - - worker_command: list[str] = [ - sys.executable, - "-m", - "celery", - "-A", - "support.worker_shutdown_process_app:celery_app", - "worker", - "--pool=gevent", - f"--concurrency={concurrency}", - "--loglevel=INFO", - "--hostname=shutdown-contract@%h", - "-Q", - _PROCESS_CONTRACT_QUEUE_NAME, - "--without-gossip", - "--without-mingle", - "--without-heartbeat", - ] - return subprocess.Popen( - worker_command, - env=process_environment, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - ) - - -def _create_shutdown_contract_producer(broker_directory: Path) -> Celery: - producer_app: Celery = Celery( - "worker-shutdown-contract-producer", - broker="filesystem://", - ) - producer_app.conf.broker_transport_options = { - "data_folder_in": str(broker_directory), - "data_folder_out": str(broker_directory), - "control_folder": str(broker_directory / "control"), - "store_processed": False, - } - return producer_app - - -def _wait_for_worker_exit( - worker_process: subprocess.Popen[str], - repeat_signal: bool = False, -) -> tuple[float, str]: - shutdown_started_at: float = time.monotonic() - os.kill(worker_process.pid, signal.SIGTERM) - if repeat_signal: - time.sleep(0.05) - os.kill(worker_process.pid, signal.SIGTERM) - - try: - worker_process.wait(timeout=_PROCESS_CONTRACT_EXIT_DEADLINE_SECONDS) - except subprocess.TimeoutExpired as exc: - worker_output: str = _stop_process(worker_process) - raise AssertionError( - "warm SIGTERM did not cancel the active gevent task before " - "the scaled ECS stop deadline\n" - f"{worker_output}" - ) from exc - - shutdown_elapsed_seconds: float = time.monotonic() - shutdown_started_at - return shutdown_elapsed_seconds, _stop_process(worker_process) - - -def test_should_cancel_an_active_gevent_task_before_the_ecs_stop_deadline( - tmp_path: Path, - worker_contract_environment: None, -) -> None: - broker_directory: Path = tmp_path / "broker" - broker_directory.mkdir() - started_marker_path: Path = tmp_path / "task-started" - worker_process: subprocess.Popen[str] = _start_shutdown_contract_worker( - repository_root=Path(__file__).resolve().parents[4], - broker_directory=broker_directory, - started_marker_path=started_marker_path, - heartbeat_path=tmp_path / "worker-heartbeat", - concurrency=1, - shutdown_timeout_seconds=_PROCESS_CONTRACT_SOFT_TIMEOUT_SECONDS, - fail_result_backend=False, - ) - producer_app: Celery = _create_shutdown_contract_producer(broker_directory) - - try: - producer_app.send_task( - _PROCESS_CONTRACT_TASK_NAME, - args=[1, 30.0], - queue=_PROCESS_CONTRACT_QUEUE_NAME, - ) - assert _wait_for_started_tasks( - started_marker_path, - worker_process, - task_count=1, - timeout_seconds=15, - ), _stop_process(worker_process) - - shutdown_elapsed_seconds, _ = _wait_for_worker_exit(worker_process) - assert ( - shutdown_elapsed_seconds < _PROCESS_CONTRACT_EXIT_DEADLINE_SECONDS - ) - assert worker_process.returncode == 0 - finally: - producer_app.close() - _stop_process(worker_process) - - -def test_should_continue_cancelling_tasks_when_result_backend_recording_fails( - tmp_path: Path, - worker_contract_environment: None, -) -> None: - broker_directory: Path = tmp_path / "broker" - broker_directory.mkdir() - started_marker_path: Path = tmp_path / "task-started" - worker_process: subprocess.Popen[str] = _start_shutdown_contract_worker( - repository_root=Path(__file__).resolve().parents[4], - broker_directory=broker_directory, - started_marker_path=started_marker_path, - heartbeat_path=tmp_path / "worker-heartbeat", - concurrency=2, - shutdown_timeout_seconds=_PROCESS_CONTRACT_SOFT_TIMEOUT_SECONDS, - fail_result_backend=True, - ) - producer_app: Celery = _create_shutdown_contract_producer(broker_directory) - - try: - for task_number in (1, 2): - producer_app.send_task( - _PROCESS_CONTRACT_TASK_NAME, - args=[task_number, 30.0], - queue=_PROCESS_CONTRACT_QUEUE_NAME, - ) - assert _wait_for_started_tasks( - started_marker_path, - worker_process, - task_count=2, - timeout_seconds=15, - ), _stop_process(worker_process) - - shutdown_elapsed_seconds, worker_output = _wait_for_worker_exit( - worker_process - ) - - assert worker_process.returncode == 0 - assert ( - shutdown_elapsed_seconds < _PROCESS_CONTRACT_EXIT_DEADLINE_SECONDS - ) - assert worker_output.count("could not record shutdown retry") == 2, worker_output - finally: - producer_app.close() - _stop_process(worker_process) - - -def test_should_not_cancel_a_task_that_finishes_before_shutdown_timeout( - tmp_path: Path, - worker_contract_environment: None, -) -> None: - broker_directory: Path = tmp_path / "broker" - broker_directory.mkdir() - started_marker_path: Path = tmp_path / "task-started" - worker_process: subprocess.Popen[str] = _start_shutdown_contract_worker( - repository_root=Path(__file__).resolve().parents[4], - broker_directory=broker_directory, - started_marker_path=started_marker_path, - heartbeat_path=tmp_path / "worker-heartbeat", - concurrency=1, - shutdown_timeout_seconds=0.75, - fail_result_backend=False, - ) - producer_app: Celery = _create_shutdown_contract_producer(broker_directory) - - try: - producer_app.send_task( - _PROCESS_CONTRACT_TASK_NAME, - args=[1, 0.05], - queue=_PROCESS_CONTRACT_QUEUE_NAME, - ) - assert _wait_for_started_tasks( - started_marker_path, - worker_process, - task_count=1, - timeout_seconds=15, - ), _stop_process(worker_process) - time.sleep(0.2) - - _, worker_output = _wait_for_worker_exit(worker_process) - - assert worker_process.returncode == 0 - assert "cancelling 1 unacknowledged active task" not in worker_output - finally: - producer_app.close() - _stop_process(worker_process) - - -def test_should_schedule_only_one_shutdown_timer_for_repeated_sigterm( - tmp_path: Path, - worker_contract_environment: None, -) -> None: - broker_directory: Path = tmp_path / "broker" - broker_directory.mkdir() - started_marker_path: Path = tmp_path / "task-started" - worker_process: subprocess.Popen[str] = _start_shutdown_contract_worker( - repository_root=Path(__file__).resolve().parents[4], - broker_directory=broker_directory, - started_marker_path=started_marker_path, - heartbeat_path=tmp_path / "worker-heartbeat", - concurrency=1, - shutdown_timeout_seconds=_PROCESS_CONTRACT_SOFT_TIMEOUT_SECONDS, - fail_result_backend=False, - ) - producer_app: Celery = _create_shutdown_contract_producer(broker_directory) - - try: - producer_app.send_task( - _PROCESS_CONTRACT_TASK_NAME, - args=[1, 30.0], - queue=_PROCESS_CONTRACT_QUEUE_NAME, - ) - assert _wait_for_started_tasks( - started_marker_path, - worker_process, - task_count=1, - timeout_seconds=15, - ), _stop_process(worker_process) - - _, worker_output = _wait_for_worker_exit( - worker_process, - repeat_signal=True, - ) - - assert worker_process.returncode == 0 - assert worker_output.count( - "Scheduled bounded warm shutdown cancellation" - ) == 1 - finally: - producer_app.close() - _stop_process(worker_process) - - -def test_should_keep_gevent_pool_cancellation_disabled_for_broker_reconnects( - monkeypatch: pytest.MonkeyPatch, - worker_contract_environment: None, -) -> None: - import gevent - from app.core.gevent_worker_shutdown import GeventWorkerShutdownController - from celery.concurrency.gevent import TaskPool as GeventTaskPool - from celery.worker import WorkController - - original_terminate_job: Callable[ - [GeventTaskPool, int, int | None], None - ] = cast( - Callable[[GeventTaskPool, int, int | None], None], - getattr(GeventTaskPool, "_original_terminate_job", None), - ) - if not callable(original_terminate_job): - original_terminate_job = GeventTaskPool.terminate_job - - pool: GeventTaskPool = GeventTaskPool(1) - pool.start() - controller = GeventWorkerShutdownController( - worker=cast( - WorkController, - SimpleNamespace(pool=pool), - ), - timeout_seconds=1, - ) - - def wait_for_reconnect_contract() -> None: - gevent.sleep(10) - - running_greenlet = gevent.spawn(wait_for_reconnect_contract) - pool._pool_map[id(running_greenlet)] = running_greenlet - - try: - pool.terminate_job(id(running_greenlet)) - gevent.sleep(0) - assert running_greenlet.dead is False - finally: - controller.close() - original_terminate_job(pool, id(running_greenlet), None) - gevent.sleep(0) - pool.stop() - monkeypatch.setattr( - GeventTaskPool, - "terminate_job", - original_terminate_job, - ) - def test_should_preserve_fargate_worker_sigterm_shutdown_contract( worker_contract_environment: None, ) -> None: from shared.core.celery_app import celery_app - from app.core import worker_bootstrap repository_root: Path = Path(__file__).resolve().parents[4] task_definition_path: Path = ( @@ -411,15 +36,6 @@ def test_should_preserve_fargate_worker_sigterm_shutdown_contract( assert celery_app.conf.worker_enable_soft_shutdown_on_idle is True assert "REMAP_SIGTERM" not in environment_values assert worker_container["stopTimeout"] == 120 - maximum_child_cleanup_seconds: float = 2 * ( - worker_bootstrap._CHILD_PROCESS_TERM_TIMEOUT_SECONDS - + worker_bootstrap._CHILD_PROCESS_KILL_TIMEOUT_SECONDS - ) - assert ( - celery_app.conf.worker_soft_shutdown_timeout - + maximum_child_cleanup_seconds - < worker_container["stopTimeout"] - ) def test_should_redeliver_interrupted_tasks_before_processing_jobs_expire( diff --git a/apps/worker/tests/support/worker_shutdown_process_app.py b/apps/worker/tests/support/worker_shutdown_process_app.py deleted file mode 100644 index 104c0cebd..000000000 --- a/apps/worker/tests/support/worker_shutdown_process_app.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Minimal Celery application for the worker shutdown process contract.""" - -from __future__ import annotations - -# This process must reproduce the production worker's cooperative runtime. -# Patching after Celery or Redis imports would leave blocking sockets in place. -import gevent.monkey - -gevent.monkey.patch_all() - -import importlib -import os -from pathlib import Path - -import gevent -from celery.worker.request import Request -from kombu import Queue - -# Importing the production bootstrap registers its Celery lifecycle receivers. -# The contract deliberately exercises those receivers instead of a test-only hook. -importlib.import_module("app.core.worker_bootstrap") -from shared.core.celery_app import celery_app - -_BROKER_DIRECTORY_ENVIRONMENT_VARIABLE: str = ( - "WORKER_SHUTDOWN_CONTRACT_BROKER_DIRECTORY" -) -_STARTED_MARKER_ENVIRONMENT_VARIABLE: str = ( - "WORKER_SHUTDOWN_CONTRACT_STARTED_MARKER" -) -_SHUTDOWN_TIMEOUT_ENVIRONMENT_VARIABLE: str = ( - "WORKER_SHUTDOWN_CONTRACT_TIMEOUT_SECONDS" -) -_FAIL_RESULT_BACKEND_ENVIRONMENT_VARIABLE: str = ( - "WORKER_SHUTDOWN_CONTRACT_FAIL_RESULT_BACKEND" -) -_QUEUE_NAME: str = "worker_shutdown_contract" -_TASK_NAME: str = "worker_shutdown_contract.run_blocking_task" - -broker_directory: Path = Path( - os.environ[_BROKER_DIRECTORY_ENVIRONMENT_VARIABLE] -) -started_marker_path: Path = Path( - os.environ[_STARTED_MARKER_ENVIRONMENT_VARIABLE] -) -shutdown_timeout_seconds: float = float( - os.environ[_SHUTDOWN_TIMEOUT_ENVIRONMENT_VARIABLE] -) - -celery_app.conf.broker_url = "filesystem://" -celery_app.conf.result_backend = "cache+memory://" -celery_app.conf.broker_transport_options = { - "data_folder_in": str(broker_directory), - "data_folder_out": str(broker_directory), - "control_folder": str(broker_directory / "control"), - "store_processed": False, -} -celery_app.conf.task_default_queue = _QUEUE_NAME -celery_app.conf.task_queues = (Queue(_QUEUE_NAME),) -celery_app.conf.task_routes = {_TASK_NAME: {"queue": _QUEUE_NAME}} -celery_app.conf.worker_soft_shutdown_timeout = shutdown_timeout_seconds -celery_app.conf.worker_enable_soft_shutdown_on_idle = True - - -def _raise_result_backend_failure(*args: object, **kwargs: object) -> None: - raise RuntimeError("shutdown contract result backend unavailable") - - -if os.getenv(_FAIL_RESULT_BACKEND_ENVIRONMENT_VARIABLE) == "1": - # Request.cancel kills the greenlet before it records retry state. Force a - # bookkeeping failure after that call so the contract proves later tasks - # are still canceled when the result backend is unavailable. - celery_app.backend.mark_as_retry = _raise_result_backend_failure - _original_request_cancel = Request.cancel - - def _cancel_then_fail( - request: Request, - pool: object, - signal: int | None = None, - ) -> None: - _original_request_cancel(request, pool, signal) - raise RuntimeError("shutdown contract retry bookkeeping failed") - - Request.cancel = _cancel_then_fail - - -@celery_app.task(name=_TASK_NAME, acks_late=True) -def run_blocking_task(task_number: int, duration_seconds: float) -> None: - """Remain active until the shutdown contract cancels this greenlet.""" - with started_marker_path.open("a", encoding="utf-8") as marker_file: - marker_file.write(f"{task_number}\n") - marker_file.flush() - gevent.sleep(duration_seconds) From aab58ddf73c7dc8415a3c58e852c0848604bd75d Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 14 Aug 2026 22:49:10 +0800 Subject: [PATCH 2/2] perf(worker): cap Fargate concurrency at 10 (#22) --- deploy/ecs/task-definition-worker.staging.json | 1 + deploy/ecs/test_render_task_definitions.py | 1 + 2 files changed, 2 insertions(+) diff --git a/deploy/ecs/task-definition-worker.staging.json b/deploy/ecs/task-definition-worker.staging.json index 1dae56648..f9ef7d6c3 100644 --- a/deploy/ecs/task-definition-worker.staging.json +++ b/deploy/ecs/task-definition-worker.staging.json @@ -23,6 +23,7 @@ {"name": "DB_SSL_MODE", "value": "require"}, {"name": "DB_SYNC_POOL_SIZE", "value": "2"}, {"name": "DB_SYNC_MAX_OVERFLOW", "value": "2"}, + {"name": "WORKER_CONCURRENCY", "value": "10"}, {"name": "TMP_PATH", "value": "/tmp/aismart_bid"}, {"name": "S3_TYPE", "value": "s3"}, {"name": "S3_BUCKET_NAME", "value": "knowhere-storage-staging"}, diff --git a/deploy/ecs/test_render_task_definitions.py b/deploy/ecs/test_render_task_definitions.py index 7d4bb39ac..30248561c 100644 --- a/deploy/ecs/test_render_task_definitions.py +++ b/deploy/ecs/test_render_task_definitions.py @@ -53,6 +53,7 @@ "PDF_PROFILE_TOC_ENABLED": "true", "REDIS_SYNC_MAX_CONNECTIONS": "200", "SUMMARY_LLM_MAX_CONCURRENT": "8", + "WORKER_CONCURRENCY": "10", }