🎉 feat(cgns writter): add callback function to save_to_disk

[tmolcard]

Checklist

• Typing enforced
• Documentation updated
• Changelog updated
• Tests and Example updates
• Coverage should be 100%

📋 Description

Add callback function to cgns storage save_to_disk method.

Context:
We need to read datasets, process them and save them to an external object storage.
Current implementation only enables us to save the whole dataset to disk before processing it. However, when using cgns, once written, a sample remains untouched and could be processed.

This feature sends a callback after each samples is written. It lets callers process samples out (e.g. upload to object storage then delete locally) sample per sample instead of processing the whole dataset.
Currently supported for the cgns backend with num_proc == 1.

[CLAassistant]

All committers have signed the CLA.

[casenave]

Hi @tmolcard, thanks you for your contribution!

Can you provide more elements to you use case. For what I understood, I'm not sure you need is really writing a plaid dataset, since you seem to want to act on the sample or modify it rather than write it to disk. The write_to_disk function needs a sample_constructor callable, and I would guess that you can iterate over your data and apply the function of your callabacl on the output of sample_constructor directly ?

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[tmolcard]

Hi @tmolcard, thanks you for your contribution!

Can you provide more elements to you use case. For what I understood, I'm not sure you need is really writing a plaid dataset, since you seem to want to act on the sample or modify it rather than write it to disk. The write_to_disk function needs a sample_constructor callable, and I would guess that you can iterate over your data and apply the function of your callabacl on the output of sample_constructor directly ?

Hi @casenave, thank you very much for your quick review!

Regarding our use case: yes, we are indeed writing a dataset. We are currently designing a very simple data pipeline that writes our data from various formats to a unified PLAID format, and in particular we do this on remote object storage. The pipeline is still very simple and early-stage:
Read simulations from S3 -> Process them into in-memory PLAID objects (through sample_constructor) -> Write the PLAID dataset to disk -> Upload it back to the same object storage.
We'd like to keep using init_from_disk and save_to_disk to benefit from PLAID's convenient API. However, in its current state save_to_disk only hands control back once the full dataset has been written. Our datasets can reach several hundred GB, if not TB, so storing the whole thing on disk isn't very convenient. Using save_to_disk saves us from reimplementing the sample directory convention, and the point of the hook isn't to manipulate the PLAID object once it's created, but rather to trigger an "upload and delete" process on each written sample.

That said, I understand this might be too narrow a use case to justify changing PLAID's core. We could also work with samples and metadata using save_to_dir, save_problem_definitions_to_disk and save_infos_to_disk separately, as a workaround.

Let me know what you think!

[xroynard]

Thanks @tmolcard, that use case makes it crystal clear — and it's a legitimate one. Reading from S3 → building PLAID samples in memory → writing → uploading back, at hundreds of GB / TB scale, is exactly the case where you can't afford to materialize the whole dataset before getting control back. And you're right that save_to_disk is worth keeping precisely so you don't have to reimplement the sample_XXXXXXXXX directory convention and the infos/pb_defs assembly.

To be clear on @casenave's earlier point: applying a function to sample_constructor's output wouldn't solve this — you need the on-disk sample directory (to upload-then-delete), which only exists inside the writer. So a post-write hook is the right shape. I'd also gently push back on the save_to_dir + save_infos_to_disk + save_problem_definitions_to_disk workaround: it works, but it forces you to re-own the orchestration and couples your code to PLAID's internal layout — which is the thing you're trying to avoid.

Two suggestions:

1. Extend the hook to num_proc > 1. At TB scale you'll almost certainly want parallel writes, and that's currently rejected with NotImplementedError — i.e. the hook is unavailable in the very regime where it matters most. Concretely: invoke the callback inside each worker right after save_to_dir, and document the relaxed contract:

   • callback must be picklable and process-safe (separate processes, concurrent uploads),
   • no global ordering guarantee across workers,
   • index is per-shard unless you surface the deterministic global start index the writer already computes (prefix sums) to keep it contiguous.

2. Longer term, consider writing straight to remote storage. Since your source and destination are S3, an fsspec-aware output_folder (s3://…) would remove the write→upload→delete dance entirely and benefit every backend. That's a bigger change (CGNS/HDF5 remote writes usually need a local buffer), so it's follow-up material — but it's where this pipeline naturally wants to go.

Happy to keep this PR minimal (cgns + mono-proc) and track the parallel variant as a follow-up if you'd rather land something small first.

Also, a small design note: exposing sample_path effectively turns PLAID's internal data/{split}/sample_{i:09d} naming into a public guarantee — once callers start parsing that path (e.g. to derive an S3 key), we can no longer change the layout without silently breaking them. Worth deciding explicitly and documenting it one way or the other: either "this path format is stable, you can rely on it", or "treat the path as opaque, don't parse its name" (e.g. by relying on the split_name and index the callback already passes separately).

[xroynard]

Regarding suggestion 2 (writing straight to remote storage): I've opened #479 to track it, so this PR can stay focused on the minimal local hook (cgns + mono-proc). The longer-term fsspec-aware output_folder direction — which would eventually make the write→upload→delete workaround unnecessary — will be followed up there.

[casenave]

Hi @tmolcard,

My view is that the proposed changes are too specific to be incorporated into write_to_disk, as they target a single backend and impose constraints on num_proc.

Your suggested workaround—using save_to_dir, save_problem_definitions_to_disk, and save_infos_to_disk separately for samples and metadata—seems spot-on. @xroynard has opened issue #479 to discuss a more general evolution, and you are very welcome to contribute there.

It would also be valuable to document this workaround in the user documentation in another PR. With your agreement, I suggest that we close this PR.

Please be assured that your interest and contribution are greatly appreciated, regardless of the outcome of this PR. Thank you!

[tmolcard]

Thank you @xroynard for your feedback and the opening of the issue.

I agree with your points :

• A long-term feature would be to enable remote storage natively, but as you mentioned I wanted to keep it as simple as possible in a first step.
• Locking it to num_proc == 1 is too restrictive and I just implemented the generalisation.
• On the sample_path question: in my case I only mirror the relative path into the object key. I never parse the name, treating it as opaque works for me, I can document it that way.

Also about your last message @casenave, I don't know if the extension to num_proc >= 1 would be enough for you to keep considering the PR. I do understand this might remain too specific, it seems that it could be generalized to zarr backend, but I think it can't be to hf backend.

Concerning the documentation of the workaround, it pushes PLAID's sample-directory convention onto the caller. I'd hardcode data/{split}/sample_{i:09d} myself. That's the public layout coupling @xroynard flagged. The post-write hook keeps that ownership inside PLAID. I'd be a bit hesitant to document the workaround as a recommended pattern i would rather point users to #479.

That said, if it still feels too specific I'm ok with closing the PR, using the work around and contributing to the issue !

Thank you again for your time !