initial implementation of DiskStore

[quaquel]

summary from Claude, will write this out in more readable form later.

DiskStore: durable, resumable storage for parameter sweeps

Part 4 of the experimentation framework PR sequence. Adds a persistent Store
implementation backed by Arrow IPC streams and JSON manifests, designed for
HPC use: shared filesystems (Lustre/GPFS), SLURM timeouts, and multi-node
execution.

What this PR adds

Three new modules under mesa/experimental/scenarios/:

• store_metadata.py — store-agnostic durable metadata: the store
  manifest (store.json), the scenario manifest (scenarios.json), and the
  status log (status.log). Reusable by any future persistent store including
  persisting InMemoryStore.
• disk_writer.py — the worker-side DiskStreamWriter, which appends
  each completed run as one Arrow record batch to per-worker stream files.
• disk_store.py — the root-side DiskStore, implementing the Store
  protocol: manifest writes, status recording, and the read path that
  assembles outcomes back from the worker files.

Plus a small change to runner.py: RunConfiguration.__call__ now enforces a
postcondition that every extracted output has at least one column (a
columnless frame is recorded as an EXTRACTING failure).

On-disk layout

    store_dir/
    ├── store.json          # format version, sessions, provenance (written once)
    ├── scenarios.json      # authoritative, seed-complete (written once)
    ├── status.log          # append-only, root-written, torn-tail-tolerant
    └── outputs/
        └── {output_name}/
            └── worker-{session}-{host}-{pid}.arrow

Key design decisions and their rationale

Data is written where it is produced

Workers write their own outcomes to disk; only a key-only DiskReference
(just a RunId) crosses back to the root. Nothing is funneled through the
root process, so result size does not bottleneck on the executor's return
pipe.

Arrow IPC stream format, one batch per run, no buffering

Each completed run is appended immediately as one record batch. The stream
format is append-native and every flushed batch is durable at once — a worker
killed by a SLURM timeout loses at most the batch being written.

The IPC file format is deliberately rejected: its footer is written at close, so a killed worker
would leave nothing readable. Reads in this implementation tolerate torn tails by keeping every
complete batch before the tear.

One file per worker per output

Bounds file count at workers × outputs regardless of run count — few metadata
operations, which is what shared parallel filesystems want. The filename
disambiguates along every axis that can collide: pid within a node, host
across nodes (pid spaces are per-node; relevant for the MPI backend),
and session across invocations. A resumed sweep opens new files rather
than appending: an IPC stream cannot be reopened once its end-of-stream marker
is written, and repairing a torn file would mean read-modify-write on a shared
filesystem.

Stateless writer + per-process stream registry

The writer is re-pickled with every job, but we need open streams on each worker. So open streams live in a module-level registry keyed by (session, output), reused across a worker's runs. The other design alternative is to use an initializer function but this makes using executors more complicated. The user has to specify the right function depending on what store she wants to use. The implementation here hides the complexity.

Per-run atomic pre-write validation

to_reference validates every output (safe name, Arrow-convertible, schema
stable within the worker) before writing any, so a run contributes a full
batch set or nothing. Output names must match a strict pattern (they become
directory names); unsafe names are rejected rather than encoded. Within-worker
schema deviation is a recorded WRITING failure; cross-worker deviation is
deferred to the reader. A future PR will expand RunConfiguration to include an
explicit schema instead of deriving it from the first run on each worker.

Status: append-only log, replayed

The root appends one JSON line per terminal transition (SUCCEEDED / FAILED /
ABORTED, with structured FailureInfo); PENDING is the absence of a line.
Reading back this log is last-write-wins and tolerates exactly one torn final line (the only
damage the write pattern can produce); a malformed earlier line raises as
corruption. In-process, status lives on the RunRecord — the log is its
write-through durable shadow, so query methods stay simple dict comprehensions
matching InMemoryStore.

Scenario manifest round-trips seeds exactly

scenarios.json stores each scenario's parameters plus SeedSequence
entropy, spawn key, and bit-generator class — enough to reconstruct RNGs
bit-exactly (required for cross-process resume correctness). No pickle
anywhere. Numpy scalar parameters are narrowed to Python natives on write.

Read path: fan-in with schema reconciliation

Reading an output concatenates all its worker files
(promote_options="permissive" null-fills columns absent in some workers),
gated by an on_schema_conflict knob ("warn" | "raise").
retrieve_output(run_id) is status-gated like InMemoryStore and returns
every sweep output as a key; a zero-row frame means the run wrote a valid
empty batch or never wrote to that output — indistinguishable by design,
since the reference is key-only. Reads are uncached (whole-sweep reads
dominate; a RunId-list signature is noted as a future option).

Two construction paths

• DiskStore(store_dir, ...) creates: writes store.json exclusively
  (open("x") — pointing a new sweep at an existing store fails loudly
  instead of clobbering it), mints a session token.
• DiskStore.from_directory(store_dir, scenario_class) attaches for
  read-back: replays manifests and log, gates on the store format version,
  and has no session — writer() raises.

Provenance (informational, non-gating)

store.json records Mesa version, creation time, and a best-effort git
commit + dirty flag for the tree containing the user's model class (located
via inspect.getfile), Never gates any operation;
anything undeterminable is omitted.

Explicitly out of scope (later PRs)

• Resume: re-dispatching pending runs, new-session minting on an
  existing store, and reconciling the window where a worker's batch landed but
  the root died before logging success. The format commitments resume needs
  (session tokens in filenames, ordered session list, format version) are
  already in place, since they cannot be retrofitted without a format bump.
• Declared output schemas: validation will live in
  RunConfiguration so InMemoryStore benefits too.

Dependencies

Adds pyarrow>=14 (floor set by both Python 3.12 wheel availability and
concat_tables(promote_options=...); current stable is 24.x). Packaging as
an optional extra with guarded imports is handled in this PR / a follow-up
(TODO before merge).

[github-actions]

Performance benchmarks:

Model	Size	Init time [95% CI]	Run time [95% CI]
BoltzmannWealth	small	🔵 +0.3% [-0.3%, +1.0%]	🔵 +0.1% [-0.2%, +0.4%]
BoltzmannWealth	large	🔵 -0.9% [-1.9%, -0.1%]	🔵 +4.3% [+2.8%, +5.7%]
Schelling	small	🔵 +1.9% [+1.5%, +2.5%]	🔵 +2.3% [+1.5%, +3.2%]
Schelling	large	🔵 +1.7% [+1.0%, +2.3%]	🔴 +5.4% [+4.2%, +6.6%]
WolfSheep	small	🔵 +1.0% [+0.5%, +1.5%]	🔵 -1.6% [-5.8%, +2.5%]
WolfSheep	large	🔵 -0.0% [-0.7%, +0.5%]	🔵 -0.3% [-1.6%, +1.0%]
SugarscapeG1mt	small	🔵 +1.0% [+0.6%, +1.3%]	🔵 +1.3% [-0.4%, +2.9%]
SugarscapeG1mt	large	🔵 -0.0% [-1.0%, +1.1%]	🔵 +1.2% [-0.7%, +3.0%]
BoidFlockers	small	🔵 +1.8% [+1.4%, +2.3%]	🔵 +1.5% [+0.9%, +2.0%]
BoidFlockers	large	🔵 +2.2% [+1.7%, +2.7%]	🔵 +1.2% [+1.0%, +1.4%]

[codebreaker32]

I've gone through the DiskStreamWriter implementation for now and left some feedback regarding a few I/O edge cases. Fantastic work so far. I'll take a closer look at the rest of the PR soon.

[codebreaker32]

I think using _PID is dangerous for massive HPC sweeps because operating systems recycle PIDs.

If a worker crashes (e.g., OOM spike or SLURM timeout) and the executor respawns it, the OS can assign the replacement a previously used PID. Because _open_stream initializes the file using pa.OSFile(str(path), "wb"), it explicitly truncates existing files. The new worker will instantly and silently wipe out all record batches written by the dead worker.

swapping it for a UUID might be a better choice

[codebreaker32]

In the docstring, it states:

the check is atomic across the whole outcome, so a run either contributes a full batch to each of its streams or nothing at all, never a partial set.

While the validation is indeed atomic, the disk write phase is not. If "agents" appends successfully, but "model" throws an OSError (e.g., disk full, network drive blip, permission drop), the "agents" data is already permanently flushed to the Arrow stream.

The worker's _safe_call will correctly catch this and mark the run as FAILED. However, if the user later resumes the sweep to retry failed runs, the replacement run will execute and append the "agents" data again, resulting in silent duplicate data in that specific output file.

[quaquel]

This is an important corner case to consider in the future PR where I built the resume logic. I'll update the docs here and at a fixme as a reminder for this corner case.

In the resume situation, the best solution is to first clear out any data related to failed or abandoned runs.

[codebreaker32]

To actually clear out the failed data, the resume script would have to:

1. Open the .arrow file.
2. Read the entire thing into memory.
3. Find the rows belonging to the failed RunId and drop them.
4. Rewrite a brand new .arrow file to the disk.
   Doing this across hundreds of files before a sweep can resume is slow, memory-intensive, and risks corrupting the good data if the resume script crashes halfway through!

Instead of mutating the disk, we let the partial/failed data sit there. When the user eventually calls retrieve_output() to get their final Pandas dataframe, the reader simply looks at the status. If a RunId doesn't say SUCCEEDED, the reader just drops those rows instantly in memory using PyArrow or Pandas. It takes milliseconds and requires zero risky disk rewrites.

[codebreaker32]

In the runner.py diff, you explicitly allow a "columned zero-row frame" to pass extraction. While logically correct, this creates a dangerous schema lock-in trap when PyArrow gets involved.

If a worker's very first scenario execution results in a 0-row DataFrame, pandas often defaults those empty columns to generic float64 or object dtypes (unless the user explicitly cast them). PyArrow converts this zero-row frame into a RecordBatch and permanently locks that generic inferred schema into _SCHEMAS[key] during _open_stream.

When that same worker processes its second scenario, if that run actually produces rows populated with real data (e.g., int64 or specific categorical strings), PyArrow will generate a new schema with the actual types. _check_schema will compare this new accurate schema against the generic one locked in from the empty run, see a mismatch, and raise a ValueError, failing a perfectly valid run.

[quaquel]

You are right about this. It was something I considered as well. The only solution to this is to have the user provide an explicit schema. I plan to add support for an optional schema to RunConfiguration with the current behavior as the fallback if no schema is provided.

I'll also ensure to document this explicitly.

However, it is a real corner case, so to enforce people to always pass a schema just to cover this case seems a bit of an overkill to me.

[codebreaker32]

For the documentation, simply mentioning that users can explicitly type their columns if they anticipate a 0-row output on their first run is more than enough to defuse the trap for now!

[codebreaker32]

One more thing worth noting is we currently utilize Python's ProcessPoolExecutor to spin up a worker for each core.
However, PyArrow uses its own underlying C++ thread pool for I/O and compression operations. By default, PyArrow allocates threads equal to the total number of system cores( see here.)

When these two interact on high-core machines, every single Python worker process spins up its own full-sized PyArrow C++ thread pool. On a 64-core node, this results in 4,096 active threads fighting for CPU time. This may cause thrashing.

We can fix this globally by:

pa.set_cpu_count(1)
pa.set_io_thread_count(1)

src : https://arrow.apache.org/docs/python/generated/pyarrow.set_cpu_count.html https://arrow.apache.org/docs/python/generated/pyarrow.set_io_thread_count.html#pyarrow.set_io_thread_count