bmb.Model Progress Bar Freezes on DataBricks

Hello,

I’m trying to run the model, but after around 6 minutes, the progress bar tends to completely freeze and it stops updating the progress (as a Databricks widget).

Could there be a reason for this? Is it related to the Databricks socket connections and widgets?
Are there any other ways to display progress bars of the bmb model fitting?

bmb.Model(
    ".... ",
    pandas_df,
    family="bernoulli"
)

Sometimes, progress bar does not show anything at all and shows the error:

Loading the widget is taking longer than expected. 

I’d really appreciate any help with the progress bar! :blush:

The PyMC progress bar uses the rich.progress package. I don’t know anything about databricks, but if they don’t play nice, that gives you something to search for?

External samplers (nutpie, numpyro, blackjax) all use their own progress bars (numpyro uses fastprogress, nutpie has a custom one), so you might be able to just switch the sampler to fix it. The bambi docs have details on choosing a sampling backend.