Is there a way to halt a model if it is taking too long?

Looking over the variational api docs, there’s an example of a custom function being used to stop an advi fitting procedure if too little iterations/sec are being made (the model is going too slowly).

def stop_after_10(approx, loss_history, i):
    if (i > 0) and (i % 10) == 0:
        raise StopIteration('I was slow, sorry')

with model:
    advifit = pm.fit(callbacks=[stop_after_10])

I have two questions. Is there a way to way to define a similar function for pm.sample(), since it does not seem to have a callbacks parameter, and is there a way to make the function instead read the time left for the model to complete, and halt the function after a certain time if the reported time left is too high? For me, at least, the time left of a model is a better constraint to optimize against than raw iterations/sec.

Thank you for your time.

Unfortunately there is no way to do it for pm.sample() - contribution welcome :slight_smile:

So, in that case I have two questions I hope you could answer so I can work this out.

Could you explain the logic behind if (i > 0) and (i % 10) == 0 ?

I imagine i % 10 checks the it/sec every 10 seconds, but how does i > 0 check if the it/sec is too low?

Secondly, going over how advi (github) implements model stopping, and what needs to be accomplished in sample (github), it seems like advi does it like this:

  1. a user-defined callback raises StopIteration if too little iterations occurring,
  2. StopIteration is followed by a progress.close (lines: 83, 173, 177, 230, 250), which does not exist in pm.sample,
  3. pm.sample, however, has sampling.close, strace.close, and trace.close in many places, which I also assume roughly shuts the model down (lines: 551, 643-652, 905-915, 977-995).

So is my analysis right so far?