Databricks Notebook运行pymc-marketing模型时单元格假死问题求助

问题现象 #

在Databricks Notebook中运行pymc-marketing拟合模型时出现异常：

• 执行模型拟合的单元格显示假死状态，进度条停止更新
• 实际代码已执行完成，所有输出文件均已写入指定文件夹
• 点击Interrupt按钮无响应，需手动分离Notebook才能结束会话
• 查看集群日志确认整个运行流程完全正常

上下文信息 #

数据 #

• 样本量：22周数据（小样本用于测试模型运行）
• 输入变量：11个

模型规格 #

{
"meida_channel": [
    "x1",
    "x2",
    "x3",
    "x4",
    "x5",
    "x6",
    "x7"
],
"control": [
    "c1",
    "c2",
    "c3",
    "c4"
],
"response_var": "units",
"date_var": "end_date",
"beta_channel": {
    "dist": "TruncatedNormal",
    "kwargs":{
        "mu": [stuff here],
        "sigma": 0.001,
        "lower": 0
    }
},
"alpha": {
    "dist": "Beta",
    "kwargs": {
        "alpha": 1,
        "beta": 3
    }
},
"lam": {
    "dist": "Gamma",
    "kwargs": {
        "alpha": 2,
        "beta": 1
    }
},
"gamma_control": {
    "dist": "Normal",
    "kwargs": {
        "mu": [stuff here],
        "sigma": 0.001
    }
}}

采样参数 #

sampler_kwargs = {
    "draws": 8_000,
    "target_accept": 0.95,
    "chains": 4,
    "random_seed": rng,
    # "nuts": {"max_treedepth": 15},
}

排查思路 #

• 前端与后端通信异常：Databricks Notebook进度条依赖WebSocket同步状态，模型完成后通信中断会导致前端未收到结束信号。可在代码末尾添加显式输出（如print("模型拟合完成，输出已写入指定路径")），强制触发前端更新。
• 进度回调兼容性问题：pymc的采样进度条在Databricks环境下可能存在输出流阻塞。尝试在采样时关闭进度条，添加progressbar=False参数到pm.sample()调用中，观察是否仍出现假死。
• 后台进程未正常退出：模型跑完后，采样子进程或资源未完全释放，导致单元格无法结束。可在代码末尾添加资源清理逻辑：del model; del sampler; import gc; gc.collect()，强制释放资源。
• Notebook会话缓存异常：个别会话可能存在缓存问题，尝试新建Notebook重新运行代码，或重启集群后测试。
• 采样参数影响：当前4链8000次采样的配置，可能在截断分布采样结束后有收尾操作未被前端捕获。临时降低draws（如1000）或chains（如2），测试是否仍复现假死，缩小问题范围。

内容的提问来源于stack exchange，提问作者Timothy Mcwilliams