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) 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", }