Allow for pymc native samplers to resume sampling from ZarrTrace

[lucianopaz]

Description

Big PR approaching! This finishes adding the ability of pymc native step methods to resume sampling from an existing trace (as long as it's a ZarrTrace!). This means that you can now continue tuning or sampling from a pre-existing sample run. For example

with model:
    # First tuning run
    pm.sample(tune=400, draws=0, trace=trace)

    # Do whatever to decide if you want to continue tuning   
    pm.sample(tune=800, draws=0, trace=trace)

    # Switch to sampling
    pm.sample(tune=800, draws=1000, trace=trace)

Another thing is that the chunks_per_draw from ZarrTrace along with its persistent storage backends (like ZipStore or DirectoryStore) makes the sampling store the results and final sampling state periodically, so in case of a crash during sampling, you can use the existing store to load the trace using ZarrTrace.from_store and then resume sampling from there.

The only thing that I haven't tested for yet is to add an Op that makes pm.sample crash to see if I can reload the partial results from the store and resume sampling. @ricardoV94 gave me some pointers to that, but I won't be working on this for the rest of the month and I thought it best to open a draft PR to kick off any discussion you have or collect feedback

Related Issue

• Closes ENH: Add checkpoints during sampling #7503
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

📚 Documentation preview 📚: https://pymc--7687.org.readthedocs.build/en/7687/

[fonnesbeck]

Maybe just InitializedTrace? Seems a little verbose!

[ricardoV94]

Sounds fine to me, it's an internal thing

[ricardoV94]

I still don't like this abstraction leaking elsewhere, just provide a default to the Ndarray backend that makes it work for either method. In that case I suppose start everything at zero

[ricardoV94]

Same here all this logic including the old link_stepper can have a sensible default in the base trace class so you don't need to worry about what kind of trace you have here. Just make link_stepper a no op and stored_draw_idx to be zero by default?

[ricardoV94]

Duplicated logic, should be a property of the backend object?

[ricardoV94]

abstraction leaking

[ricardoV94]

I like the new functionality, I am deeply against all the if isinstance(..., ZarrTrace) in the codebase. Either our code is supposed to allow different trace backends or it is not, this suggests you want to drop the Ndarray altogether, which fine if you do.

Otherwise all these cases seem like they could be handled by the BaseTrace having sensible default for these methods. We used to have continuation of traces in the past with Ndarray, I don't see anything that fundamentally needs ZarrTrace other than dev interest in it? So just make it raise NotImplementedErrors or make them no-ops and adjust the external code appropriately

I stopped half-way so it was not an extensive review. I think this is a bigger design point that needs decision before settling on the details of the PR.