orbax
orbax copied to clipboard
google
CheckpointManager randomly crashes the kernel of the JupyterLab
The following code in multiple cells in the JupyterLab crashes randomly when executed:
import jax
import flax
from flax import nnx
import orbax.checkpoint as ocp
print(jax.__version__)
print(flax.__version__)
print(ocp.__version__)
path = ocp.test_utils.erase_and_create_empty("/tmp/repro-crash-model/")
options = ocp.CheckpointManagerOptions(max_to_keep=2)
mngr = ocp.CheckpointManager(path, options=options)
def save_model(epoch):
state = nnx.state(model)
# We should convert PRNGKeyArray to the old format for Dropout layers
# https://github.com/google/flax/issues/4231
def get_key_data(x):
if isinstance(x, jax._src.prng.PRNGKeyArray):
if isinstance(x.dtype, jax._src.prng.KeyTy):
return jax.random.key_data(x)
return x
serializable_state = jax.tree.map(get_key_data, state)
mngr.save(epoch, args=ocp.args.StandardSave(serializable_state))
mngr.wait_until_finished()
class TestModel(nnx.Module):
def __init__(self, n1, n2, n3, *, rngs = nnx.Rngs(0)):
self.m1 = nnx.Linear(n1, n2, rngs=rngs)
self.dropout = nnx.Dropout(0.1, rngs=rngs)
self.m1 = nnx.Linear(n2, n3, rngs=rngs)
def __call__(self, x):
x = self.m1(x)
x = self.dropout(x)
x = self.m2(x)
return x
Here is the ipynb to reproduce the crash: https://gist.github.com/vfdev-5/da7f01b9c0c2948046b227a32db2cb2b
Can you please attach additional details like env and error stack.
May be generic debugging tips for JupyterLab can be helpful. https://stackoverflow.com/questions/74154123/how-to-debug-jupyter-kernel-crashes
Env info:
- docker container: pytorch/pytorch:2.4.0-cuda12.4-cudnn9-devel,
cat /etc/lsb-release
DISTRIB_ID=Ubuntu
DISTRIB_RELEASE=22.04
DISTRIB_CODENAME=jammy
DISTRIB_DESCRIPTION="Ubuntu 22.04.4 LTS"
pip list
Package Version
---------------------------- -------------------
absl-py 2.1.0
aiohappyeyeballs 2.4.3
aiohttp 3.10.10
aiosignal 1.3.1
albucore 0.0.20
albumentations 1.4.21
anaconda-anon-usage 0.4.4
annotated-types 0.7.0
anyio 4.6.2.post1
archspec 0.2.3
argon2-cffi 23.1.0
argon2-cffi-bindings 21.2.0
array_record 0.5.1
arrow 1.3.0
asttokens 2.0.5
astunparse 1.6.3
async-lru 2.0.4
attrs 23.1.0
babel 2.16.0
beautifulsoup4 4.12.3
bleach 6.1.0
boltons 23.0.0
Brotli 1.0.9
cachetools 5.5.0
certifi 2024.7.4
cffi 1.16.0
cfgv 3.4.0
chardet 4.0.0
charset-normalizer 2.0.4
chex 0.1.87
click 8.1.7
cloudpickle 3.1.0
comm 0.2.2
conda 24.5.0
conda-build 24.5.1
conda-content-trust 0.2.0
conda_index 0.5.0
conda-libmamba-solver 24.1.0
conda-package-handling 2.3.0
conda_package_streaming 0.10.0
contourpy 1.3.0
cryptography 42.0.5
cycler 0.12.1
datasets 3.1.0
debugpy 1.8.7
decorator 5.1.1
defusedxml 0.7.1
dill 0.3.8
distlib 0.3.9
distro 1.9.0
dm-tree 0.1.8
dnspython 2.6.1
docstring_parser 0.16
einops 0.8.0
etils 1.9.4
eval_type_backport 0.2.0
executing 0.8.3
expecttest 0.2.1
fastjsonschema 2.20.0
filelock 3.13.1
flatbuffers 24.3.25
flax 0.10.1
fonttools 4.54.1
fqdn 1.5.1
frozendict 2.4.2
frozenlist 1.5.0
fsspec 2024.9.0
gast 0.6.0
gcsfs 2024.10.0
gmpy2 2.1.2
google-api-core 2.22.0
google-auth 2.35.0
google-auth-oauthlib 1.2.1
google-cloud-core 2.4.1
google-cloud-storage 2.18.2
google-crc32c 1.6.0
google-pasta 0.2.0
google-resumable-media 2.7.2
googleapis-common-protos 1.65.0
grain 0.2.2
grpcio 1.67.1
h11 0.14.0
h5py 3.12.1
httpcore 1.0.6
httpx 0.27.2
huggingface-hub 0.26.2
humanize 4.11.0
hypothesis 6.108.4
identify 2.6.1
idna 3.7
immutabledict 4.2.0
importlib_resources 6.4.5
ipykernel 6.29.5
ipython 8.25.0
ipywidgets 8.1.5
isoduration 20.11.0
jax 0.4.34
jax-cuda12-pjrt 0.4.34
jax-cuda12-plugin 0.4.34
jaxlib 0.4.34
jaxtyping 0.2.34
jedi 0.19.1
Jinja2 3.1.4
json5 0.9.25
jsonpatch 1.33
jsonpointer 2.1
jsonschema 4.19.2
jsonschema-specifications 2023.7.1
jupyter_client 8.6.3
jupyter_core 5.7.2
jupyter-events 0.10.0
jupyter-lsp 2.2.5
jupyter_server 2.14.2
jupyter_server_terminals 0.5.3
jupyterlab 4.2.5
jupyterlab_pygments 0.3.0
jupyterlab_server 2.27.3
jupyterlab_widgets 3.0.13
jupytext 1.16.4
keras 3.6.0
keras-nightly 3.7.0.dev2024110103
kiwisolver 1.4.7
libarchive-c 2.9
libclang 18.1.1
libmambapy 1.5.8
lintrunner 0.12.5
Markdown 3.7
markdown-it-py 3.0.0
MarkupSafe 2.1.3
matplotlib 3.9.2
matplotlib-inline 0.1.6
mdit-py-plugins 0.4.2
mdurl 0.1.2
menuinst 2.1.1
mistune 3.0.2
mkl-fft 1.3.8
mkl-random 1.2.4
mkl-service 2.4.0
ml-dtypes 0.4.1
monai 1.4.0
more-itertools 10.1.0
mpmath 1.3.0
msgpack 1.1.0
multidict 6.1.0
multiprocess 0.70.16
namex 0.0.8
nbclient 0.10.0
nbconvert 7.16.4
nbformat 5.10.4
nest-asyncio 1.6.0
networkx 3.3
ninja 1.11.1.1
nodeenv 1.9.1
notebook_shim 0.2.4
numpy 1.26.4
nvidia-cublas-cu12 12.4.5.8
nvidia-cuda-cupti-cu12 12.4.127
nvidia-cuda-nvcc-cu12 12.6.77
nvidia-cuda-nvrtc-cu12 12.4.127
nvidia-cuda-runtime-cu12 12.4.127
nvidia-cudnn-cu12 9.1.0.70
nvidia-cufft-cu12 11.2.1.3
nvidia-curand-cu12 10.3.5.147
nvidia-cusolver-cu12 11.6.1.9
nvidia-cusparse-cu12 12.3.1.170
nvidia-nccl-cu12 2.21.5
nvidia-nvjitlink-cu12 12.4.127
nvidia-nvtx-cu12 12.4.127
oauthlib 3.2.2
opencv-python-headless 4.10.0.84
opt_einsum 3.4.0
optax 0.2.4
optree 0.12.1
orbax-checkpoint 0.9.1
overrides 7.7.0
packaging 24.1
pandas 2.2.3
pandocfilters 1.5.1
parso 0.8.3
pexpect 4.8.0
pillow 11.0.0
pip 24.0
pkginfo 1.10.0
platformdirs 3.10.0
pluggy 1.0.0
pre_commit 4.0.1
prometheus_client 0.21.0
promise 2.3
prompt-toolkit 3.0.43
propcache 0.2.0
proto-plus 1.25.0
protobuf 5.28.2
psutil 5.9.0
ptyprocess 0.7.0
pure-eval 0.2.2
pyarrow 18.0.0
pyasn1 0.6.1
pyasn1_modules 0.4.1
pycosat 0.6.6
pycparser 2.21
pydantic 2.9.2
pydantic_core 2.23.4
Pygments 2.15.1
pyparsing 3.2.0
PySocks 1.7.1
python-dateutil 2.9.0.post0
python-etcd 0.4.5
python-json-logger 2.0.7
pytz 2024.1
PyYAML 6.0.1
pyzmq 26.2.0
referencing 0.30.2
regex 2024.11.6
requests 2.32.3
requests-oauthlib 2.0.0
rfc3339-validator 0.1.4
rfc3986-validator 0.1.1
rich 13.9.2
rpds-py 0.10.6
rsa 4.9
ruamel.yaml 0.17.21
safetensors 0.4.5
scipy 1.14.1
seaborn 0.13.2
Send2Trash 1.8.3
setuptools 69.5.1
simple-parsing 0.1.6
simsimd 6.0.1
six 1.16.0
sniffio 1.3.1
sortedcontainers 2.4.0
soupsieve 2.5
stack-data 0.2.0
stringzilla 3.10.8
sympy 1.13.1
tb-nightly 2.19.0a20241106
tensorboard 2.18.0
tensorboard-data-server 0.7.2
tensorflow 2.18.0
tensorflow-datasets 4.9.7
tensorflow-io-gcs-filesystem 0.37.1
tensorflow-metadata 1.16.1
tensorstore 0.1.66
termcolor 2.5.0
terminado 0.18.1
tiktoken 0.8.0
tinycss2 1.3.0
tokenizers 0.20.3
toml 0.10.2
toolz 1.0.0
torch 2.5.1
torchaudio 2.4.0
torchelastic 0.2.2
torchvision 0.20.1
tornado 6.4.1
tqdm 4.66.4
traitlets 5.14.3
transformers 4.46.3
treescope 0.1.5
triton 3.1.0
truststore 0.8.0
typeguard 2.13.3
types-dataclasses 0.6.6
types-python-dateutil 2.9.0.20241003
typing_extensions 4.11.0
tzdata 2024.2
uri-template 1.3.0
urllib3 2.2.2
virtualenv 20.26.6
wcwidth 0.2.5
webcolors 24.8.0
webencodings 0.5.1
websocket-client 1.8.0
Werkzeug 3.1.2
wheel 0.43.0
widgetsnbextension 4.0.13
wrapt 1.16.0
xxhash 3.5.0
yarl 1.17.1
zipp 3.20.2
zstandard 0.22.0
JupyterLab console shows the following before the crash:
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/conda/lib/python3.11/site-packages/ipykernel_launcher.py", line 18, in <module>
app.launch_new_instance()
File "/opt/conda/lib/python3.11/site-packages/traitlets/config/application.py", line 1075, in launch_instance
app.start()
File "/opt/conda/lib/python3.11/site-packages/ipykernel/kernelapp.py", line 739, in start
self.io_loop.start()
File "/opt/conda/lib/python3.11/site-packages/tornado/platform/asyncio.py", line 205, in start
self.asyncio_loop.run_forever()
File "/opt/conda/lib/python3.11/asyncio/base_events.py", line 608, in run_forever
self._run_once()
File "/opt/conda/lib/python3.11/asyncio/base_events.py", line 1921, in _run_once
handle = self._ready.popleft()
^^^^^^^^^^^^^^^^^^^^^
IndexError: pop from an empty deque
Hope this helps.
PS: few observations about the issue:
- repro code in the single cell does not crash (at least I've never seen)
Output of `jupyter troubleshoot`:
$PATH:
/usr/local/nvidia/bin
/usr/local/cuda/bin
/opt/conda/bin
/usr/local/nvidia/bin
/usr/local/cuda/bin
/usr/local/sbin
/usr/local/bin
/usr/sbin
/usr/bin
/sbin
/bin
sys.path:
/opt/conda/bin
/opt/conda/lib/python311.zip
/opt/conda/lib/python3.11
/opt/conda/lib/python3.11/lib-dynload
/opt/conda/lib/python3.11/site-packages
sys.executable:
/opt/conda/bin/python
sys.version:
3.11.9 (main, Apr 19 2024, 16:48:06) [GCC 11.2.0]
platform.platform():
Linux-5.15.0-124-generic-x86_64-with-glibc2.35
which -a jupyter:
/opt/conda/bin/jupyter
pip list:
Package Version
---------------------------- -------------------
absl-py 2.1.0
aiohappyeyeballs 2.4.3
aiohttp 3.10.10
aiosignal 1.3.1
albucore 0.0.20
albumentations 1.4.21
anaconda-anon-usage 0.4.4
annotated-types 0.7.0
anyio 4.6.2.post1
archspec 0.2.3
argon2-cffi 23.1.0
argon2-cffi-bindings 21.2.0
array_record 0.5.1
arrow 1.3.0
asttokens 2.0.5
astunparse 1.6.3
async-lru 2.0.4
attrs 23.1.0
babel 2.16.0
beautifulsoup4 4.12.3
bleach 6.1.0
boltons 23.0.0
Brotli 1.0.9
cachetools 5.5.0
certifi 2024.7.4
cffi 1.16.0
cfgv 3.4.0
chardet 4.0.0
charset-normalizer 2.0.4
chex 0.1.87
click 8.1.7
cloudpickle 3.1.0
comm 0.2.2
conda 24.5.0
conda-build 24.5.1
conda-content-trust 0.2.0
conda_index 0.5.0
conda-libmamba-solver 24.1.0
conda-package-handling 2.3.0
conda_package_streaming 0.10.0
contourpy 1.3.0
cryptography 42.0.5
cycler 0.12.1
datasets 3.1.0
debugpy 1.8.7
decorator 5.1.1
defusedxml 0.7.1
dill 0.3.8
distlib 0.3.9
distro 1.9.0
dm-tree 0.1.8
dnspython 2.6.1
docstring_parser 0.16
einops 0.8.0
etils 1.9.4
eval_type_backport 0.2.0
executing 0.8.3
expecttest 0.2.1
fastjsonschema 2.20.0
filelock 3.13.1
flatbuffers 24.3.25
flax 0.10.1
fonttools 4.54.1
fqdn 1.5.1
frozendict 2.4.2
frozenlist 1.5.0
fsspec 2024.9.0
gast 0.6.0
gcsfs 2024.10.0
gmpy2 2.1.2
google-api-core 2.22.0
google-auth 2.35.0
google-auth-oauthlib 1.2.1
google-cloud-core 2.4.1
google-cloud-storage 2.18.2
google-crc32c 1.6.0
google-pasta 0.2.0
google-resumable-media 2.7.2
googleapis-common-protos 1.65.0
grain 0.2.2
grpcio 1.67.1
h11 0.14.0
h5py 3.12.1
httpcore 1.0.6
httpx 0.27.2
huggingface-hub 0.26.2
humanize 4.11.0
hypothesis 6.108.4
identify 2.6.1
idna 3.7
immutabledict 4.2.0
importlib_resources 6.4.5
ipykernel 6.29.5
ipython 8.25.0
ipywidgets 8.1.5
isoduration 20.11.0
jax 0.4.34
jax-cuda12-pjrt 0.4.34
jax-cuda12-plugin 0.4.34
jaxlib 0.4.34
jaxtyping 0.2.34
jedi 0.19.1
Jinja2 3.1.4
json5 0.9.25
jsonpatch 1.33
jsonpointer 2.1
jsonschema 4.19.2
jsonschema-specifications 2023.7.1
jupyter_client 8.6.3
jupyter_core 5.7.2
jupyter-events 0.10.0
jupyter-lsp 2.2.5
jupyter_server 2.14.2
jupyter_server_terminals 0.5.3
jupyterlab 4.3.1
jupyterlab_pygments 0.3.0
jupyterlab_server 2.27.3
jupyterlab_widgets 3.0.13
jupytext 1.16.4
keras 3.6.0
keras-nightly 3.7.0.dev2024110103
kiwisolver 1.4.7
libarchive-c 2.9
libclang 18.1.1
libmambapy 1.5.8
lintrunner 0.12.5
Markdown 3.7
markdown-it-py 3.0.0
MarkupSafe 2.1.3
matplotlib 3.9.2
matplotlib-inline 0.1.6
mdit-py-plugins 0.4.2
mdurl 0.1.2
menuinst 2.1.1
mistune 3.0.2
mkl-fft 1.3.8
mkl-random 1.2.4
mkl-service 2.4.0
ml-dtypes 0.4.1
monai 1.4.0
more-itertools 10.1.0
mpmath 1.3.0
msgpack 1.1.0
multidict 6.1.0
multiprocess 0.70.16
namex 0.0.8
nbclient 0.10.0
nbconvert 7.16.4
nbformat 5.10.4
nest-asyncio 1.6.0
networkx 3.3
ninja 1.11.1.1
nodeenv 1.9.1
notebook_shim 0.2.4
numpy 1.26.4
nvidia-cublas-cu12 12.4.5.8
nvidia-cuda-cupti-cu12 12.4.127
nvidia-cuda-nvcc-cu12 12.6.77
nvidia-cuda-nvrtc-cu12 12.4.127
nvidia-cuda-runtime-cu12 12.4.127
nvidia-cudnn-cu12 9.1.0.70
nvidia-cufft-cu12 11.2.1.3
nvidia-curand-cu12 10.3.5.147
nvidia-cusolver-cu12 11.6.1.9
nvidia-cusparse-cu12 12.3.1.170
nvidia-nccl-cu12 2.21.5
nvidia-nvjitlink-cu12 12.4.127
nvidia-nvtx-cu12 12.4.127
oauthlib 3.2.2
opencv-python-headless 4.10.0.84
opt_einsum 3.4.0
optax 0.2.4
optree 0.12.1
orbax-checkpoint 0.9.1
overrides 7.7.0
packaging 24.1
pandas 2.2.3
pandocfilters 1.5.1
parso 0.8.3
pexpect 4.8.0
pillow 11.0.0
pip 24.0
pkginfo 1.10.0
platformdirs 3.10.0
pluggy 1.0.0
pre_commit 4.0.1
prometheus_client 0.21.0
promise 2.3
prompt-toolkit 3.0.43
propcache 0.2.0
proto-plus 1.25.0
protobuf 5.28.2
psutil 5.9.0
ptyprocess 0.7.0
pure-eval 0.2.2
pyarrow 18.0.0
pyasn1 0.6.1
pyasn1_modules 0.4.1
pycosat 0.6.6
pycparser 2.21
pydantic 2.9.2
pydantic_core 2.23.4
Pygments 2.15.1
pyparsing 3.2.0
PySocks 1.7.1
python-dateutil 2.9.0.post0
python-etcd 0.4.5
python-json-logger 2.0.7
pytz 2024.1
PyYAML 6.0.1
pyzmq 26.2.0
referencing 0.30.2
regex 2024.11.6
requests 2.32.3
requests-oauthlib 2.0.0
rfc3339-validator 0.1.4
rfc3986-validator 0.1.1
rich 13.9.2
rpds-py 0.10.6
rsa 4.9
ruamel.yaml 0.17.21
safetensors 0.4.5
scipy 1.14.1
seaborn 0.13.2
Send2Trash 1.8.3
setuptools 69.5.1
simple-parsing 0.1.6
simsimd 6.0.1
six 1.16.0
sniffio 1.3.1
sortedcontainers 2.4.0
soupsieve 2.5
stack-data 0.2.0
stringzilla 3.10.8
sympy 1.13.1
tb-nightly 2.19.0a20241106
tensorboard 2.18.0
tensorboard-data-server 0.7.2
tensorflow 2.18.0
tensorflow-datasets 4.9.7
tensorflow-io-gcs-filesystem 0.37.1
tensorflow-metadata 1.16.1
tensorstore 0.1.66
termcolor 2.5.0
terminado 0.18.1
tiktoken 0.8.0
tinycss2 1.3.0
tokenizers 0.20.3
toml 0.10.2
toolz 1.0.0
torch 2.5.1
torchaudio 2.4.0
torchelastic 0.2.2
torchvision 0.20.1
tornado 6.4.1
tqdm 4.66.4
traitlets 5.14.3
transformers 4.46.3
treescope 0.1.5
triton 3.1.0
truststore 0.8.0
typeguard 2.13.3
types-dataclasses 0.6.6
types-python-dateutil 2.9.0.20241003
typing_extensions 4.11.0
tzdata 2024.2
uri-template 1.3.0
urllib3 2.2.2
virtualenv 20.26.6
wcwidth 0.2.5
webcolors 24.8.0
webencodings 0.5.1
websocket-client 1.8.0
Werkzeug 3.1.2
wheel 0.43.0
widgetsnbextension 4.0.13
wrapt 1.16.0
xxhash 3.5.0
yarl 1.17.1
zipp 3.20.2
zstandard 0.22.0
conda list:
# packages in environment at /opt/conda:
#
# Name Version Build Channel
_libgcc_mutex 0.1 main
_openmp_mutex 5.1 1_gnu
absl-py 2.1.0 pypi_0 pypi
aiohappyeyeballs 2.4.3 pypi_0 pypi
aiohttp 3.10.10 pypi_0 pypi
aiosignal 1.3.1 pypi_0 pypi
albucore 0.0.20 pypi_0 pypi
albumentations 1.4.21 pypi_0 pypi
anaconda-anon-usage 0.4.4 py311hfc0e8ea_100
annotated-types 0.7.0 pypi_0 pypi
anyio 4.6.2.post1 pypi_0 pypi
archspec 0.2.3 pyhd3eb1b0_0
argon2-cffi 23.1.0 pypi_0 pypi
argon2-cffi-bindings 21.2.0 pypi_0 pypi
array-record 0.5.1 pypi_0 pypi
arrow 1.3.0 pypi_0 pypi
asttokens 2.0.5 pyhd3eb1b0_0
astunparse 1.6.3 pypi_0 pypi
async-lru 2.0.4 pypi_0 pypi
attrs 23.1.0 py311h06a4308_0
babel 2.16.0 pypi_0 pypi
beautifulsoup4 4.12.3 py311h06a4308_0
blas 1.0 mkl
bleach 6.1.0 pypi_0 pypi
boltons 23.0.0 py311h06a4308_0
brotli-python 1.0.9 py311h6a678d5_8
bzip2 1.0.8 h5eee18b_6
c-ares 1.19.1 h5eee18b_0
ca-certificates 2024.7.2 h06a4308_0
cachetools 5.5.0 pypi_0 pypi
certifi 2024.7.4 py311h06a4308_0
cffi 1.16.0 py311h5eee18b_1
cfgv 3.4.0 pypi_0 pypi
chardet 4.0.0 py311h06a4308_1003
charset-normalizer 2.0.4 pyhd3eb1b0_0
chex 0.1.87 pypi_0 pypi
click 8.1.7 py311h06a4308_0
cloudpickle 3.1.0 pypi_0 pypi
cmake 3.26.4 h96355d8_0
comm 0.2.2 pypi_0 pypi
conda 24.5.0 py311h06a4308_0
conda-build 24.5.1 py311h06a4308_0
conda-content-trust 0.2.0 py311h06a4308_1
conda-index 0.5.0 py311h06a4308_0
conda-libmamba-solver 24.1.0 pyhd3eb1b0_0
conda-package-handling 2.3.0 py311h06a4308_0
conda-package-streaming 0.10.0 py311h06a4308_0
contourpy 1.3.0 pypi_0 pypi
cryptography 42.0.5 py311hdda0065_1
cuda-cudart 12.4.127 0 nvidia
cuda-cupti 12.4.127 0 nvidia
cuda-libraries 12.4.0 0 nvidia
cuda-nvrtc 12.4.127 0 nvidia
cuda-nvtx 12.4.127 0 nvidia
cuda-opencl 12.5.39 0 nvidia
cuda-runtime 12.4.0 0 nvidia
cuda-version 12.5 3 nvidia
cycler 0.12.1 pypi_0 pypi
datasets 3.1.0 pypi_0 pypi
debugpy 1.8.7 pypi_0 pypi
decorator 5.1.1 pyhd3eb1b0_0
defusedxml 0.7.1 pypi_0 pypi
dill 0.3.8 pypi_0 pypi
distlib 0.3.9 pypi_0 pypi
distro 1.9.0 py311h06a4308_0
dm-tree 0.1.8 pypi_0 pypi
dnspython 2.6.1 pypi_0 pypi
docstring-parser 0.16 pypi_0 pypi
einops 0.8.0 pypi_0 pypi
etils 1.9.4 pypi_0 pypi
eval-type-backport 0.2.0 pypi_0 pypi
executing 0.8.3 pyhd3eb1b0_0
expat 2.6.2 h6a678d5_0
expecttest 0.2.1 pypi_0 pypi
fastjsonschema 2.20.0 pypi_0 pypi
ffmpeg 4.3 hf484d3e_0 pytorch
filelock 3.13.1 py311h06a4308_0
flatbuffers 24.3.25 pypi_0 pypi
flax 0.10.1 pypi_0 pypi
fmt 9.1.0 hdb19cb5_1
fonttools 4.54.1 pypi_0 pypi
fqdn 1.5.1 pypi_0 pypi
freetype 2.12.1 h4a9f257_0
frozendict 2.4.2 py311h06a4308_0
frozenlist 1.5.0 pypi_0 pypi
fsspec 2024.9.0 pypi_0 pypi
gast 0.6.0 pypi_0 pypi
gcsfs 2024.10.0 pypi_0 pypi
gmp 6.2.1 h295c915_3
gmpy2 2.1.2 py311hc9b5ff0_0
gnutls 3.6.15 he1e5248_0
google-api-core 2.22.0 pypi_0 pypi
google-auth 2.35.0 pypi_0 pypi
google-auth-oauthlib 1.2.1 pypi_0 pypi
google-cloud-core 2.4.1 pypi_0 pypi
google-cloud-storage 2.18.2 pypi_0 pypi
google-crc32c 1.6.0 pypi_0 pypi
google-pasta 0.2.0 pypi_0 pypi
google-resumable-media 2.7.2 pypi_0 pypi
googleapis-common-protos 1.65.0 pypi_0 pypi
grain 0.2.2 pypi_0 pypi
grpcio 1.67.1 pypi_0 pypi
h11 0.14.0 pypi_0 pypi
h5py 3.12.1 pypi_0 pypi
httpcore 1.0.6 pypi_0 pypi
httpx 0.27.2 pypi_0 pypi
huggingface-hub 0.26.2 pypi_0 pypi
humanize 4.11.0 pypi_0 pypi
hypothesis 6.108.4 pypi_0 pypi
icu 73.1 h6a678d5_0
identify 2.6.1 pypi_0 pypi
idna 3.7 py311h06a4308_0
immutabledict 4.2.0 pypi_0 pypi
importlib-resources 6.4.5 pypi_0 pypi
intel-openmp 2023.1.0 hdb19cb5_46306
ipykernel 6.29.5 pypi_0 pypi
ipython 8.25.0 py311h06a4308_0
ipywidgets 8.1.5 pypi_0 pypi
isoduration 20.11.0 pypi_0 pypi
jax 0.4.34 pypi_0 pypi
jax-cuda12-pjrt 0.4.34 pypi_0 pypi
jax-cuda12-plugin 0.4.34 pypi_0 pypi
jaxlib 0.4.34 pypi_0 pypi
jaxtyping 0.2.34 pypi_0 pypi
jedi 0.19.1 py311h06a4308_0
jinja2 3.1.4 py311h06a4308_0
jpeg 9e h5eee18b_2
json5 0.9.25 pypi_0 pypi
jsonpatch 1.33 py311h06a4308_1
jsonpointer 2.1 pyhd3eb1b0_0
jsonschema 4.19.2 py311h06a4308_0
jsonschema-specifications 2023.7.1 py311h06a4308_0
jupyter-client 8.6.3 pypi_0 pypi
jupyter-core 5.7.2 pypi_0 pypi
jupyter-events 0.10.0 pypi_0 pypi
jupyter-lsp 2.2.5 pypi_0 pypi
jupyter-server 2.14.2 pypi_0 pypi
jupyter-server-terminals 0.5.3 pypi_0 pypi
jupyterlab 4.3.1 pypi_0 pypi
jupyterlab-pygments 0.3.0 pypi_0 pypi
jupyterlab-server 2.27.3 pypi_0 pypi
jupyterlab-widgets 3.0.13 pypi_0 pypi
jupytext 1.16.4 pypi_0 pypi
keras 3.6.0 pypi_0 pypi
keras-nightly 3.7.0.dev2024110103 pypi_0 pypi
kiwisolver 1.4.7 pypi_0 pypi
krb5 1.20.1 h143b758_1
lame 3.100 h7b6447c_0
lcms2 2.12 h3be6417_0
ld_impl_linux-64 2.38 h1181459_1
lerc 3.0 h295c915_0
libarchive 3.6.2 h6ac8c49_3
libclang 18.1.1 pypi_0 pypi
libcublas 12.4.2.65 0 nvidia
libcufft 11.2.0.44 0 nvidia
libcufile 1.10.1.7 0 nvidia
libcurand 10.3.6.82 0 nvidia
libcurl 8.7.1 h251f7ec_0
libcusolver 11.6.0.99 0 nvidia
libcusparse 12.3.0.142 0 nvidia
libdeflate 1.17 h5eee18b_1
libedit 3.1.20230828 h5eee18b_0
libev 4.33 h7f8727e_1
libffi 3.4.4 h6a678d5_1
libgcc-ng 11.2.0 h1234567_1
libgomp 11.2.0 h1234567_1
libiconv 1.16 h5eee18b_3
libidn2 2.3.4 h5eee18b_0
libjpeg-turbo 2.0.0 h9bf148f_0 pytorch
liblief 0.12.3 h6a678d5_0
libmamba 1.5.8 hfe524e5_2
libmambapy 1.5.8 py311h2dafd23_2
libnghttp2 1.57.0 h2d74bed_0
libnpp 12.2.5.2 0 nvidia
libnvfatbin 12.5.82 0 nvidia
libnvjitlink 12.4.99 0 nvidia
libnvjpeg 12.3.1.89 0 nvidia
libpng 1.6.39 h5eee18b_0
libsolv 0.7.24 he621ea3_1
libssh2 1.11.0 h251f7ec_0
libstdcxx-ng 11.2.0 h1234567_1
libtasn1 4.19.0 h5eee18b_0
libtiff 4.5.1 h6a678d5_0
libunistring 0.9.10 h27cfd23_0
libuuid 1.41.5 h5eee18b_0
libuv 1.48.0 h5eee18b_0
libwebp-base 1.3.2 h5eee18b_0
libxml2 2.10.4 hfdd30dd_2
lintrunner 0.12.5 pypi_0 pypi
llvm-openmp 14.0.6 h9e868ea_0
lz4-c 1.9.4 h6a678d5_1
markdown 3.7 pypi_0 pypi
markdown-it-py 3.0.0 pypi_0 pypi
markupsafe 2.1.3 py311h5eee18b_0
matplotlib 3.9.2 pypi_0 pypi
matplotlib-inline 0.1.6 py311h06a4308_0
mdit-py-plugins 0.4.2 pypi_0 pypi
mdurl 0.1.2 pypi_0 pypi
menuinst 2.1.1 py311h06a4308_0
mistune 3.0.2 pypi_0 pypi
mkl 2023.1.0 h213fc3f_46344
mkl-service 2.4.0 py311h5eee18b_1
mkl_fft 1.3.8 py311h5eee18b_0
mkl_random 1.2.4 py311hdb19cb5_0
ml-dtypes 0.4.1 pypi_0 pypi
monai 1.4.0 pypi_0 pypi
more-itertools 10.1.0 py311h06a4308_0
mpc 1.1.0 h10f8cd9_1
mpfr 4.0.2 hb69a4c5_1
mpmath 1.3.0 py311h06a4308_0
msgpack 1.1.0 pypi_0 pypi
multidict 6.1.0 pypi_0 pypi
multiprocess 0.70.16 pypi_0 pypi
namex 0.0.8 pypi_0 pypi
nbclient 0.10.0 pypi_0 pypi
nbconvert 7.16.4 pypi_0 pypi
nbformat 5.10.4 pypi_0 pypi
ncurses 6.4 h6a678d5_0
nest-asyncio 1.6.0 pypi_0 pypi
nettle 3.7.3 hbbd107a_1
networkx 3.3 py311h06a4308_0
ninja 1.11.1.1 pypi_0 pypi
nodeenv 1.9.1 pypi_0 pypi
notebook-shim 0.2.4 pypi_0 pypi
numpy 1.26.4 py311h08b1b3b_0
numpy-base 1.26.4 py311hf175353_0
nvidia-cublas-cu12 12.4.5.8 pypi_0 pypi
nvidia-cuda-cupti-cu12 12.4.127 pypi_0 pypi
nvidia-cuda-nvcc-cu12 12.6.77 pypi_0 pypi
nvidia-cuda-nvrtc-cu12 12.4.127 pypi_0 pypi
nvidia-cuda-runtime-cu12 12.4.127 pypi_0 pypi
nvidia-cudnn-cu12 9.1.0.70 pypi_0 pypi
nvidia-cufft-cu12 11.2.1.3 pypi_0 pypi
nvidia-curand-cu12 10.3.5.147 pypi_0 pypi
nvidia-cusolver-cu12 11.6.1.9 pypi_0 pypi
nvidia-cusparse-cu12 12.3.1.170 pypi_0 pypi
nvidia-nccl-cu12 2.21.5 pypi_0 pypi
nvidia-nvjitlink-cu12 12.4.127 pypi_0 pypi
nvidia-nvtx-cu12 12.4.127 pypi_0 pypi
oauthlib 3.2.2 pypi_0 pypi
opencv-python-headless 4.10.0.84 pypi_0 pypi
openh264 2.1.1 h4ff587b_0
openjpeg 2.4.0 h9ca470c_2
openssl 3.0.14 h5eee18b_0
opt-einsum 3.4.0 pypi_0 pypi
optax 0.2.4 pypi_0 pypi
optree 0.12.1 pypi_0 pypi
orbax-checkpoint 0.9.1 pypi_0 pypi
overrides 7.7.0 pypi_0 pypi
packaging 24.1 py311h06a4308_0
pandas 2.2.3 pypi_0 pypi
pandocfilters 1.5.1 pypi_0 pypi
parso 0.8.3 pyhd3eb1b0_0
patch 2.7.6 h7b6447c_1001
patchelf 0.17.2 h6a678d5_0
pcre2 10.42 hebb0a14_1
pexpect 4.8.0 pyhd3eb1b0_3
pillow 11.0.0 pypi_0 pypi
pip 24.0 py311h06a4308_0
pkginfo 1.10.0 py311h06a4308_0
platformdirs 3.10.0 py311h06a4308_0
pluggy 1.0.0 py311h06a4308_1
pre-commit 4.0.1 pypi_0 pypi
prometheus-client 0.21.0 pypi_0 pypi
promise 2.3 pypi_0 pypi
prompt-toolkit 3.0.43 py311h06a4308_0
prompt_toolkit 3.0.43 hd3eb1b0_0
propcache 0.2.0 pypi_0 pypi
proto-plus 1.25.0 pypi_0 pypi
protobuf 5.28.2 pypi_0 pypi
psutil 5.9.0 py311h5eee18b_0
ptyprocess 0.7.0 pyhd3eb1b0_2
pure_eval 0.2.2 pyhd3eb1b0_0
py-lief 0.12.3 py311h6a678d5_0
pyarrow 18.0.0 pypi_0 pypi
pyasn1 0.6.1 pypi_0 pypi
pyasn1-modules 0.4.1 pypi_0 pypi
pybind11-abi 4 hd3eb1b0_1
pycosat 0.6.6 py311h5eee18b_1
pycparser 2.21 pyhd3eb1b0_0
pydantic 2.9.2 pypi_0 pypi
pydantic-core 2.23.4 pypi_0 pypi
pygments 2.15.1 py311h06a4308_1
pyparsing 3.2.0 pypi_0 pypi
pysocks 1.7.1 py311h06a4308_0
python 3.11.9 h955ad1f_0
python-dateutil 2.9.0.post0 pypi_0 pypi
python-etcd 0.4.5 pypi_0 pypi
python-json-logger 2.0.7 pypi_0 pypi
python-libarchive-c 2.9 pyhd3eb1b0_1
pytorch-cuda 12.4 hc786d27_6 pytorch
pytorch-mutex 1.0 cuda pytorch
pytz 2024.1 py311h06a4308_0
pyyaml 6.0.1 py311h5eee18b_0
pyzmq 26.2.0 pypi_0 pypi
readline 8.2 h5eee18b_0
referencing 0.30.2 py311h06a4308_0
regex 2024.11.6 pypi_0 pypi
reproc 14.2.4 h6a678d5_2
reproc-cpp 14.2.4 h6a678d5_2
requests 2.32.3 py311h06a4308_0
requests-oauthlib 2.0.0 pypi_0 pypi
rfc3339-validator 0.1.4 pypi_0 pypi
rfc3986-validator 0.1.1 pypi_0 pypi
rhash 1.4.3 hdbd6064_0
rich 13.9.2 pypi_0 pypi
rpds-py 0.10.6 py311hb02cf49_0
rsa 4.9 pypi_0 pypi
ruamel.yaml 0.17.21 py311h5eee18b_0
safetensors 0.4.5 pypi_0 pypi
scipy 1.14.1 pypi_0 pypi
seaborn 0.13.2 pypi_0 pypi
send2trash 1.8.3 pypi_0 pypi
setuptools 69.5.1 py311h06a4308_0
simple-parsing 0.1.6 pypi_0 pypi
simsimd 6.0.1 pypi_0 pypi
six 1.16.0 pyhd3eb1b0_1
sniffio 1.3.1 pypi_0 pypi
sortedcontainers 2.4.0 pypi_0 pypi
soupsieve 2.5 py311h06a4308_0
sqlite 3.45.3 h5eee18b_0
stack_data 0.2.0 pyhd3eb1b0_0
stringzilla 3.10.8 pypi_0 pypi
sympy 1.13.1 pypi_0 pypi
tb-nightly 2.19.0a20241106 pypi_0 pypi
tbb 2021.8.0 hdb19cb5_0
tensorboard 2.18.0 pypi_0 pypi
tensorboard-data-server 0.7.2 pypi_0 pypi
tensorflow 2.18.0 pypi_0 pypi
tensorflow-datasets 4.9.7 pypi_0 pypi
tensorflow-io-gcs-filesystem 0.37.1 pypi_0 pypi
tensorflow-metadata 1.16.1 pypi_0 pypi
tensorstore 0.1.66 pypi_0 pypi
termcolor 2.5.0 pypi_0 pypi
terminado 0.18.1 pypi_0 pypi
tiktoken 0.8.0 pypi_0 pypi
tinycss2 1.3.0 pypi_0 pypi
tk 8.6.14 h39e8969_0
tokenizers 0.20.3 pypi_0 pypi
toml 0.10.2 pypi_0 pypi
toolz 1.0.0 pypi_0 pypi
torch 2.5.1 pypi_0 pypi
torchaudio 2.4.0 py311_cu124 pytorch
torchelastic 0.2.2 pypi_0 pypi
torchvision 0.20.1 pypi_0 pypi
tornado 6.4.1 pypi_0 pypi
tqdm 4.66.4 py311h92b7b1e_0
traitlets 5.14.3 py311h06a4308_0
transformers 4.46.3 pypi_0 pypi
treescope 0.1.5 pypi_0 pypi
triton 3.1.0 pypi_0 pypi
truststore 0.8.0 py311h06a4308_0
typeguard 2.13.3 pypi_0 pypi
types-dataclasses 0.6.6 pypi_0 pypi
types-python-dateutil 2.9.0.20241003 pypi_0 pypi
typing_extensions 4.11.0 py311h06a4308_0
tzdata 2024.2 pypi_0 pypi
uri-template 1.3.0 pypi_0 pypi
urllib3 2.2.2 py311h06a4308_0
virtualenv 20.26.6 pypi_0 pypi
wcwidth 0.2.5 pyhd3eb1b0_0
webcolors 24.8.0 pypi_0 pypi
webencodings 0.5.1 pypi_0 pypi
websocket-client 1.8.0 pypi_0 pypi
werkzeug 3.1.2 pypi_0 pypi
wheel 0.43.0 py311h06a4308_0
widgetsnbextension 4.0.13 pypi_0 pypi
wrapt 1.16.0 pypi_0 pypi
xxhash 3.5.0 pypi_0 pypi
xz 5.4.6 h5eee18b_1
yaml 0.2.5 h7b6447c_0
yaml-cpp 0.8.0 h6a678d5_1
yarl 1.17.1 pypi_0 pypi
zipp 3.20.2 pypi_0 pypi
zlib 1.2.13 h5eee18b_1
zstandard 0.22.0 py311h2c38b39_0
zstd 1.5.5 hc292b87_2
conda env:
name: base
channels:
- pytorch
- nvidia
- defaults
dependencies:
- _libgcc_mutex=0.1=main
- _openmp_mutex=5.1=1_gnu
- anaconda-anon-usage=0.4.4=py311hfc0e8ea_100
- archspec=0.2.3=pyhd3eb1b0_0
- asttokens=2.0.5=pyhd3eb1b0_0
- attrs=23.1.0=py311h06a4308_0
- beautifulsoup4=4.12.3=py311h06a4308_0
- blas=1.0=mkl
- boltons=23.0.0=py311h06a4308_0
- brotli-python=1.0.9=py311h6a678d5_8
- bzip2=1.0.8=h5eee18b_6
- c-ares=1.19.1=h5eee18b_0
- ca-certificates=2024.7.2=h06a4308_0
- certifi=2024.7.4=py311h06a4308_0
- cffi=1.16.0=py311h5eee18b_1
- chardet=4.0.0=py311h06a4308_1003
- charset-normalizer=2.0.4=pyhd3eb1b0_0
- click=8.1.7=py311h06a4308_0
- cmake=3.26.4=h96355d8_0
- conda=24.5.0=py311h06a4308_0
- conda-build=24.5.1=py311h06a4308_0
- conda-content-trust=0.2.0=py311h06a4308_1
- conda-index=0.5.0=py311h06a4308_0
- conda-libmamba-solver=24.1.0=pyhd3eb1b0_0
- conda-package-handling=2.3.0=py311h06a4308_0
- conda-package-streaming=0.10.0=py311h06a4308_0
- cryptography=42.0.5=py311hdda0065_1
- cuda-cudart=12.4.127=0
- cuda-cupti=12.4.127=0
- cuda-libraries=12.4.0=0
- cuda-nvrtc=12.4.127=0
- cuda-nvtx=12.4.127=0
- cuda-opencl=12.5.39=0
- cuda-runtime=12.4.0=0
- cuda-version=12.5=3
- decorator=5.1.1=pyhd3eb1b0_0
- distro=1.9.0=py311h06a4308_0
- executing=0.8.3=pyhd3eb1b0_0
- expat=2.6.2=h6a678d5_0
- ffmpeg=4.3=hf484d3e_0
- filelock=3.13.1=py311h06a4308_0
- fmt=9.1.0=hdb19cb5_1
- freetype=2.12.1=h4a9f257_0
- frozendict=2.4.2=py311h06a4308_0
- gmp=6.2.1=h295c915_3
- gmpy2=2.1.2=py311hc9b5ff0_0
- gnutls=3.6.15=he1e5248_0
- icu=73.1=h6a678d5_0
- idna=3.7=py311h06a4308_0
- intel-openmp=2023.1.0=hdb19cb5_46306
- ipython=8.25.0=py311h06a4308_0
- jedi=0.19.1=py311h06a4308_0
- jinja2=3.1.4=py311h06a4308_0
- jpeg=9e=h5eee18b_2
- jsonpatch=1.33=py311h06a4308_1
- jsonpointer=2.1=pyhd3eb1b0_0
- jsonschema=4.19.2=py311h06a4308_0
- jsonschema-specifications=2023.7.1=py311h06a4308_0
- krb5=1.20.1=h143b758_1
- lame=3.100=h7b6447c_0
- lcms2=2.12=h3be6417_0
- ld_impl_linux-64=2.38=h1181459_1
- lerc=3.0=h295c915_0
- libarchive=3.6.2=h6ac8c49_3
- libcublas=12.4.2.65=0
- libcufft=11.2.0.44=0
- libcufile=1.10.1.7=0
- libcurand=10.3.6.82=0
- libcurl=8.7.1=h251f7ec_0
- libcusolver=11.6.0.99=0
- libcusparse=12.3.0.142=0
- libdeflate=1.17=h5eee18b_1
- libedit=3.1.20230828=h5eee18b_0
- libev=4.33=h7f8727e_1
- libffi=3.4.4=h6a678d5_1
- libgcc-ng=11.2.0=h1234567_1
- libgomp=11.2.0=h1234567_1
- libiconv=1.16=h5eee18b_3
- libidn2=2.3.4=h5eee18b_0
- libjpeg-turbo=2.0.0=h9bf148f_0
- liblief=0.12.3=h6a678d5_0
- libmamba=1.5.8=hfe524e5_2
- libmambapy=1.5.8=py311h2dafd23_2
- libnghttp2=1.57.0=h2d74bed_0
- libnpp=12.2.5.2=0
- libnvfatbin=12.5.82=0
- libnvjitlink=12.4.99=0
- libnvjpeg=12.3.1.89=0
- libpng=1.6.39=h5eee18b_0
- libsolv=0.7.24=he621ea3_1
- libssh2=1.11.0=h251f7ec_0
- libstdcxx-ng=11.2.0=h1234567_1
- libtasn1=4.19.0=h5eee18b_0
- libtiff=4.5.1=h6a678d5_0
- libunistring=0.9.10=h27cfd23_0
- libuuid=1.41.5=h5eee18b_0
- libuv=1.48.0=h5eee18b_0
- libwebp-base=1.3.2=h5eee18b_0
- libxml2=2.10.4=hfdd30dd_2
- llvm-openmp=14.0.6=h9e868ea_0
- lz4-c=1.9.4=h6a678d5_1
- markupsafe=2.1.3=py311h5eee18b_0
- matplotlib-inline=0.1.6=py311h06a4308_0
- menuinst=2.1.1=py311h06a4308_0
- mkl=2023.1.0=h213fc3f_46344
- mkl-service=2.4.0=py311h5eee18b_1
- mkl_fft=1.3.8=py311h5eee18b_0
- mkl_random=1.2.4=py311hdb19cb5_0
- more-itertools=10.1.0=py311h06a4308_0
- mpc=1.1.0=h10f8cd9_1
- mpfr=4.0.2=hb69a4c5_1
- mpmath=1.3.0=py311h06a4308_0
- ncurses=6.4=h6a678d5_0
- nettle=3.7.3=hbbd107a_1
- networkx=3.3=py311h06a4308_0
- numpy=1.26.4=py311h08b1b3b_0
- numpy-base=1.26.4=py311hf175353_0
- openh264=2.1.1=h4ff587b_0
- openjpeg=2.4.0=h9ca470c_2
- openssl=3.0.14=h5eee18b_0
- packaging=24.1=py311h06a4308_0
- parso=0.8.3=pyhd3eb1b0_0
- patch=2.7.6=h7b6447c_1001
- patchelf=0.17.2=h6a678d5_0
- pcre2=10.42=hebb0a14_1
- pexpect=4.8.0=pyhd3eb1b0_3
- pip=24.0=py311h06a4308_0
- pkginfo=1.10.0=py311h06a4308_0
- platformdirs=3.10.0=py311h06a4308_0
- pluggy=1.0.0=py311h06a4308_1
- prompt-toolkit=3.0.43=py311h06a4308_0
- prompt_toolkit=3.0.43=hd3eb1b0_0
- psutil=5.9.0=py311h5eee18b_0
- ptyprocess=0.7.0=pyhd3eb1b0_2
- pure_eval=0.2.2=pyhd3eb1b0_0
- py-lief=0.12.3=py311h6a678d5_0
- pybind11-abi=4=hd3eb1b0_1
- pycosat=0.6.6=py311h5eee18b_1
- pycparser=2.21=pyhd3eb1b0_0
- pygments=2.15.1=py311h06a4308_1
- pysocks=1.7.1=py311h06a4308_0
- python=3.11.9=h955ad1f_0
- python-libarchive-c=2.9=pyhd3eb1b0_1
- pytorch-cuda=12.4=hc786d27_6
- pytorch-mutex=1.0=cuda
- pytz=2024.1=py311h06a4308_0
- pyyaml=6.0.1=py311h5eee18b_0
- readline=8.2=h5eee18b_0
- referencing=0.30.2=py311h06a4308_0
- reproc=14.2.4=h6a678d5_2
- reproc-cpp=14.2.4=h6a678d5_2
- requests=2.32.3=py311h06a4308_0
- rhash=1.4.3=hdbd6064_0
- rpds-py=0.10.6=py311hb02cf49_0
- ruamel.yaml=0.17.21=py311h5eee18b_0
- setuptools=69.5.1=py311h06a4308_0
- six=1.16.0=pyhd3eb1b0_1
- soupsieve=2.5=py311h06a4308_0
- sqlite=3.45.3=h5eee18b_0
- stack_data=0.2.0=pyhd3eb1b0_0
- tbb=2021.8.0=hdb19cb5_0
- tk=8.6.14=h39e8969_0
- torchaudio=2.4.0=py311_cu124
- tqdm=4.66.4=py311h92b7b1e_0
- traitlets=5.14.3=py311h06a4308_0
- truststore=0.8.0=py311h06a4308_0
- typing_extensions=4.11.0=py311h06a4308_0
- urllib3=2.2.2=py311h06a4308_0
- wcwidth=0.2.5=pyhd3eb1b0_0
- wheel=0.43.0=py311h06a4308_0
- xz=5.4.6=h5eee18b_1
- yaml=0.2.5=h7b6447c_0
- yaml-cpp=0.8.0=h6a678d5_1
- zlib=1.2.13=h5eee18b_1
- zstandard=0.22.0=py311h2c38b39_0
- zstd=1.5.5=hc292b87_2
- pip:
- absl-py==2.1.0
- aiohappyeyeballs==2.4.3
- aiohttp==3.10.10
- aiosignal==1.3.1
- albucore==0.0.20
- albumentations==1.4.21
- annotated-types==0.7.0
- anyio==4.6.2.post1
- argon2-cffi==23.1.0
- argon2-cffi-bindings==21.2.0
- array-record==0.5.1
- arrow==1.3.0
- astunparse==1.6.3
- async-lru==2.0.4
- babel==2.16.0
- bleach==6.1.0
- cachetools==5.5.0
- cfgv==3.4.0
- chex==0.1.87
- cloudpickle==3.1.0
- comm==0.2.2
- contourpy==1.3.0
- cycler==0.12.1
- datasets==3.1.0
- debugpy==1.8.7
- defusedxml==0.7.1
- dill==0.3.8
- distlib==0.3.9
- dm-tree==0.1.8
- dnspython==2.6.1
- docstring-parser==0.16
- einops==0.8.0
- etils==1.9.4
- eval-type-backport==0.2.0
- expecttest==0.2.1
- fastjsonschema==2.20.0
- flatbuffers==24.3.25
- flax==0.10.1
- fonttools==4.54.1
- fqdn==1.5.1
- frozenlist==1.5.0
- fsspec==2024.9.0
- gast==0.6.0
- gcsfs==2024.10.0
- google-api-core==2.22.0
- google-auth==2.35.0
- google-auth-oauthlib==1.2.1
- google-cloud-core==2.4.1
- google-cloud-storage==2.18.2
- google-crc32c==1.6.0
- google-pasta==0.2.0
- google-resumable-media==2.7.2
- googleapis-common-protos==1.65.0
- grain==0.2.2
- grpcio==1.67.1
- h11==0.14.0
- h5py==3.12.1
- httpcore==1.0.6
- httpx==0.27.2
- huggingface-hub==0.26.2
- humanize==4.11.0
- hypothesis==6.108.4
- identify==2.6.1
- immutabledict==4.2.0
- importlib-resources==6.4.5
- ipykernel==6.29.5
- ipywidgets==8.1.5
- isoduration==20.11.0
- jax==0.4.34
- jax-cuda12-pjrt==0.4.34
- jax-cuda12-plugin==0.4.34
- jaxlib==0.4.34
- jaxtyping==0.2.34
- json5==0.9.25
- jupyter-client==8.6.3
- jupyter-core==5.7.2
- jupyter-events==0.10.0
- jupyter-lsp==2.2.5
- jupyter-server==2.14.2
- jupyter-server-terminals==0.5.3
- jupyterlab==4.3.1
- jupyterlab-pygments==0.3.0
- jupyterlab-server==2.27.3
- jupyterlab-widgets==3.0.13
- jupytext==1.16.4
- keras==3.6.0
- keras-nightly==3.7.0.dev2024110103
- kiwisolver==1.4.7
- libclang==18.1.1
- lintrunner==0.12.5
- markdown==3.7
- markdown-it-py==3.0.0
- matplotlib==3.9.2
- mdit-py-plugins==0.4.2
- mdurl==0.1.2
- mistune==3.0.2
- ml-dtypes==0.4.1
- monai==1.4.0
- msgpack==1.1.0
- multidict==6.1.0
- multiprocess==0.70.16
- namex==0.0.8
- nbclient==0.10.0
- nbconvert==7.16.4
- nbformat==5.10.4
- nest-asyncio==1.6.0
- ninja==1.11.1.1
- nodeenv==1.9.1
- notebook-shim==0.2.4
- nvidia-cublas-cu12==12.4.5.8
- nvidia-cuda-cupti-cu12==12.4.127
- nvidia-cuda-nvcc-cu12==12.6.77
- nvidia-cuda-nvrtc-cu12==12.4.127
- nvidia-cuda-runtime-cu12==12.4.127
- nvidia-cudnn-cu12==9.1.0.70
- nvidia-cufft-cu12==11.2.1.3
- nvidia-curand-cu12==10.3.5.147
- nvidia-cusolver-cu12==11.6.1.9
- nvidia-cusparse-cu12==12.3.1.170
- nvidia-nccl-cu12==2.21.5
- nvidia-nvjitlink-cu12==12.4.127
- nvidia-nvtx-cu12==12.4.127
- oauthlib==3.2.2
- opencv-python-headless==4.10.0.84
- opt-einsum==3.4.0
- optax==0.2.4
- optree==0.12.1
- orbax-checkpoint==0.9.1
- overrides==7.7.0
- pandas==2.2.3
- pandocfilters==1.5.1
- pillow==11.0.0
- pre-commit==4.0.1
- prometheus-client==0.21.0
- promise==2.3
- propcache==0.2.0
- proto-plus==1.25.0
- protobuf==5.28.2
- pyarrow==18.0.0
- pyasn1==0.6.1
- pyasn1-modules==0.4.1
- pydantic==2.9.2
- pydantic-core==2.23.4
- pyparsing==3.2.0
- python-dateutil==2.9.0.post0
- python-etcd==0.4.5
- python-json-logger==2.0.7
- pyzmq==26.2.0
- regex==2024.11.6
- requests-oauthlib==2.0.0
- rfc3339-validator==0.1.4
- rfc3986-validator==0.1.1
- rich==13.9.2
- rsa==4.9
- safetensors==0.4.5
- scipy==1.14.1
- seaborn==0.13.2
- send2trash==1.8.3
- simple-parsing==0.1.6
- simsimd==6.0.1
- sniffio==1.3.1
- sortedcontainers==2.4.0
- stringzilla==3.10.8
- sympy==1.13.1
- tb-nightly==2.19.0a20241106
- tensorboard==2.18.0
- tensorboard-data-server==0.7.2
- tensorflow==2.18.0
- tensorflow-datasets==4.9.7
- tensorflow-io-gcs-filesystem==0.37.1
- tensorflow-metadata==1.16.1
- tensorstore==0.1.66
- termcolor==2.5.0
- terminado==0.18.1
- tiktoken==0.8.0
- tinycss2==1.3.0
- tokenizers==0.20.3
- toml==0.10.2
- toolz==1.0.0
- torch==2.5.1
- torchelastic==0.2.2
- torchvision==0.20.1
- tornado==6.4.1
- transformers==4.46.3
- treescope==0.1.5
- triton==3.1.0
- typeguard==2.13.3
- types-dataclasses==0.6.6
- types-python-dateutil==2.9.0.20241003
- tzdata==2024.2
- uri-template==1.3.0
- virtualenv==20.26.6
- webcolors==24.8.0
- webencodings==0.5.1
- websocket-client==1.8.0
- werkzeug==3.1.2
- widgetsnbextension==4.0.13
- wrapt==1.16.0
- xxhash==3.5.0
- yarl==1.17.1
- zipp==3.20.2
prefix: /opt/conda
repro code in the single cell does not crash (at least I've never seen)
Just to be clear, here by cell you meant the notebook cell. Correct?
repro code in the single cell does not crash (at least I've never seen)
Just to be clear, here by
cellyou meant the notebook cell. Correct?
Yes, if I merge all notebook code cells into a single cell, I can't repro the crash. When the code is distributed into multiple cells as in the ipynb, than I can get notebook kernel crash rather frequently.
Hi @vfdev-5
I could not reproduce the issue when I tested it on JupyterLab locally on MAC OS and on Colab CPU with Orbax-Checkpoint 0.11.0, JAX 0.5.0 and Flax 0.10.2. Please find the attached screenshot of JupyterLab for reference.
Could please verify if issue still persists in your setup?
Thank you.
Hi @rajasekharporeddy , I retried on my setup and on few others and yes I could repro the crash. Here are detailed repro steps on a linux machine with docker:
- Run docker container using python:3.13-slim-bookworm image:
docker run -it --rm --name=jax-ai-test -v $PWD:/ws -w /ws --ipc=host --network=host python:3.13-slim-bookworm /bin/bash
- Install dependencies in the container:
pip install jupyterlab
pip install jax[cpu] flax orbax
- Start jupyterlab and open it in the browser
jupyter lab --port 8844 --allow-root --no-browser
- Run the following multiple times: restart jupyter kernel and execute cells. On the 2nd, 3rd tries it crashes with
Traceback (most recent call last):
File "<frozen runpy>", line 198, in _run_module_as_main
File "<frozen runpy>", line 88, in _run_code
File "/usr/local/lib/python3.13/site-packages/ipykernel_launcher.py", line 18, in <module>
app.launch_new_instance()
~~~~~~~~~~~~~~~~~~~~~~~^^
File "/usr/local/lib/python3.13/site-packages/traitlets/config/application.py", line 1075, in launch_instance
app.start()
~~~~~~~~~^^
File "/usr/local/lib/python3.13/site-packages/ipykernel/kernelapp.py", line 739, in start
self.io_loop.start()
~~~~~~~~~~~~~~~~~~^^
File "/usr/local/lib/python3.13/site-packages/tornado/platform/asyncio.py", line 205, in start
self.asyncio_loop.run_forever()
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^
File "/usr/local/lib/python3.13/asyncio/base_events.py", line 678, in run_forever
self._run_once()
~~~~~~~~~~~~~~^^
File "/usr/local/lib/python3.13/asyncio/base_events.py", line 2018, in _run_once
handle = self._ready.popleft()
IndexError: pop from an empty deque
[I 2025-01-29 08:47:59.223 ServerApp] AsyncIOLoopKernelRestarter: restarting kernel (1/5), keep random ports
[W 2025-01-29 08:47:59.223 ServerApp] kernel cb35e8f3-f9f9-42e5-8da9-139047ef2e73 restarted
Asking around, I was suggested to try using the alpha version of ipykernel which has a substantial rewrite: pip install -U --pre 'ipykernel>=v7.0.0a0'
Using that I can't reproduce the crash. I think we can close the issue in orbax as mostly unrelated to orbax code.
Same issue here, observe with orbax-checkpoint 0.11.9 and both ipykernel 6.30.0 as-well as ipykernel 7.0.0.a1. With the former, I get the same traceback as @vfdev-5 :
Traceback (most recent call last):
File "<frozen runpy>", line 198, in _run_module_as_main
File "<frozen runpy>", line 88, in _run_code
File "/Users/yves/.pyenv/versions/3.13.2/Library/Frameworks/Python.framework/Versions/3.13/lib/python3.13/site-packages/ipykernel_launcher.py", line 18, in <module>
app.launch_new_instance()
~~~~~~~~~~~~~~~~~~~~~~~^^
File "/Users/yves/.pyenv/versions/3.13.2/Library/Frameworks/Python.framework/Versions/3.13/lib/python3.13/site-packages/traitlets/config/application.py", line 1075, in launch_instance
app.start()
~~~~~~~~~^^
File "/Users/yves/.pyenv/versions/3.13.2/Library/Frameworks/Python.framework/Versions/3.13/lib/python3.13/site-packages/ipykernel/kernelapp.py", line 739, in start
self.io_loop.start()
~~~~~~~~~~~~~~~~~~^^
File "/Users/yves/.pyenv/versions/3.13.2/Library/Frameworks/Python.framework/Versions/3.13/lib/python3.13/site-packages/tornado/platform/asyncio.py", line 205, in start
self.asyncio_loop.run_forever()
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^
File "/Users/yves/.pyenv/versions/3.13.2/Library/Frameworks/Python.framework/Versions/3.13/lib/python3.13/asyncio/base_events.py", line 683, in run_forever
self._run_once()
~~~~~~~~~~~~~~^^
File "/Users/yves/.pyenv/versions/3.13.2/Library/Frameworks/Python.framework/Versions/3.13/lib/python3.13/asyncio/base_events.py", line 2025, in _run_once
handle = self._ready.popleft()
IndexError: pop from an empty deque
With ipykernel 7.0.0a1, I don't see any traceback on the Jupyter console, though it still crashes the kernel. The crash is definitely related to orbax - commenting out just the orbax stuff makes it go away.
I do have to say, the fact that I see asyncio in that traceback worries me a lot, and makes me questioning the design of orbax: Do you somehow mess with the asyncio event loop of Jupyter, despite no async visible on any API? That would be a big red flag for me.
asyncio is used a lot under the hood, but the final layer the user sees wraps everything in a future-like interface. We use nest_asyncio to allow it to run on colabs and Jupyter notebooks - this may be the likeliest spot to look for a culprit.
Hm. In this case, I would have preferred an clean, structured async API instead. There is a reason why event loops don't like to be nested. On the other hand, with structured concurrency, things like wait_until_finished wouldn't be needed (nor would they be allowed: structured concurrency forbids implicit lifetimes of async tasks). That's a real bummer that there is no way to use orbax without it stomping over the main event loop.
I just noticed that etils.epath.Path accessing a file on GCS crashes my notebook (with nothing else in there). The crash was fixable by setting ETILS_USE_TF=false in the environment, i.e. disabling their default choice of using a TF backend when they see TF is installed. Could this be related to this bug here?