`` that updates in-place via
+ ``IPython.display``.
+
+ Args:
+ stats (dict[str, str]): Mapping of label to formatted value.
+ extra_prefix (str): Optional text shown above the table.
+ """
+ from IPython.display import HTML, display
+
+ groups = _group_stats(stats, self._groups)
+ # All columns rendered side-by-side in a single table row
+ max_rows = max(len(items) for items in groups.values())
+
+ col_htmls = []
+ for group_name, items in groups.items():
+ col = (
+ f"| "
+ f"{group_name} | "
+ )
+ rows = []
+ for k, v in items:
+ rows.append(
+ f""
+ f"{k}"
+ f" "
+ f"{v}"
+ f" | "
+ )
+ # Pad shorter columns
+ for _ in range(max_rows - len(items)):
+ rows.append(" | ")
+ col_htmls.append((col, rows))
+
+ # Assemble: header row, then data rows
+ header_row = "" + "".join(c[0] for c in col_htmls) + "
"
+ data_rows = ""
+ for i in range(max_rows):
+ data_rows += "" + "".join(c[1][i] for c in col_htmls) + "
"
+
+ prefix_html = (
+ f""
+ f"{extra_prefix}
"
+ if extra_prefix
+ else ""
+ )
+ html = (
+ f"{prefix_html}"
+ f""
+ f"{header_row}{data_rows}
"
+ )
+
+ if self._handle is None:
+ self._handle = display(HTML(html), display_id=True)
+ else:
+ self._handle.update(HTML(html))
+
+ def _update_terminal(self, stats: dict[str, str], extra_prefix: str) -> None:
+ """Render stats as grouped text lines in a terminal.
+
+ Uses ANSI escape codes to erase the previous output and overwrite it
+ with the updated stats, one line per group.
+
+ Args:
+ stats (dict[str, str]): Mapping of label to formatted value.
+ extra_prefix (str): Optional text shown on the first line.
+ """
+ # Erase previous output
+ if self._last_line_count > 0:
+ sys.stderr.write(f"\033[{self._last_line_count}A\033[J")
+
+ groups = _group_stats(stats, self._groups)
+ lines = []
+ if extra_prefix:
+ lines.append(f" {extra_prefix}")
+ for group_name, items in groups.items():
+ values = " ".join(f"{k}={v}" for k, v in items)
+ lines.append(f" {group_name}: {values}")
+
+ output = "\n".join(lines)
+ sys.stderr.write(output + "\n")
+ sys.stderr.flush()
+ self._last_line_count = len(lines)
+
+ def close(self) -> None:
+ """Clean up the display.
+
+ In terminal mode, erases the stats block using ANSI escape codes.
+ In notebook mode, the HTML output remains visible.
+ """
+ if self._disabled:
+ return
+ # In terminal, erase the stats block
+ if not self._is_notebook and self._last_line_count > 0:
+ sys.stderr.write(f"\033[{self._last_line_count}A\033[J")
+ sys.stderr.flush()
+ self._last_line_count = 0
diff --git a/llmeter/results.py b/llmeter/results.py
index 45a1c00..1652d0f 100644
--- a/llmeter/results.py
+++ b/llmeter/results.py
@@ -6,7 +6,6 @@
import os
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
-from functools import cached_property
from numbers import Number
from typing import Any, Sequence
@@ -44,9 +43,9 @@ class Result:
"""Results of a test run."""
responses: list[InvocationResponse]
- total_requests: int
- clients: int
- n_requests: int
+ total_requests: int | None = None
+ clients: int = 1
+ n_requests: int | None = None
total_test_time: float | None = None
model_id: str | None = None
output_path: os.PathLike | None = None
@@ -169,8 +168,8 @@ def load_responses(self) -> list[InvocationResponse]:
InvocationResponse(**json.loads(line)) for line in f if line
]
logger.info("Loaded %d responses from %s", len(self.responses), responses_path)
- # Invalidate cached stats so they are recomputed with the loaded responses
- self.__dict__.pop("_builtin_stats", None)
+ # Recompute stats from the freshly loaded responses
+ self._preloaded_stats = self._compute_stats(self)
return self.responses
@classmethod
@@ -241,9 +240,9 @@ def load(
result = cls(responses=responses, **summary)
- # When skipping responses, load pre-computed stats from stats.json if available
- # so that result.stats works without needing the responses
+ # Load or compute stats
if not load_responses:
+ # Use pre-computed stats from disk when responses aren't loaded
stats_path = result_path / "stats.json"
if stats_path.exists():
with stats_path.open("r") as s:
@@ -260,78 +259,97 @@ def load(
pass
else:
result._preloaded_stats = None
+ else:
+ # Compute stats from the loaded responses, but also merge any
+ # contributed stats that were persisted in stats.json so they
+ # survive a save/load round-trip.
+ result._preloaded_stats = cls._compute_stats(result)
+ stats_path = result_path / "stats.json"
+ if stats_path.exists():
+ with stats_path.open("r") as s:
+ saved_stats = json.loads(s.read())
+ # Contributed stats are any keys in the saved file that are
+ # not produced by _compute_stats (i.e. they came from callbacks).
+ for key, value in saved_stats.items():
+ if key not in result._preloaded_stats:
+ result._preloaded_stats[key] = value
return result
- @cached_property
- def _builtin_stats(self) -> dict:
- """
- Default run metrics and aggregated statistics provided by LLMeter core
+ @classmethod
+ def _compute_stats(cls, result: "Result") -> dict:
+ """Compute stats from in-memory responses.
- Users should generally refer to the `.stats` property instead, which combines this data
- with any additional values contributed by callbacks or other extensions.
+ This is the fallback used when ``_preloaded_stats`` is not available — for
+ example when a ``Result`` is constructed manually or after
+ :meth:`load_responses` reloads data from disk.
- This is a read-only and `@cached_property`, which means the result is computed once and
- then cached for subsequent accesses - improving performance.
+ Args:
+ result: A ``Result`` instance whose ``responses`` list is populated.
Returns:
- stats: A dictionary containing all computed statistics. The keys are:
- - All key-value pairs from the Result's dictionary representation
- - Test-specific statistics
- - Aggregated statistics with keys in the format "{stat_name}-{aggregation_type}"
- where stat_name is one of the four metrics listed above, and
- aggregation_type includes measures like mean, median, etc.
- """
+ A flat dictionary matching the ``Result.stats`` schema, containing
+ run-level metrics (``failed_requests``, ``requests_per_minute``, …)
+ and per-metric aggregations (``time_to_first_token-p50``, …).
+ Example::
+
+ result = Result(responses=my_responses, total_requests=100, ...)
+ stats = Result._compute_stats(result)
+ stats["time_to_first_token-p90"] # 0.485
+ """
aggregation_metrics = [
"time_to_last_token",
"time_to_first_token",
"num_tokens_output",
"num_tokens_input",
]
-
- results_stats = _get_stats_from_results(
- self,
- aggregation_metrics,
- )
+ results_stats = _get_stats_from_results(result, aggregation_metrics)
return {
- **self.to_dict(),
- **_get_run_stats(self),
+ **result.to_dict(),
+ **_get_run_stats(result),
**{f"{k}-{j}": v for k, o in results_stats.items() for j, v in o.items()},
}
@property
def stats(self) -> dict:
+ """Run metrics and aggregated statistics over the individual requests.
+
+ Returns a flat dictionary combining:
+
+ * Basic run information (from ``to_dict()``).
+ * Aggregated statistics (``average``, ``p50``, ``p90``, ``p99``) for
+ ``time_to_last_token``, ``time_to_first_token``, ``num_tokens_output``,
+ and ``num_tokens_input``. Keys use the format
+ ``"{metric}-{aggregation}"``.
+ * Run-level throughput metrics (``requests_per_minute``,
+ ``total_input_tokens``, etc.).
+ * Any additional stats contributed by callbacks via
+ :meth:`_update_contributed_stats`.
+
+ During a live run, stats are computed incrementally by
+ :class:`~llmeter.utils.RunningStats` and stored in ``_preloaded_stats``.
+ When loading from disk with ``load_responses=False``, pre-computed stats
+ from ``stats.json`` are used. As a fallback (e.g. manually constructed
+ ``Result``), stats are computed on the fly from ``self.responses``.
+
+ Returns:
+ A new shallow copy of the stats dictionary on each access.
+
+ Example::
+
+ result = await runner.run(payload=my_payload, clients=5)
+ result.stats["time_to_first_token-p50"] # 0.312
+ result.stats["requests_per_minute"] # 141.2
+ result.stats["failed_requests"] # 0
"""
- Run metrics and aggregated statistics over the individual requests
-
- This combined view includes:
- - Basic information about the run (from the Result's dictionary representation)
- - Aggregated statistics ('average', 'p50', 'p90', 'p99') for:
- - Time to last token
- - Time to first token
- - Number of tokens output
- - Number of tokens input
-
- Aggregated statistics are keyed in the format "{stat_name}-{aggregation_type}"
-
- This property is read-only and returns a new shallow copy of the data on each access.
- Default stats provided by LLMeter are calculated on first access and then cached. Callbacks
- Callbacks or other mechanisms needing to augment stats should use the
- `_update_contributed_stats()` method.
-
- When the Result was loaded with ``load_responses=False``, pre-computed stats from
- ``stats.json`` are returned if available. Call ``load_responses()`` to load the
- individual responses and recompute stats from the raw data.
- """
- # Use preloaded stats when responses were not loaded
- if not self.responses and self._preloaded_stats is not None:
+ if self._preloaded_stats is not None:
+ stats = self._preloaded_stats.copy()
+ else:
+ # Fallback: compute from responses (e.g. Result constructed manually)
+ # Cache so subsequent accesses don't recompute.
+ self._preloaded_stats = self._compute_stats(self)
stats = self._preloaded_stats.copy()
- if self._contributed_stats:
- stats.update(self._contributed_stats)
- return stats
-
- stats = self._builtin_stats.copy()
if self._contributed_stats:
stats.update(self._contributed_stats)
diff --git a/llmeter/runner.py b/llmeter/runner.py
index 47626e3..7922868 100644
--- a/llmeter/runner.py
+++ b/llmeter/runner.py
@@ -12,7 +12,7 @@
from concurrent.futures import ThreadPoolExecutor
from copy import deepcopy
from dataclasses import InitVar, asdict, dataclass, fields, replace
-from datetime import datetime
+from datetime import datetime, timedelta
from itertools import cycle
from typing import TYPE_CHECKING, Any
from uuid import uuid4
@@ -20,7 +20,8 @@
from tqdm.auto import tqdm, trange
from upath import UPath as Path
-from llmeter.utils import now_utc
+from llmeter.live_display import LiveStatsDisplay
+from llmeter.utils import RunningStats, now_utc
if TYPE_CHECKING:
# Avoid circular import: We only need typing for Callback
@@ -56,11 +57,14 @@ class _RunConfig:
tokenizer: Tokenizer | Any | None = None
clients: int = 1
n_requests: int | None = None
+ run_duration: int | float | timedelta | None = None
payload: dict | list[dict] | os.PathLike | str | None = None
run_name: str | None = None
run_description: str | None = None
timeout: int | float = 60
callbacks: list[Callback] | None = None
+ low_memory: bool = False
+ progress_bar_stats: dict[str, str | tuple[str, str]] | None = None
disable_per_client_progress_bar: InitVar[bool] = True
disable_clients_progress_bar: InitVar[bool] = True
@@ -69,9 +73,14 @@ def __post_init__(self, disable_client_progress_bar, disable_clients_progress_ba
self._disable_clients_progress_bar = disable_clients_progress_bar
self._random_seed = 0
- if self.n_requests is not None:
+ if self.n_requests is not None and self.run_duration is None:
assert self.n_requests > 0, "Number of requests must be a positive integer"
+ if self.run_duration is not None:
+ if isinstance(self.run_duration, timedelta):
+ self.run_duration = self.run_duration.total_seconds()
+ assert self.run_duration > 0, "Run duration must be a positive number"
+
assert self.clients > 0, "Number of clients must be a positive integer"
if self.run_name is not None:
@@ -149,30 +158,74 @@ class _Run(_RunConfig):
"""
def __post_init__(self, disable_client_progress_bar, disable_clients_progress_bar):
- assert (
- self.run_name is not None
- ), "Test Run must be created with an explicit run_name"
+ assert self.run_name is not None, (
+ "Test Run must be created with an explicit run_name"
+ )
super().__post_init__(disable_client_progress_bar, disable_clients_progress_bar)
- assert (
- self.endpoint is not None
- ), "Test Run must be created with an explicit Endpoint"
+ assert self.endpoint is not None, (
+ "Test Run must be created with an explicit Endpoint"
+ )
self._validate_and_prepare_payload()
self._responses = []
+ if self.low_memory:
+ assert self.output_path is not None, (
+ "output_path is required when low_memory=True "
+ "(responses must be written to disk)"
+ )
+
+ self._running_stats = RunningStats(
+ metrics=[
+ "time_to_last_token",
+ "time_to_first_token",
+ "time_per_output_token",
+ "num_tokens_output",
+ "num_tokens_input",
+ ]
+ )
+
def _validate_and_prepare_payload(self):
- """Validate and prepare the payload for the test run and update n_requests
+ """Validate and prepare the payload for the test run.
+
+ Normalizes the payload into a list of dicts, validates that ``n_requests``
+ and ``run_duration`` are not both set, and sets ``_time_bound`` and
+ ``n_requests`` accordingly.
- This method ensures that the payload is valid and prepared for the test run.
+ For count-bound runs, ``n_requests`` defaults to the number of payloads
+ when not explicitly provided. For time-bound runs, ``n_requests`` is set
+ to 0 since the actual count is unknown upfront.
+
+ Raises:
+ AssertionError: If no payload is provided.
+ ValueError: If both ``n_requests`` and ``run_duration`` are set.
+ FileNotFoundError: If the payload path does not exist.
"""
assert self.payload, "No payload provided"
if isinstance(self.payload, (os.PathLike, str)):
self.payload = list(load_payloads(self.payload))
if isinstance(self.payload, dict):
self.payload = [self.payload]
- self._n_requests = self.n_requests or len(self.payload)
+
+ if (
+ self.run_duration is not None
+ and self.n_requests is not None
+ and self.n_requests != 0
+ ):
+ raise ValueError(
+ "Cannot set both n_requests and run_duration. "
+ "Use n_requests for request-bound runs or run_duration for time-bound runs."
+ )
+
+ self._time_bound = self.run_duration is not None
+ if self._time_bound:
+ # For time-bound runs, n_requests is unknown upfront; set to 0
+ # and update to the actual count after the run completes.
+ self.n_requests = 0
+ else:
+ self.n_requests = self.n_requests or len(self.payload)
@staticmethod
async def _compute_time_per_output_token(response: InvocationResponse):
@@ -251,10 +304,23 @@ async def _process_results_from_q(self, output_path: Path | None = None):
if self.callbacks is not None:
[await cb.after_invoke(response) for cb in self.callbacks]
- self._responses.append(response)
- if self._progress_bar:
+ if self.low_memory and self._running_stats is not None:
+ self._running_stats.update(response.to_dict())
+ else:
+ self._responses.append(response)
+ self._running_stats.update(response.to_dict())
+
+ if self._progress_bar is not None and not self._time_bound:
self._progress_bar.update(1)
+ if self._stats_display is not None:
+ raw = self._running_stats.to_stats()
+ if raw:
+ prefix = (
+ f"reqs={self._running_stats._count}" if self._time_bound else ""
+ )
+ self._stats_display.update(raw, extra_prefix=prefix)
+
if output_path:
output_path.parent.mkdir(parents=True, exist_ok=True)
with output_path.open("a") as f:
@@ -266,25 +332,32 @@ def _invoke_n_no_wait(
self,
payload: list[dict],
n: int | None = None,
+ duration: float | None = None,
shuffle_order=True,
) -> list[InvocationResponse]:
- """
- Generate multiple invocations for the given payload.
+ """Generate invocations synchronously for a single client.
+
+ Terminates when either *n* requests have been sent or *duration* seconds
+ have elapsed, whichever is specified. Exactly one of *n* or *duration*
+ must be provided.
- This method generates `n` invocations for the given payload(s) by sending
- requests to the endpoint in a loop. If a sequence of payloads is provided,
- the payloads are cycled through until `n` invocations are generated. If a
- single payload is provided, it is used for all `n` invocations.
+ Cycles through *payload*, sending each request to the endpoint and
+ pushing the response onto ``self._queue`` for async token-counting and
+ stats collection.
Args:
- payload: The input payload to generate invocations for.
- n (int|None, optional): The number of invocations to generate.
- If not specified, every element in the payload is used once.
+ payload (list[dict]): The input payloads to cycle through.
+ n (int | None, optional): The number of invocations to generate.
+ If not specified, every element in the payload is used once
+ (only when *duration* is also ``None``).
+ duration (float | None, optional): Maximum wall-clock seconds to
+ keep sending requests. When set, requests are sent continuously
+ until the deadline.
shuffle_order (bool, optional): Whether to shuffle the order of payloads
before generating invocations. Defaults to True.
Returns:
- List[EndpointResponse]: A list of response objects.
+ list[InvocationResponse]: A list of response objects.
"""
# ToDo: replace with an async method to prepare payloads, including possible callbacks,
@@ -294,22 +367,46 @@ def _invoke_n_no_wait(
random.seed(0)
payload = random.sample(payload, k=len(payload))
- responses = []
- if n is None:
- n = len(payload)
- for p, _ in zip(
- cycle(payload),
+ responses: list[InvocationResponse] = []
+ if not payload:
+ return responses
+
+ time_bound = duration is not None
+ if time_bound:
+ deadline = time.perf_counter() + duration
+ else:
+ if n is None:
+ n = len(payload)
+
+ payload_iter = cycle(payload)
+
+ # Count-bound runs get a trange progress bar; time-bound runs use a
+ # separate _tick_time_bar task so we skip the per-client bar here.
+ pbar = (
trange(
n,
leave=False,
desc="Requests",
disable=_disable_tqdm or self._disable_per_client_progress_bar,
- ),
- ):
+ )
+ if not time_bound
+ else None
+ )
+
+ sent = 0
+ while True:
+ if time_bound:
+ if time.perf_counter() >= deadline:
+ break
+ else:
+ if sent >= n:
+ break
+
+ p = next(payload_iter)
try:
p = asyncio.run(process_before_invoke_callbacks(self.callbacks, p))
+ self._running_stats.record_send()
response = self._endpoint.invoke(p)
-
except Exception as e:
logger.exception(f"Error with invocation with payload {p}: {e}")
response = InvocationResponse.error_output(
@@ -322,26 +419,33 @@ def _invoke_n_no_wait(
self._queue._loop.call_soon_threadsafe( # type: ignore
self._queue.put_nowait, response
)
+ sent += 1
+ if pbar is not None:
+ pbar.update(1)
+
+ if pbar is not None:
+ pbar.close()
return responses
- async def _invoke_n(
+ async def _invoke_client(
self,
payload: list[dict],
n: int | None = None,
+ duration: float | None = None,
add_start_jitter=True,
shuffle_order=True,
) -> list[InvocationResponse]:
- """
- Asynchronously generate multiple invocations for the given payload.
+ """Asynchronously generate invocations for a single client.
- This method generates `n` invocations for the given payload(s) by sending
- requests to the endpoint asynchronously. If a sequence of payloads is provided,
- the payloads are cycled through until `n` invocations are generated. If a
- single payload is provided, it is used for all `n` invocations.
+ Wraps :meth:`_invoke_n_no_wait` in a thread. For count-bound runs an
+ overall timeout of ``self.timeout * n`` is applied; time-bound runs
+ have no extra timeout (the duration itself is the limit).
Args:
- payload (Dict[str, str] | Sequence[Dict[str, str]]): The input payload(s) to generate invocations for.
- n (int | None, optional): The number of invocations to generate. Defaults to None.
+ payload (list[dict]): The input payload(s) to generate invocations for.
+ n (int | None, optional): The number of invocations to generate.
+ Defaults to None (one per payload element).
+ duration (float | None, optional): Maximum wall-clock seconds.
add_start_jitter (bool, optional): Whether to add a random delay before
starting the invocations loop to avoid batch bunching when using
multiple clients. Defaults to True.
@@ -349,7 +453,8 @@ async def _invoke_n(
before generating invocations. Defaults to True.
Returns:
- List[EndpointResponse]: A list of response objects.
+ list[InvocationResponse]: A list of response objects. Returns an empty
+ list if the overall timeout is exceeded (count-bound only).
"""
if add_start_jitter:
@@ -358,72 +463,117 @@ async def _invoke_n(
if shuffle_order:
self._random_seed = random.randint(0, 2**16 - 1)
+ coro = asyncio.to_thread(
+ self._invoke_n_no_wait, payload, n, duration, shuffle_order
+ )
+
+ if duration is not None:
+ # Time-bound: no extra timeout — the duration is the limit
+ return await coro
+
try:
- response = await asyncio.wait_for(
- asyncio.to_thread(self._invoke_n_no_wait, payload, n, shuffle_order),
+ return await asyncio.wait_for(
+ coro,
timeout=self.timeout * (n or len(payload)),
)
except asyncio.TimeoutError:
logger.error("client timeout!")
return []
- return response
-
- async def _invoke_n_c(
+ async def _invoke_clients(
self,
payload: list[dict],
n_requests: int | None = None,
+ duration: float | None = None,
clients: int = 1,
) -> tuple[float, float, float]:
- """
- Asynchronously generates multiple invocations for a given payload.
+ """Spawn *clients* concurrent invocation loops.
+
+ Each client generates invocations by delegating to
+ :meth:`_invoke_client`. All clients run concurrently and the method
+ waits for all of them to finish before signalling the token-counting
+ queue to stop.
Args:
- payload (dict): The input data for generating invocations.
- queue (asyncio.Queue): The queue to store the generated responses.
- n_requests (int | None, optional): The number of invocations to generate per connection. Defaults to None.
- clients (int, optional): The number of concurrent connections to generate invocations. Defaults to 1.
+ payload (list[dict]): The input payloads to send.
+ n_requests (int | None, optional): The number of invocations to
+ generate per client (count-bound). Defaults to None.
+ duration (float | None, optional): Maximum wall-clock seconds per
+ client (time-bound). Defaults to None.
+ clients (int, optional): The number of concurrent client connections.
+ Defaults to 1.
Returns:
- None
-
- Raises:
- None
+ tuple[float, float, float]: A ``(total_test_time, start_t, end_t)``
+ tuple of ``time.perf_counter`` values.
"""
- logger.info(
- f"Generating {clients} connections with {n_requests} invocations each"
- )
+ if duration is not None:
+ logger.info(f"Generating {clients} connections for {duration}s each")
+ else:
+ logger.info(
+ f"Generating {clients} connections with {n_requests} invocations each"
+ )
start_t = time.perf_counter()
await tqdm.gather(
- *[self._invoke_n(payload, n_requests) for _ in range(clients)],
+ *[
+ self._invoke_client(payload, n=n_requests, duration=duration)
+ for _ in range(clients)
+ ],
leave=False,
desc="Clients",
disable=_disable_tqdm or self._disable_clients_progress_bar,
)
end_t = time.perf_counter()
total_test_time = end_t - start_t
- logger.info(
- f"Generated {clients} connections with {n_requests} invocations each in {total_test_time*1000:.2f} seconds"
- )
- # Signal the token counting task to exit
+ if duration is not None:
+ logger.info(
+ f"Completed {clients} clients x {duration}s in "
+ f"{total_test_time * 1000:.2f}ms"
+ )
+ else:
+ logger.info(
+ f"Completed {clients} clients x {n_requests} requests in "
+ f"{total_test_time * 1000:.2f}ms"
+ )
+
if self._queue:
await self._queue.put(None)
logger.debug("Signaling token counting task to exit")
return total_test_time, start_t, end_t
+ async def _tick_time_bar(self):
+ """Advance ``_progress_bar`` every 0.5 s until ``run_duration`` is reached.
+
+ Designed to run as a concurrent task alongside the invocation loops so
+ the user sees a smooth time-based progress bar.
+ """
+ start = time.perf_counter()
+ duration = self.run_duration
+ prev = 0
+ while True:
+ await asyncio.sleep(0.5)
+ elapsed = time.perf_counter() - start
+ tick = min(int(elapsed), int(duration)) - prev
+ if tick > 0 and self._progress_bar is not None:
+ self._progress_bar.update(tick)
+ prev += tick
+ if elapsed >= duration:
+ break
+
async def _run(self):
"""Run the test with the given configuration
This method is expected to be called *exactly once* after the _Run object is created.
Attempting to re-use a _Run object may result in undefined behavior.
"""
+ # For time-bound runs, total_requests is unknown upfront
result = Result(
responses=[],
total_test_time=None,
- total_requests=self._n_requests * self.clients,
+ total_requests=0 if self._time_bound else self.n_requests * self.clients,
clients=self.clients,
- n_requests=self._n_requests,
+ n_requests=self.n_requests,
output_path=self.output_path, # type: ignore
model_id=self._endpoint.model_id,
provider=self._endpoint.provider,
@@ -444,27 +594,67 @@ async def _run(self):
loop.set_default_executor(ThreadPoolExecutor(max_workers=self.clients + 5))
logger.info("Starting test")
self._queue = asyncio.Queue()
- self._progress_bar = tqdm(
- total=self.clients * self._n_requests,
- leave=False,
- desc="Total requests",
- disable=_disable_tqdm,
+
+ if self._time_bound:
+ # Time-bound: progress bar shows elapsed seconds
+ self._progress_bar = tqdm(
+ total=int(self.run_duration),
+ leave=False,
+ desc="Elapsed",
+ unit="s",
+ bar_format="{desc}: {bar}| {n:.0f}/{total:.0f}s [{elapsed}]",
+ disable=_disable_tqdm,
+ )
+ else:
+ # Count-bound: progress bar shows completed requests
+ self._progress_bar = tqdm(
+ total=self.clients * self.n_requests,
+ leave=False,
+ desc="Total requests",
+ disable=_disable_tqdm,
+ )
+
+ # Live stats display — renders as an HTML table in notebooks, multi-line in terminals
+ self._stats_display = LiveStatsDisplay(
+ disabled=_disable_tqdm,
+ display_stats=self.progress_bar_stats,
)
+ # Show the table layout immediately with placeholder values
+ prefix = "reqs=0" if self._time_bound else ""
+ self._stats_display.update({}, extra_prefix=prefix)
+
try:
run_start_time = now_utc()
- _, (total_test_time, start_time, end_time) = await asyncio.gather(
- self._process_results_from_q(
- output_path=Path(self.output_path) / "responses.jsonl"
- if self.output_path
- else None,
- ),
- self._invoke_n_c(
+ if self._time_bound:
+ invoke_coro = self._invoke_clients(
payload=self.payload, # type: ignore
- n_requests=self._n_requests,
+ duration=self.run_duration,
clients=self.clients,
- ),
- )
+ )
+ _, (total_test_time, start_time, end_time), _ = await asyncio.gather(
+ self._process_results_from_q(
+ output_path=Path(self.output_path) / "responses.jsonl"
+ if self.output_path
+ else None,
+ ),
+ invoke_coro,
+ self._tick_time_bar(),
+ )
+ else:
+ invoke_coro = self._invoke_clients(
+ payload=self.payload, # type: ignore
+ n_requests=self.n_requests,
+ clients=self.clients,
+ )
+ _, (total_test_time, start_time, end_time) = await asyncio.gather(
+ self._process_results_from_q(
+ output_path=Path(self.output_path) / "responses.jsonl"
+ if self.output_path
+ else None,
+ ),
+ invoke_coro,
+ )
run_end_time = now_utc()
except asyncio.CancelledError:
@@ -474,16 +664,40 @@ async def _run(self):
return result
self._progress_bar.close()
- logger.info(f"Test completed in {total_test_time*1000:.2f} seconds.")
+ if self._stats_display is not None:
+ self._stats_display.close()
+ logger.info(f"Test completed in {total_test_time * 1000:.2f} seconds.")
+
+ actual_total = self._running_stats._count
result = replace(
result,
responses=self._responses,
total_test_time=total_test_time,
+ total_requests=actual_total,
+ n_requests=actual_total // max(self.clients, 1)
+ if self._time_bound
+ else self.n_requests,
start_time=run_start_time,
end_time=run_end_time,
)
+ # Compute stats from the running accumulators
+ result._preloaded_stats = self._running_stats.to_stats(
+ total_requests=result.total_requests,
+ total_test_time=total_test_time,
+ result_dict=result.to_dict(),
+ )
+ result._preloaded_stats["start_time"] = run_start_time
+ result._preloaded_stats["end_time"] = run_end_time
+ result._preloaded_stats["total_test_time"] = total_test_time
+
+ if self.low_memory:
+ logger.info(
+ "Low-memory mode: responses not stored in memory. "
+ "Use result.load_responses() to load from disk."
+ )
+
if self.callbacks is not None:
[await cb.after_run(result) for cb in self.callbacks]
@@ -537,7 +751,12 @@ class Runner(_RunConfig):
`DummyTokenizer` will be used if needed.
clients (int): The number of concurrent clients to use for sending requests. Defaults to 1.
n_requests (int | None): The number of LLM invocations to generate *per client*. By
- default, each request in `payload` will be sent once by each client.
+ default, each request in `payload` will be sent once by each client. Mutually
+ exclusive with ``run_duration``.
+ run_duration (int | float | timedelta | None): Run each client for this many seconds instead of a
+ fixed request count. Clients send requests continuously until the duration expires.
+ Mutually exclusive with ``n_requests``. Defaults to ``None`` (count-bound mode).
+ Accepts a number of seconds or a ``timedelta``.
payload (dict | list[dict] | os.PathLike | str | None): The request data to send to the
endpoint under test. You can provide a single JSON payload (dict), a list of payloads
(list[dict]), or a path to one or more JSON/JSON-Lines files to be loaded by
@@ -554,6 +773,15 @@ class Runner(_RunConfig):
endpoint. Defaults to 60 seconds.
callbacks (list[Callback] | None): Optional callbacks to enable during the test Run. See
`llmeter.callbacks` for more information.
+ low_memory (bool): When ``True``, responses are written to disk but not kept in memory
+ during the run. Stats are computed incrementally via
+ :class:`~llmeter.utils.RunningStats`. Requires ``output_path`` to be set. Use
+ ``result.load_responses()`` to load responses from disk after the run. Defaults to
+ ``False``.
+ progress_bar_stats (dict | None): Controls which live stats appear on the progress bar.
+ Maps short display labels to canonical stat keys — see
+ :data:`~llmeter.live_display.DEFAULT_DISPLAY_STATS` for the format and defaults. Pass ``{}``
+ to disable live stats entirely. Defaults to ``None`` (use built-in defaults).
disable_per_client_progress_bar (bool): Set `True` to disable per-client progress bars
from showing during the run. Default `False` (each client's progress will be shown).
disable_clients_progress_bar (bool): Set `True` to disable overall progress bar from
@@ -595,11 +823,14 @@ async def run(
tokenizer: Tokenizer | Any | None = None,
clients: int | None = None,
n_requests: int | None = None,
+ run_duration: int | float | timedelta | None = None,
payload: dict | list[dict] | os.PathLike | str | None = None,
run_name: str | None = None,
run_description: str | None = None,
timeout: int | float | None = None,
callbacks: list[Callback] | None = None,
+ low_memory: bool | None = None,
+ progress_bar_stats: dict[str, str | tuple[str, str]] | None = None,
disable_per_client_progress_bar: bool | None = None,
disable_clients_progress_bar: bool | None = None,
) -> Result:
@@ -623,6 +854,18 @@ async def run(
output token counts for endpoints that don't report exact information.
clients (int): The number of concurrent clients to use for sending requests.
n_requests (int | None): The number of LLM invocations to generate *per client*.
+ Mutually exclusive with ``run_duration``.
+ run_duration (int | float | timedelta | None): Run each client for this many seconds
+ instead of a fixed request count. Clients send requests continuously
+ until the duration expires. Mutually exclusive with ``n_requests``.
+ Accepts a number of seconds or a ``timedelta``.
+
+ Example::
+
+ # Run for 60 seconds with 5 concurrent clients:
+ result = await runner.run(run_duration=60, clients=5)
+ result.total_requests # actual count completed
+
payload (dict | list[dict] | os.PathLike | str | None): The request data to send to the
endpoint under test. You can provide a single JSON payload (dict), a list of
payloads (list[dict]), or a path to one or more JSON/JSON-Lines files to be loaded
@@ -635,6 +878,36 @@ async def run(
endpoint.
callbacks (list[Callback] | None): Optional callbacks to enable during the test Run. See
`llmeter.callbacks` for more information.
+ low_memory (bool): When ``True``, responses are written to disk but not
+ kept in memory. Stats are computed incrementally via
+ :class:`~llmeter.utils.RunningStats`. Requires ``output_path``.
+ Use ``result.load_responses()`` to access responses after the run.
+
+ Example::
+
+ result = await runner.run(
+ output_path="/tmp/my_run",
+ low_memory=True,
+ )
+ result.stats # works (computed incrementally)
+ result.responses # [] (empty)
+ result.load_responses() # loads from disk
+
+ progress_bar_stats (dict): Controls which live stats appear on the
+ progress bar. Maps short display labels to canonical stat keys — see
+ :data:`~llmeter.live_display.DEFAULT_DISPLAY_STATS` for the format and
+ defaults. Pass ``{}`` to disable live stats entirely.
+
+ Example::
+
+ # Show only p99 latency and tokens per second:
+ result = await runner.run(
+ progress_bar_stats={
+ "p99_ttlt": "time_to_last_token-p99",
+ "tps": ("time_per_output_token-p50", "inv"),
+ "fail": "failed_requests",
+ },
+ )
disable_per_client_progress_bar (bool): Set `True` to disable per-client progress bars
from showing during the run.
disable_clients_progress_bar (bool): Set `True` to disable overall progress bar from
@@ -662,11 +935,14 @@ async def run(
tokenizer=tokenizer,
clients=clients,
n_requests=n_requests,
+ run_duration=run_duration,
payload=payload,
run_name=run_name,
run_description=run_description,
timeout=timeout,
callbacks=callbacks,
+ low_memory=low_memory,
+ progress_bar_stats=progress_bar_stats,
disable_per_client_progress_bar=disable_per_client_progress_bar,
disable_clients_progress_bar=disable_clients_progress_bar,
)
diff --git a/llmeter/utils.py b/llmeter/utils.py
index d072e58..74198ac 100644
--- a/llmeter/utils.py
+++ b/llmeter/utils.py
@@ -1,5 +1,7 @@
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
+import bisect
+import time
from datetime import datetime, timezone
from itertools import filterfalse
from math import isnan
@@ -83,6 +85,184 @@ def summary_stats_from_list(
return {}
+class RunningStats:
+ """Accumulate summary statistics incrementally from individual responses.
+
+ Maintains sorted value lists per metric so that percentiles (p50, p90, p99),
+ averages, and sums can be computed at any point — both mid-run (for live
+ progress-bar display via :meth:`snapshot`) and at the end of a run (for the
+ final :class:`~llmeter.results.Result` stats via :meth:`to_stats`).
+
+ Args:
+ metrics: Names of numeric response fields to track (e.g.
+ ``"time_to_first_token"``, ``"num_tokens_output"``).
+
+ Example::
+
+ rs = RunningStats(metrics=["time_to_first_token", "time_to_last_token"])
+ rs.update({"time_to_first_token": 0.3, "time_to_last_token": 0.8})
+ rs.update({"time_to_first_token": 0.5, "time_to_last_token": 1.2, "error": None})
+ rs.to_stats()
+ # {'failed_requests': 0, ..., 'time_to_first_token-p50': 0.4, ...}
+ """
+
+ def __init__(self, metrics: Sequence[str]):
+ self._metrics = list(metrics)
+ self._count = 0
+ self._failed = 0
+ self._sends = 0
+ self._first_send_time: float | None = None
+ self._last_send_time: float | None = None
+ self._sums: dict[str, float] = {m: 0.0 for m in metrics}
+ self._values: dict[str, list[float]] = {m: [] for m in metrics}
+
+ def record_send(self) -> None:
+ """Record that a request was dispatched to the endpoint.
+
+ Call this from the invocation loop each time a request is sent, *before*
+ waiting for the response. This tracks the send-side time window used for
+ accurate RPM and throughput calculations.
+
+ The send window (``_first_send_time`` to ``_last_send_time``) excludes
+ the tail latency of the final response, giving a more accurate picture
+ of the request dispatch rate.
+ """
+ now = time.perf_counter()
+ self._sends += 1
+ if self._first_send_time is None:
+ self._first_send_time = now
+ self._last_send_time = now
+
+ def update(self, response_dict: dict[str, Any]) -> None:
+ """Record one response's metric values.
+
+ Call this once per :class:`~llmeter.endpoints.base.InvocationResponse`
+ (typically via ``response.to_dict()``). The method extracts each tracked
+ metric from *response_dict*, skipping ``None`` and ``NaN`` values, and
+ increments the failure counter when an ``"error"`` key is present.
+
+ Args:
+ response_dict: A flat dictionary of response fields, as returned by
+ ``InvocationResponse.to_dict()``.
+
+ Example::
+
+ rs = RunningStats(metrics=["time_to_first_token"])
+ rs.update({"time_to_first_token": 0.42, "error": None})
+ rs.update({"time_to_first_token": None, "error": "timeout"})
+ assert rs._failed == 1
+ """
+ self._count += 1
+ if response_dict.get("error") is not None:
+ self._failed += 1
+ for m in self._metrics:
+ val = response_dict.get(m)
+ if val is not None and not (isinstance(val, float) and isnan(val)):
+ self._sums[m] += val
+ bisect.insort(self._values[m], val)
+
+ def to_stats(
+ self,
+ total_requests: int | None = None,
+ total_test_time: float | None = None,
+ result_dict: dict[str, Any] | None = None,
+ ) -> dict[str, Any]:
+ """Compute all accumulated statistics as raw numeric values.
+
+ This is the single source of truth for stats computation. It is called
+ once at the end of a run (with all three optional arguments) to produce
+ the full ``Result.stats`` dict, and also called internally by
+ :meth:`snapshot` (without arguments) for mid-run progress display.
+
+ Args:
+ total_requests: Total number of requests across all clients. When
+ provided, enables ``failed_requests_rate`` and
+ ``requests_per_minute`` computation.
+ total_test_time: Wall-clock duration of the run in seconds. When
+ provided, enables throughput metrics (requests/min, tokens/min).
+ result_dict: Base key-value pairs to include in the output (typically
+ from ``Result.to_dict()``). When ``None``, only metric
+ aggregations and failure counts are returned.
+
+ Returns:
+ A flat dictionary of statistics. Keys include:
+
+ * ``failed_requests``, ``failed_requests_rate``, ``requests_per_minute``
+ * ``total_input_tokens``, ``total_output_tokens``
+ * ``average_input_tokens_per_minute``, ``average_output_tokens_per_minute``
+ * ``{metric}-{agg}`` for each tracked metric and each aggregation
+ (``average``, ``p50``, ``p90``, ``p99``).
+
+ Example::
+
+ rs = RunningStats(metrics=["time_to_first_token", "num_tokens_output"])
+ for resp in responses:
+ rs.update(resp.to_dict())
+
+ # Mid-run (no run-level context):
+ partial = rs.to_stats()
+ partial["time_to_first_token-p50"] # 0.312
+
+ # End of run (full Result.stats schema):
+ full = rs.to_stats(
+ total_requests=100,
+ total_test_time=42.5,
+ result_dict=result.to_dict(),
+ )
+ full["requests_per_minute"] # 141.2
+ """
+ stats: dict[str, Any] = {}
+ if result_dict is not None:
+ stats.update(result_dict)
+
+ # Run-level stats
+ stats["failed_requests"] = self._failed
+ stats["failed_requests_rate"] = total_requests and self._failed / total_requests
+ stats["requests_per_minute"] = (
+ total_test_time and total_requests / total_test_time * 60
+ if total_requests
+ else None
+ )
+ stats["total_input_tokens"] = self._sums.get("num_tokens_input", 0)
+ stats["total_output_tokens"] = self._sums.get("num_tokens_output", 0)
+ stats["average_input_tokens_per_minute"] = (
+ total_test_time and stats["total_input_tokens"] / total_test_time * 60
+ )
+ stats["average_output_tokens_per_minute"] = (
+ total_test_time and stats["total_output_tokens"] / total_test_time * 60
+ )
+
+ # Per-metric aggregations
+ for m in self._metrics:
+ agg = summary_stats_from_list(self._values.get(m, []))
+ for j, v in agg.items():
+ stats[f"{m}-{j}"] = v
+
+ # Send-window throughput (live RPM and output tokens/s).
+ # These use the dispatch timestamps rather than response timestamps,
+ # giving a more accurate picture of the request rate.
+ send_window = self._send_window()
+ if send_window and send_window > 0:
+ stats["rpm"] = self._count / send_window * 60
+ total_out = self._sums.get("num_tokens_output", 0)
+ stats["output_tps"] = total_out / send_window
+
+ return stats
+
+ def _send_window(self) -> float | None:
+ """Return the elapsed seconds between first and last ``record_send`` call.
+
+ Returns ``None`` when fewer than two sends have been recorded.
+ """
+ if (
+ self._first_send_time is not None
+ and self._last_send_time is not None
+ and self._last_send_time > self._first_send_time
+ ):
+ return self._last_send_time - self._first_send_time
+ return None
+
+
def now_utc() -> datetime:
"""Returns the current UTC datetime.
diff --git a/tests/unit/test_experiments.py b/tests/unit/test_experiments.py
index 13fbd1f..8996578 100644
--- a/tests/unit/test_experiments.py
+++ b/tests/unit/test_experiments.py
@@ -531,3 +531,122 @@ def test_get_n_requests_parametrized(
)
result = runner._get_n_requests(clients)
assert result == expected, f"Expected {expected}, but got {result}"
+
+
+# ── LoadTest with run_duration, low_memory, progress_bar_stats ───────────────
+
+
+class TestLoadTestTimeBound:
+ def test_load_test_with_run_duration(self, mock_endpoint):
+ """run_duration should be stored on the LoadTest instance."""
+ lt = LoadTest(
+ endpoint=mock_endpoint,
+ payload={"input": "test"},
+ sequence_of_clients=[1, 2],
+ run_duration=30,
+ )
+ assert lt.run_duration == 30
+
+ def test_load_test_with_low_memory(self, mock_endpoint):
+ """low_memory should be stored on the LoadTest instance."""
+ lt = LoadTest(
+ endpoint=mock_endpoint,
+ payload={"input": "test"},
+ sequence_of_clients=[1],
+ low_memory=True,
+ )
+ assert lt.low_memory is True
+
+ def test_load_test_with_progress_bar_stats(self, mock_endpoint):
+ """progress_bar_stats should be stored on the LoadTest instance."""
+ custom_stats = {"rpm": "rpm", "fail": "failed_requests"}
+ lt = LoadTest(
+ endpoint=mock_endpoint,
+ payload={"input": "test"},
+ sequence_of_clients=[1],
+ progress_bar_stats=custom_stats,
+ )
+ assert lt.progress_bar_stats == custom_stats
+
+ @pytest.mark.asyncio
+ async def test_run_duration_passed_to_runner(self, mock_endpoint):
+ """When run_duration is set, runner.run() should receive it."""
+ mock_runner_instance = AsyncMock(spec=Runner)
+ mock_runner_instance.run.return_value = MagicMock(
+ spec=Result, clients=1, total_requests=10
+ )
+
+ with patch("llmeter.experiments.Runner", return_value=mock_runner_instance):
+ lt = LoadTest(
+ endpoint=mock_endpoint,
+ payload={"input": "test"},
+ sequence_of_clients=[1, 3],
+ run_duration=15,
+ )
+ await lt.run()
+
+ # Check that run_duration was passed in each call
+ for call in mock_runner_instance.run.call_args_list:
+ assert call.kwargs["run_duration"] == 15
+ assert "n_requests" not in call.kwargs
+
+ @pytest.mark.asyncio
+ async def test_count_bound_does_not_pass_run_duration(self, mock_endpoint):
+ """When run_duration is None, runner.run() should receive n_requests."""
+ mock_runner_instance = AsyncMock(spec=Runner)
+ mock_runner_instance.run.return_value = MagicMock(
+ spec=Result, clients=1, total_requests=10
+ )
+
+ with patch("llmeter.experiments.Runner", return_value=mock_runner_instance):
+ lt = LoadTest(
+ endpoint=mock_endpoint,
+ payload={"input": "test"},
+ sequence_of_clients=[1],
+ )
+ await lt.run()
+
+ call_kwargs = mock_runner_instance.run.call_args_list[0].kwargs
+ assert "n_requests" in call_kwargs
+ assert "run_duration" not in call_kwargs
+
+ @pytest.mark.asyncio
+ async def test_low_memory_passed_to_runner(self, mock_endpoint):
+ """low_memory should be forwarded to each runner.run() call."""
+ mock_runner_instance = AsyncMock(spec=Runner)
+ mock_runner_instance.run.return_value = MagicMock(
+ spec=Result, clients=1, total_requests=10
+ )
+
+ with patch("llmeter.experiments.Runner", return_value=mock_runner_instance):
+ lt = LoadTest(
+ endpoint=mock_endpoint,
+ payload={"input": "test"},
+ sequence_of_clients=[1, 2],
+ low_memory=True,
+ )
+ await lt.run()
+
+ for call in mock_runner_instance.run.call_args_list:
+ assert call.kwargs["low_memory"] is True
+
+ @pytest.mark.asyncio
+ async def test_progress_bar_stats_passed_to_runner(self, mock_endpoint):
+ """progress_bar_stats should be forwarded to each runner.run() call."""
+ custom_stats = {"rpm": "rpm"}
+ mock_runner_instance = AsyncMock(spec=Runner)
+ mock_runner_instance.run.return_value = MagicMock(
+ spec=Result, clients=1, total_requests=10
+ )
+
+ with patch("llmeter.experiments.Runner", return_value=mock_runner_instance):
+ lt = LoadTest(
+ endpoint=mock_endpoint,
+ payload={"input": "test"},
+ sequence_of_clients=[1],
+ progress_bar_stats=custom_stats,
+ )
+ await lt.run()
+
+ call_kwargs = mock_runner_instance.run.call_args_list[0].kwargs
+ assert call_kwargs["progress_bar_stats"] == custom_stats
diff --git a/tests/unit/test_lazy_load.py b/tests/unit/test_lazy_load.py
index ed02b86..14436ce 100644
--- a/tests/unit/test_lazy_load.py
+++ b/tests/unit/test_lazy_load.py
@@ -1,12 +1,12 @@
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
+from unittest.mock import MagicMock, patch
+
import pytest
from upath import UPath
from llmeter.endpoints.base import InvocationResponse
-from unittest.mock import MagicMock, patch
-
from llmeter.experiments import LoadTestResult
from llmeter.results import Result
@@ -129,14 +129,14 @@ def test_load_responses_returns_correct_data(self, sample_responses, saved_resul
assert orig.time_to_first_token == loaded_resp.time_to_first_token
assert orig.time_to_last_token == loaded_resp.time_to_last_token
- def test_load_responses_invalidates_cached_stats(self, saved_result):
+ def test_load_responses_recomputes_stats(self, saved_result):
loaded = Result.load(saved_result, load_responses=True)
- # Access _builtin_stats to cache it
- _ = loaded._builtin_stats
- assert "_builtin_stats" in loaded.__dict__
+ original_stats = loaded._preloaded_stats.copy()
loaded.load_responses()
- assert "_builtin_stats" not in loaded.__dict__
+ # Stats should be recomputed (same values, but a fresh dict)
+ assert loaded._preloaded_stats is not original_stats
+ assert loaded._preloaded_stats == original_stats
def test_load_responses_stats_match_full_load(self, saved_result):
full = Result.load(saved_result, load_responses=True)
@@ -153,9 +153,9 @@ def test_load_responses_stats_match_full_load(self, saved_result):
"failed_requests",
"requests_per_minute",
]:
- assert lazy_stats[key] == pytest.approx(
- full_stats[key]
- ), f"Mismatch on {key}"
+ assert lazy_stats[key] == pytest.approx(full_stats[key]), (
+ f"Mismatch on {key}"
+ )
def test_load_responses_no_output_path_raises(self):
result = Result(
diff --git a/tests/unit/test_live_display.py b/tests/unit/test_live_display.py
new file mode 100644
index 0000000..6e0fb92
--- /dev/null
+++ b/tests/unit/test_live_display.py
@@ -0,0 +1,225 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from collections import OrderedDict
+from unittest.mock import patch
+
+from llmeter.live_display import (
+ DEFAULT_DISPLAY_STATS,
+ LiveStatsDisplay,
+ _classify,
+ _format_stat,
+ _group_stats,
+ _in_notebook,
+)
+
+
+# ── _classify ────────────────────────────────────────────────────────────────
+
+
+class TestClassify:
+ def test_rpm_goes_to_throughput(self):
+ assert _classify("rpm") == "Throughput"
+
+ def test_tps_goes_to_throughput(self):
+ assert _classify("p50_tps") == "Throughput"
+ assert _classify("output_tps") == "Throughput"
+
+ def test_ttft_goes_to_ttft(self):
+ assert _classify("p50_ttft") == "TTFT"
+ assert _classify("p90_ttft") == "TTFT"
+
+ def test_ttlt_goes_to_ttlt(self):
+ assert _classify("p50_ttlt") == "TTLT"
+ assert _classify("p90_ttlt") == "TTLT"
+
+ def test_token_goes_to_tokens(self):
+ assert _classify("input_tokens") == "Tokens"
+ assert _classify("output_tokens") == "Tokens"
+
+ def test_fail_goes_to_errors(self):
+ assert _classify("fail") == "Errors"
+
+ def test_unknown_goes_to_other(self):
+ assert _classify("custom_metric") == "Other"
+
+ def test_case_insensitive(self):
+ assert _classify("RPM") == "Throughput"
+ assert _classify("P50_TTFT") == "TTFT"
+
+
+# ── _group_stats ─────────────────────────────────────────────────────────────
+
+
+class TestGroupStats:
+ def test_groups_by_category(self):
+ stats = {
+ "rpm": "185.9",
+ "p50_ttft": "0.312s",
+ "p90_ttlt": "1.203s",
+ "input_tokens": "12540",
+ "fail": "0",
+ }
+ groups = _group_stats(stats)
+ assert "Throughput" in groups
+ assert "TTFT" in groups
+ assert "TTLT" in groups
+ assert "Tokens" in groups
+ assert "Errors" in groups
+
+ def test_preserves_order(self):
+ stats = OrderedDict(
+ [
+ ("rpm", "185.9"),
+ ("p50_ttft", "0.312s"),
+ ("p50_ttlt", "0.847s"),
+ ("fail", "0"),
+ ]
+ )
+ groups = _group_stats(stats)
+ group_names = list(groups.keys())
+ assert group_names == ["Throughput", "TTFT", "TTLT", "Errors"]
+
+ def test_unknown_keys_go_to_other(self):
+ stats = {"custom": "42"}
+ groups = _group_stats(stats)
+ assert "Other" in groups
+ assert groups["Other"] == [("custom", "42")]
+
+ def test_empty_stats(self):
+ groups = _group_stats({})
+ assert len(groups) == 0
+
+
+# ── _in_notebook ─────────────────────────────────────────────────────────────
+
+
+class TestInNotebook:
+ def test_returns_false_outside_notebook(self):
+ assert _in_notebook() is False
+
+ def test_returns_false_for_terminal_ipython(self):
+ mock_shell = type("TerminalInteractiveShell", (), {})()
+ with patch("IPython.get_ipython", return_value=mock_shell):
+ assert _in_notebook() is False
+
+ def test_returns_true_for_zmq_shell(self):
+ mock_shell = type("ZMQInteractiveShell", (), {})()
+ with patch("IPython.get_ipython", return_value=mock_shell):
+ assert _in_notebook() is True
+
+ def test_returns_false_for_none(self):
+ with patch("IPython.get_ipython", return_value=None):
+ assert _in_notebook() is False
+
+
+# ── _format_stat ─────────────────────────────────────────────────────────────
+
+
+class TestFormatStat:
+ def test_time_metric(self):
+ assert _format_stat("time_to_first_token-p50", 0.312) == "0.312s"
+
+ def test_rpm_metric(self):
+ assert _format_stat("rpm", 185.9) == "185.9"
+
+ def test_tps_metric(self):
+ assert _format_stat("output_tps", 80.0) == "80.0 tok/s"
+
+ def test_inverse(self):
+ result = _format_stat("time_per_output_token-p50", 0.04, invert=True)
+ assert "tok/s" in result
+ assert "25.0" in result
+
+ def test_integer_value(self):
+ assert _format_stat("failed_requests", 3) == "3"
+
+ def test_float_that_is_whole(self):
+ assert _format_stat("failed_requests", 0.0) == "0"
+
+
+# ── LiveStatsDisplay ─────────────────────────────────────────────────────────
+
+
+class TestLiveStatsDisplay:
+ def test_disabled_does_nothing(self):
+ display = LiveStatsDisplay(disabled=True)
+ display.update({"rpm": 100})
+ display.close()
+
+ def test_format_stats_with_empty_raw(self):
+ display = LiveStatsDisplay()
+ result = display.format_stats({})
+ assert all(v == "—" for v in result.values())
+ assert "rpm" in result
+ assert "fail" in result
+
+ def test_format_stats_with_data(self):
+ display = LiveStatsDisplay(
+ display_stats={
+ "rpm": "rpm",
+ "fail": "failed_requests",
+ "p50_ttft": "time_to_first_token-p50",
+ }
+ )
+ raw = {
+ "rpm": 185.9,
+ "failed_requests": 0,
+ "time_to_first_token-p50": 0.312,
+ }
+ result = display.format_stats(raw)
+ assert result["rpm"] == "185.9"
+ assert result["fail"] == "0"
+ assert result["p50_ttft"] == "0.312s"
+
+ def test_format_stats_inverse(self):
+ display = LiveStatsDisplay(
+ display_stats={"tps": ("time_per_output_token-p50", "inv")}
+ )
+ raw = {"time_per_output_token-p50": 0.04}
+ result = display.format_stats(raw)
+ assert "tok/s" in result["tps"]
+
+ def test_format_stats_missing_key_shows_placeholder(self):
+ display = LiveStatsDisplay(
+ display_stats={"rpm": "rpm", "missing": "nonexistent_key"}
+ )
+ result = display.format_stats({"rpm": 100.0})
+ assert result["rpm"] == "100.0"
+ assert result["missing"] == "—"
+
+ def test_custom_display_stats(self):
+ custom = {"latency": "time_to_last_token-p99", "errors": "failed_requests"}
+ display = LiveStatsDisplay(display_stats=custom)
+ assert display._display_stats == custom
+
+ def test_default_display_stats_used(self):
+ display = LiveStatsDisplay()
+ assert display._display_stats is DEFAULT_DISPLAY_STATS
+
+ def test_terminal_output(self):
+ display = LiveStatsDisplay(
+ disabled=False,
+ display_stats={"rpm": "rpm", "fail": "failed_requests"},
+ )
+ display._is_notebook = False
+ display.update({"rpm": 100.0, "failed_requests": 0})
+ assert display._last_line_count > 0
+ display.close()
+ assert display._last_line_count == 0
+
+ def test_terminal_with_prefix(self):
+ display = LiveStatsDisplay(disabled=False, display_stats={"rpm": "rpm"})
+ display._is_notebook = False
+ display.update({"rpm": 100.0}, extra_prefix="reqs=42")
+ assert display._last_line_count >= 2 # prefix line + stats line
+ display.close()
+
+ def test_terminal_overwrites_previous(self):
+ display = LiveStatsDisplay(disabled=False, display_stats={"rpm": "rpm"})
+ display._is_notebook = False
+ display.update({"rpm": 100.0})
+ first_count = display._last_line_count
+ display.update({"rpm": 200.0})
+ assert display._last_line_count == first_count
+ display.close()
diff --git a/tests/unit/test_results.py b/tests/unit/test_results.py
index 9262d16..e949b04 100644
--- a/tests/unit/test_results.py
+++ b/tests/unit/test_results.py
@@ -192,7 +192,9 @@ def test_stats_property(sample_result: Result):
assert key in stats
# Test caching returns same object for built-in stats:
- assert sample_result._builtin_stats is sample_result._builtin_stats
+ assert sample_result._preloaded_stats is None or isinstance(
+ sample_result._preloaded_stats, dict
+ )
def test_stats_property_empty_result():
@@ -306,3 +308,189 @@ def test_save_method_existing_responses(sample_result: Result, temp_dir: UPath):
responses = [json.loads(line) for line in f]
assert len(responses) == 6 # 5 original + 1 extra
assert responses[-1]["id"] == "extra_response"
+
+
+# ── Contributed stats round-trip ─────────────────────────────────────────────
+
+
+class TestContributedStatsRoundTrip:
+ """Verify that callback-contributed stats survive save → load cycles."""
+
+ @pytest.fixture
+ def result_with_contributed_stats(self):
+ responses = [
+ InvocationResponse(
+ id=f"r{i}",
+ response_text=f"resp {i}",
+ input_prompt=f"prompt {i}",
+ time_to_first_token=0.1 * i,
+ time_to_last_token=0.2 * i,
+ num_tokens_output=10 * i,
+ num_tokens_input=5 * i,
+ )
+ for i in range(1, 4)
+ ]
+ result = Result(
+ responses=responses,
+ total_requests=3,
+ clients=1,
+ n_requests=3,
+ total_test_time=1.0,
+ )
+ result._update_contributed_stats(
+ {"custom_metric_a": 42.0, "custom_metric_b": 99.5}
+ )
+ return result
+
+ def test_contributed_stats_appear_in_stats(self, result_with_contributed_stats):
+ stats = result_with_contributed_stats.stats
+ assert stats["custom_metric_a"] == 42.0
+ assert stats["custom_metric_b"] == 99.5
+
+ def test_contributed_stats_written_to_stats_json(
+ self, result_with_contributed_stats, tmp_path
+ ):
+ output = UPath(tmp_path / "out")
+ result_with_contributed_stats.save(output)
+
+ with (output / "stats.json").open() as f:
+ saved = json.load(f)
+ assert saved["custom_metric_a"] == 42.0
+ assert saved["custom_metric_b"] == 99.5
+
+ def test_load_with_responses_preserves_contributed_stats(
+ self, result_with_contributed_stats, tmp_path
+ ):
+ output = UPath(tmp_path / "out")
+ result_with_contributed_stats.save(output)
+
+ loaded = Result.load(output, load_responses=True)
+
+ assert loaded.stats["custom_metric_a"] == 42.0
+ assert loaded.stats["custom_metric_b"] == 99.5
+
+ def test_load_without_responses_preserves_contributed_stats(
+ self, result_with_contributed_stats, tmp_path
+ ):
+ output = UPath(tmp_path / "out")
+ result_with_contributed_stats.save(output)
+
+ loaded = Result.load(output, load_responses=False)
+
+ assert loaded.stats["custom_metric_a"] == 42.0
+ assert loaded.stats["custom_metric_b"] == 99.5
+
+ def test_contributed_stats_do_not_clobber_builtin_stats(
+ self, result_with_contributed_stats, tmp_path
+ ):
+ output = UPath(tmp_path / "out")
+ result_with_contributed_stats.save(output)
+
+ loaded = Result.load(output, load_responses=True)
+
+ # Builtin stats must still be present and correct
+ assert "failed_requests" in loaded.stats
+ assert loaded.stats["total_requests"] == 3
+ assert "time_to_first_token-p50" in loaded.stats
+
+ def test_builtin_stats_not_overwritten_by_stale_saved_values(self, tmp_path):
+ """If a builtin key exists in stats.json with a stale value, the freshly
+ computed value from responses should win."""
+ responses = [
+ InvocationResponse(
+ id="x",
+ response_text="r",
+ input_prompt="p",
+ time_to_first_token=0.5,
+ time_to_last_token=1.0,
+ num_tokens_output=10,
+ num_tokens_input=5,
+ )
+ ]
+ result = Result(
+ responses=responses,
+ total_requests=1,
+ clients=1,
+ n_requests=1,
+ total_test_time=2.0,
+ )
+ output = UPath(tmp_path / "out")
+ result.save(output)
+
+ # Tamper with stats.json: set a wrong value for a builtin key
+ stats_path = output / "stats.json"
+ with stats_path.open() as f:
+ saved = json.load(f)
+ saved["failed_requests"] = 999
+ with stats_path.open("w") as f:
+ json.dump(saved, f)
+
+ loaded = Result.load(output, load_responses=True)
+
+ # The freshly computed value (0 failures) should win over the tampered 999
+ assert loaded.stats["failed_requests"] == 0
+
+ def test_load_responses_recomputes_but_keeps_contributed(self, tmp_path):
+ """After load(load_responses=False) + load_responses(), contributed
+ stats from stats.json should still be accessible via _preloaded_stats
+ even though responses were reloaded."""
+ responses = [
+ InvocationResponse(
+ id="z",
+ response_text="r",
+ input_prompt="p",
+ time_to_first_token=0.3,
+ time_to_last_token=0.6,
+ num_tokens_output=8,
+ num_tokens_input=4,
+ )
+ ]
+ result = Result(
+ responses=responses,
+ total_requests=1,
+ clients=1,
+ n_requests=1,
+ total_test_time=1.0,
+ )
+ result._update_contributed_stats({"cb_stat": 7.0})
+ output = UPath(tmp_path / "out")
+ result.save(output)
+
+ loaded = Result.load(output, load_responses=False)
+ assert loaded.stats["cb_stat"] == 7.0
+
+ # Now reload responses — _preloaded_stats gets recomputed from
+ # responses only, so cb_stat won't be in _preloaded_stats anymore,
+ # but it was never in _contributed_stats on the loaded instance either.
+ loaded.load_responses()
+ # After recompute, builtin stats should be correct
+ assert loaded.stats["failed_requests"] == 0
+ assert "time_to_first_token-p50" in loaded.stats
+
+ def test_multiple_contributed_stats_updates_merge(self, tmp_path):
+ responses = [
+ InvocationResponse(
+ id="m",
+ response_text="r",
+ input_prompt="p",
+ num_tokens_output=5,
+ num_tokens_input=3,
+ )
+ ]
+ result = Result(
+ responses=responses,
+ total_requests=1,
+ clients=1,
+ n_requests=1,
+ total_test_time=0.5,
+ )
+ result._update_contributed_stats({"stat_a": 1.0})
+ result._update_contributed_stats({"stat_b": 2.0})
+ result._update_contributed_stats({"stat_a": 10.0}) # overwrite
+
+ output = UPath(tmp_path / "out")
+ result.save(output)
+
+ loaded = Result.load(output, load_responses=True)
+ assert loaded.stats["stat_a"] == 10.0
+ assert loaded.stats["stat_b"] == 2.0
diff --git a/tests/unit/test_runner.py b/tests/unit/test_runner.py
index 4150dee..616988e 100644
--- a/tests/unit/test_runner.py
+++ b/tests/unit/test_runner.py
@@ -70,11 +70,11 @@ def run(mock_endpoint: MagicMock, mock_tokenizer: MagicMock):
mock_run._queue = AsyncMock()
mock_run._queue.task_done = MagicMock()
- # Mock the _invoke_n_c method to return a simple result
- async def mock_invoke_n_c(payload, n_requests, clients):
+ # Mock the _invoke_clients method to return a simple result
+ async def mock_invoke_clients(payload, n_requests=None, duration=None, clients=1):
return 1.0, [], []
- mock_run._invoke_n_c = mock_invoke_n_c
+ mock_run._invoke_clients = mock_invoke_clients
# Mock the _process_results_from_q method
async def mock_process_results_from_q(output_path=None):
@@ -126,7 +126,7 @@ async def test_invoke_n(run: _Run):
]
)
- result = await run._invoke_n(
+ result = await run._invoke_client(
payload=[{"prompt": "test1"}, {"prompt": "test2"}], n=2
)
@@ -180,7 +180,7 @@ async def test_invoke_n_no_wait(run: _Run):
@pytest.mark.asyncio
async def test_invoke_n_c(run: _Run):
# Remove the fixture override and create a proper mock
- async def mock_invoke_n_c(payload, n_requests, clients):
+ async def mock_invoke_clients(payload, n_requests=None, duration=None, clients=1):
# Simulate the actual behavior
responses = [
InvocationResponse(id="1", input_prompt="test1", response_text="response1"),
@@ -189,9 +189,9 @@ async def mock_invoke_n_c(payload, n_requests, clients):
return 1.5, responses, [] # total_time, responses, errors
# Replace the fixture mock with our test-specific mock
- run._invoke_n_c = mock_invoke_n_c
+ run._invoke_clients = mock_invoke_clients
- total_test_time, responses, _ = await run._invoke_n_c(
+ total_test_time, responses, _ = await run._invoke_clients(
payload=[{"prompt": "test"}], n_requests=2, clients=1
)
@@ -238,7 +238,7 @@ async def test_run_with_output_path(runner: Runner, tmp_path: Path):
@pytest.mark.asyncio
async def test_run_error_handling(run: _Run):
- run._invoke_n_c = AsyncMock(side_effect=Exception("Test error"))
+ run._invoke_clients = AsyncMock(side_effect=Exception("Test error"))
run._process_results_from_q = AsyncMock()
with pytest.raises(Exception, match="Test error"):
@@ -432,7 +432,7 @@ def test_run_output_path(runner: Runner, tmp_path: Path):
@pytest.mark.asyncio
async def test_invoke_n_edge_cases(run: _Run):
# Test with empty payload
- result = await run._invoke_n(payload=[], n=5)
+ result = await run._invoke_client(payload=[], n=5)
assert not result
# Test with n=None (should use all payloads)
@@ -442,7 +442,7 @@ async def test_invoke_n_edge_cases(run: _Run):
InvocationResponse(id="2", input_prompt="test2", response_text="response2"),
]
)
- result = await run._invoke_n(
+ result = await run._invoke_client(
payload=[{"prompt": "test1"}, {"prompt": "test2"}], n=None
)
assert len(result) == 2
@@ -535,7 +535,9 @@ def test_prepare_run_combinations(
)
assert isinstance(run.payload, list)
- assert run.n_requests == n_requests
+ # When n_requests is None, it defaults to len(payload)
+ expected_n = n_requests if n_requests is not None else len(run.payload)
+ assert run.n_requests == expected_n
assert run.clients == clients
assert run.output_path == (Path(output_path) if output_path else None)
assert run.run_name is not None
@@ -566,7 +568,9 @@ async def test_run_with_different_payloads(
@pytest.mark.asyncio
async def test_invoke_n_c_concurrent_execution(run: _Run):
- async def mock_invoke_n(payload, n, add_start_jitter=True, shuffle_order=True):
+ async def mock_invoke_client(
+ payload, n=None, duration=None, add_start_jitter=True, shuffle_order=True
+ ):
await asyncio.sleep(0.1) # Simulate some processing time
return [
InvocationResponse(
@@ -575,10 +579,10 @@ async def mock_invoke_n(payload, n, add_start_jitter=True, shuffle_order=True):
for i in range(n)
]
- run._invoke_n = mock_invoke_n # type: ignore
+ run._invoke_client = mock_invoke_client # type: ignore
start_time = time.perf_counter()
- total_test_time, _, _ = await run._invoke_n_c(
+ total_test_time, _, _ = await run._invoke_clients(
payload=[{"prompt": "test"}], n_requests=5, clients=3
)
end_time = time.perf_counter()
@@ -757,7 +761,9 @@ def test_prepare_run_more_edge_cases(
)
assert isinstance(run_config.payload, list)
- assert run_config.n_requests == n_requests
+ # When n_requests is None, it defaults to len(payload)
+ expected_n = n_requests if n_requests is not None else len(run_config.payload)
+ assert run_config.n_requests == expected_n
assert run_config.clients == clients if clients is not None else 1
assert run_config.output_path == (Path(output_path) if output_path else None)
assert run_config.run_name is not None
@@ -806,7 +812,9 @@ async def test_run_with_optional_parameters(
async def test_invoke_n_c_with_different_clients(
run: _Run, clients: Literal[1] | Literal[3] | Literal[5] | Literal[10]
):
- async def mock_invoke_n(payload, n, add_start_jitter=True, shuffle_order=True):
+ async def mock_invoke_client(
+ payload, n=None, duration=None, add_start_jitter=True, shuffle_order=True
+ ):
await asyncio.sleep(0.1) # Simulate some processing time
return [
InvocationResponse(
@@ -815,10 +823,10 @@ async def mock_invoke_n(payload, n, add_start_jitter=True, shuffle_order=True):
for i in range(n)
]
- run._invoke_n = mock_invoke_n # type: ignore
+ run._invoke_client = mock_invoke_client # type: ignore
start_time = time.perf_counter()
- total_test_time, _, _ = await run._invoke_n_c(
+ total_test_time, _, _ = await run._invoke_clients(
payload=[{"prompt": "test"}], n_requests=5, clients=clients
)
end_time = time.perf_counter()
@@ -956,7 +964,7 @@ async def test_invoke_n_with_different_options(
]
)
- result = await run._invoke_n(
+ result = await run._invoke_client(
payload=[{"prompt": "test1"}, {"prompt": "test2"}],
n=2,
shuffle_order=shuffle_order,
@@ -965,7 +973,7 @@ async def test_invoke_n_with_different_options(
assert len(result) == 2
run._invoke_n_no_wait.assert_called_once_with(
- [{"prompt": "test1"}, {"prompt": "test2"}], 2, shuffle_order
+ [{"prompt": "test1"}, {"prompt": "test2"}], 2, None, shuffle_order
)
@@ -997,3 +1005,208 @@ async def test_count_tokens_from_q_with_custom_output_path(run: _Run, tmp_path:
# Add more tests for edge cases and other methods as needed
+
+
+# ── Time-bound (run_duration) tests ──────────────────────────────────────────
+
+
+def test_run_duration_and_n_requests_mutually_exclusive(
+ mock_endpoint: Endpoint, mock_tokenizer: MagicMock
+):
+ """Setting both n_requests and run_duration should raise ValueError."""
+ with pytest.raises(ValueError, match="Cannot set both"):
+ _Run(
+ endpoint=mock_endpoint,
+ tokenizer=mock_tokenizer,
+ payload=[{"prompt": "test"}],
+ n_requests=10,
+ run_duration=5,
+ clients=1,
+ run_name="test_run",
+ )
+
+
+def test_run_duration_sets_time_bound_flag(
+ mock_endpoint: Endpoint, mock_tokenizer: MagicMock
+):
+ """When run_duration is set, _time_bound should be True and n_requests 0."""
+ run = _Run(
+ endpoint=mock_endpoint,
+ tokenizer=mock_tokenizer,
+ payload=[{"prompt": "test"}],
+ run_duration=5,
+ clients=1,
+ run_name="test_run",
+ )
+ assert run._time_bound is True
+ assert run.n_requests == 0
+
+
+def test_n_requests_sets_count_bound(
+ mock_endpoint: Endpoint, mock_tokenizer: MagicMock
+):
+ """When n_requests is set (no run_duration), _time_bound should be False."""
+ run = _Run(
+ endpoint=mock_endpoint,
+ tokenizer=mock_tokenizer,
+ payload=[{"prompt": "test"}],
+ n_requests=10,
+ clients=1,
+ run_name="test_run",
+ )
+ assert run._time_bound is False
+ assert run.n_requests == 10
+
+
+def test_run_duration_must_be_positive(
+ mock_endpoint: Endpoint, mock_tokenizer: MagicMock
+):
+ """run_duration must be > 0."""
+ with pytest.raises(AssertionError, match="positive"):
+ _Run(
+ endpoint=mock_endpoint,
+ tokenizer=mock_tokenizer,
+ payload=[{"prompt": "test"}],
+ run_duration=-1,
+ clients=1,
+ run_name="test_run",
+ )
+
+
+def test_invoke_for_duration_respects_deadline(
+ mock_endpoint: Endpoint, mock_tokenizer: MagicMock
+):
+ """_invoke_n_no_wait with duration should stop after the specified duration."""
+ run = _Run(
+ endpoint=mock_endpoint,
+ tokenizer=mock_tokenizer,
+ payload=[{"prompt": "test"}],
+ run_duration=0.5,
+ clients=1,
+ run_name="test_run",
+ )
+ run.callbacks = None
+ run._queue = MagicMock()
+ run._queue._loop.call_soon_threadsafe = MagicMock()
+
+ # Make invoke take ~100ms so we get a handful of requests
+ def slow_invoke(payload):
+ time.sleep(0.1)
+ return InvocationResponse(id="1", input_prompt="test", response_text="response")
+
+ run._endpoint.invoke.side_effect = slow_invoke
+
+ start = time.perf_counter()
+ responses = run._invoke_n_no_wait(payload=[{"prompt": "test"}], duration=0.5)
+ elapsed = time.perf_counter() - start
+
+ assert len(responses) > 0
+ assert elapsed < 1.0 # Should not overshoot by much
+ assert all(isinstance(r, InvocationResponse) for r in responses)
+
+
+def test_invoke_for_duration_cycles_payloads(
+ mock_endpoint: Endpoint, mock_tokenizer: MagicMock
+):
+ """_invoke_n_no_wait with duration should cycle through payloads."""
+ run = _Run(
+ endpoint=mock_endpoint,
+ tokenizer=mock_tokenizer,
+ payload=[{"prompt": "a"}, {"prompt": "b"}],
+ run_duration=0.3,
+ clients=1,
+ run_name="test_run",
+ )
+ run.callbacks = None
+ run._queue = MagicMock()
+ run._queue._loop.call_soon_threadsafe = MagicMock()
+
+ payloads_seen = []
+
+ def tracking_invoke(payload):
+ payloads_seen.append(payload)
+ return InvocationResponse(id="1", input_prompt=str(payload), response_text="ok")
+
+ run._endpoint.invoke.side_effect = tracking_invoke
+
+ responses = run._invoke_n_no_wait(
+ payload=[{"prompt": "a"}, {"prompt": "b"}],
+ duration=0.3,
+ shuffle_order=False,
+ )
+
+ assert len(responses) >= 2
+ # Should see both payloads used (cycling)
+ prompts = [p.get("prompt") for p in payloads_seen]
+ assert "a" in prompts
+ assert "b" in prompts
+
+
+@pytest.mark.asyncio
+async def test_run_with_duration(runner: Runner):
+ """Full run() with run_duration should complete and report actual counts."""
+ result = await runner.run(
+ payload={"prompt": "test"},
+ run_duration=0.3,
+ clients=1,
+ )
+
+ assert result.total_requests > 0
+ assert result.n_requests > 0
+ assert result.total_test_time is not None
+ assert result.total_test_time > 0
+ assert result.stats["total_requests"] == result.total_requests
+
+
+@pytest.mark.asyncio
+async def test_run_with_duration_multiple_clients(runner: Runner):
+ """Time-bound run with multiple clients should aggregate counts."""
+ result = await runner.run(
+ payload={"prompt": "test"},
+ run_duration=0.3,
+ clients=3,
+ )
+
+ assert result.total_requests > 0
+ assert result.clients == 3
+ assert result.total_test_time is not None
+
+
+@pytest.mark.asyncio
+async def test_run_with_duration_and_output_path(runner: Runner, tmp_path: Path):
+ """Time-bound run with output_path should save results to disk."""
+ result = await runner.run(
+ payload={"prompt": "test"},
+ run_duration=0.3,
+ clients=1,
+ output_path=tmp_path / "duration_run",
+ run_name="dur_test",
+ )
+
+ assert result.output_path is not None
+ assert (tmp_path / "duration_run" / "responses.jsonl").exists()
+ assert (tmp_path / "duration_run" / "summary.json").exists()
+ assert (tmp_path / "duration_run" / "stats.json").exists()
+
+
+def test_prepare_run_with_duration(runner: Runner):
+ """_prepare_run should pass run_duration through to _Run."""
+ run = runner._prepare_run(
+ payload={"prompt": "test"},
+ run_duration=30,
+ clients=2,
+ )
+ assert run._time_bound is True
+ assert run.run_duration == 30
+ assert run.n_requests == 0
+
+
+def test_prepare_run_duration_and_n_requests_conflict(runner: Runner):
+ """_prepare_run should raise when both are set."""
+ with pytest.raises(ValueError, match="Cannot set both"):
+ runner._prepare_run(
+ payload={"prompt": "test"},
+ n_requests=10,
+ run_duration=30,
+ clients=2,
+ )
diff --git a/tests/unit/test_running_stats.py b/tests/unit/test_running_stats.py
new file mode 100644
index 0000000..634483a
--- /dev/null
+++ b/tests/unit/test_running_stats.py
@@ -0,0 +1,198 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import time
+
+import pytest
+
+from llmeter.utils import RunningStats
+
+
+@pytest.fixture
+def rs():
+ return RunningStats(
+ metrics=[
+ "time_to_first_token",
+ "time_to_last_token",
+ "time_per_output_token",
+ "num_tokens_input",
+ "num_tokens_output",
+ ]
+ )
+
+
+@pytest.fixture
+def populated_rs(rs):
+ """A RunningStats with 3 responses recorded."""
+ responses = [
+ {
+ "time_to_first_token": 0.3,
+ "time_to_last_token": 0.8,
+ "time_per_output_token": 0.02,
+ "num_tokens_input": 100,
+ "num_tokens_output": 25,
+ "error": None,
+ },
+ {
+ "time_to_first_token": 0.5,
+ "time_to_last_token": 1.2,
+ "time_per_output_token": 0.03,
+ "num_tokens_input": 120,
+ "num_tokens_output": 30,
+ "error": None,
+ },
+ {
+ "time_to_first_token": 0.4,
+ "time_to_last_token": 1.0,
+ "time_per_output_token": 0.025,
+ "num_tokens_input": 110,
+ "num_tokens_output": 28,
+ "error": "timeout",
+ },
+ ]
+ for r in responses:
+ rs.record_send()
+ rs.update(r)
+ return rs
+
+
+# ── record_send ──────────────────────────────────────────────────────────────
+
+
+class TestRecordSend:
+ def test_first_send_sets_first_time(self, rs):
+ assert rs._first_send_time is None
+ rs.record_send()
+ assert rs._first_send_time is not None
+ assert rs._last_send_time is not None
+ assert rs._sends == 1
+
+ def test_subsequent_sends_update_last_time(self, rs):
+ rs.record_send()
+ first = rs._first_send_time
+ time.sleep(0.01)
+ rs.record_send()
+ assert rs._first_send_time == first
+ assert rs._last_send_time > first
+ assert rs._sends == 2
+
+ def test_send_count_increments(self, rs):
+ for _ in range(5):
+ rs.record_send()
+ assert rs._sends == 5
+
+
+# ── update ───────────────────────────────────────────────────────────────────
+
+
+class TestUpdate:
+ def test_count_increments(self, rs):
+ rs.update({"time_to_first_token": 0.3, "error": None})
+ assert rs._count == 1
+ rs.update({"time_to_first_token": 0.5, "error": None})
+ assert rs._count == 2
+
+ def test_failed_count(self, rs):
+ rs.update({"error": "timeout"})
+ rs.update({"error": None})
+ rs.update({"error": "connection refused"})
+ assert rs._failed == 2
+
+ def test_none_values_skipped(self, rs):
+ rs.update({"time_to_first_token": None, "error": None})
+ assert len(rs._values["time_to_first_token"]) == 0
+
+ def test_nan_values_skipped(self, rs):
+ rs.update({"time_to_first_token": float("nan"), "error": None})
+ assert len(rs._values["time_to_first_token"]) == 0
+
+ def test_sums_accumulated(self, rs):
+ rs.update({"num_tokens_output": 10, "error": None})
+ rs.update({"num_tokens_output": 20, "error": None})
+ assert rs._sums["num_tokens_output"] == 30
+
+ def test_values_sorted(self, rs):
+ rs.update({"time_to_first_token": 0.5, "error": None})
+ rs.update({"time_to_first_token": 0.1, "error": None})
+ rs.update({"time_to_first_token": 0.3, "error": None})
+ assert rs._values["time_to_first_token"] == [0.1, 0.3, 0.5]
+
+
+# ── to_stats ─────────────────────────────────────────────────────────────────
+
+
+class TestToStats:
+ def test_basic_stats(self, populated_rs):
+ stats = populated_rs.to_stats()
+ assert stats["failed_requests"] == 1
+ assert "time_to_first_token-p50" in stats
+ assert "time_to_last_token-average" in stats
+ assert "num_tokens_output-p90" in stats
+
+ def test_with_run_context(self, populated_rs):
+ stats = populated_rs.to_stats(
+ total_requests=3,
+ total_test_time=10.0,
+ result_dict={"model_id": "test"},
+ )
+ assert stats["model_id"] == "test"
+ assert stats["requests_per_minute"] == pytest.approx(18.0)
+ assert stats["failed_requests_rate"] == pytest.approx(1 / 3)
+ assert stats["total_output_tokens"] == 83
+
+ def test_without_run_context(self, populated_rs):
+ stats = populated_rs.to_stats()
+ assert stats["failed_requests"] == 1
+ assert stats["total_input_tokens"] == 330
+ assert stats["total_output_tokens"] == 83
+
+ def test_empty_stats(self, rs):
+ stats = rs.to_stats()
+ assert stats["failed_requests"] == 0
+ assert stats["total_input_tokens"] == 0
+
+
+# ── send-window throughput in to_stats ────────────────────────────────────────
+
+
+class TestSendWindowStats:
+ def test_rpm_uses_send_window(self, rs):
+ rs._first_send_time = 100.0
+ rs._last_send_time = 110.0 # 10 second window
+ rs.update({"error": None})
+ rs.update({"error": None})
+ rs.update({"error": None})
+ stats = rs.to_stats()
+ # 3 responses / 10 seconds * 60 = 18.0 rpm
+ assert stats["rpm"] == pytest.approx(18.0)
+
+ def test_output_tps_uses_send_window(self, rs):
+ rs._first_send_time = 100.0
+ rs._last_send_time = 110.0 # 10 second window
+ rs.update({"num_tokens_output": 500, "error": None})
+ rs.update({"num_tokens_output": 300, "error": None})
+ stats = rs.to_stats()
+ # 800 tokens / 10 seconds = 80.0 tok/s
+ assert stats["output_tps"] == pytest.approx(80.0)
+
+ def test_no_send_window_when_single_send(self, rs):
+ """With only one send, first == last, no window to compute RPM."""
+ rs._first_send_time = 100.0
+ rs._last_send_time = 100.0
+ rs.update({"error": None})
+ stats = rs.to_stats()
+ assert "rpm" not in stats
+ assert "output_tps" not in stats
+
+ def test_no_send_window_when_no_sends(self, rs):
+ stats = rs.to_stats()
+ assert "rpm" not in stats
+ assert "output_tps" not in stats
+
+ def test_send_window_helper(self, rs):
+ assert rs._send_window() is None
+ rs._first_send_time = 10.0
+ rs._last_send_time = 10.0
+ assert rs._send_window() is None
+ rs._last_send_time = 20.0
+ assert rs._send_window() == pytest.approx(10.0)