Flow gets stuck when sending logs to UI from some libraries

[williamjamir]

Bug summary

I'm having trouble using Prefect with the pymc library. When I send logs to the UI, the flow gets stuck and doesn't finish, so I have to cancel it manually. This problem happens with both Prefect versions 2 and 3.

I've encountered similar issues with other situations, like with Dask and other integrations, but I couldn't figure out the problem. I think solving this issue will also help with other similar unnoticed problems.

Here's a minimal reproducible example. When it's working properly (without sending logs to the UI), it should run in 5 to 6 seconds.

logging.yaml

version: 1
disable_existing_loggers: False
handlers:
  console:
    level: WARNING
    class: prefect.logging.handlers.PrefectConsoleHandler

  api:
    class: prefect.logging.handlers.APILogHandler

loggers:
    pymc:
        level: INFO
        handlers: [ api ]
        propagate: false

root:
  level: INFO
  handlers: [ console ]

acme.py

from pathlib import Path
from prefect import flow
from pymc_marketing.mmm import MMM, GeometricAdstock, LogisticSaturation
import pandas as pd 

@flow()
def acme():
    if Path("data.csv").exists():
        data = pd.read_csv('data.csv')
    else:
        data_url = 'https://raw.githubusercontent.com/pymc-labs/pymc-marketing/main/data/mmm_example.csv'
        data = pd.read_csv(data_url, parse_dates=['date_week'])
        data.to_csv('data.csv')
    
    mmm = MMM(
        adstock=GeometricAdstock(l_max=1),
        saturation=LogisticSaturation(),
        date_column='date_week',
        channel_columns=['x1', 'x2'],
        control_columns=['event_1', 'event_2', 't'],
        yearly_seasonality=1,
    )
    X = data.drop('y', axis=1)
    y = data['y']
    mmm.fit(X, y)

if __name__ == '__main__':
    acme()

pip install prefect pymc-marketing

PREFECT_LOGGING_SETTINGS_PATH =<path>/logging.yml python acme.py

Version info (prefect version output)

Version:             3.0.0
API version:         0.8.4
Python version:      3.11.9
Git commit:          c40d069d
Built:               Tue, Sep 3, 2024 11:13 AM
OS/Arch:             darwin/arm64
Profile:             local
Server type:         server
Pydantic version:    2.9.2

Additional context

I already tried disabling the progress bar from Pymc, but the issue is still persisting (flow gets stuck)

    mmm.fit(X, y,  progressbar=False)

Also, this is the stack trace when interrupting the execution on Prefect 3

python -m acme
Using NumPy C-API based implementation for BLAS functions.
Sampling 4 chains, 1 divergences ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╺━━  92% 0:00:01 / 0:00:39
^C^CCrash detected! Execution was aborted by an interrupt signal.
Finished in state Crashed('Execution was aborted by an interrupt signal.')
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/opt/opensource/dummy_folder/acme.py", line 31, in <module>
    hi()
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/prefect/flows.py", line 1334, in __call__
    return run_flow(
           ^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/prefect/flow_engine.py", line 810, in run_flow
    return run_flow_sync(**kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/prefect/flow_engine.py", line 688, in run_flow_sync
    engine.call_flow_fn()
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/prefect/flow_engine.py", line 667, in call_flow_fn
    result = call_with_parameters(self.flow.fn, self.parameters)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/prefect/utilities/callables.py", line 206, in call_with_parameters
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/opt/opensource/dummy_folder/acme.py", line 27, in hi
    mmm.fit(X, y)
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pymc_marketing/model_builder.py", line 606, in fit
    idata = pm.sample(**sampler_kwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pymc/sampling/mcmc.py", line 870, in sample
    return _sample_return(
           ^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pymc/sampling/mcmc.py", line 938, in _sample_return
    idata = pm.to_inference_data(mtrace, **ikwargs)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pymc/backends/arviz.py", line 520, in to_inference_data
    return InferenceDataConverter(
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pymc/backends/arviz.py", line 208, in __init__
    self.posterior_trace, self.warmup_trace = self.split_trace()
                                              ^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pymc/backends/arviz.py", line 253, in split_trace
    trace_posterior = self.trace[self.ntune :]
                      ~~~~~~~~~~^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pymc/backends/base.py", line 362, in __getitem__
    return self._slice(idx)
           ^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pymc/backends/base.py", line 529, in _slice
    new_traces = [trace._slice(slice) for trace in self._straces.values()]
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pymc/backends/base.py", line 529, in <listcomp>
    new_traces = [trace._slice(slice) for trace in self._straces.values()]
                  ^^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pymc/backends/ndarray.py", line 168, in _slice
    sliced = NDArray(model=self.model, vars=self.vars)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pymc/backends/ndarray.py", line 44, in __init__
    super().__init__(name, model, vars, test_point)
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pymc/backends/base.py", line 158, in __init__
    self.fn = model.compile_fn(vars, inputs=model.value_vars, on_unused_input="ignore")
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pymc/model/core.py", line 1648, in compile_fn
    fn = compile_pymc(
         ^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pymc/pytensorf.py", line 1040, in compile_pymc
    pytensor_function = pytensor.function(
                        ^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pytensor/compile/function/__init__.py", line 315, in function
    fn = pfunc(
         ^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pytensor/compile/function/pfunc.py", line 465, in pfunc
    return orig_function(
           ^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pytensor/compile/function/types.py", line 1750, in orig_function
    m = Maker(
        ^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pytensor/compile/function/types.py", line 1523, in __init__
    self.prepare_fgraph(inputs, outputs, found_updates, fgraph, mode, profile)
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pytensor/compile/function/types.py", line 1411, in prepare_fgraph
    rewriter_profile = rewriter(fgraph)
                       ^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 125, in __call__
    return self.rewrite(fgraph)
           ^^^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 121, in rewrite
    return self.apply(fgraph, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 291, in apply
    sub_prof = rewriter.apply(fgraph)
               ^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 2427, in apply
    node_rewriter_change = self.process_node(
                           ^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pytensor/graph/rewriting/basic.py", line 1965, in process_node
    fgraph.replace_all_validate_remove(  # type: ignore
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pytensor/graph/features.py", line 628, in replace_all_validate_remove
    chk = fgraph.replace_all_validate(replacements, reason=reason, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pytensor/graph/features.py", line 603, in replace_all_validate
    fgraph.validate()
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pytensor/graph/features.py", line 478, in validate_
    ret = fgraph.execute_callbacks("validate")
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/william/.pyenv/versions/3.11.9/envs/dummy_env/lib/python3.11/site-packages/pytensor/graph/fg.py", line 712, in execute_callbacks
    fn = getattr(feature, name)
         ^^^^^^^^^^^^^^^^^^^^^^
KeyboardInterrupt
make: *** [run] Interrupt: 2