9db64ecda3
* chore: run benchmarks with reflect mode * chore: run benchmarks with reflect mode * fixes * new mm * bunch of fixes * initial commit * fixes * fixes * fixes * fix: sometimes memories gets extracted in the wrong language
404 lines
15 KiB
Python
404 lines
15 KiB
Python
"""
|
|
Worker poller for distributed task execution.
|
|
|
|
Polls PostgreSQL for pending tasks and executes them using
|
|
FOR UPDATE SKIP LOCKED for safe concurrent claiming.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import time
|
|
import traceback
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
if TYPE_CHECKING:
|
|
import asyncpg
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Progress logging interval in seconds
|
|
PROGRESS_LOG_INTERVAL = 30
|
|
|
|
|
|
def fq_table(table: str, schema: str | None = None) -> str:
|
|
"""Get fully-qualified table name with optional schema prefix."""
|
|
if schema:
|
|
return f'"{schema}".{table}'
|
|
return table
|
|
|
|
|
|
class WorkerPoller:
|
|
"""
|
|
Polls PostgreSQL for pending tasks and executes them.
|
|
|
|
Uses FOR UPDATE SKIP LOCKED for safe distributed claiming,
|
|
allowing multiple workers to process tasks without conflicts.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
pool: "asyncpg.Pool",
|
|
worker_id: str,
|
|
executor: Callable[[dict[str, Any]], Awaitable[None]],
|
|
poll_interval_ms: int = 500,
|
|
batch_size: int = 10,
|
|
max_retries: int = 3,
|
|
schema: str | None = None,
|
|
):
|
|
"""
|
|
Initialize the worker poller.
|
|
|
|
Args:
|
|
pool: asyncpg connection pool
|
|
worker_id: Unique identifier for this worker
|
|
executor: Async function to execute tasks (typically MemoryEngine.execute_task)
|
|
poll_interval_ms: Interval between polls when no tasks found (milliseconds)
|
|
batch_size: Maximum number of tasks to claim per poll cycle
|
|
max_retries: Maximum retry attempts before marking task as failed
|
|
schema: Database schema for multi-tenant support (optional)
|
|
"""
|
|
self._pool = pool
|
|
self._worker_id = worker_id
|
|
self._executor = executor
|
|
self._poll_interval_ms = poll_interval_ms
|
|
self._batch_size = batch_size
|
|
self._max_retries = max_retries
|
|
self._schema = schema
|
|
self._shutdown = asyncio.Event()
|
|
self._current_tasks: set[asyncio.Task] = set()
|
|
self._in_flight_count = 0
|
|
self._in_flight_lock = asyncio.Lock()
|
|
self._last_progress_log = 0.0
|
|
self._tasks_completed_since_log = 0
|
|
self._active_banks: set[str] = set()
|
|
|
|
async def claim_batch(self) -> list[tuple[str, dict[str, Any]]]:
|
|
"""
|
|
Claim up to batch_size pending tasks atomically.
|
|
|
|
Uses FOR UPDATE SKIP LOCKED to ensure no conflicts with other workers.
|
|
|
|
For consolidation tasks specifically, skips pending tasks if there's already
|
|
a processing consolidation for the same bank (to avoid duplicate work).
|
|
|
|
Returns:
|
|
List of tuples (operation_id, task_dict)
|
|
"""
|
|
table = fq_table("async_operations", self._schema)
|
|
|
|
async with self._pool.acquire() as conn:
|
|
async with conn.transaction():
|
|
# Select and lock pending tasks
|
|
# For consolidation: skip if same bank already has one processing
|
|
rows = await conn.fetch(
|
|
f"""
|
|
SELECT operation_id, task_payload
|
|
FROM {table} AS pending
|
|
WHERE status = 'pending' AND task_payload IS NOT NULL
|
|
AND (
|
|
-- Non-consolidation tasks: always claimable
|
|
operation_type != 'consolidation'
|
|
OR
|
|
-- Consolidation: only if no other consolidation processing for same bank
|
|
NOT EXISTS (
|
|
SELECT 1 FROM {table} AS processing
|
|
WHERE processing.bank_id = pending.bank_id
|
|
AND processing.operation_type = 'consolidation'
|
|
AND processing.status = 'processing'
|
|
)
|
|
)
|
|
ORDER BY created_at
|
|
LIMIT $1
|
|
FOR UPDATE SKIP LOCKED
|
|
""",
|
|
self._batch_size,
|
|
)
|
|
|
|
if not rows:
|
|
return []
|
|
|
|
# Claim the tasks by updating status and worker_id
|
|
operation_ids = [row["operation_id"] for row in rows]
|
|
await conn.execute(
|
|
f"""
|
|
UPDATE {table}
|
|
SET status = 'processing', worker_id = $1, claimed_at = now(), updated_at = now()
|
|
WHERE operation_id = ANY($2)
|
|
""",
|
|
self._worker_id,
|
|
operation_ids,
|
|
)
|
|
|
|
# Parse and return task payloads
|
|
return [(str(row["operation_id"]), json.loads(row["task_payload"])) for row in rows]
|
|
|
|
async def _mark_completed(self, operation_id: str):
|
|
"""Mark a task as completed."""
|
|
table = fq_table("async_operations", self._schema)
|
|
await self._pool.execute(
|
|
f"""
|
|
UPDATE {table}
|
|
SET status = 'completed', completed_at = now(), updated_at = now()
|
|
WHERE operation_id = $1
|
|
""",
|
|
operation_id,
|
|
)
|
|
|
|
async def _mark_failed(self, operation_id: str, error_message: str):
|
|
"""Mark a task as failed with error message."""
|
|
table = fq_table("async_operations", self._schema)
|
|
# Truncate error message if too long (max 5000 chars in schema)
|
|
error_message = error_message[:5000] if len(error_message) > 5000 else error_message
|
|
await self._pool.execute(
|
|
f"""
|
|
UPDATE {table}
|
|
SET status = 'failed', error_message = $2, completed_at = now(), updated_at = now()
|
|
WHERE operation_id = $1
|
|
""",
|
|
operation_id,
|
|
error_message,
|
|
)
|
|
|
|
async def _retry_or_fail(self, operation_id: str, error_message: str):
|
|
"""Increment retry count or mark as failed if max retries exceeded."""
|
|
table = fq_table("async_operations", self._schema)
|
|
|
|
# Get current retry count
|
|
row = await self._pool.fetchrow(
|
|
f"SELECT retry_count FROM {table} WHERE operation_id = $1",
|
|
operation_id,
|
|
)
|
|
|
|
if row is None:
|
|
logger.warning(f"Operation {operation_id} not found, cannot retry")
|
|
return
|
|
|
|
retry_count = row["retry_count"]
|
|
|
|
if retry_count >= self._max_retries:
|
|
# Max retries exceeded, mark as failed
|
|
await self._mark_failed(
|
|
operation_id, f"Max retries ({self._max_retries}) exceeded. Last error: {error_message}"
|
|
)
|
|
logger.error(f"Task {operation_id} failed after {retry_count} retries")
|
|
else:
|
|
# Increment retry and reset to pending
|
|
await self._pool.execute(
|
|
f"""
|
|
UPDATE {table}
|
|
SET status = 'pending', worker_id = NULL, claimed_at = NULL,
|
|
retry_count = retry_count + 1, updated_at = now()
|
|
WHERE operation_id = $1
|
|
""",
|
|
operation_id,
|
|
)
|
|
logger.warning(f"Task {operation_id} failed, will retry (attempt {retry_count + 1}/{self._max_retries})")
|
|
|
|
async def execute_task(self, operation_id: str, task_dict: dict[str, Any]):
|
|
"""Execute a single task and update its status."""
|
|
task_type = task_dict.get("type", "unknown")
|
|
bank_id = task_dict.get("bank_id", "unknown")
|
|
|
|
try:
|
|
logger.debug(f"Executing task {operation_id} (type={task_type}, bank={bank_id})")
|
|
await self._executor(task_dict)
|
|
await self._mark_completed(operation_id)
|
|
logger.debug(f"Task {operation_id} completed successfully")
|
|
except Exception as e:
|
|
error_msg = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
|
|
logger.error(f"Task {operation_id} failed: {e}")
|
|
await self._retry_or_fail(operation_id, error_msg)
|
|
|
|
async def recover_own_tasks(self) -> int:
|
|
"""
|
|
Recover tasks that were assigned to this worker but not completed.
|
|
|
|
This handles the case where a worker crashes while processing tasks.
|
|
On startup, we reset any tasks stuck in 'processing' for this worker_id
|
|
back to 'pending' so they can be picked up again.
|
|
|
|
Returns:
|
|
Number of tasks recovered
|
|
"""
|
|
table = fq_table("async_operations", self._schema)
|
|
|
|
result = await self._pool.execute(
|
|
f"""
|
|
UPDATE {table}
|
|
SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now()
|
|
WHERE status = 'processing' AND worker_id = $1
|
|
""",
|
|
self._worker_id,
|
|
)
|
|
|
|
# Parse "UPDATE N" to get count
|
|
count = int(result.split()[-1]) if result else 0
|
|
if count > 0:
|
|
logger.info(f"Worker {self._worker_id} recovered {count} stale tasks from previous run")
|
|
return count
|
|
|
|
async def run(self):
|
|
"""
|
|
Main polling loop.
|
|
|
|
Continuously polls for pending tasks, claims them, and executes them
|
|
until shutdown is signaled.
|
|
"""
|
|
# Recover any tasks from a previous crash before starting
|
|
await self.recover_own_tasks()
|
|
|
|
logger.info(f"Worker {self._worker_id} starting polling loop")
|
|
|
|
while not self._shutdown.is_set():
|
|
try:
|
|
# Claim a batch of tasks
|
|
tasks = await self.claim_batch()
|
|
|
|
if tasks:
|
|
# Log batch info
|
|
task_types = {}
|
|
for _, task_dict in tasks:
|
|
t = task_dict.get("type", "unknown")
|
|
task_types[t] = task_types.get(t, 0) + 1
|
|
types_str = ", ".join(f"{k}:{v}" for k, v in task_types.items())
|
|
logger.info(f"Worker {self._worker_id} claimed {len(tasks)} tasks: {types_str}")
|
|
|
|
# Track in-flight tasks
|
|
async with self._in_flight_lock:
|
|
self._in_flight_count += len(tasks)
|
|
|
|
# Execute tasks concurrently
|
|
try:
|
|
await asyncio.gather(
|
|
*[self.execute_task(op_id, task_dict) for op_id, task_dict in tasks],
|
|
return_exceptions=True,
|
|
)
|
|
finally:
|
|
async with self._in_flight_lock:
|
|
self._in_flight_count -= len(tasks)
|
|
else:
|
|
# No tasks found, wait before polling again
|
|
try:
|
|
await asyncio.wait_for(
|
|
self._shutdown.wait(),
|
|
timeout=self._poll_interval_ms / 1000,
|
|
)
|
|
except asyncio.TimeoutError:
|
|
pass # Normal timeout, continue polling
|
|
|
|
# Log progress stats periodically
|
|
await self._log_progress_if_due()
|
|
|
|
except asyncio.CancelledError:
|
|
logger.info(f"Worker {self._worker_id} polling loop cancelled")
|
|
break
|
|
except Exception as e:
|
|
logger.error(f"Worker {self._worker_id} error in polling loop: {e}")
|
|
traceback.print_exc()
|
|
# Backoff on error
|
|
await asyncio.sleep(1)
|
|
|
|
logger.info(f"Worker {self._worker_id} polling loop stopped")
|
|
|
|
async def shutdown_graceful(self, timeout: float = 30.0):
|
|
"""
|
|
Signal shutdown and wait for current tasks to complete.
|
|
|
|
Args:
|
|
timeout: Maximum time to wait for in-flight tasks (seconds)
|
|
"""
|
|
logger.info(f"Worker {self._worker_id} initiating graceful shutdown")
|
|
self._shutdown.set()
|
|
|
|
# Wait for in-flight tasks to complete
|
|
start_time = asyncio.get_event_loop().time()
|
|
while asyncio.get_event_loop().time() - start_time < timeout:
|
|
async with self._in_flight_lock:
|
|
in_flight = self._in_flight_count
|
|
|
|
if in_flight == 0:
|
|
logger.info(f"Worker {self._worker_id} graceful shutdown complete")
|
|
return
|
|
|
|
logger.info(f"Worker {self._worker_id} waiting for {in_flight} in-flight tasks")
|
|
await asyncio.sleep(0.5)
|
|
|
|
logger.warning(f"Worker {self._worker_id} shutdown timeout after {timeout}s")
|
|
|
|
async def _log_progress_if_due(self):
|
|
"""Log progress stats every PROGRESS_LOG_INTERVAL seconds."""
|
|
now = time.time()
|
|
if now - self._last_progress_log < PROGRESS_LOG_INTERVAL:
|
|
return
|
|
|
|
self._last_progress_log = now
|
|
|
|
try:
|
|
table = fq_table("async_operations", self._schema)
|
|
async with self._pool.acquire() as conn:
|
|
# Get global stats by status
|
|
stats = await conn.fetch(
|
|
f"""
|
|
SELECT status, COUNT(*) as count
|
|
FROM {table}
|
|
WHERE created_at > now() - interval '24 hours'
|
|
GROUP BY status
|
|
"""
|
|
)
|
|
|
|
# Get currently processing tasks grouped by type and bank
|
|
processing = await conn.fetch(
|
|
f"""
|
|
SELECT operation_type, bank_id, COUNT(*) as count
|
|
FROM {table}
|
|
WHERE status = 'processing'
|
|
GROUP BY operation_type, bank_id
|
|
"""
|
|
)
|
|
|
|
# Build stats dict
|
|
status_counts = {row["status"]: row["count"] for row in stats}
|
|
pending = status_counts.get("pending", 0)
|
|
processing_count = status_counts.get("processing", 0)
|
|
completed = status_counts.get("completed", 0)
|
|
failed = status_counts.get("failed", 0)
|
|
|
|
# Build processing breakdown
|
|
processing_info = []
|
|
banks_working = set()
|
|
for row in processing:
|
|
op_type = row["operation_type"]
|
|
bank_id = row["bank_id"]
|
|
count = row["count"]
|
|
banks_working.add(bank_id)
|
|
processing_info.append(f"{op_type}:{bank_id}({count})")
|
|
|
|
# Format log
|
|
async with self._in_flight_lock:
|
|
in_flight = self._in_flight_count
|
|
|
|
processing_str = ", ".join(processing_info[:10]) if processing_info else "none"
|
|
if len(processing_info) > 10:
|
|
processing_str += f" +{len(processing_info) - 10} more"
|
|
|
|
logger.info(
|
|
f"[WORKER_STATS] worker={self._worker_id} in_flight={in_flight} | "
|
|
f"global: pending={pending} processing={processing_count} "
|
|
f"completed_24h={completed} failed_24h={failed} | "
|
|
f"active: {processing_str}"
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.debug(f"Failed to log progress stats: {e}")
|
|
|
|
@property
|
|
def worker_id(self) -> str:
|
|
"""Get the worker ID."""
|
|
return self._worker_id
|
|
|
|
@property
|
|
def is_shutdown(self) -> bool:
|
|
"""Check if shutdown has been signaled."""
|
|
return self._shutdown.is_set()
|