Skip to content

Unload datasets by default - #587

Draft
gtrevisan wants to merge 5 commits into
devfrom
glt/unload
Draft

Unload datasets by default#587
gtrevisan wants to merge 5 commits into
devfrom
glt/unload

Conversation

@gtrevisan

Copy link
Copy Markdown
Member

we want DictOutputSetting, which is our under-the-hood de facto base class, to be memory-efficient and therefore more reliable for large-scale workflows.

  • once a Pool-processed dataset is fed to output_shot:
    • temporarily save/shard it to disk,
    • unload it from memory,
    • replace it with a temporary lazy-loaded dataset,
  • once all results are written to disk:
    • move the corresponding shards to their final path,
    • replace each dataset with its final lazy-loaded version.

furthermore, for the (abstract) subclass SingleOutputSetting:

  • when getting the results:
    • explicit read before concatenation,
  • when saving to disk:
    • removed superseded shards.

as a consequence:

  • all workflows collect partial shards in the run's temporary folder, even in fatal cases (eg OOM killed, SLURM time limits, any framework-level unhandled exception),
  • DictOutputSetting should complete the workflow with a low memory footprint, namely just enough to have lazy-loaded datatasets as dictionary values,
  • by design, as before, DictOutputSetting requires explicit downstream concatenation.
  • SingleOutputSetting should work effectively as before, although the new logic is to unload each process results to disk and to load them up again from disk just before concatenation. this results in a minor performance hit which is negligible due to the no-encoding and no-compression defaults of xarray's to_netcdf, and which is overshadowed by the newly-introduced impossibility of losing data during the concatenation stage,
  • new debug statements are now tracking the maximum RSS memory (with a high-watermark level!) throughout the workflow.

example runs for 1,000 C-MOD shots:

  • DatasetOutputSetting:
Starting workflow: MaxRSS = 201.0 MB
Completed workflow: MaxRSS = 729.0 MB
Read shots: MaxRSS = 1,231.0 MB
Concatenated shots: MaxRSS = 2,555.4 MB
Saved to disk: MaxRSS = 2,555.4 MB
  • DictOutputSetting:
Starting workflow: MaxRSS = 197.0 MB
Completed workflow: MaxRSS = 719.0 MB
Saved to disk: MaxRSS = 731.0 MB
# custom post-processed load step to compare final footprint
Loaded: MaxRSS = 1,227.0 MB

I already have a plan to develop a new OutputSetting to efficiently concat large-scale workflow results.
@samc24 thoughts on the architecture? for sure we could change many things, including fully revamping the multiprocessing data exchange and all the output settings, but this felt a great bang-for-the-buck improvement to me.
@nbarbour13 tests should already pass, but feel free to play around with the various output settings to make sure the results still end up where they are supposed to, and no shards get left behind.

@gtrevisan
gtrevisan requested review from nbarbour13 and samc24 and a balanced review from Copilot August 18, 2026 15:53

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

Improves output memory efficiency by sharding per-shot datasets and retaining lazy-loaded references.

Changes:

  • Writes temporary NetCDF shards and moves them to final paths.
  • Loads shards before single-output concatenation and removes them afterward.
  • Adds MaxRSS workflow logging.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File Description
disruption_py/workflow.py Adds memory watermark logging.
disruption_py/settings/output_setting.py Implements dataset sharding, lazy loading, finalization, and cleanup.
Suppressed comments (3)

disruption_py/workflow.py:229

  • This MaxRSS conversion is incorrect on macOS, where ru_maxrss is bytes; the logged MB value will be 1,024× too large. Apply the same platform-aware conversion here.
        mem=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024,

disruption_py/workflow.py:237

  • This also mislabels macOS's byte-valued ru_maxrss after saving, producing an MB value 1,024× too large. Normalize the unit by platform here as well.
        mem=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024,

disruption_py/settings/output_setting.py:301

  • This second SingleOutputSetting watermark has the same macOS unit error: bytes divided by 1,024 are not MB. Use the platform-aware conversion here too.
            mem=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024,

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread disruption_py/workflow.py Outdated
Comment thread disruption_py/settings/output_setting.py
Comment thread disruption_py/settings/output_setting.py Outdated
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