Segmentation fault on import trax

[tobigue]

Description

I'm getting a segmentation fault when trying to import trax.
Has anyone encountered the same problem?

Environment information

OS: Ubuntu 18.04.3 LTS
Docker image: tensorflow/tensorflow:2.1.0-gpu-py3

$ pip freeze | grep tensor
mesh-tensorflow==0.1.12
tensor2tensor==1.15.4
tensorboard==2.1.1
tensorflow-datasets==2.1.0
tensorflow-estimator==2.1.0
tensorflow-gan==2.0.0
tensorflow-gpu==2.1.0
tensorflow-hub==0.7.0
tensorflow-metadata==0.21.1
tensorflow-probability==0.7.0

$ pip freeze | grep jax
jax==0.1.59
jaxlib==0.1.40

$ pip freeze | grep trax
trax==1.2.3

$ pip freeze | grep matplotlib
matplotlib==3.2.0
# also tried matplotlib==2.2.5 with same results

$ python -V
Python 3.6.9

$ lshw -C display
  *-display                 
       description: VGA compatible controller
       product: GP104 [GeForce GTX 1080]
       vendor: NVIDIA Corporation
       physical id: 0
       bus info: pci@0000:01:00.0
       version: a1
       width: 64 bits
       clock: 33MHz
       capabilities: vga_controller bus_master cap_list rom
       configuration: driver=nvidia latency=0
       resources: irq:126 memory:ee000000-eeffffff memory:d0000000-dfffffff memory:e0000000-e1ffffff ioport:e000(size=128) memory:ef000000-ef07ffff

For bugs: reproduction and error logs

# Steps to reproduce:
Python 3.6.9 (default, Nov  7 2019, 10:44:02) 
[GCC 8.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorflow as tf
>>> tf.__version__
'2.1.0'
>>> print(tf.config.list_physical_devices('GPU'))
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
>>> import trax
Segmentation fault (core dumped)

Python 3.6.9 (default, Nov  7 2019, 10:44:02)                     
[GCC 8.3.0] on linux                                                       
Type "help", "copyright", "credits" or "license" for more information.
>>> import faulthandler                                                      
>>> faulthandler.enable()                                            
>>> import trax                                                                             
Fatal Python error: Segmentation fault                                       
                                                                        
Thread 0x00007f77d6bb5740 (most recent call first):               
  File "/tmp/env/lib/python3.6/site-packages/matplotlib/font_manager.py", line 1007 in addfont
  File "/tmp/env/lib/python3.6/site-packages/matplotlib/font_manager.py", line 991 in __init__
  File "/tmp/env/lib/python3.6/site-packages/matplotlib/font_manager.py", line 1334 in _rebuild
  File "/tmp/env/lib/python3.6/site-packages/matplotlib/font_manager.py", line 1343 in <module>
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed                 
  File "<frozen importlib._bootstrap_external>", line 678 in exec_module     
  File "<frozen importlib._bootstrap>", line 665 in _load_unlocked      
  File "<frozen importlib._bootstrap>", line 955 in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 971 in _find_and_load         
  File "/tmp/env/lib/python3.6/site-packages/matplotlib/contour.py", line 16 in <module>
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed                      
  File "<frozen importlib._bootstrap_external>", line 678 in exec_module     
  File "<frozen importlib._bootstrap>", line 665 in _load_unlocked      
  File "<frozen importlib._bootstrap>", line 955 in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 971 in _find_and_load         
  File "/tmp/env/lib/python3.6/site-packages/matplotlib/colorbar.py", line 31 in <module>
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed
  File "<frozen importlib._bootstrap_external>", line 678 in exec_module
  File "<frozen importlib._bootstrap>", line 665 in _load_unlocked                    
  File "<frozen importlib._bootstrap>", line 955 in _find_and_load_unlocked  
  File "<frozen importlib._bootstrap>", line 971 in _find_and_load      
  File "/tmp/env/lib/python3.6/site-packages/matplotlib/pyplot.py", line 32 in <module>
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed
  File "<frozen importlib._bootstrap_external>", line 678 in exec_module
  File "<frozen importlib._bootstrap>", line 665 in _load_unlocked                       
  File "<frozen importlib._bootstrap>", line 955 in _find_and_load_unlocked  
  File "<frozen importlib._bootstrap>", line 971 in _find_and_load      
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed
  File "<frozen importlib._bootstrap>", line 1023 in _handle_fromlist      
  File "/tmp/env/lib/python3.6/site-packages/matplotlib/__init__.py", line 1258 in use
  File "/tmp/env/lib/python3.6/site-packages/matplotlib/cbook/deprecation.py", line 358 in wrapper
  File "/tmp/env/lib/python3.6/site-packages/matplotlib/cbook/deprecation.py", line 296 in wrapper
  File "/tmp/env/lib/python3.6/site-packages/tensor2tensor/data_generators/video_generated.py", line 35 in <module>
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed
  File "<frozen importlib._bootstrap_external>", line 678 in exec_module
  File "<frozen importlib._bootstrap>", line 665 in _load_unlocked
  File "<frozen importlib._bootstrap>", line 955 in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 971 in _find_and_load
  File "<frozen importlib._bootstrap>", line 994 in _gcd_import
  File "/usr/lib/python3.6/importlib/__init__.py", line 126 in import_module
  File "/tmp/env/lib/python3.6/site-packages/tensor2tensor/data_generators/all_problems.py", line 140 in import_modules
  File "/tmp/env/lib/python3.6/site-packages/tensor2tensor/problems_colab.py", line 36 in <module>
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed
  File "<frozen importlib._bootstrap_external>", line 678 in exec_module
  File "<frozen importlib._bootstrap>", line 665 in _load_unlocked
  File "<frozen importlib._bootstrap>", line 955 in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 971 in _find_and_load
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed
  File "<frozen importlib._bootstrap>", line 1023 in _handle_fromlist
  File "/tmp/env/lib/python3.6/site-packages/trax/supervised/inputs.py", line 31 in <module>
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed
  File "<frozen importlib._bootstrap_external>", line 678 in exec_module
  File "<frozen importlib._bootstrap>", line 665 in _load_unlocked
  File "<frozen importlib._bootstrap>", line 955 in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 971 in _find_and_load
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed
  File "<frozen importlib._bootstrap>", line 1023 in _handle_fromlist
  File "/tmp/env/lib/python3.6/site-packages/trax/supervised/__init__.py", line 18 in <module>
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed
  File "<frozen importlib._bootstrap_external>", line 678 in exec_module
  File "<frozen importlib._bootstrap>", line 665 in _load_unlocked
  File "<frozen importlib._bootstrap>", line 955 in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 971 in _find_and_load
  File "/tmp/env/lib/python3.6/site-packages/trax/rl/simulated_env_problem.py", line 35 in <module>
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed
  File "<frozen importlib._bootstrap_external>", line 678 in exec_module
  File "<frozen importlib._bootstrap>", line 665 in _load_unlocked
  File "<frozen importlib._bootstrap>", line 955 in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 971 in _find_and_load
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed
  File "<frozen importlib._bootstrap>", line 1023 in _handle_fromlist
  File "/tmp/env/lib/python3.6/site-packages/trax/rl/__init__.py", line 24 in <module>
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed
  File "<frozen importlib._bootstrap_external>", line 678 in exec_module
  File "<frozen importlib._bootstrap>", line 665 in _load_unlocked
  File "<frozen importlib._bootstrap>", line 955 in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 971 in _find_and_load
  File "/tmp/env/lib/python3.6/site-packages/trax/learning_rate.py", line 294 in <module>
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed
  File "<frozen importlib._bootstrap_external>", line 678 in exec_module
  File "<frozen importlib._bootstrap>", line 665 in _load_unlocked
  File "<frozen importlib._bootstrap>", line 955 in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 971 in _find_and_load
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed
  File "<frozen importlib._bootstrap>", line 1023 in _handle_fromlist
  File "/tmp/env/lib/python3.6/site-packages/trax/__init__.py", line 19 in <module>
  File "<frozen importlib._bootstrap>", line 219 in _call_with_frames_removed
  File "<frozen importlib._bootstrap_external>", line 678 in exec_module
  File "<frozen importlib._bootstrap>", line 665 in _load_unlocked
  File "<frozen importlib._bootstrap>", line 955 in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 971 in _find_and_load
  File "<stdin>", line 1 in <module>
Segmentation fault (core dumped)

[lukaszkaiser]

Sorry, this looks more like a python bug than anything we can handle.

[devjeetr]

I had the exact same issue on a cloud VM. It turns out that this was because matplotlib uses XWindows display backend by default, which is not present on non-gui OSes (not entirely sure about this just a guess). You can override this behavior and set it to use another backend that doesn't need XWindows. Refer to this stackoverflow thread.