diff --git a/tests/bench/beam.py b/tests/bench/beam.py index 0567802e..b61c0e81 100644 --- a/tests/bench/beam.py +++ b/tests/bench/beam.py @@ -171,7 +171,9 @@ class BEAMRunner(BaseRunner[ConversationResult]): user_peer = ctx.peers["user"] assistant_peer = ctx.peers["assistant"] - ctx.session = await ctx.honcho_client.aio.session(id=ctx.session_id) + ctx.session = await ctx.honcho_client.aio.session( + id=ctx.session_id, configuration=self._get_session_configuration() + ) await ctx.session.aio.add_peers( [ diff --git a/tests/bench/locomo.py b/tests/bench/locomo.py index f91a1285..f9400088 100644 --- a/tests/bench/locomo.py +++ b/tests/bench/locomo.py @@ -256,7 +256,9 @@ class LoCoMoRunner(BaseRunner[ConversationResult]): peer_a = ctx.peers["speaker_a"] peer_b = ctx.peers["speaker_b"] - ctx.session = await ctx.honcho_client.aio.session(id=ctx.session_id) + ctx.session = await ctx.honcho_client.aio.session( + id=ctx.session_id, configuration=self._get_session_configuration() + ) # Observe both peers since questions ask about both speakers await ctx.session.aio.add_peers( diff --git a/tests/bench/longmem.py b/tests/bench/longmem.py index 7271e189..09a417c8 100644 --- a/tests/bench/longmem.py +++ b/tests/bench/longmem.py @@ -206,7 +206,9 @@ class LongMemEvalRunner(BaseRunner[TestResult]): if self.merge_sessions: # Create a single merged session - ctx.session = await ctx.honcho_client.aio.session(id=ctx.session_id) + ctx.session = await ctx.honcho_client.aio.session( + id=ctx.session_id, configuration=self._get_session_configuration() + ) if is_assistant_type: await ctx.session.aio.add_peers( @@ -301,7 +303,9 @@ class LongMemEvalRunner(BaseRunner[TestResult]): for session_date, session_id, session_messages in zip( parsed_dates, haystack_session_ids, haystack_sessions, strict=True ): - session = await ctx.honcho_client.aio.session(id=session_id) + session = await ctx.honcho_client.aio.session( + id=session_id, configuration=self._get_session_configuration() + ) if is_assistant_type: await session.aio.add_peers( diff --git a/tests/bench/run_tests.py b/tests/bench/run_tests.py index 13ed113b..fcf141de 100644 --- a/tests/bench/run_tests.py +++ b/tests/bench/run_tests.py @@ -23,11 +23,15 @@ import tiktoken from anthropic import AsyncAnthropic from dotenv import load_dotenv from honcho import Honcho +from honcho.api_types import SessionConfiguration, SummaryConfiguration from honcho.session import SessionPeerConfig from typing_extensions import TypedDict load_dotenv() +# Default session configuration with summaries disabled for tests +_TEST_SESSION_CONFIG = SessionConfiguration(summary=SummaryConfiguration(enabled=False)) + class SessionResult(TypedDict): """Type definition for session creation results.""" @@ -319,8 +323,10 @@ Evaluate whether the actual response contains the core correct information from peers[peer_name] = await honcho_client.aio.peer(id=peer_name) for session_name, session_data in sessions.items(): - # Create session - session = await honcho_client.aio.session(id=str(session_name)) + # Create session with summaries disabled + session = await honcho_client.aio.session( + id=str(session_name), configuration=_TEST_SESSION_CONFIG + ) output_lines.append(f"\n session: {session_name}") @@ -507,7 +513,9 @@ Evaluate whether the actual response contains the core correct information from session_name = str(get_context_call["session"]) summary = get_context_call["summary"] max_tokens: int | None = get_context_call.get("max_tokens") - session = await honcho_client.aio.session(id=session_name) + session = await honcho_client.aio.session( + id=session_name, configuration=_TEST_SESSION_CONFIG + ) # Wait for deriver queue to be empty for this session # TODO implement this differently! diff --git a/tests/bench/runner_common.py b/tests/bench/runner_common.py index a12b4834..dea8e544 100644 --- a/tests/bench/runner_common.py +++ b/tests/bench/runner_common.py @@ -20,6 +20,7 @@ from typing import Any, Generic, Literal, TypeVar import redis.asyncio as aioredis from anthropic import AsyncAnthropic from honcho import Honcho +from honcho.api_types import SessionConfiguration, SummaryConfiguration from openai import AsyncOpenAI from redis.asyncio.client import Redis @@ -49,6 +50,7 @@ class RunnerConfig: api_key: str | None = None skip_dream: bool = False json_output: Path | None = None + max_concurrent: int | None = None # None means no limit (use batch_size) @classmethod def from_args( @@ -70,6 +72,7 @@ class RunnerConfig: api_key=args.api_key, skip_dream=args.skip_dream, json_output=args.json_output, + max_concurrent=args.max_concurrent, ) @@ -173,6 +176,13 @@ def add_common_arguments(parser: argparse.ArgumentParser) -> None: help="Skip the dream consolidation step (default: False)", ) + parser.add_argument( + "--max-concurrent", + type=int, + default=None, + help="Maximum concurrent items executing at once (default: unlimited, use for rate-limited remote instances)", + ) + def validate_common_arguments(args: argparse.Namespace) -> str | None: """ @@ -190,6 +200,9 @@ def validate_common_arguments(args: argparse.Namespace) -> str | None: if args.pool_size <= 0: return f"Error: Pool size must be positive, got {args.pool_size}" + if args.max_concurrent is not None and args.max_concurrent <= 0: + return f"Error: Max concurrent must be positive, got {args.max_concurrent}" + return None @@ -309,6 +322,10 @@ class BaseRunner(ABC, Generic[ResultT]): f"{self.get_metrics_prefix()}_{datetime.now().strftime('%Y%m%d_%H%M%S')}" ) self.logger: Logger = configure_logging() + # Semaphore for rate limiting concurrent item execution + self._concurrency_semaphore: asyncio.Semaphore | None = ( + asyncio.Semaphore(config.max_concurrent) if config.max_concurrent else None + ) # ------------------------------------------------------------------------- # Abstract methods - must be implemented by subclasses @@ -410,6 +427,8 @@ class BaseRunner(ABC, Generic[ResultT]): f"Distributing across {self.config.pool_size} Honcho instances " + f"(ports {self.config.base_api_port}-{self.config.base_api_port + self.config.pool_size - 1})" ) + if self.config.max_concurrent: + print(f"Limiting to {self.config.max_concurrent} concurrent item(s)") overall_start = time.time() all_results: list[ResultT] = [] @@ -425,10 +444,10 @@ class BaseRunner(ABC, Generic[ResultT]): print(f"Processing batch {batch_num}/{total_batches} ({len(batch)} items)") print(f"{'=' * 60}") - # Run items in batch concurrently + # Run items in batch concurrently (with optional rate limiting) batch_results = await asyncio.gather( *[ - self.execute_item(item, self._get_honcho_url(i + idx)) + self._execute_item_with_limit(item, self._get_honcho_url(i + idx)) for idx, item in enumerate(batch) ] ) @@ -442,6 +461,13 @@ class BaseRunner(ABC, Generic[ResultT]): return all_results, overall_duration + async def _execute_item_with_limit(self, item: Any, honcho_url: str) -> ResultT: + """Wrapper that applies concurrency limiting if configured.""" + if self._concurrency_semaphore: + async with self._concurrency_semaphore: + return await self.execute_item(item, honcho_url) + return await self.execute_item(item, honcho_url) + async def execute_item(self, item: Any, honcho_url: str) -> ResultT: """ Execute a single benchmark item. @@ -591,6 +617,10 @@ class BaseRunner(ABC, Generic[ResultT]): api_key=self.config.api_key, ) + def _get_session_configuration(self) -> SessionConfiguration: + """Get default session configuration with summaries disabled.""" + return SessionConfiguration(summary=SummaryConfiguration(enabled=False)) + async def _flush_deriver_queue(self) -> None: """Enable deriver flush mode to bypass batch token threshold.""" if self.config.base_url: