Skip to content

add memory preflight check#168

Open
mike-ferguson wants to merge 1 commit intomainfrom
add_mem_preflight
Open

add memory preflight check#168
mike-ferguson wants to merge 1 commit intomainfrom
add_mem_preflight

Conversation

@mike-ferguson
Copy link
Copy Markdown
Member

Add preallocate_memory hook and score_benchmark helper to Benchmark

Problem

Neural benchmarks can run for 6-7 hours and crash with an OOM error at the very end
when trying to allocate the full activation array. There was no way to detect this early.

Changes

brainscore_core/benchmarks/__init__.py

  • Added Benchmark.preallocate_memory(candidate) — a no-op by default. Domain-specific
    subclasses (e.g. NeuralBenchmark in brainscore_vision) can override this to run a
    pre-flight memory check before scoring begins.
  • Added score_benchmark(benchmark, candidate) — the recommended call site instead of
    benchmark(candidate) directly. Calls preallocate_memory first, then scores.

tests/test_plugin_management/test_benchmarks.py — new unit test file (13 tests)
covering the no-op default, call ordering, MemoryError propagation, and env var bypass.

Design notes

  • The default preallocate_memory is a no-op, so all existing benchmarks and plugins
    are completely unaffected.
  • score_benchmark is the hook point — domain libraries call this instead of
    benchmark(candidate) and get the check for free.
  • The actual probe logic (running 1 stimulus through the model, estimating GB) lives in
    brainscore_vision, not here. Core only defines the interface.
  • Language benchmarks can adopt the same pattern by overriding preallocate_memory when
    ready — no further core changes needed.

Companion PR (vision)

brainscore_vision will add the probe implementation in benchmark_helpers/memory.py and
wires NeuralBenchmark.preallocate_memory to call it. This core PR must be merged and
released before the vision PR can be merged.

@mike-ferguson mike-ferguson requested a review from KartikP April 9, 2026 15:44
Copy link
Copy Markdown
Contributor

@KartikP KartikP left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! Lots of tests, but comprehensive

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.

2 participants