RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Track worker ownership (#216) 

* fix: simplify; use debug logs

* fix: random type error

* fix: CR comments

* fix: move limit check before opening db connection

* fix: use single SQL call to verify ownership and fetch message batch

* fix: CR comments 2
This commit is contained in:
Rajat Ahuja 2025-09-25 15:14:48 -04:00 committed by GitHub
parent 1ae2c4599a
commit 6e01af2005
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 215 additions and 81 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

192
src/deriver/queue_manager.py
View File

@ -4,9 +4,11 @@ from asyncio import Task
from collections.abc import Sequence
from datetime import datetime, timedelta, timezone
from logging import getLogger
from typing import NamedTuple
import sentry_sdk
from dotenv import load_dotenv
from nanoid import generate as generate_nanoid
from sentry_sdk.integrations.asyncio import AsyncioIntegration
from sqlalchemy import BigInteger, delete, select, update
from sqlalchemy.dialects.postgresql import insert
@ -25,11 +27,18 @@ logger = getLogger(__name__)
load_dotenv(override=True)
class WorkerOwnership(NamedTuple):
"""Represents the instance of a work unit that a worker is processing."""
work_unit_key: str
aqs_id: str # The ID of the ActiveQueueSession that the worker is processing
class QueueManager:
def __init__(self):
self.shutdown_event: asyncio.Event = asyncio.Event()
self.active_tasks: set[asyncio.Task[None]] = set()
self.owned_work_units: set[str] = set()
self.worker_ownership: dict[str, WorkerOwnership] = {}
self.queue_empty_flag: asyncio.Event = asyncio.Event()
# Initialize from settings
@ -53,13 +62,25 @@ class QueueManager:
self.active_tasks.add(task)
task.add_done_callback(self.active_tasks.discard)
def track_work_unit(self, work_unit_key: str) -> None:
"""Track a new work unit owned by this process"""
self.owned_work_units.add(work_unit_key)
def track_worker_work_unit(
self, worker_id: str, work_unit_key: str, aqs_id: str
) -> None:
"""Track a work unit owned by a specific worker"""
self.worker_ownership[worker_id] = WorkerOwnership(work_unit_key, aqs_id)
def untrack_work_unit(self, work_unit_key: str) -> None:
"""Remove a work unit from tracking"""
self.owned_work_units.discard(work_unit_key)
def untrack_worker_work_unit(self, worker_id: str, work_unit_key: str) -> None:
"""Remove a work unit from worker tracking"""
ownership = self.worker_ownership.get(worker_id)
if ownership and ownership.work_unit_key == work_unit_key:
del self.worker_ownership[worker_id]
def create_worker_id(self) -> str:
"""Generate a unique worker ID for this processing task"""
return generate_nanoid()
def get_total_owned_work_units(self) -> int:
"""Get the total number of work units owned by all workers"""
return len(self.worker_ownership)
async def initialize(self) -> None:
"""Setup signal handlers, initialize client, and start the main polling loop"""
@ -94,15 +115,19 @@ class QueueManager:
async def cleanup(self) -> None:
"""Clean up owned work units"""
if self.owned_work_units:
logger.info(f"Cleaning up {len(self.owned_work_units)} owned work units...")
total_work_units = self.get_total_owned_work_units()
if total_work_units > 0:
logger.debug(f"Cleaning up {total_work_units} owned work units...")
try:
# Use the tracked_db dependency for transaction safety
async with tracked_db("queue_cleanup") as db:
for work_unit_key in self.owned_work_units:
aqs_ids = [
ownership.aqs_id for ownership in self.worker_ownership.values()
]
if aqs_ids:
await db.execute(
delete(models.ActiveQueueSession).where(
models.ActiveQueueSession.work_unit_key == work_unit_key
models.ActiveQueueSession.id.in_(aqs_ids)
)
)
await db.commit()
@ -111,6 +136,8 @@ class QueueManager:
logger.error(f"Error during cleanup: {str(e)}")
if settings.SENTRY.ENABLED:
sentry_sdk.capture_exception(e)
finally:
self.worker_ownership.clear()
##########################
# Polling and Scheduling #
@ -145,18 +172,17 @@ class QueueManager:
)
await db.commit()
async def get_and_claim_work_units(self) -> Sequence[str]:
async def get_and_claim_work_units(self) -> dict[str, str]:
"""
Get available work units that aren't being processed.
Returns a list of work unit keys.
Returns a dict mapping work_unit_key to aqs_id.
"""
claimed_units: list[str] = []
limit: int = max(0, self.workers - self.get_total_owned_work_units())
if limit == 0:
return {}
async with tracked_db(
"get_available_work_units"
) as db: # Get number of available workers
limit: int = max(0, self.workers - len(self.owned_work_units))
query = (
select(models.QueueItem.work_unit_key)
.outerjoin(
@ -175,32 +201,38 @@ class QueueManager:
available_units = result.scalars().all()
if not available_units:
await db.commit()
return []
return {}
claimed_units = await self.claim_work_units(db, available_units)
claimed_mapping = await self.claim_work_units(db, available_units)
await db.commit()
for work_unit in claimed_units:
self.track_work_unit(work_unit)
return claimed_units
return claimed_mapping
async def claim_work_units(
self, db: AsyncSession, work_unit_keys: Sequence[str]
) -> list[str]:
) -> dict[str, str]:
"""
Claim work units and return a mapping of work_unit_key to aqs_id.
Returns only the work units that were successfully claimed.
"""
values = [{"work_unit_key": key} for key in work_unit_keys]
stmt = (
insert(models.ActiveQueueSession)
.values(values)
.on_conflict_do_nothing()
.returning(models.ActiveQueueSession.work_unit_key)
.returning(
models.ActiveQueueSession.work_unit_key, models.ActiveQueueSession.id
)
)
result = await db.execute(stmt)
claimed_units = result.scalars().all()
logger.debug(f"Claimed {len(claimed_units)} work units")
return list(claimed_units)
claimed_rows = result.all()
claimed_mapping = {row[0]: row[1] for row in claimed_rows}
logger.debug(
f"Claimed {len(claimed_mapping)} work units: {list(claimed_mapping.keys())}"
)
return claimed_mapping
async def polling_loop(self) -> None:
"""Main polling loop to find and process new work units"""
@ -221,13 +253,19 @@ class QueueManager:
try:
await self.cleanup_stale_work_units()
new_work_units = await self.get_and_claim_work_units()
if new_work_units:
for work_unit in new_work_units:
claimed_work_units = await self.get_and_claim_work_units()
if claimed_work_units:
for work_unit_key, aqs_id in claimed_work_units.items():
# Create a new task for processing this work unit
if not self.shutdown_event.is_set():
# Track worker ownership
worker_id = self.create_worker_id()
self.track_worker_work_unit(
worker_id, work_unit_key, aqs_id
)
task: Task[None] = asyncio.create_task(
self.process_work_unit(work_unit)
self.process_work_unit(work_unit_key, worker_id)
)
self.add_task(task)
else:
@ -249,9 +287,11 @@ class QueueManager:
######################
@sentry_sdk.trace
async def process_work_unit(self, work_unit_key: str):
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."""
logger.debug(f"Starting to process work unit {work_unit_key}")
logger.debug(
f"Worker {worker_id} starting to process work unit {work_unit_key}"
)
async with self.semaphore:
message_count = 0
try:
@ -259,15 +299,25 @@ class QueueManager:
task_type = parsed_key["task_type"]
while not self.shutdown_event.is_set():
messages_to_process: list[QueueItem] = await self.get_message_batch(
work_unit_key,
task_type,
)
if not messages_to_process:
logger.debug(f"No more messages for work unit {work_unit_key}")
# Get worker ownership info for verification
ownership = self.worker_ownership.get(worker_id)
if not ownership or ownership.work_unit_key != work_unit_key:
logger.warning(
f"Worker {worker_id} lost ownership of work unit {work_unit_key}, stopping processing {work_unit_key}"
)
break
# Process the batch/single item
messages_to_process: list[QueueItem] = await self.get_message_batch(
task_type, work_unit_key, ownership.aqs_id
)
logger.debug(
f"Worker {worker_id} retrieved {len(messages_to_process)} messages for work unit {work_unit_key} (AQS ID: {ownership.aqs_id})"
)
if not messages_to_process:
logger.debug(
f"No more messages to process for work unit {work_unit_key} for worker {worker_id}"
)
break
try:
payloads = [msg.payload for msg in messages_to_process]
await process_items(task_type, payloads)
@ -294,9 +344,15 @@ class QueueManager:
finally:
# Remove work unit from active_queue_sessions when done
logger.debug(f"Removing work unit {work_unit_key} from active sessions")
removed = await self._cleanup_work_unit(work_unit_key)
ownership: WorkerOwnership | None = self.worker_ownership.get(worker_id)
if ownership and ownership.work_unit_key == work_unit_key:
removed = await self._cleanup_work_unit(
ownership.aqs_id, work_unit_key
)
else:
removed = False
self.untrack_worker_work_unit(worker_id, work_unit_key)
if removed and message_count > 0:
# Only publish webhook if we actually removed an active session
try:
@ -307,7 +363,7 @@ class QueueManager:
parsed_key = parse_work_unit_key(work_unit_key)
if parsed_key["task_type"] in ["representation", "summary"]:
logger.info(
logger.debug(
f"Publishing queue.empty event for {work_unit_key}"
)
await publish_webhook_event(
@ -330,24 +386,42 @@ class QueueManager:
f"Work unit {work_unit_key} already cleaned up by another worker, skipping webhook"
)
self.untrack_work_unit(work_unit_key)
@sentry_sdk.trace
async def get_message_batch(
self, work_unit_key: str, task_type: str
self, task_type: str, work_unit_key: str, aqs_id: str
) -> list[QueueItem]:
"""
Get a batch of unprocessed messages for a specific work unit ordered by id.
For representation tasks, this will be a batch of messages up to REPRESENTATION_BATCH_MAX_TOKENS.
For other tasks, it will be a single message.
Args:
task_type: The type of task to process
work_unit_key: The key of the work unit to process
aqs_id: The ID of the active queue session to process
Returns:
A list of QueueItem objects
"""
async with tracked_db("get_message_batch") as db:
# ActiveQueueSession conditions for worker ownership verification
aqs_conditions = [
models.ActiveQueueSession.work_unit_key == work_unit_key,
models.ActiveQueueSession.id == aqs_id,
]
if task_type != "representation":
# For non-representation tasks, just get the next single message.
query = (
select(models.QueueItem)
.join(
models.ActiveQueueSession,
models.QueueItem.work_unit_key
== models.ActiveQueueSession.work_unit_key,
)
.where(models.QueueItem.work_unit_key == work_unit_key)
.where(~models.QueueItem.processed)
.where(*aqs_conditions)
.order_by(models.QueueItem.id)
.limit(1)
)
@ -389,8 +463,14 @@ class QueueManager:
# Select messages where either:
# 1. It's the first message (row_num = 1), OR
# 2. The cumulative token count is within the limit
# Also ensure worker ownership verification by joining with ActiveQueueSession
query = (
select(models.QueueItem)
.join(
models.ActiveQueueSession,
models.QueueItem.work_unit_key
== models.ActiveQueueSession.work_unit_key,
)
.where(
models.QueueItem.id.in_(
select(cte.c.id).where(
@ -402,6 +482,7 @@ class QueueManager:
)
)
)
.where(*aqs_conditions)
.order_by(models.QueueItem.id)
)
@ -415,7 +496,7 @@ class QueueManager:
async def mark_messages_as_processed(
self, messages: list[QueueItem], work_unit_key: str
):
) -> None:
if not messages:
return
async with tracked_db("process_message_batch") as db:
@ -432,12 +513,19 @@ class QueueManager:
)
await db.commit()
async def _cleanup_work_unit(self, work_unit_key: str) -> bool:
async def _cleanup_work_unit(
self,
aqs_id: str,
work_unit_key: str,
) -> bool:
"""
Clean up a specific work unit session by both work_unit_key and AQS ID.
"""
async with tracked_db("cleanup_work_unit") as db:
result = await db.execute(
delete(models.ActiveQueueSession).where(
models.ActiveQueueSession.work_unit_key == work_unit_key
)
delete(models.ActiveQueueSession)
.where(models.ActiveQueueSession.id == aqs_id)
.where(models.ActiveQueueSession.work_unit_key == work_unit_key)
)
await db.commit()
return result.rowcount > 0

3
tests/bench/harness.py
View File

@ -563,7 +563,8 @@ except Exception as e:
# Create an empty .env file in temp directory to satisfy docker-compose
# (even though the database service doesn't actually use it)
(self.temp_dir / ".env").touch()
if self.temp_dir:
(self.temp_dir / ".env").touch()
# Start database
self.start_database()

101
tests/deriver/test_queue_processing.py
View File

@ -2,11 +2,12 @@ from collections.abc import Callable
from typing import Any
import pytest
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from src import models
from src.config import settings
from src.deriver.queue_manager import QueueManager
from src.deriver.queue_manager import QueueManager, WorkerOwnership
@pytest.mark.asyncio
@ -58,9 +59,8 @@ class TestQueueProcessing:
assert len(work_units) > 0
# The API already claimed returned units; verify it's tracked and not returned again
from sqlalchemy import select
work_unit = work_units[0]
work_unit = next(iter(work_units.keys()))
tracked = (
await db_session.execute(
select(models.ActiveQueueSession).where(
@ -99,11 +99,11 @@ class TestQueueProcessing:
queue_manager = QueueManager()
claimed = await queue_manager.get_and_claim_work_units()
assert len(claimed) > 0
key = claimed[0]
key = list(claimed.keys())[0]
# Trying to claim the same key again via the API should return empty list
# Trying to claim the same key again via the API should return empty dict
claimed_again = await queue_manager.claim_work_units(db_session, [key])
assert claimed_again == []
assert claimed_again == {}
@pytest.mark.asyncio
async def test_get_next_message_orders_and_filters_simple(
@ -113,8 +113,6 @@ class TestQueueProcessing:
create_queue_payload: Callable[..., Any],
add_queue_items: Callable[..., Any],
) -> None:
from sqlalchemy import select
session, peers = sample_session_with_peers
peer = peers[0]
@ -163,9 +161,18 @@ class TestQueueProcessing:
first, second = ordered[0], ordered[1]
qm = QueueManager()
aqs = models.ActiveQueueSession(
work_unit_key=first.work_unit_key,
)
db_session.add(aqs)
await db_session.commit()
await db_session.refresh(aqs)
batch = await qm.get_message_batch(
first.work_unit_key,
task_type="representation",
work_unit_key=first.work_unit_key,
aqs_id=aqs.id,
)
nxt = batch[0] if batch else None
assert nxt is not None and nxt.id == first.id
@ -174,8 +181,9 @@ class TestQueueProcessing:
first.processed = True
await db_session.commit()
batch2 = await qm.get_message_batch(
first.work_unit_key,
task_type="representation",
work_unit_key=first.work_unit_key,
aqs_id=aqs.id,
)
nxt2 = batch2[0] if batch2 else None
assert nxt2 is not None and nxt2.id == second.id
@ -186,14 +194,13 @@ class TestQueueProcessing:
sample_queue_items: list[models.QueueItem], # noqa: ARG001 # pyright: ignore[reportUnusedParameter]
db_session: AsyncSession,
) -> None:
from sqlalchemy import select
qm = QueueManager()
claimed = await qm.get_and_claim_work_units()
assert len(claimed) > 0
key = claimed[0]
key = list(claimed.keys())[0]
aqs_id = claimed[key]
removed = await qm._cleanup_work_unit(key) # pyright: ignore[reportPrivateUsage]
removed = await qm._cleanup_work_unit(aqs_id, key) # pyright: ignore[reportPrivateUsage]
assert removed is True
remaining = (
@ -225,7 +232,7 @@ class TestQueueProcessing:
work_units = await queue_manager.get_and_claim_work_units()
# This test ensures the cleanup logic doesn't break, though we don't have stale entries yet
assert isinstance(work_units, list)
assert isinstance(work_units, dict)
async def test_work_unit_key_format(
self, sample_session_with_peers: tuple[models.Session, list[models.Peer]]
@ -340,10 +347,21 @@ class TestQueueProcessing:
# Process work unit and verify batching
qm = QueueManager()
work_unit_key = queue_items[0].work_unit_key
worker_id = "test_worker"
# Manually claim and assign ownership
claimed_units = await qm.claim_work_units(db_session, [work_unit_key])
aqs_id = claimed_units[work_unit_key]
qm.worker_ownership[worker_id] = WorkerOwnership(
work_unit_key=work_unit_key, aqs_id=aqs_id
)
with patch(
"src.deriver.queue_manager.process_items", side_effect=mock_process_items
"src.deriver.queue_manager.process_items",
side_effect=mock_process_items,
):
await qm.process_work_unit(queue_items[0].work_unit_key)
await qm.process_work_unit(work_unit_key, worker_id)
# Should create 2 batches due to token limits
assert len(processed_batches) == 2
@ -418,19 +436,26 @@ class TestQueueProcessing:
qm = QueueManager()
work_unit_key = queue_items[0].work_unit_key
worker_id = "test_worker"
# Manually claim and assign ownership
claimed_units = await qm.claim_work_units(db_session, [work_unit_key])
aqs_id = claimed_units[work_unit_key]
qm.worker_ownership[worker_id] = WorkerOwnership(
work_unit_key=work_unit_key, aqs_id=aqs_id
)
with patch(
"src.deriver.queue_manager.process_items", side_effect=mock_process_items
"src.deriver.queue_manager.process_items",
side_effect=mock_process_items,
):
await qm.process_work_unit(work_unit_key)
await qm.process_work_unit(work_unit_key, worker_id)
# Verify both messages were processed in separate batches
assert len(processed_batches) == 2
assert all(batch["task_type"] == "summary" for batch in processed_batches)
assert all(batch["payload_count"] == 1 for batch in processed_batches)
# Verify the corresponding DB records are marked as processed
from sqlalchemy import select
# Query for the summary queue items that were processed
processed_items = (
(
@ -540,12 +565,22 @@ class TestQueueProcessing:
}
)
# Process work unit and verify batching
qm = QueueManager()
work_unit_key = queue_items[0].work_unit_key
worker_id = "test_worker"
# Manually claim and assign ownership
claimed_units = await qm.claim_work_units(db_session, [work_unit_key])
aqs_id = claimed_units[work_unit_key]
qm.worker_ownership[worker_id] = WorkerOwnership(
work_unit_key=work_unit_key, aqs_id=aqs_id
)
with patch(
"src.deriver.queue_manager.process_items", side_effect=mock_process_items
"src.deriver.queue_manager.process_items",
side_effect=mock_process_items,
):
await qm.process_work_unit(queue_items[0].work_unit_key)
await qm.process_work_unit(work_unit_key, worker_id)
# Should create 2 batches: first large message alone, then second and third together
assert len(processed_batches) == 2
@ -641,12 +676,22 @@ class TestQueueProcessing:
}
)
# Process work unit and verify batching
qm = QueueManager()
work_unit_key = queue_items[0].work_unit_key
worker_id = "test_worker"
# Manually claim and assign ownership
claimed_units = await qm.claim_work_units(db_session, [work_unit_key])
aqs_id = claimed_units[work_unit_key]
qm.worker_ownership[worker_id] = WorkerOwnership(
work_unit_key=work_unit_key, aqs_id=aqs_id
)
with patch(
"src.deriver.queue_manager.process_items", side_effect=mock_process_items
"src.deriver.queue_manager.process_items",
side_effect=mock_process_items,
):
await qm.process_work_unit(queue_items[0].work_unit_key)
await qm.process_work_unit(work_unit_key, worker_id)
# Should create 2 batches: first two messages together (exactly at limit), third alone
assert len(processed_batches) == 2