Kernel crashes during pathfinder fitting

[fonnesbeck]

For several models (not sure what characteristics yet cause if to fail), fitting with pathfinder causes the Python kernel to crash. I've tried this both on CPU and GPU. The log is below:

info 12:18:18.397: Disposing kernel .jvsc74a57bd0a0954a124d2bf69a9421ff2dbb2519042f2244cade5040e3d13151bd4c002865./Users/cfonnesbeck/mambaforge/envs/pie/python./Users/cfonnesbeck/mambaforge/envs/pie/python.-m#ipykernel_launcher for notebook Interactive-1.interactive due to selection of another kernel or closing of the notebook
info 12:18:18.397: Dispose Kernel 'Interactive-1.interactive' associated with '/Users/cfonnesbeck/phillies/pie/research/projections/pitchers/pitcher_proj.py'
info 12:18:18.406: Dispose Kernel 'Interactive-1.interactive' associated with '/Users/cfonnesbeck/phillies/pie/research/projections/pitchers/pitcher_proj.py'
[I 12:18:18.433 NotebookApp] Kernel shutdown: db435396-0e9c-48cf-8b9d-7ac31bccd3da
[I 12:18:19.402 NotebookApp] Starting buffering for db435396-0e9c-48cf-8b9d-7ac31bccd3da:05dcd16b-b821-4492-bbbf-5224df1874c6
info 12:18:23.589: Starting interactive window for resource '/Users/cfonnesbeck/phillies/pie/research/projections/pitchers/pitcher_proj.py' with controller '.jvsc74a57bd0a0954a124d2bf69a9421ff2dbb2519042f2244cade5040e3d13151bd4c002865./Users/cfonnesbeck/mambaforge/envs/pie/python./Users/cfonnesbeck/mambaforge/envs/pie/python.-m#ipykernel_launcher (Interactive)'
info 12:18:23.636: Attempting to start a server because of preload conditions ...
info 12:18:23.661: Process Execution: > ~/mambaforge/envs/pie/bin/python -m pip list
> ~/mambaforge/envs/pie/bin/python -m pip list
info 12:18:24.433: Got env vars with python /Users/cfonnesbeck/mambaforge/envs/pie/bin/python, with env var count 104 and custom env var count 93 in 780ms
info 12:18:24.694: Starting Jupyter Session startUsingPythonInterpreter, .jvsc74a57bd0a0954a124d2bf69a9421ff2dbb2519042f2244cade5040e3d13151bd4c002865./Users/cfonnesbeck/mambaforge/envs/pie/python./Users/cfonnesbeck/mambaforge/envs/pie/python.-m#ipykernel_launcher (Python Path: /Users/cfonnesbeck/mambaforge/envs/pie, EnvType: Conda, EnvName: 'pie', Version: 3.9.15 | packaged by conda-forge | (main, Nov 22 2022, 08:48:25) 
[Clang 14.0.6 ]) for 'Interactive-1.interactive' (disableUI=false)
[I 12:18:24.699 NotebookApp] Creating new notebook in /research/projections/pitchers
info 12:18:24.709: installMissingDependencies /Users/cfonnesbeck/mambaforge/envs/pie/bin/python, ui.disabled=false for resource '/Users/cfonnesbeck/phillies/pie/research/projections/pitchers/pitcher_proj.py'
info 12:18:24.710: Got env vars with python /Users/cfonnesbeck/mambaforge/envs/pie/bin/python, with env var count 104 and custom env var count 93 in 0ms
info 12:18:24.711: Process Execution: > ~/mambaforge/envs/pie/bin/python -c "import ipykernel"
> ~/mambaforge/envs/pie/bin/python -c "import ipykernel"
info 12:18:24.923: Spec argv[0] updated from '/Users/cfonnesbeck/mambaforge/envs/pie/bin/python' to '/Users/cfonnesbeck/mambaforge/envs/pie/bin/python'
info 12:18:24.934: Got env vars with python /Users/cfonnesbeck/mambaforge/envs/pie/bin/python, with env var count 104 and custom env var count 93 in 797ms
[I 12:18:24.952 NotebookApp] Kernel started: abd00f78-2946-4cf5-a842-d3e4b08ea45e, name: pythonjvsc74a57bd0a0954a124d2bf69a9421ff2dbb2519042f2244cade5040e3d13151bd4c002865
[W 12:18:24.964 NotebookApp] delete /research/projections/pitchers/pitcher_proj.py-jvsc-1a415af0-45bf-4b4b-a0c6-7ef772010b09afc69f4d-cad7-4334-ad4e-a60d17a4e4cf.ipynb
info 12:18:25.640: Started session for kernel startUsingPythonInterpreter:.jvsc74a57bd0a0954a124d2bf69a9421ff2dbb2519042f2244cade5040e3d13151bd4c002865./Users/cfonnesbeck/mambaforge/envs/pie/python./Users/cfonnesbeck/mambaforge/envs/pie/python.-m#ipykernel_launcher
info 12:18:25.642: UpdateWorkingDirectoryAndPath in Kernel
info 12:18:25.671: Generated code for 1 = <ipython-input-1-e8e064250631> with 30 lines
info 12:18:25.748: Got env vars with python /Users/cfonnesbeck/mambaforge/envs/pie/bin/python, with env var count 104 and custom env var count 93 in 825ms
info 12:18:28.377: Generated code for 2 = <ipython-input-2-7fcc05f8245f> with 97 lines
info 12:18:28.406: Generated code for 3 = <ipython-input-3-2e36476b3ed3> with 217 lines
info 12:18:28.442: Generated code for 4 = <ipython-input-4-1d2b8ff62d64> with 228 lines
info 12:18:28.482: Generated code for 5 = <ipython-input-5-7f1077aeacc4> with 27 lines
info 12:18:28.506: Generated code for 6 = <ipython-input-6-f7c4cace876c> with 96 lines
info 12:18:36.647: Cancel all remaining cells true || Error || undefined
info 12:22:22.264: Cancel all remaining cells true || Error || undefined
info 12:22:50.255: Disposing kernel .jvsc74a57bd0a0954a124d2bf69a9421ff2dbb2519042f2244cade5040e3d13151bd4c002865./Users/cfonnesbeck/mambaforge/envs/pie/python./Users/cfonnesbeck/mambaforge/envs/pie/python.-m#ipykernel_launcher for notebook Interactive-1.interactive due to selection of another kernel or closing of the notebook
info 12:22:50.255: Dispose Kernel 'Interactive-1.interactive' associated with '/Users/cfonnesbeck/phillies/pie/research/projections/pitchers/pitcher_proj.py'
info 12:22:50.259: Dispose Kernel 'Interactive-1.interactive' associated with '/Users/cfonnesbeck/phillies/pie/research/projections/pitchers/pitcher_proj.py'
[I 12:22:50.262 NotebookApp] Kernel shutdown: abd00f78-2946-4cf5-a842-d3e4b08ea45e
info 12:22:55.344: Starting interactive window for resource '/Users/cfonnesbeck/phillies/pie/research/projections/pitchers/pitcher_proj.py' with controller '.jvsc74a57bd0a0954a124d2bf69a9421ff2dbb2519042f2244cade5040e3d13151bd4c002865./Users/cfonnesbeck/mambaforge/envs/pie/python./Users/cfonnesbeck/mambaforge/envs/pie/python.-m#ipykernel_launcher (Interactive)'
info 12:22:55.391: Attempting to start a server because of preload conditions ...
info 12:22:55.453: Starting Jupyter Session startUsingPythonInterpreter, .jvsc74a57bd0a0954a124d2bf69a9421ff2dbb2519042f2244cade5040e3d13151bd4c002865./Users/cfonnesbeck/mambaforge/envs/pie/python./Users/cfonnesbeck/mambaforge/envs/pie/python.-m#ipykernel_launcher (Python Path: /Users/cfonnesbeck/mambaforge/envs/pie, EnvType: Conda, EnvName: 'pie', Version: 3.9.15 | packaged by conda-forge | (main, Nov 22 2022, 08:48:25) 
[Clang 14.0.6 ]) for 'Interactive-1.interactive' (disableUI=false)
[I 12:22:55.457 NotebookApp] Creating new notebook in /research/projections/pitchers
info 12:22:55.465: installMissingDependencies /Users/cfonnesbeck/mambaforge/envs/pie/bin/python, ui.disabled=false for resource '/Users/cfonnesbeck/phillies/pie/research/projections/pitchers/pitcher_proj.py'
info 12:22:55.466: Got env vars with python /Users/cfonnesbeck/mambaforge/envs/pie/bin/python, with env var count 104 and custom env var count 93 in 1ms
info 12:22:55.467: Process Execution: > ~/mambaforge/envs/pie/bin/python -c "import ipykernel"
> ~/mambaforge/envs/pie/bin/python -c "import ipykernel"
info 12:22:55.654: Spec argv[0] updated from '/Users/cfonnesbeck/mambaforge/envs/pie/bin/python' to '/Users/cfonnesbeck/mambaforge/envs/pie/bin/python'
info 12:22:55.655: Got env vars with python /Users/cfonnesbeck/mambaforge/envs/pie/bin/python, with env var count 104 and custom env var count 93 in 1ms
[I 12:22:55.678 NotebookApp] Kernel started: 00b2089f-89e5-4996-bf18-f5be75bb5cd3, name: pythonjvsc74a57bd0a0954a124d2bf69a9421ff2dbb2519042f2244cade5040e3d13151bd4c002865
[W 12:22:55.685 NotebookApp] delete /research/projections/pitchers/pitcher_proj.py-jvsc-a9a4b23d-ab0d-47a1-8cc3-d7a0922f8fd0e9568bc4-19d4-4ea4-9442-b172f24c2401.ipynb
info 12:22:55.894: Got env vars with python /Users/cfonnesbeck/mambaforge/envs/pie/bin/python, with env var count 104 and custom env var count 93 in 0ms
info 12:22:56.150: Started session for kernel startUsingPythonInterpreter:.jvsc74a57bd0a0954a124d2bf69a9421ff2dbb2519042f2244cade5040e3d13151bd4c002865./Users/cfonnesbeck/mambaforge/envs/pie/python./Users/cfonnesbeck/mambaforge/envs/pie/python.-m#ipykernel_launcher
info 12:22:56.151: UpdateWorkingDirectoryAndPath in Kernel
info 12:22:56.173: Generated code for 1 = <ipython-input-1-e8e064250631> with 30 lines
info 12:22:58.356: Generated code for 2 = <ipython-input-2-7fcc05f8245f> with 97 lines
info 12:22:58.388: Generated code for 3 = <ipython-input-3-2e36476b3ed3> with 217 lines
info 12:22:58.428: Generated code for 4 = <ipython-input-4-1d2b8ff62d64> with 228 lines
info 12:22:58.465: Generated code for 5 = <ipython-input-5-7f1077aeacc4> with 27 lines
info 12:22:58.494: Generated code for 6 = <ipython-input-6-f7c4cace876c> with 96 lines
info 12:23:19.337: Cancel all remaining cells true || Error || undefined
info 12:25:01.080: Cancel all remaining cells true || Error || undefined
[I 12:27:01.670 NotebookApp] KernelRestarter: restarting kernel (1/5), keep random ports
error 12:27:01.677: Error in waiting for cell to complete Error: Canceled future for execute_request message before replies were done
    at t.KernelShellFutureHandler.dispose (/Users/cfonnesbeck/.vscode/extensions/ms-toolsai.jupyter-2022.11.1003412109/out/extension.node.js:2:32353)
    at /Users/cfonnesbeck/.vscode/extensions/ms-toolsai.jupyter-2022.11.1003412109/out/extension.node.js:2:26572
    at Map.forEach (<anonymous>)
    at v._clearKernelState (/Users/cfonnesbeck/.vscode/extensions/ms-toolsai.jupyter-2022.11.1003412109/out/extension.node.js:2:26557)
    at /Users/cfonnesbeck/.vscode/extensions/ms-toolsai.jupyter-2022.11.1003412109/out/extension.node.js:2:29000
    at processTicksAndRejections (node:internal/process/task_queues:96:5)
warn 12:27:01.679: Cell completed with errors {
  message: 'Canceled future for execute_request message before replies were done'
}
info 12:27:01.680: Cancel all remaining cells true || Error || undefined
[I 12:27:01.677 NotebookApp] Starting buffering for 00b2089f-89e5-4996-bf18-f5be75bb5cd3:8468a364-83c1-4bc6-bc88-d1fcfce66644
[I 12:27:01.688 NotebookApp] Restoring connection for 00b2089f-89e5-4996-bf18-f5be75bb5cd3:8468a364-83c1-4bc6-bc88-d1fcfce66644

[fonnesbeck]

Initially I thought it may have something to do with running under Jupyter, but failure occurs from the Python console as well:

/Users/cfonnesbeck/mambaforge/envs/pie/lib/python3.9/site-packages/pymc/sampling/jax.py:37: UserWarning: This module is experimental.
  warnings.warn("This module is experimental.")
2022-12-12 12:34:47 INFO jax._src.lib.xla_bridge: Remote TPU is not linked into jax; skipping remote TPU.
2022-12-12 12:34:47 INFO jax._src.lib.xla_bridge: Unable to initialize backend 'tpu_driver': Could not initialize backend 'tpu_driver'
2022-12-12 12:34:47 INFO jax._src.lib.xla_bridge: Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
2022-12-12 12:34:47 INFO jax._src.lib.xla_bridge: Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
2022-12-12 12:34:47 INFO jax._src.lib.xla_bridge: Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'
2022-12-12 12:34:47 INFO jax._src.lib.xla_bridge: Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.
Killed: 9
/Users/cfonnesbeck/mambaforge/envs/pie/lib/python3.9/multiprocessing/resource_tracker.py:216: UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '

[ricardoV94]

Two problems with our implementation:

• Initializing the model with an isometric MvNormal draw is silly and consumes a lot of RAM on my machine (killing the kernel).
• We are returning the whole path?

I'll fix those and see if it helps. My model is still too big that the JAX Scan (I think) tries to preallocate 30GB of RAM, but I think that's outside of PyMC control at that point (maybe still tweakable in blackjax - CC @junpenglao?)

[ricardoV94]

The API changed in blackjax, but no release since last October :(