From 414e31c960d57623f39d9f10892d74d8dda0e9d9 Mon Sep 17 00:00:00 2001 From: Harish Kukreja Date: Mon, 22 Jun 2026 17:49:08 -0400 Subject: [PATCH] feat(deriver): age-flush stalled representation batches (#826) Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> --- .env.template | 1 + CHANGELOG.md | 4 + config.toml.example | 1 + docs/v3/contributing/configuration.mdx | 1 + docs/v3/contributing/troubleshooting.mdx | 2 +- .../core-concepts/design-patterns.mdx | 2 +- .../guides/recipes/unified-memory-setup.mdx | 11 +- src/config.py | 5 + src/deriver/queue_manager.py | 56 ++++- tests/deriver/test_queue_processing.py | 227 ++++++++++++++++++ tests/test_config.py | 14 ++ 11 files changed, 309 insertions(+), 15 deletions(-) diff --git a/.env.template b/.env.template index 0b4a6f99..451def53 100644 --- a/.env.template +++ b/.env.template @@ -129,6 +129,7 @@ LLM_OPENAI_API_KEY=your-api-key-here # DERIVER_MAX_CUSTOM_INSTRUCTIONS_TOKENS=2000 # DERIVER_WORKING_REPRESENTATION_MAX_OBSERVATIONS=100 # DERIVER_REPRESENTATION_BATCH_MAX_TOKENS=1024 +# DERIVER_REPRESENTATION_BATCH_MAX_AGE_SECONDS=1800 # DERIVER_FLUSH_ENABLED=false # Bypass batch token threshold, process work immediately # DERIVER_MODEL_CONFIG__FALLBACK__MODEL= # DERIVER_MODEL_CONFIG__FALLBACK__TRANSPORT= diff --git a/CHANGELOG.md b/CHANGELOG.md index d31b9d72..6bc93dc8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](http://semver.org/). ## [Unreleased] +### Added + +- `DERIVER_REPRESENTATION_BATCH_MAX_AGE_SECONDS` (default 1800s) lets sub-threshold representation work units flush once their oldest unprocessed queue item ages out. Set it to `0` to keep the legacy behavior where sub-threshold tails wait indefinitely unless `DERIVER_FLUSH_ENABLED=true`. + ### Changed - Peer-scoped JWTs now get read-only access to the sessions their peer is an active member of (session context, summaries, peers, their own per-session config, search, and message reads). Session-scoped JWTs remain confined to their session and cannot reach peer routes. diff --git a/config.toml.example b/config.toml.example index 3aacfbb7..9e0e9c11 100644 --- a/config.toml.example +++ b/config.toml.example @@ -109,6 +109,7 @@ MAX_INPUT_TOKENS = 25000 MAX_CUSTOM_INSTRUCTIONS_TOKENS = 2000 WORKING_REPRESENTATION_MAX_OBSERVATIONS = 100 REPRESENTATION_BATCH_MAX_TOKENS = 1024 +REPRESENTATION_BATCH_MAX_AGE_SECONDS = 1800 FLUSH_ENABLED = false # Bypass batch token threshold, process work immediately [deriver.model_config] diff --git a/docs/v3/contributing/configuration.mdx b/docs/v3/contributing/configuration.mdx index 06add0d0..cd74a829 100644 --- a/docs/v3/contributing/configuration.mdx +++ b/docs/v3/contributing/configuration.mdx @@ -402,6 +402,7 @@ DERIVER_DEDUPLICATE=true DERIVER_LOG_OBSERVATIONS=false DERIVER_WORKING_REPRESENTATION_MAX_OBSERVATIONS=100 DERIVER_REPRESENTATION_BATCH_MAX_TOKENS=1024 +DERIVER_REPRESENTATION_BATCH_MAX_AGE_SECONDS=1800 ``` **Peer Card:** diff --git a/docs/v3/contributing/troubleshooting.mdx b/docs/v3/contributing/troubleshooting.mdx index bb71a475..7bfc6290 100644 --- a/docs/v3/contributing/troubleshooting.mdx +++ b/docs/v3/contributing/troubleshooting.mdx @@ -109,7 +109,7 @@ Messages are stored but no observations, summaries, or representations are being ```bash DERIVER_WORKERS=4 ``` -5. **Representation Batch Max** — By default the deriver is set to buffer its operations until there are enough tokens for a given representation in a session. This is set via the `REPRESENTATION_BATCH_MAX_TOKENS` environment variable. If you aren't seeing tasks continue it may be that the batch size is set too high or enough data hasn't flowed into to the session yet. See [token batching](/v3/documentation/core-concepts/reasoning#token-batching) for more details +5. **Representation Batch Max** — By default the deriver buffers representation work until a session has enough tokens for that representation, set via `DERIVER_REPRESENTATION_BATCH_MAX_TOKENS`. Sub-threshold tails become eligible after `DERIVER_REPRESENTATION_BATCH_MAX_AGE_SECONDS` (default 1800 seconds), so quiet sessions eventually flush without disabling batching globally. Set the age to `0` for legacy behavior where sub-threshold tails wait indefinitely. See [token batching](/v3/documentation/core-concepts/reasoning#token-batching) for more details ## Alternative Provider Issues diff --git a/docs/v3/documentation/core-concepts/design-patterns.mdx b/docs/v3/documentation/core-concepts/design-patterns.mdx index 06582d80..9893d3f4 100644 --- a/docs/v3/documentation/core-concepts/design-patterns.mdx +++ b/docs/v3/documentation/core-concepts/design-patterns.mdx @@ -63,7 +63,7 @@ Sessions define the temporal boundaries of an interaction. How you scope them af Create a **new** session when context resets (new conversation, new day, new topic); **reuse** one when context should keep accumulating (ongoing channel, persistent thread). -**Don't scope sessions too thin.** Honcho only reasons over a peer once it accumulates ~1,000 tokens *within a single session* ([token batching](/v3/documentation/core-concepts/reasoning#token-batching)). Many tiny sessions each stall below that threshold, so low-volume or trickle inputs should append to one ongoing session rather than fragment across many (nothing is lost — it just waits). +**Don't scope sessions too thin.** Honcho batches reasoning until a peer accumulates ~1,000 tokens *within a single session*, with a default age-based flush for quiet tails ([token batching](/v3/documentation/core-concepts/reasoning#token-batching)). Low-volume or trickle inputs should still append to one ongoing session rather than fragment across many, so reasoning runs with useful context instead of many small delayed batches. **How cross-session reasoning works** diff --git a/docs/v3/guides/recipes/unified-memory-setup.mdx b/docs/v3/guides/recipes/unified-memory-setup.mdx index 1d547db2..ea921cc3 100644 --- a/docs/v3/guides/recipes/unified-memory-setup.mdx +++ b/docs/v3/guides/recipes/unified-memory-setup.mdx @@ -131,16 +131,17 @@ for i in range(0, len(messages), 100): session.add_messages(messages[i:i + 100]) ``` -Honcho only reasons over a peer once it accumulates ~1,000 tokens *within a single session* -([token batching](/v3/documentation/core-concepts/reasoning#token-batching)). Scope -the session to the volume you ingest: +Honcho batches reasoning until a peer accumulates ~1,000 tokens *within a single session*, +with a default age-based flush for quiet tails +([token batching](/v3/documentation/core-concepts/reasoning#token-batching)). Scope the +session to the volume you ingest: - **High-volume runs** (a day of emails, a CRM export) clear the threshold easily — a per-run session like `email-import-{date}` is fine. - **Low-volume or trickle imports** (a few short records at a time) should append to one **ongoing per-source session** (e.g. `email-import-gmail`), so content - accumulates across runs instead of fragmenting into thin sessions that each stall - below the threshold (nothing is lost — it just waits). + accumulates across runs instead of fragmenting into thin sessions that each flush + later with little context. The [Gmail](/v3/guides/gmail) and [Granola](/v3/guides/granola) guides are related import examples. diff --git a/src/config.py b/src/config.py index 267a8154..3c20b34f 100644 --- a/src/config.py +++ b/src/config.py @@ -830,6 +830,11 @@ class DeriverSettings(HonchoSettings): int, Field(default=1024, ge=128, le=16_384), ] = 1024 + # Sub-threshold work units become eligible once their oldest unprocessed + # item exceeds this age. 0 disables age-based flushing. + REPRESENTATION_BATCH_MAX_AGE_SECONDS: Annotated[int, Field(default=1800, ge=0)] = ( + 1800 + ) # When enabled, bypasses the batch token threshold and processes work immediately FLUSH_ENABLED: bool = False diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index e97e67ed..93bed1a1 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -330,8 +330,10 @@ class QueueManager: async def get_and_claim_work_units(self) -> dict[str, str]: """ Get available work units that aren't being processed. - For representation tasks, only returns work units with accumulated tokens - >= REPRESENTATION_BATCH_MAX_TOKENS (forced batching), unless FLUSH_ENABLED is True. + For representation tasks, only returns work units whose accumulated + tokens reach REPRESENTATION_BATCH_MAX_TOKENS or whose oldest pending + item exceeds REPRESENTATION_BATCH_MAX_AGE_SECONDS, unless + FLUSH_ENABLED is True. Returns a dict mapping work_unit_key to aqs_id. """ limit: int = max(0, self.workers - self.get_total_owned_work_units()) @@ -346,6 +348,7 @@ class QueueManager: select( models.QueueItem.work_unit_key, func.sum(models.Message.token_count).label("total_tokens"), + func.min(models.QueueItem.created_at).label("oldest_created_at"), ) .join( models.Message, @@ -358,15 +361,21 @@ class QueueManager: ) work_units_subq = ( - select(models.QueueItem.work_unit_key) + select( + models.QueueItem.work_unit_key, + func.min(models.QueueItem.created_at).label("oldest_created_at"), + ) .where(~models.QueueItem.processed) .group_by(models.QueueItem.work_unit_key) .subquery() ) query = ( - select(work_units_subq.c.work_unit_key) - .limit(limit) + select( + work_units_subq.c.work_unit_key, + token_stats_subq.c.total_tokens, + token_stats_subq.c.oldest_created_at, + ) .outerjoin( token_stats_subq, work_units_subq.c.work_unit_key == token_stats_subq.c.work_unit_key, @@ -379,22 +388,53 @@ class QueueManager: ) .exists() ) + .order_by( + work_units_subq.c.oldest_created_at.asc(), + work_units_subq.c.work_unit_key.asc(), + ) + .limit(limit) ) # Apply batch threshold filter (skip if FLUSH_ENABLED is True) if not settings.DERIVER.FLUSH_ENABLED and batch_max_tokens > 0: + max_age_seconds = settings.DERIVER.REPRESENTATION_BATCH_MAX_AGE_SECONDS + threshold_clause = ( + func.coalesce(token_stats_subq.c.total_tokens, 0) + >= batch_max_tokens + ) + if max_age_seconds > 0: + threshold_clause = or_( + threshold_clause, + token_stats_subq.c.oldest_created_at + <= func.now() - timedelta(seconds=max_age_seconds), + ) query = query.where( or_( ~work_units_subq.c.work_unit_key.startswith( representation_prefix ), - func.coalesce(token_stats_subq.c.total_tokens, 0) - >= batch_max_tokens, + threshold_clause, ) ) result = await db.execute(query) - available_units = result.scalars().all() + available_rows = result.all() + available_units: list[str] = [] + for work_unit_key, total_tokens, oldest_created_at in available_rows: + available_units.append(work_unit_key) + if ( + not settings.DERIVER.FLUSH_ENABLED + and settings.DERIVER.REPRESENTATION_BATCH_MAX_AGE_SECONDS > 0 + and work_unit_key.startswith(representation_prefix) + and int(total_tokens or 0) < batch_max_tokens + ): + logger.info( + "age-flushing work unit %s (tokens=%s < %s, oldest=%s)", + work_unit_key, + total_tokens or 0, + batch_max_tokens, + oldest_created_at, + ) if not available_units: await db.commit() return {} diff --git a/tests/deriver/test_queue_processing.py b/tests/deriver/test_queue_processing.py index 4c658d97..68d4852c 100644 --- a/tests/deriver/test_queue_processing.py +++ b/tests/deriver/test_queue_processing.py @@ -1,5 +1,6 @@ import asyncio from collections.abc import Callable +from datetime import datetime, timedelta, timezone from typing import Any from unittest.mock import patch @@ -18,6 +19,71 @@ from src.utils.work_unit import construct_work_unit_key class TestQueueProcessing: """Test suite for queue processing functionality""" + async def _add_representation_work_unit( + self, + *, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + token_counts: list[int], + created_ats: list[datetime] | None = None, + ) -> tuple[str, list[models.QueueItem]]: + session, peers = sample_session_with_peers + peer = peers[0] + + messages: list[models.Message] = [] + for index, token_count in enumerate(token_counts): + message = models.Message( + session_name=session.name, + workspace_name=session.workspace_name, + peer_name=peer.name, + content=f"Message {index}", + token_count=token_count, + seq_in_session=index + 1, + ) + db_session.add(message) + messages.append(message) + + await db_session.commit() + for message in messages: + await db_session.refresh(message) + + work_unit_key = "" + queue_items: list[models.QueueItem] = [] + for index, message in enumerate(messages): + payload = create_queue_payload( + message=message, + task_type="representation", + observed=peer.name, + observer=peer.name, + ) + work_unit_key = work_unit_key or construct_work_unit_key( + session.workspace_name, payload + ) + + queue_item_kwargs: dict[str, Any] = {} + if created_ats: + queue_item_kwargs["created_at"] = created_ats[index] + + queue_item = models.QueueItem( + session_id=session.id, + task_type="representation", + work_unit_key=work_unit_key, + payload=payload, + processed=False, + workspace_name=session.workspace_name, + message_id=message.id, + **queue_item_kwargs, + ) + db_session.add(queue_item) + queue_items.append(queue_item) + + await db_session.commit() + for queue_item in queue_items: + await db_session.refresh(queue_item) + + return work_unit_key, queue_items + async def test_get_and_claim_work_units( self, db_session: AsyncSession, @@ -1413,6 +1479,167 @@ class TestQueueProcessing: claimed2 = await qm.get_and_claim_work_units() assert rep_work_unit_key in claimed2 + @pytest.mark.asyncio + async def test_age_flush_waits_for_fresh_sub_threshold_items( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(settings.DERIVER, "FLUSH_ENABLED", False) + monkeypatch.setattr( + settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 1800 + ) + + work_unit_key, _queue_items = await self._add_representation_work_unit( + db_session=db_session, + sample_session_with_peers=sample_session_with_peers, + create_queue_payload=create_queue_payload, + token_counts=[100, 100, 100], + ) + + claimed = await QueueManager().get_and_claim_work_units() + + assert work_unit_key not in claimed + + @pytest.mark.asyncio + async def test_age_flush_claims_old_sub_threshold_items_and_fetches_tail( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(settings.DERIVER, "FLUSH_ENABLED", False) + monkeypatch.setattr( + settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 1800 + ) + old_timestamp = datetime.now(timezone.utc) - timedelta(hours=2) + + work_unit_key, queue_items = await self._add_representation_work_unit( + db_session=db_session, + sample_session_with_peers=sample_session_with_peers, + create_queue_payload=create_queue_payload, + token_counts=[100, 100, 100], + created_ats=[old_timestamp, old_timestamp, old_timestamp], + ) + + qm = QueueManager() + claimed = await qm.get_and_claim_work_units() + + assert work_unit_key in claimed + batch = await qm.get_queue_item_batch( + task_type="representation", + work_unit_key=work_unit_key, + aqs_id=claimed[work_unit_key], + ) + assert [item.id for item in batch.items_to_process] == [ + item.id for item in queue_items + ] + + @pytest.mark.asyncio + async def test_age_flush_zero_preserves_legacy_wait_for_old_items( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(settings.DERIVER, "FLUSH_ENABLED", False) + monkeypatch.setattr(settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 0) + old_timestamp = datetime.now(timezone.utc) - timedelta(hours=2) + + work_unit_key, _queue_items = await self._add_representation_work_unit( + db_session=db_session, + sample_session_with_peers=sample_session_with_peers, + create_queue_payload=create_queue_payload, + token_counts=[100, 100], + created_ats=[old_timestamp, old_timestamp], + ) + + claimed = await QueueManager().get_and_claim_work_units() + + assert work_unit_key not in claimed + + @pytest.mark.asyncio + async def test_flush_enabled_bypasses_age_and_token_thresholds( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(settings.DERIVER, "FLUSH_ENABLED", True) + monkeypatch.setattr( + settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 1800 + ) + + work_unit_key, _queue_items = await self._add_representation_work_unit( + db_session=db_session, + sample_session_with_peers=sample_session_with_peers, + create_queue_payload=create_queue_payload, + token_counts=[100, 100], + ) + + claimed = await QueueManager().get_and_claim_work_units() + + assert work_unit_key in claimed + + @pytest.mark.asyncio + async def test_age_flush_uses_oldest_unprocessed_item( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(settings.DERIVER, "FLUSH_ENABLED", False) + monkeypatch.setattr( + settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 1800 + ) + now = datetime.now(timezone.utc) + + work_unit_key, _queue_items = await self._add_representation_work_unit( + db_session=db_session, + sample_session_with_peers=sample_session_with_peers, + create_queue_payload=create_queue_payload, + token_counts=[100, 100], + created_ats=[now - timedelta(hours=2), now], + ) + + claimed = await QueueManager().get_and_claim_work_units() + + assert work_unit_key in claimed + + @pytest.mark.asyncio + async def test_age_flush_ignores_old_processed_items( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr(settings.DERIVER, "FLUSH_ENABLED", False) + monkeypatch.setattr( + settings.DERIVER, "REPRESENTATION_BATCH_MAX_AGE_SECONDS", 1800 + ) + now = datetime.now(timezone.utc) + + work_unit_key, queue_items = await self._add_representation_work_unit( + db_session=db_session, + sample_session_with_peers=sample_session_with_peers, + create_queue_payload=create_queue_payload, + token_counts=[100, 100], + created_ats=[now - timedelta(hours=2), now], + ) + queue_items[0].processed = True + await db_session.commit() + + claimed = await QueueManager().get_and_claim_work_units() + + assert work_unit_key not in claimed + @pytest.mark.asyncio async def test_forced_batching_single_large_message( self, diff --git a/tests/test_config.py b/tests/test_config.py index 86ea0994..28b58d81 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -8,6 +8,7 @@ def _make_deriver_settings( MAX_INPUT_TOKENS: int = 25000, MAX_CUSTOM_INSTRUCTIONS_TOKENS: int = 2000, REPRESENTATION_BATCH_MAX_TOKENS: int = 1024, + REPRESENTATION_BATCH_MAX_AGE_SECONDS: int = 1800, ) -> DeriverSettings: return DeriverSettings( MODEL_CONFIG=ConfiguredModelSettings( @@ -17,6 +18,7 @@ def _make_deriver_settings( MAX_INPUT_TOKENS=MAX_INPUT_TOKENS, MAX_CUSTOM_INSTRUCTIONS_TOKENS=MAX_CUSTOM_INSTRUCTIONS_TOKENS, REPRESENTATION_BATCH_MAX_TOKENS=REPRESENTATION_BATCH_MAX_TOKENS, + REPRESENTATION_BATCH_MAX_AGE_SECONDS=REPRESENTATION_BATCH_MAX_AGE_SECONDS, ) @@ -25,6 +27,7 @@ def test_deriver_defaults_enable_custom_instructions_at_supported_cap() -> None: assert settings.MAX_INPUT_TOKENS == 25000 assert settings.MAX_CUSTOM_INSTRUCTIONS_TOKENS == 2000 + assert settings.REPRESENTATION_BATCH_MAX_AGE_SECONDS == 1800 def test_custom_instructions_tokens_can_be_disabled_with_zero() -> None: @@ -36,3 +39,14 @@ def test_custom_instructions_tokens_can_be_disabled_with_zero() -> None: def test_custom_instructions_tokens_cannot_exceed_supported_cap() -> None: with pytest.raises(ValueError, match="less than or equal to 2000"): _make_deriver_settings(MAX_CUSTOM_INSTRUCTIONS_TOKENS=2001) + + +def test_representation_batch_age_can_be_disabled_with_zero() -> None: + settings = _make_deriver_settings(REPRESENTATION_BATCH_MAX_AGE_SECONDS=0) + + assert settings.REPRESENTATION_BATCH_MAX_AGE_SECONDS == 0 + + +def test_representation_batch_age_rejects_negative_values() -> None: + with pytest.raises(ValueError, match="greater than or equal to 0"): + _make_deriver_settings(REPRESENTATION_BATCH_MAX_AGE_SECONDS=-1)