Expose callback during sampling

[wd60622]

I want to be able to run a python callback on each sample. The signature would be different than this:

https://github.com/wd60622/nutpie/blob/5b881638d658fb8236d4b627209d7f034ee02050/python/nutpie/sample.py#L361-L363

[aseyboldt]

There isn't currently any mechanism to call a function for each draw. The progress callback is called once every 50ms or so.
I think that might be better for the mlflow usecase as well? That way we don't slow things down a lot if the sampler is fast?

Either option wouldn't be too difficult to add though.
What info would you want to have for each update call?

[wd60622]

That would do the trick! Yeah, no need to slow it down.

So I'm hearing that the progressbar html is updated every 50ms or so. Is there much overhead in that?

What info would you want to have for each update call?

Can you point me to what would be possible to get? Are there some objects already that could be used?

[aseyboldt]

If we want to push progress info to mlflow, I think once every second or so would probably be enough? We don't want to spam the poor thing.
I can easily expose anything here for each chain https://github.com/pymc-devs/nuts-rs/blob/main/src/sampler.rs#L386

[wd60622]

Thanks for sharing that link.

Totally. Would that be up to the user though? Something that could be defined in the callback:

def create_callback(update_frequency): 
    # Psuedo code
    step = None
    def callback(state):
        if step is None: 
            step = time.time()

        if time.time() - step > update_frequency:
            step = time.time()

            # do something with state

    return callback