[mj023]

Implements #591: ExecutionConfig(value_storage=...) keeps solved values in pinned host memory or on disk, so the device holds only the resident window instead of every solved period.

Stacked on #630 (codex/570-singleton-initial). The diff also carries the #590 commits that the stack does not have yet.

Summary

ExecutionConfig(value_storage="device")      # default, unchanged
ExecutionConfig(value_storage="host")        # pinned host memory
ExecutionConfig(value_storage=Path("dir"))   # one archive per solve, SolutionResult.save format

Solve, period-major, per period t:

 run kernels of t
-keep V[t] on the device until solve returns
+copy V[t] and the EGM-family policies of t to pinned host memory, without blocking
+  (under a directory, a writer thread then writes them to solution-<id>.h5)
+NaN/Inf flags of t become host booleans; at the first NaN, validate_V runs now,
+  while V[t] and V[t+1] are still on the device
+wait for the copy of t+1, so at most 3 periods are on the device

• Planner: budget admission charges only the resident window. Under "device", a refusal caused by stored values names value_storage="host".
• Simulation:
  • With several chunks it runs period by period, chunks inside.
  • Reads go through a caching device view, so each period's values and policies are copied once for every chunk.
  • The next period's declared reads start copying before a period's chunks run.
  • Under a budget the planner reserves every subject's carry and that read set.
• Block-major: each code is solved with host storage. Under a directory, all codes' blocks go into one archive.
• Failures: a failed copy or write raises ValueStorageError, naming the period, the regime and the target. Under a directory the file stays as solution-<id>.incomplete, recording fingerprints, the pylcm version and commit, the dtype and the last completed period.

Evidence

Every outline test in #591 is in tests/solution/test_value_storage.py and tests/solution/test_value_storage_sharded.py:

• Bitwise equality: device, host and directory storage give bitwise-equal values, EGM policies and panels. Checked on the stochastic 4-period model, with one chunk and with several, and on NB-EGM.

• Resident window:

  live device values at each kernel, three-regime model
  "device": every solved period stays  -> window check fails, as intended
  "host":   only {t, t+1, t+2}         -> window check passes
  

• Round trips: a disk archive is solved in one process and loaded and simulated in another: one device to one, two sharded to one, one to two sharded.

• Failures: an injected write failure names the period, the regime and the .h5.partial path. Before this change the raw OSError came through, and the record said period 0 instead of 2, because the writer kept writing after the failure.

• Fewer copies:

  • With several chunks, policy copies now equal one chunk's. Before: 19 vs 11 under "host", and 8 vs 4 disk reads.
  • A NaN under host or disk storage reads no value back. Before, 2 values were read back.

• Blocking references: blocked host storage matches the unblocked solve, and block-major in every mode matches period-major, within 2 ULP. That is the gap blocking itself causes; nothing else shrank.

• Local runs (after merging Honor finite floors in paired maximum reductions #630 at 1ed7559d2):

  Where	fp64	fp32
  CPU	431 passed, 16 skipped (15 value-storage, archive, persistence, simulation and diagnostics files)	95/95
  CPU, four devices	16/16 (test_distributed_simulation_values.py, the 16 skipped above)
  Local RTX 3050	89/89 (value-storage and archive files, before the merge)	89/89

CI

CI ran the full suite on 2ca4c2c05: 41 jobs pass, 23 fail. Every failure comes from the base #630:

• All 35 failing tests fail with the same errors on Honor finite floors in paired maximum reductions #630's CI at 1ed7559d2; which shard they land in differs.
• This PR changes none of the failing test files. The value-storage tests pass in every job, including macOS, Windows, fp32 and the four- and eight-device jobs.
• The failures, grouped by cause:
  • NaN values in solved value functions: test_fgp_model, test_economic_validation, test_egm_passive_asset_row, test_process_simulation_laws, test_unbudgeted_subject_width[9000].
  • The simulate_draws= → simulated_draws= rename: test_density_derivative_accuracy still passes the old name.
  • Memory and size limits:
    • limits exceeded: test_regime_mixture_reduction, test_invariant_assembly_reservation;
    • budget refusals: test_current_iid_grid_search, test_policy_lookup_devices[8].
  • An XLA FAILED_PRECONDITION error on the GPU: test_public_exit_queries in the bounded GPU job.
  • Crashed test workers: test_participation_pricing_depth::test_a_saved_solution_… and test_collective_household_example. The first passes locally in 290 s.
  • Single failures: test_collective_extended_real, test_participation_read_alias_process_code[solve], test_outer_candidate_drop_report, test_block_major_lifetime, and the regression output for precautionary_savings_health.
• The GPU benchmark job also fails on Honor finite floors in paired maximum reductions #630.

These are left to #630 and should clear here once it is fixed and merged in.

Device memory, measured

Setup:

• RTX 3050, fp64, default budget.
• One grid-search regime: 40,000 wealth points × 25 preference types × 40 consumption points, then a terminal regime.
• Each run in a fresh process; the figures are the device allocator's peak_bytes_in_use during the solve and bytes_in_use after it.

Periods	Values	Storage	Peak during solve	Held after solve	Solve time
60	491 MB	"device"	1786 MB	492 MB	7.4 s
60	491 MB	"host"	1312 MB	0.6 MB	7.4 s
60	491 MB	directory	1312 MB	0.6 MB	19.7 s
120	990 MB	"device"	2285 MB	991 MB	15.0 s
120	990 MB	"host"	1312 MB	0.6 MB	14.6 s

• Peak: under host or disk storage the peak stays flat as periods grow. Under "device" it grows by every period's values: −474 MB at 60 periods and −973 MB at 120, about all values but the resident window. What remains is the period's workspace.
• Time: copying to pinned host memory costs no solve time here. Writing to disk costs about 2.7× on this machine.

Merge Danger

Door: two-way

"device" stays the default, and with it nothing changes except one thing every mode shares: the NaN check now runs at each period's barrier instead of after the loop. The disk archive is in the existing SolutionResult.save format.

Blast Radius: moderate

• Every mode:
  • The diagnostics change touches all solves.
  • The engine view's policies are now a caching view under host or disk storage.
  • Artifact leaves in pinned host memory are copied onto the device when read.
• Known limits:
  • The early load runs only when a simulation has several chunks.
  • Disk storage gathers each value on one host before writing it.
  • Component jobs accept "device" only.
  • Budgeted runs of the NB-EGM toy and the blocked test model are refused on a GPU under any storage. That predates this PR, so off-device policies are untested under a budget.
• Certificate: pins and seals are updated after the merge, and verify.py passes.
• Not included: the issue comment's re-check of only the refused cell after a narrowing is left to its own change.

🤖 Generated with Claude Code

[read-the-docs-community]

Documentation build overview

📚 pylcm | 🛠️ Build #35697200 | 📁 Comparing e9c483a against latest (f57269f)

  🔍 Preview build  

38 files changed · + 1 added · ± 37 modified

+ Added

• shocks/index.html

± Modified

• age-specialized/index.html
• architecture/index.html
• architecture-transition-ledger/index.html
• benchmarking/index.html
• beta-delta/index.html
• categories-with-characteristics/index.html
• certification/index.html
• choosing-a-solver/index.html
• collective-regimes/index.html
• collective-regimes-1/index.html
• and 27 more...

[codecov]

Codecov Report

❌ Patch coverage is 97.55382% with 25 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.15%. Comparing base (a831dc1) to head (e9c483a).

Files with missing lines	Patch %	Lines
src/_lcm/persistence/solution.py	95.30%	18 Missing ⚠️
src/_lcm/solution/host_values.py	97.88%	4 Missing ⚠️
src/_lcm/solution/backward_induction.py	95.74%	2 Missing ⚠️
src/_lcm/simulation/simulate.py	98.36%	1 Missing ⚠️

Additional details and impacted files

@@                       Coverage Diff                       @@
##           codex/570-singleton-initial     #633      +/-   ##
===============================================================
+ Coverage                        94.12%   94.15%   +0.02%     
===============================================================
  Files                              660      662       +2     
  Lines                           111076   111997     +921     
===============================================================
+ Hits                            104554   105449     +895     
- Misses                            6522     6548      +26     

Flag	Coverage Δ
cpu-python	94.15% <97.55%> (+0.02%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Benchmark comparison (main → HEAD)

Comparing 0fb55673 (main) → e9c483ac (HEAD)

Benchmark	Statistic	before	after	Ratio	Alert
ACA (reduced)	warm solve + simulate (reuses compiled code)	105.722 s	245.611 s	2.32	❌
	cold solve + simulate (first run, includes compilation)	808.56 s	1695.49 s	2.10	❌
	peak CPU mem	6.16 GB	9.17 GB	1.49	❌
Mahler-Yum	warm solve + simulate (reuses compiled code)	6.054 s	9.222 s	1.52	❌
	cold solve + simulate (first run, includes compilation)	19.02 s	36.19 s	1.90	❌
	peak CPU mem	1.66 GB	1.78 GB	1.07
	peak GPU mem: automatic solve + simulate	1.61 GB	714 MB	0.44
	peak GPU mem: load saved solution + simulate	1.89 GB	837 MB	0.44
	peak GPU mem: solve + save all persistable artifacts	571 MB	714 MB	1.25	❌
Precautionary Savings - Solve	execution time	84.0 ms	89.1 ms	1.06
	first call (including compilation)	1.29 s	2.00 s	1.55	❌
	peak CPU mem	1.44 GB	1.45 GB	1.00
Precautionary Savings - Simulate	execution time	901.9 ms	1.199 s	1.33	❌
	peak GPU mem	324 MB	690 MB	2.13	❌
	first call (including compilation)	3.40 s	5.01 s	1.47	❌
	peak CPU mem	1.57 GB	1.58 GB	1.01
Precautionary Savings - Solve & Simulate	execution time	661.7 ms	1.007 s	1.52	❌
	peak GPU mem	1.51 GB	4.24 GB	2.81	❌
	first call (including compilation)	4.28 s	6.83 s	1.60	❌
	peak CPU mem	1.53 GB	1.55 GB	1.01
Precautionary Savings - Solve & Simulate (irreg)	execution time	802.4 ms	1.110 s	1.38	❌
	peak GPU mem	2.21 GB	4.24 GB	1.92	❌
	first call (including compilation)	4.45 s	7.40 s	1.66	❌
	peak CPU mem	1.53 GB	1.55 GB	1.01
Iskhakov et al. (2017) - GridSearch Solve	execution time	195.2 ms	226.2 ms	1.16	❌
	peak GPU mem	120 MB	120 MB	1.00
	first call (including compilation)	2.43 s	3.30 s	1.36	❌
	peak CPU mem	1.63 GB	1.64 GB	1.01
Iskhakov et al. (2017) - DC-EGM Solve	execution time	826.2 ms	1.077 s	1.30	❌
	first call (including compilation)	6.87 s	8.58 s	1.25	❌
	peak CPU mem	1.99 GB	1.97 GB	0.99
Iskhakov et al. (2017) - GridSearch Simulate	execution time	918.0 ms	1.215 s	1.32	❌
	peak GPU mem	891 MB	890 MB	1.00
	first call (including compilation)	7.02 s	8.79 s	1.25	❌
	peak CPU mem	1.76 GB	1.77 GB	1.00
Iskhakov et al. (2017) - DC-EGM Simulate	execution time	1.281 s	1.850 s	1.44	❌
	peak GPU mem	891 MB	891 MB	1.00
	first call (including compilation)	9.25 s	11.42 s	1.24	❌
	peak CPU mem	2.03 GB	2.08 GB	1.02
Collective Household - Construct	execution time	1.547 s	1.805 s	1.17	❌
	peak CPU mem	1.43 GB	1.43 GB	1.00
Collective Household - Solve	execution time	5.453 s	6.216 s	1.14	❌
	first call (including compilation)	55.91 s	55.95 s	1.00
	peak CPU mem	4.60 GB	2.74 GB	0.60
Collective Household - Simulate (1000)	execution time	32.017 s	42.573 s	1.33	❌
	first call (including compilation)	505.24 s	736.88 s	1.46	❌
	peak CPU mem	25.19 GB	28.84 GB	1.15	❌
Collective Household - Simulate (100000)	execution time	nan µs	nan µs	nan
	peak GPU mem	nan MB	nan MB	nan
	first call (including compilation)	nan s	nan s	nan
	peak CPU mem	nan MB	nan MB	nan
Reference Chain - Solve (1)	execution time	39.8 ms	58.0 ms	1.46	❌
	first call (including compilation)	1.48 s	2.11 s	1.43	❌
	peak CPU mem	1.16 GB	1.19 GB	1.02
Reference Chain - Solve (8)	execution time	253.5 ms	371.6 ms	1.47	❌
	first call (including compilation)	5.35 s	7.70 s	1.44	❌
	peak CPU mem	1.29 GB	1.32 GB	1.02
SimulationDispatch ('dissolution', 'off')	warm simulation time per period and regime (ms)	78.72 ms	99.75 ms	1.27	❌
	backend compile requests on second simulation call		0
SimulationDispatch ('dissolution', 'progress')	warm simulation time per period and regime (ms)	83.99 ms	102.79 ms	1.22	❌
	backend compile requests on second simulation call		0
SimulationDispatch ('multi_regime', 'off')	warm simulation time per period and regime (ms)	56.41 ms	72.21 ms	1.28	❌
	backend compile requests on second simulation call		0
SimulationDispatch ('multi_regime', 'progress')	warm simulation time per period and regime (ms)	50.58 ms	80.29 ms	1.59	❌
	backend compile requests on second simulation call		0

[hmgaudecker]

Review summary

Thank you, the host and disk storage design reads well. The read views hand out device arrays, and the writer thread, the failure latch and the per-solve archive names are easy to follow.

One thing I fixed on the branch, with a test in 53cbc91 and an inline comment on 1a2d304 explaining why. One thing I would like changed before merging: the disk archive's failure handling. It is built on a try/except in Model.solve that this PR introduces, and the three ways it misses its own contract all come from that choice. The archive can record its progress as it goes and then needs no handler. Details below. Line numbers refer to 1a2d304. Where I say something ran, it ran in the branch's locked tests-cpu environment on CPU with x64.

Disk archive: record progress as it happens, drop the handler

Model.solve wraps the solve in try/except BaseException and calls abandon with a record built in the handler (src/lcm/model.py:2778-2820). This PR introduces that handler. I would rather it did not exist, and I think it does not have to.

Every field of the record is known before the handler runs. dtype, params_fingerprint, pylcm_version, pylcm_commit and the model_fingerprint that describe collects at model.py:2860 are known before the first period. last_completed_period and completed_codes are known the moment each period or code finishes, and the archive already tracks them in memory. If the archive writes the static fields into the .h5.partial file when it is created and updates the progress fields as each period lands, then the documented record is on disk at every moment of the solve. A failure anywhere leaves the file with the record already in it. There is nothing to do on the way out, so there is no handler. Completion writes the manifest, flushes, fsyncs and renames, as _close does now. The finalizer's rename of a leftover .partial to .incomplete at solution.py:3729-3740 then completes the picture instead of producing a file with no record.

The three gaps I found are consequences of the handler and go away with it:

• Model.simulate(solution=None) on the block-major route creates its own BlockedSolutionArchive at model.py:4916-4924 and completes it at model.py:5151-5152, but never abandons it. A NaN value function leaves a .partial that the finalizer renames with no record inside. With the archive owning its record, this route gets the same record as solve for free.
• abandon writes through the same HDF5 handle that just failed (solution.py:3250-3264), and the call is unprotected. On a full volume the cleanup error replaces the ValueStorageError that names the period and regime. A record written before the failure cannot fail after it.
• _close closes the handle before fsync and the rename (solution.py:3308-3316). If either fails, abandon calls create_dataset on a closed handle; with h5py 3.16.0 that raises ValueError: Unable to synchronously create dataset (invalid identifier type to function), which replaces the original error. The file at that point holds a complete manifest and payloads, and the loader checks contents (solution.py:606-609), so load_solution would accept it even after the finalizer renames it. I derived this from the code and did not run it. Without a handler there is no second failure, and the question becomes only whether a file with a complete manifest that was never renamed should load. I would say it may, since the payload checksums are verified, but please decide and write it down.

For tests: a block-major simulate(solution=None) with a NaN in the first and in a later code, checked before the exception is released; a solve where the period write fails on a full volume, which must raise the ValueStorageError and leave the record with the last completed period; a solve where fsync fails in _close, which must raise that error and leave a file whose fate is the one you decided above.

Lesser notes

• The archive is created before the solve starts (model.py:2765-2777), and a read-only directory passes _fail_if_value_storage_invalid (execution.py:522-530). solve then raises h5py's PermissionError for the .h5.partial path rather than a ValueStorageError. I ran this; the directory is left empty, but the error does not name the storage setting.
• The planning descriptors drop weak_type (host_values.py:549-553 and the two device_stand_in methods in solution.py near lines 3090 and 3183), while the device path keeps it in _placed_abstract (forward_program_profiles.py:870-876). A weakly typed stored leaf would be planned with a different signature than it is dispatched with. I did not check whether any stored leaf is weakly typed; 1a2d304 does not change this.
• DiskValueEntry.materialize and DiskPolicyEntry.materialize use a plain jax.device_put (solution.py:3111-3116, 3149-3166), while the saved-archive loader refuses implicit x64 narrowing in _to_jax_without_narrowing (solution.py:2237-2248). A change of the x64 flag between solve and read narrows silently on the live path. Please consider the same guard.
• PeriodValueWriter.submit queues the new period before it waits (solution.py:3640-3655), so a slow writer can hold three unwritten periods in host memory, not two. The wording at execution.py:354-356 is correct but easy to misread; please say three.
• submit waits on the oldest write with no timeout (solution.py:3650-3653). A write that never returns, on a stalled network volume say, holds the solve forever. If that is intended, please say so in the docstring.
• The per-period NaN and Inf flags now go to the host in _materialize_period_feasibility_counts (diagnostics.py:216-253). The merge base already blocked on the running flags with .item() each period, so I see no new synchronisation, only more flags per transfer. A warm timing against the merge base would settle it.

[hmgaudecker]

I pushed this commit, so here is why. A budgeted simulation of host- or disk-stored finite NNBEGM policies failed at 628c38b with a BeartypeCallHintParamViolation in ReplayPayload.from_policy. The chunk planner receives planning_stand_in descriptors instead of the stored policies (chunk_admission.py:249, host_values.py:541-568), and the finite NNBEGM planning step still insisted on concrete arrays: this annotation, the guard at forward_program_profiles.py:661-666, and the leaf annotation of _abstract_policy_leaf at forward_program_profiles.py:709. Those are type assumptions from #440, from when every stored policy was a device array.

Nothing in planning needs the values. _profile_finite_decision maps every leaf through _abstract_policy_leaf, which keeps shape, dtype, weak type and placement, and hands that to _prepare_program. So the commit admits jax.ShapeDtypeStruct at the three sites, plus a fourth: replay_payload_reads at replay_inputs.py:64 enumerates the leaves a replay acquires, and without counting a descriptor there _policy_read_sources at chunk_profiles.py:1508-1509 zips four stand-in leaves against zero reads. The corridor pins in the candidate certificate are re-anchored in the same commit.

The route is public: a same-instance solve with value_storage="host" or a directory, plus a positive device_memory_bytes, on CPU. The budget suite runs it with device storage only, and the NNBEGM storage tests remove the budget (tests/solution/test_value_storage.py:463-477), so neither suite simulates finite NNBEGM policies from storage under a budget. The DC-EGM policy-first read at forward_program_profiles.py:466-476 goes through the same _abstract_policy_leaf and failed the same way.

53cbc91 adds tests/simulation/test_budgeted_stored_policies.py, which covers both routes: host- and disk-stored panels must equal the device panel exactly. All four cases fail on the sources of 628c38b and pass on 1a2d304. With the fix, tests/simulation/test_finite_policy_budget.py and tests/solution/test_value_storage.py pass as well, 94 cases.

[hmgaudecker]

@mj023 We hit this refusal while running the EKL 2019 model (lcm-zoo) at production scale. The message itself is good, but it arrives about 4.5 hours too late.

Setup. pylcm 7d0ec63, the #639 stack, which already carries this refusal text. Source5 grid, float64, one A100, device_memory_bytes=51_539_607_552, default value_storage="device".

What happened. Every kernel compiled. Then public_solve ended status=error without solving a single period. Nearly all of the time went into four compilation_waves, with workspace narrowings in between:

run	query_cell	waves (s)	public_solve until refusal
A	64	8461, 3827, 2813, 2700	17844 s
B	32	8282, 2820, 2625, 2281	16049 s

The refusal at single_f_school period 1 was 74,361,084,696 bytes: a 37,282,115,852-byte compiler reservation plus 37,078,968,844 resident value bytes, against a 51,539,607,552-byte budget.

Suggestions for this PR:

1. Refuse before compiling. Under "device", the resident bytes of stored values are known before any compile. When they alone exceed the budget, or exceed it together with the smallest compiler reservation already seen, stop right after planning (or after the first wave) instead of narrowing through further waves.
2. Say whether "host" would fit. The planner knows the host-storage resident window, so the message can say whether switching to host storage makes the plan fit, not only suggest trying it. For example: "Resident solved values reach 37.1 GB on device 0 by single_f_school period 1; with the 37.3 GB compiler reservation this exceeds the 48.0 GiB budget. value_storage='host' keeps at most 3 periods resident; with it, this plan fits." (only when that is true; we have not verified it for this model yet)

A rerun of the same solve with value_storage="host" is queued. I'll add its outcome here, since it also exercises this PR's storage path on a large model.

[hmgaudecker]

Follow-up on the rerun I mentioned above. Same pin (7d0ec63), same EKL source5 solve on one A100, width 64, with ExecutionConfig(value_storage="host").

Outcome: refused at the same cell after a similar compile time.

• _CellBudgetError at single_f_school period 1, core _chunks.
• The smallest total is 74,361,459,472 bytes: a 37,282,490,628-byte compiler reservation plus 37,078,968,844 resident bytes. The resident figure is identical to the value_storage="device" run.
• public_solve ended with an error after 18,078.7 s, almost all of it in four compilation_waves. No period was solved.

The resident split the engine logs explains why host storage has no effect here. At single_m_school period 1, the busiest device holds:

• 25,194,240,000 bytes of regime_value:married_school. That is exactly 12 periods of 2,099,520,000 bytes, and it shrinks by one period at each later school period.
• 6,561,000,000 bytes of married_postfertile values and 5,248,800,000 bytes of married_fertile values.

My reading of the code at 7d0ec63 follows. I have not instrumented it.

• Under host storage, _build_planned_input_liveness keeps a regime value live only through its own period and the next one.
• The ledger never makes pinned_artifacts eligible for release.
• undeclared_read_pins pins the same-period reference values and the next-period values of any dispatch with a program that declares no value reads.
• So if one of EKL's school cores declares no reads, every married period it can reach stays pinned for the whole solve, under either storage mode.

Questions:

1. Should a pin from an undeclared read end after the reading period under value_storage="host", the same way declared reads do?
2. If that is not safe, could the refusal name the pinned owners and the program that caused the pin? At the moment it suggests value_storage="host", which cannot help when the resident bytes are pins.

The message-side requests are collected in #643.