Skip to content

feat(model): per-attempt ModelEvent retry accounting and timing#4526

Open
sjawhar wants to merge 1 commit into
UKGovernmentBEIS:mainfrom
trajectory-labs-pbc:feat/model-event-retry-accounting
Open

feat(model): per-attempt ModelEvent retry accounting and timing#4526
sjawhar wants to merge 1 commit into
UKGovernmentBEIS:mainfrom
trajectory-labs-pbc:feat/model-event-retry-accounting

Conversation

@sjawhar

@sjawhar sjawhar commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

This PR contains:

  • New features

What is the current behavior?

ModelEvent is emitted once per model.generate() call after all retries complete. Several consequences:

  • No per-attempt visibility into which attempt produced an event, how many retries occurred, or how long each attempt took.
  • Timing fields conflate per-attempt latency with retry overhead. working_time on the event is the wall span of the whole call, not per-attempt working time.
  • Two distinct retry mechanisms (Tenacity-scheduled outer retries and provider/internal HTTP retries surfaced via report_http_retry()) are not separately accounted.
  • Downstream tooling (analysis dataframes, hooks, log replay) cannot reliably reconstruct per-attempt behavior from the event stream.

What is the new behavior?

model.generate() now emits one ModelEvent per attempt, sharing a per-call identifier:

  • New ModelGenerateAccounting (in inspect_ai.model._generate_accounting) is set on a ContextVar for the duration of a generate() call and registers each attempt's event.
  • ModelEvent gains:
    • call_id (shared across all attempts of one generate())
    • attempt (1-indexed)
    • call_started_at / call_completed_at (per-attempt UTC datetimes)
    • call_working_start / call_working_time (per-attempt monotonic working-time span)
    • call_retries — count of outer scheduled retries (Tenacity before_sleep)
    • http_retries — count of provider/HTTP retry signals via report_http_retry()
  • tests/_helpers/event_assertions.py and tests/_helpers/retry_provider.py expose reusable helpers for asserting the per-attempt invariants.
  • inspect_ai.analysis._dataframe.events.columns exports the new fields so dataframe-based consumers see them.
  • Hooks (ModelUsageData / emit_model_usage) carry the new accounting forward.
  • inspect-openapi.json is regenerated. Legacy logs (without these fields) round-trip via a call_id is None sentinel handled in tests/log/test_legacy_log_compat.py.
  • count_tokens and compact remain eventless; provider-returned Exception results keep the existing RuntimeError wrap with no retry.

The split between call_retries and http_retries is intentional:

  • call_retries is the count of outer scheduled retries (each one starts a new ModelEvent attempt)
  • http_retries is the count of provider-internal HTTP retries reported via report_http_retry() (these do not start a new attempt)

Does this PR introduce a breaking change?

Additive schema changes to ModelEvent. Consumers reading the OpenAPI schema or log files will see new optional fields. Existing logs continue to deserialize (no required new fields). Downstream code that introspects ModelEvent and the dataframe export should benefit but may need to opt into the new fields.

Other information:

  • Tests: 17 new test files covering retry timing, attempt grouping, terminal finalization, attempt timeout, provider-returned exceptions, retry classification, sample/model retry interaction, dataframe export, legacy-log compatibility, property-based invariants, and tool/subtask smoke.
  • New deps (dev): hypothesis, mutmut (with a scoped [tool.mutmut] for _generate_accounting.py).
  • make check (ruff + mypy) is clean; targeted retry-accounting tests pass under asyncio and trio.
  • This is one focused PR. Several adjacent fixes I made in my fork while preparing this work are coming as separate small PRs (FastAPI 0.118 OpenAPI compat, S3 conditional-write ETag handling, GenerateConfig() shared mutable default, xdist teardown + hook-entrypoint ordering, _model_info import-cycle break).

@sjawhar

sjawhar commented Jul 17, 2026

Copy link
Copy Markdown
Contributor Author

CI status note: check-schema-and-types

The check-schema-and-types check fails on this PR because it requires the ts-mono submodule's generated.ts to match what python src/inspect_ai/_view/schema.py produces. Since this PR adds new fields to ModelEvent, the regenerated generated.ts differs from what's pinned at the current ts-mono commit (c4b97ef71a).

I've opened the corresponding ts-mono PR with the regenerated types: meridianlabs-ai/ts-mono#210.

Once that merges, I'll push a commit here that bumps the ts-mono submodule pin to the new ts-mono main, and check-schema-and-types should go green.

All other CI checks pass:

  • ruff (3.10) / ruff (3.11)
  • mypy (3.10) / mypy (3.11)
  • pre-commit
  • test (3.10) / test (3.11)
  • submodule-on-main / dist-validation / check-tool-paths / Build & inspect the package.

@sjawhar

sjawhar commented Jul 17, 2026

Copy link
Copy Markdown
Contributor Author

ts-mono types PR moved: meridianlabs-ai/ts-mono#210meridianlabs-ai/ts-mono#430 (rebased onto ts-mono main, resubmitted from the trajectory-labs-pbc fork).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant