Add QueueItem cleanup logic (#228)
* feat: start queue cleanup * feat: add Queue cleanup logic * feat: mark messages as errored * fix: rename + commit * fix: rm apscheduler * fix: add default for created_at in migration * fix: consolidate error handling and only throw away first message of an errored batch * fix: move processing to within queue manager * fix: regen migration * fix: CR comment * fix: rename message -> queue_item when dealing with queue_item in the queue manager
This commit is contained in:
Rajat Ahuja
GitHub
parent
f988aae996
commit
3d0f96d339
|
|
@ -0,0 +1,94 @@
|
|||
"""add error field to QueueItem
|
||||
|
||||
Revision ID: 76ffba56fe8c
|
||||
Revises: 08894082221a
|
||||
Create Date: 2025-10-08 11:47:17.488301
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
from migrations.utils import column_exists, get_schema, index_exists
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "76ffba56fe8c"
|
||||
down_revision: str | None = "08894082221a"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
schema = get_schema()
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# Add error column
|
||||
op.add_column(
|
||||
"queue",
|
||||
sa.Column("error", sa.TEXT(), nullable=True),
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
op.add_column(
|
||||
"queue",
|
||||
sa.Column(
|
||||
"created_at",
|
||||
sa.DateTime(timezone=True),
|
||||
nullable=True,
|
||||
server_default=sa.func.now(),
|
||||
),
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# Backfill created_at for existing queue items (batched)
|
||||
op.execute(
|
||||
sa.text(
|
||||
f"""
|
||||
DO $$
|
||||
DECLARE
|
||||
rows_updated INT;
|
||||
BEGIN
|
||||
LOOP
|
||||
UPDATE "{schema}".queue
|
||||
SET created_at = NOW()
|
||||
WHERE id IN (
|
||||
SELECT id FROM "{schema}".queue
|
||||
WHERE created_at IS NULL
|
||||
LIMIT 1000
|
||||
);
|
||||
|
||||
GET DIAGNOSTICS rows_updated = ROW_COUNT;
|
||||
EXIT WHEN rows_updated = 0;
|
||||
END LOOP;
|
||||
END $$;
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
# Make created_at non-nullable
|
||||
op.alter_column("queue", "created_at", nullable=False, schema=schema)
|
||||
|
||||
# Add index on created_at
|
||||
op.create_index(
|
||||
op.f("ix_queue_created_at"),
|
||||
"queue",
|
||||
["created_at"],
|
||||
unique=False,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
inspector = sa.inspect(op.get_bind())
|
||||
|
||||
if column_exists("queue", "error", inspector):
|
||||
op.drop_column("queue", "error", schema=schema)
|
||||
|
||||
# Drop index first
|
||||
if index_exists("queue", "ix_queue_created_at", inspector):
|
||||
op.drop_index(op.f("ix_queue_created_at"), table_name="queue", schema=schema)
|
||||
|
||||
if column_exists("queue", "created_at", inspector):
|
||||
op.drop_column("queue", "created_at", schema=schema)
|
||||
|
|
@ -191,6 +191,11 @@ class DeriverSettings(HonchoSettings):
|
|||
] = 1.0
|
||||
STALE_SESSION_TIMEOUT_MINUTES: Annotated[int, Field(default=5, gt=0, le=1440)] = 5
|
||||
|
||||
# Retention window (seconds) for keeping errored items in the queue
|
||||
QUEUE_ERROR_RETENTION_SECONDS: Annotated[
|
||||
int, Field(default=30 * 24 * 3600, gt=0)
|
||||
] = 30 * 24 * 3600 # 30 days default
|
||||
|
||||
PROVIDER: SupportedProviders = "google"
|
||||
MODEL: str = "gemini-2.5-flash-lite"
|
||||
|
||||
|
|
|
|||
|
|
@ -52,6 +52,7 @@ class QueueManager:
|
|||
self.active_tasks: set[asyncio.Task[None]] = set()
|
||||
self.worker_ownership: dict[str, WorkerOwnership] = {}
|
||||
self.queue_empty_flag: asyncio.Event = asyncio.Event()
|
||||
self._maintenance_task: asyncio.Task[None] | None = None
|
||||
|
||||
# Initialize from settings
|
||||
self.workers: int = settings.DERIVER.WORKERS
|
||||
|
|
@ -115,6 +116,12 @@ class QueueManager:
|
|||
)
|
||||
logger.debug("Signal handlers registered")
|
||||
|
||||
# Start background maintenance loop
|
||||
try:
|
||||
self._maintenance_task = asyncio.create_task(self._maintenance_loop())
|
||||
except Exception:
|
||||
logger.exception("Failed to start maintenance loop")
|
||||
|
||||
# Run the polling loop directly in this task
|
||||
logger.debug("Starting polling loop directly")
|
||||
try:
|
||||
|
|
@ -162,6 +169,14 @@ class QueueManager:
|
|||
finally:
|
||||
self.worker_ownership.clear()
|
||||
|
||||
# Cancel maintenance loop if running
|
||||
if self._maintenance_task is not None:
|
||||
from contextlib import suppress
|
||||
|
||||
self._maintenance_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await self._maintenance_task
|
||||
|
||||
##########################
|
||||
# Polling and Scheduling #
|
||||
##########################
|
||||
|
|
@ -308,13 +323,99 @@ class QueueManager:
|
|||
######################
|
||||
# Queue Worker Logic #
|
||||
######################
|
||||
|
||||
async def cleanup_queue_items(self) -> None:
|
||||
"""Delete processed queue items.
|
||||
Successfully processed queue items are deleted immediately,
|
||||
while errored queue items are deleted after retention window."""
|
||||
async with tracked_db("cleanup_queue_items") as db:
|
||||
now = datetime.now(timezone.utc)
|
||||
error_cutoff = now - timedelta(
|
||||
seconds=settings.DERIVER.QUEUE_ERROR_RETENTION_SECONDS
|
||||
)
|
||||
|
||||
await db.execute(
|
||||
delete(models.QueueItem).where(
|
||||
models.QueueItem.processed
|
||||
& (
|
||||
models.QueueItem.error.is_(None)
|
||||
| (
|
||||
models.QueueItem.error.is_not(None)
|
||||
& (models.QueueItem.created_at < error_cutoff)
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def _maintenance_loop(self) -> None:
|
||||
"""Run periodic maintenance tasks on the queue."""
|
||||
try:
|
||||
while not self.shutdown_event.is_set():
|
||||
try:
|
||||
await self.cleanup_queue_items()
|
||||
except Exception:
|
||||
logger.exception("Error during maintenance cleanup")
|
||||
if settings.SENTRY.ENABLED:
|
||||
sentry_sdk.capture_exception()
|
||||
|
||||
# Sleep until interval elapses or shutdown event is set
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self.shutdown_event.wait(),
|
||||
timeout=43200, # 12 hours
|
||||
)
|
||||
break # Shutdown event set
|
||||
except asyncio.TimeoutError:
|
||||
# Timeout means it's time for next cleanup
|
||||
pass
|
||||
except asyncio.CancelledError:
|
||||
logger.debug("Maintenance loop cancelled")
|
||||
raise
|
||||
|
||||
async def _handle_processing_error(
|
||||
self,
|
||||
error: Exception,
|
||||
items: list[QueueItem],
|
||||
work_unit_key: str,
|
||||
context: str,
|
||||
) -> None:
|
||||
"""
|
||||
Handle processing errors by marking queue items as errored, logging, and forwarding to Sentry.
|
||||
We only mark the first queue item as errored so we don't potentially throw away a batch. This allows us
|
||||
to incrementally attempt to process the batch while still maintaining progress in a work unit.
|
||||
|
||||
Args:
|
||||
error: The exception that occurred
|
||||
items: The queue items that were being processed
|
||||
work_unit_key: The work unit key for the queue items
|
||||
context: Context string describing what was being processed (e.g., "processing representation batch")
|
||||
"""
|
||||
error_msg = f"{error.__class__.__name__}: {str(error)}"
|
||||
try:
|
||||
if items:
|
||||
await self.mark_queue_item_as_errored(
|
||||
items[0], work_unit_key, error_msg
|
||||
)
|
||||
except Exception as mark_error:
|
||||
logger.error(
|
||||
f"Failed to mark queue items as errored for work unit {work_unit_key}: {mark_error}",
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
logger.error(
|
||||
f"Error {context} for work unit {work_unit_key}: {error}",
|
||||
exc_info=True,
|
||||
)
|
||||
if settings.SENTRY.ENABLED:
|
||||
sentry_sdk.capture_exception(error)
|
||||
|
||||
async def process_work_unit(self, work_unit_key: str, worker_id: str) -> None:
|
||||
"""Process all messages for a specific work unit by routing to the correct handler."""
|
||||
"""Process all queue items for a specific work unit by routing to the correct handler."""
|
||||
logger.debug(f"Starting to process work unit {work_unit_key}")
|
||||
work_unit = parse_work_unit_key(work_unit_key)
|
||||
async with self.semaphore:
|
||||
message_count = 0
|
||||
messages_to_process: list[QueueItem] = []
|
||||
queue_item_count = 0
|
||||
try:
|
||||
while not self.shutdown_event.is_set():
|
||||
# Get worker ownership info for verification
|
||||
|
|
@ -329,7 +430,7 @@ class QueueManager:
|
|||
(
|
||||
messages_context,
|
||||
items_to_process,
|
||||
) = await self.get_message_batch(
|
||||
) = await self.get_queue_item_batch(
|
||||
work_unit.task_type, work_unit_key, ownership.aqs_id
|
||||
)
|
||||
logger.debug(
|
||||
|
|
@ -341,38 +442,53 @@ class QueueManager:
|
|||
)
|
||||
break
|
||||
|
||||
# Build payloads from the unique messages context window
|
||||
await process_representation_batch(
|
||||
messages_context,
|
||||
observer=work_unit.observer,
|
||||
observed=work_unit.observed,
|
||||
)
|
||||
|
||||
await self.mark_messages_as_processed(
|
||||
items_to_process, work_unit_key
|
||||
)
|
||||
message_count += len(items_to_process)
|
||||
try:
|
||||
await process_representation_batch(
|
||||
messages_context,
|
||||
observer=work_unit.observer,
|
||||
observed=work_unit.observed,
|
||||
)
|
||||
await self.mark_queue_items_as_processed(
|
||||
items_to_process, work_unit_key
|
||||
)
|
||||
queue_item_count += len(items_to_process)
|
||||
except Exception as e:
|
||||
await self._handle_processing_error(
|
||||
e,
|
||||
items_to_process,
|
||||
work_unit_key,
|
||||
"processing representation batch",
|
||||
)
|
||||
|
||||
else:
|
||||
messages_to_process = await self.get_next_message(
|
||||
queue_item = await self.get_next_queue_item(
|
||||
work_unit.task_type, work_unit_key, ownership.aqs_id
|
||||
)
|
||||
if not messages_to_process:
|
||||
if not queue_item:
|
||||
logger.debug(
|
||||
f"No more messages to process for work unit {work_unit_key} for worker {worker_id}"
|
||||
f"No more queue items to process for work unit {work_unit_key} for worker {worker_id}"
|
||||
)
|
||||
break
|
||||
await process_item(
|
||||
work_unit.task_type, messages_to_process[0].payload
|
||||
)
|
||||
await self.mark_messages_as_processed(
|
||||
messages_to_process, work_unit_key
|
||||
)
|
||||
message_count += len(messages_to_process)
|
||||
|
||||
try:
|
||||
await process_item(
|
||||
work_unit.task_type, queue_item.payload
|
||||
)
|
||||
await self.mark_queue_items_as_processed(
|
||||
[queue_item], work_unit_key
|
||||
)
|
||||
queue_item_count += 1
|
||||
except Exception as e:
|
||||
await self._handle_processing_error(
|
||||
e,
|
||||
[queue_item],
|
||||
work_unit_key,
|
||||
"processing queue item",
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error processing tasks for work unit {work_unit_key}: {e}",
|
||||
f"Error in processing loop for work unit {work_unit_key}: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
if settings.SENTRY.ENABLED:
|
||||
|
|
@ -397,7 +513,7 @@ class QueueManager:
|
|||
removed = False
|
||||
|
||||
self.untrack_worker_work_unit(worker_id, work_unit_key)
|
||||
if removed and message_count > 0:
|
||||
if removed and queue_item_count > 0:
|
||||
# Only publish webhook if we actually removed an active session
|
||||
try:
|
||||
if work_unit.task_type in ["representation", "summary"]:
|
||||
|
|
@ -425,22 +541,21 @@ class QueueManager:
|
|||
)
|
||||
|
||||
@sentry_sdk.trace
|
||||
async def get_next_message(
|
||||
async def get_next_queue_item(
|
||||
self, task_type: str, work_unit_key: str, aqs_id: str
|
||||
) -> list[QueueItem]:
|
||||
"""Get the next message to process for a specific work unit."""
|
||||
) -> QueueItem | None:
|
||||
"""Get the next queue item to process for a specific work unit."""
|
||||
if task_type == "representation":
|
||||
raise ValueError(
|
||||
"Representation tasks are not supported for get_next_message"
|
||||
"Representation tasks are not supported for get_next_queue_item"
|
||||
)
|
||||
async with tracked_db("get_next_message") as db:
|
||||
async with tracked_db("get_next_queue_item") as db:
|
||||
# ActiveQueueSession conditions for worker ownership verification
|
||||
aqs_conditions = [
|
||||
models.ActiveQueueSession.work_unit_key == work_unit_key,
|
||||
models.ActiveQueueSession.id == aqs_id,
|
||||
]
|
||||
|
||||
# For non-representation tasks, just get the next single message.
|
||||
query = (
|
||||
select(models.QueueItem)
|
||||
.join(
|
||||
|
|
@ -455,15 +570,15 @@ class QueueManager:
|
|||
.limit(1)
|
||||
)
|
||||
result = await db.execute(query)
|
||||
messages = result.scalars().all()
|
||||
queue_item = result.scalar_one_or_none()
|
||||
|
||||
# Important: commit to avoid tracked_db's rollback expiring the instance
|
||||
# We rely on expire_on_commit=False to keep attributes accessible post-close
|
||||
await db.commit()
|
||||
return list(messages)
|
||||
return queue_item
|
||||
|
||||
@sentry_sdk.trace
|
||||
async def get_message_batch(
|
||||
async def get_queue_item_batch(
|
||||
self,
|
||||
task_type: str,
|
||||
work_unit_key: str,
|
||||
|
|
@ -476,9 +591,9 @@ class QueueManager:
|
|||
"""
|
||||
if task_type != "representation":
|
||||
raise ValueError(
|
||||
"Non-representation tasks are not supported for get_message_batch"
|
||||
"Non-representation tasks are not supported for get_queue_item_batch"
|
||||
)
|
||||
async with tracked_db("get_message_batch") as db:
|
||||
async with tracked_db("get_queue_item_batch") as db:
|
||||
# For representation tasks, get a batch based on token limit.
|
||||
# Step 1: Parse work_unit_key to get session context and focused sender
|
||||
parsed_key = parse_work_unit_key(work_unit_key)
|
||||
|
|
@ -592,16 +707,16 @@ class QueueManager:
|
|||
|
||||
return messages_context, items_to_process
|
||||
|
||||
async def mark_messages_as_processed(
|
||||
self, messages: list[QueueItem], work_unit_key: str
|
||||
async def mark_queue_items_as_processed(
|
||||
self, items: list[QueueItem], work_unit_key: str
|
||||
) -> None:
|
||||
if not messages:
|
||||
if not items:
|
||||
return
|
||||
async with tracked_db("process_message_batch") as db:
|
||||
message_ids = [msg.id for msg in messages]
|
||||
async with tracked_db("process_queue_item_batch") as db:
|
||||
item_ids = [item.id for item in items]
|
||||
await db.execute(
|
||||
update(models.QueueItem)
|
||||
.where(models.QueueItem.id.in_(message_ids))
|
||||
.where(models.QueueItem.id.in_(item_ids))
|
||||
.where(models.QueueItem.work_unit_key == work_unit_key)
|
||||
.values(processed=True)
|
||||
)
|
||||
|
|
@ -612,6 +727,26 @@ class QueueManager:
|
|||
)
|
||||
await db.commit()
|
||||
|
||||
async def mark_queue_item_as_errored(
|
||||
self, item: QueueItem, work_unit_key: str, error: str
|
||||
) -> None:
|
||||
"""Mark queue item as processed with an error"""
|
||||
if not item:
|
||||
return
|
||||
async with tracked_db("mark_queue_item_as_errored") as db:
|
||||
await db.execute(
|
||||
update(models.QueueItem)
|
||||
.where(models.QueueItem.id == item.id)
|
||||
.where(models.QueueItem.work_unit_key == work_unit_key)
|
||||
.values(processed=True, error=error[:65535]) # Truncate to TEXT limit
|
||||
)
|
||||
await db.execute(
|
||||
update(models.ActiveQueueSession)
|
||||
.where(models.ActiveQueueSession.work_unit_key == work_unit_key)
|
||||
.values(last_updated=func.now())
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def _cleanup_work_unit(
|
||||
self,
|
||||
aqs_id: str,
|
||||
|
|
|
|||
|
|
@ -102,7 +102,6 @@ if SENTRY_ENABLED:
|
|||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_: FastAPI):
|
||||
# Lifespan events are now handled by the respective services
|
||||
yield
|
||||
await engine.dispose()
|
||||
|
||||
|
|
|
|||
|
|
@ -396,6 +396,10 @@ class QueueItem(Base):
|
|||
task_type: Mapped[TaskType] = mapped_column(TEXT, nullable=False)
|
||||
payload: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False)
|
||||
processed: Mapped[bool] = mapped_column(Boolean, default=False)
|
||||
error: Mapped[str | None] = mapped_column(TEXT, nullable=True)
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), index=True, default=func.now()
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"QueueItem(id={self.id}, session_id={self.session_id}, work_unit_key={self.work_unit_key}, task_type={self.task_type}, payload={self.payload}, processed={self.processed})"
|
||||
|
|
|
|||
|
|
@ -172,7 +172,7 @@ class TestQueueProcessing:
|
|||
await db_session.commit()
|
||||
await db_session.refresh(aqs)
|
||||
|
||||
_, items_to_process = await qm.get_message_batch(
|
||||
_, items_to_process = await qm.get_queue_item_batch(
|
||||
task_type="representation",
|
||||
work_unit_key=first.work_unit_key,
|
||||
aqs_id=aqs.id,
|
||||
|
|
@ -183,7 +183,7 @@ class TestQueueProcessing:
|
|||
# Mark first processed, next should be the second
|
||||
first.processed = True
|
||||
await db_session.commit()
|
||||
_, items_to_process2 = await qm.get_message_batch(
|
||||
_, items_to_process2 = await qm.get_queue_item_batch(
|
||||
task_type="representation",
|
||||
work_unit_key=first.work_unit_key,
|
||||
aqs_id=aqs.id,
|
||||
|
|
@ -461,7 +461,7 @@ class TestQueueProcessing:
|
|||
await db_session.commit()
|
||||
await db_session.refresh(alice_aqs)
|
||||
|
||||
alice_messages, alice_items = await qm.get_message_batch(
|
||||
alice_messages, alice_items = await qm.get_queue_item_batch(
|
||||
task_type="representation",
|
||||
work_unit_key=alice_work_unit_key,
|
||||
aqs_id=alice_aqs.id,
|
||||
|
|
@ -489,7 +489,7 @@ class TestQueueProcessing:
|
|||
await db_session.commit()
|
||||
await db_session.refresh(bob_aqs)
|
||||
|
||||
bob_messages, bob_items = await qm.get_message_batch(
|
||||
bob_messages, bob_items = await qm.get_queue_item_batch(
|
||||
task_type="representation",
|
||||
work_unit_key=bob_work_unit_key,
|
||||
aqs_id=bob_aqs.id,
|
||||
|
|
@ -515,7 +515,7 @@ class TestQueueProcessing:
|
|||
await db_session.commit()
|
||||
await db_session.refresh(steve_aqs)
|
||||
|
||||
steve_messages, steve_items = await qm.get_message_batch(
|
||||
steve_messages, steve_items = await qm.get_queue_item_batch(
|
||||
task_type="representation",
|
||||
work_unit_key=steve_work_unit_key,
|
||||
aqs_id=steve_aqs.id,
|
||||
|
|
@ -628,7 +628,7 @@ class TestQueueProcessing:
|
|||
await db_session.commit()
|
||||
await db_session.refresh(alice_aqs)
|
||||
|
||||
alice_messages2, _ = await qm.get_message_batch(
|
||||
alice_messages2, _ = await qm.get_queue_item_batch(
|
||||
task_type="representation",
|
||||
work_unit_key=alice_work_unit_key,
|
||||
aqs_id=alice_aqs.id,
|
||||
|
|
@ -651,7 +651,7 @@ class TestQueueProcessing:
|
|||
await db_session.commit()
|
||||
await db_session.refresh(bob_aqs)
|
||||
|
||||
bob_messages2, _ = await qm.get_message_batch(
|
||||
bob_messages2, _ = await qm.get_queue_item_batch(
|
||||
task_type="representation",
|
||||
work_unit_key=bob_work_unit_key,
|
||||
aqs_id=bob_aqs.id,
|
||||
|
|
@ -670,7 +670,7 @@ class TestQueueProcessing:
|
|||
await db_session.commit()
|
||||
await db_session.refresh(steve_aqs)
|
||||
|
||||
steve_messages2, _ = await qm.get_message_batch(
|
||||
steve_messages2, _ = await qm.get_queue_item_batch(
|
||||
task_type="representation",
|
||||
work_unit_key=steve_work_unit_key,
|
||||
aqs_id=steve_aqs.id,
|
||||
|
|
|
|||
Loading…
Reference in New Issue