8GPU工作站同时运行4个Jupyter Notebook训练TF模型报错求助

我有一台用于TensorFlow深度学习的8GPU工作站，配置如下：

• Intel Xeon Gold x2
• NVIDIA Quadro RTXA6000 x8
• 1TB RAM

使用环境：

• Python 3.8
• UBUNTU 20.04
• Tensorflow 2.8
• CUDA 11.2
• cuDNN 8.1

我想同时运行4个Jupyter Notebook训练4个模型（认为单模型用8GPU效率不高），但训练过程中Notebook崩溃，报错信息如下：

[E 13:32:28.388 NotebookApp] Uncaught exception in ZMQStream callback
    Traceback (most recent call last):
      File "/home/super/researchvenv/lib/python3.8/site-packages/zmq/eventloop/zmqstream.py", line 584, in _run_callback
        f = callback(*args, **kwargs)
      File "/home/super/researchvenv/lib/python3.8/site-packages/zmq/eventloop/zmqstream.py", line 308, in stream_callback
        return callback(self, msg)
      File "/home/super/researchvenv/lib/python3.8/site-packages/notebook/services/kernels/handlers.py", line 572, in _on_zmq_reply
        super()._on_zmq_reply(stream, msg)
      File "/home/super/researchvenv/lib/python3.8/site-packages/notebook/base/zmqhandlers.py", line 256, in _on_zmq_reply
        self.write_message(msg, binary=isinstance(msg, bytes))
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/websocket.py", line 339, in write_message
        return self.ws_connection.write_message(message, binary=binary)
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/websocket.py", line 1086, in write_message
        fut = self._write_frame(True, opcode, message, flags=flags)
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/websocket.py", line 1061, in _write_frame
        return self.stream.write(frame)
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/iostream.py", line 546, in write
        self._handle_write()
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/iostream.py", line 976, in _handle_write
        self._write_buffer.advance(num_bytes)
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/iostream.py", line 182, in advance
        assert 0 < size <= self._size
    AssertionError

[E 13:32:28.388 NotebookApp] Uncaught exception in zmqstream callback
    Traceback (most recent call last):
      File "/home/super/researchvenv/lib/python3.8/site-packages/zmq/eventloop/zmqstream.py", line 621, in _handle_events
        self._handle_recv()
      File "/home/super/researchvenv/lib/python3.8/site-packages/zmq/eventloop/zmqstream.py", line 650, in _handle_recv
        self._run_callback(callback, msg)
      File "/home/super/researchvenv/lib/python3.8/site-packages/zmq/eventloop/zmqstream.py", line 584, in _run_callback
        f = callback(*args, **kwargs)
      File "/home/super/researchvenv/lib/python3.8/site-packages/zmq/eventloop/zmqstream.py", line 308, in stream_callback
        return callback(self, msg)
      File "/home/super/researchvenv/lib/python3.8/site-packages/notebook/services/kernels/handlers.py", line 572, in _on_zmq_reply
        super()._on_zmq_reply(stream, msg)
      File "/home/super/researchvenv/lib/python3.8/site-packages/notebook/base/zmqhandlers.py", line 256, in _on_zmq_reply
        self.write_message(msg, binary=isinstance(msg, bytes))
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/websocket.py", line 339, in write_message
        return self.ws_connection.write_message(message, binary=binary)
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/websocket.py", line 1086, in write_message
        fut = self._write_frame(True, opcode, message, flags=flags)
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/websocket.py", line 1061, in _write_frame
        return self.stream.write(frame)
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/iostream.py", line 546, in write
        self._handle_write()
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/iostream.py", line 976, in _handle_write
        self._write_buffer.advance(num_bytes)
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/iostream.py", line 182, in advance
        assert 0 < size <= self._size
    AssertionError
    
[E 13:32:28.389 NotebookApp] Exception in callback functools.partial(<function ZMQStream._update_handler.<locals>.<lambda> at 0x7f5da3064310>)
    Traceback (most recent call last):
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/ioloop.py", line 740, in _run_callback
        ret = callback()
      File "/home/super/researchvenv/lib/python3.8/site-packages/zmq/eventloop/zmqstream.py", line 705, in <lambda>
        self.io_loop.add_callback(lambda: self._handle_events(self.socket, 0))
      File "/home/super/researchvenv/lib/python3.8/site-packages/zmq/eventloop/zmqstream.py", line 621, in _handle_events
        self._handle_recv()
      File "/home/super/researchvenv/lib/python3.8/site-packages/zmq/eventloop/zmqstream.py", line 650, in _handle_recv
        self._run_callback(callback, msg)
      File "/home/super/researchvenv/lib/python3.8/site-packages/zmq/eventloop/zmqstream.py", line 584, in _run_callback
        f = callback(*args, **kwargs)
      File "/home/super/researchvenv/lib/python3.8/site-packages/zmq/eventloop/zmqstream.py", line 308, in stream_callback
        return callback(self, msg)
      File "/home/super/researchvenv/lib/python3.8/site-packages/notebook/services/kernels/handlers.py", line 572, in _on_zmq_reply
        super()._on_zmq_reply(stream, msg)
      File "/home/super/researchvenv/lib/python3.8/site-packages/notebook/base/zmqhandlers.py", line 256, in _on_zmq_reply
        self.write_message(msg, binary=isinstance(msg, bytes))
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/websocket.py", line 339, in write_message
        return self.ws_connection.write_message(message, binary=binary)
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/websocket.py", line 1086, in write_message
        fut = self._write_frame(True, opcode, message, flags=flags)
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/websocket.py", line 1061, in _write_frame
        return self.stream.write(frame)
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/iostream.py", line 546, in write
        self._handle_write()
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/iostream.py", line 976, in _handle_write
        self._write_buffer.advance(num_bytes)
      File "/home/super/researchvenv/lib/python3.8/site-packages/tornado/iostream.py", line 182, in advance
        assert 0 < size <= self._size
    AssertionError

请问这些错误是否与多Notebook运行有关？该如何解决？

错误关联性判断 #

这些报错属于Jupyter Notebook的通信层异常，根源是Tornado框架的IO缓冲区操作断言失败，和多Notebook运行本身没有直接因果关系，但多Notebook同时运行时，大量的日志输出、内核通信数据会加重通信层负担，从而触发这个Tornado旧版本的BUG。

具体解决步骤 #

1. 升级Tornado版本
   你当前使用的Tornado版本存在缓冲区处理的BUG，执行以下命令升级到修复该问题的稳定版本：

   pip install --upgrade tornado==6.1
   

   注：Tornado 6.1修复了多个IO缓冲区相关的断言错误，与Python 3.8、Notebook环境兼容性良好。

2. 限制Notebook的日志输出
   多模型训练时，每个Notebook会输出大量日志，加重通信层负担。可以在训练代码中调整日志级别：

   import logging
   logging.getLogger("tensorflow").setLevel(logging.WARNING)
   

   或者启动Jupyter时减少日志输出：

   jupyter notebook --log-level=WARN
   

3. 为每个Notebook分配独立GPU资源
   虽然报错不是GPU资源冲突导致，但为了避免训练过程中出现其他资源竞争问题，建议在每个Notebook开头指定专属GPU：

   import tensorflow as tf
   # 示例：第一个Notebook用GPU 0、1，第二个用2、3，以此类推
   gpus = tf.config.list_physical_devices('GPU')
   tf.config.set_visible_devices(gpus[0:2], 'GPU')
   

4. 重启Jupyter服务
   升级依赖后，务必重启Jupyter服务，确保新的Tornado版本生效：

   jupyter notebook stop
   jupyter notebook
   

内容的提问来源于stack exchange，提问作者Hyungjin Kim