Skip to content

Something wrong with the ProgressMeter and tqdm integration #962

@HarshBabla99

Description

@HarshBabla99

Not sure what exactly is causing the issue but here is a MWE.

import jax.numpy as jnp
import dynamiqs as dq
from equinox import filter_jit # for debugging

print(dq.__version__) # out = 0.3.2

####
c_dim = 20
a = dq.destroy(c_dim)
eye = dq.eye(c_dim)
zeros = dq.zeros(c_dim)

####
kappa_1 = 2
kappa_2 = 20*kappa_1    

alpha = 1.5
eng_diss = a@a - (alpha**2)*eye
c_ops = [jnp.sqrt(kappa_1) * a, jnp.sqrt(kappa_2) * eng_diss]

dq.set_progress_meter(True) # <-- set to False to see the issue go away!

H0 = zeros #no Hamiltonian
init = dq.coherent_dm(c_dim, alpha)
times = jnp.linspace(0.0, 3e3, 50)
result = dq.mesolve(H=H0, rho0=init, tsave=times, jump_ops=c_ops)

Here is the traceback. Observe that the simulation runs for a little bit before crashing.

|██████▍   |  64.8%elapsed 2.00sremaining 1.07s   ERROR:2025-05-08 18:21:47,006:jax._src.callback:102: jax.pure_callback failed
Traceback (most recent call last):
  File ".../lib/python3.13/site-packages/ipykernel/iostream.py", line 142, in _event_pipe
    event_pipe = self._local.event_pipe
                 ^^^^^^^^^^^^^^^^^^^^^^
AttributeError: '_thread._local' object has no attribute 'event_pipe'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File ".../lib/python3.13/site-packages/jax/_src/callback.py", line 100, in pure_callback_impl
    return tree_util.tree_map(np.asarray, callback(*args))
                                          ~~~~~~~~^^^^^^^
  File ".../lib/python3.13/site-packages/jax/_src/callback.py", line 77, in __call__
    return tree_util.tree_leaves(self.callback_func(*args, **kwargs))
                                 ~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File ".../lib/python3.13/site-packages/diffrax/_progress_meter.py", line 297, in _step
    step_bar(bar, np.array(_progress).reshape(()))
    ~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.13/site-packages/diffrax/_progress_meter.py", line 213, in _step_bar
    bar.refresh()
    ~~~~~~~~~~~^^
  File ".../lib/python3.13/site-packages/tqdm/std.py", line 1347, in refresh
    self.display()
    ~~~~~~~~~~~~^^
  File ".../lib/python3.13/site-packages/tqdm/std.py", line 1495, in display
    self.sp(self.__str__() if msg is None else msg)
    ~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.13/site-packages/tqdm/std.py", line 459, in print_status
    fp_write('\r' + s + (' ' * max(last_len[0] - len_s, 0)))
    ~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.13/site-packages/tqdm/std.py", line 452, in fp_write
    fp.write(str(s))
    ~~~~~~~~^^^^^^^^
  File ".../lib/python3.13/site-packages/tqdm/utils.py", line 196, in inner
    return func(*args, **kwargs)
  File ".../lib/python3.13/site-packages/IPython/core/interactiveshell.py", line 3027, in write
    result = original_write(data, *args, **kwargs)
  File ".../lib/python3.13/site-packages/ipykernel/iostream.py", line 694, in write
    self._schedule_flush()
    ~~~~~~~~~~~~~~~~~~~~^^
  File ".../lib/python3.13/site-packages/ipykernel/iostream.py", line 590, in _schedule_flush
    self.pub_thread.schedule(_schedule_in_thread)
    ~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.13/site-packages/ipykernel/iostream.py", line 267, in schedule
    self._event_pipe.send(b"")
    ^^^^^^^^^^^^^^^^
  File ".../lib/python3.13/site-packages/ipykernel/iostream.py", line 146, in _event_pipe
    event_pipe = ctx.socket(zmq.PUSH)
  File ".../lib/python3.13/site-packages/zmq/sugar/context.py", line 354, in socket
    socket_class(  # set PYTHONTRACEMALLOC=2 to get the calling frame
    ~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        self, socket_type, **kwargs
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File ".../lib/python3.13/site-packages/zmq/sugar/socket.py", line 156, in __init__
    super().__init__(
    ~~~~~~~~~~~~~~~~^
        ctx_or_socket,
        ^^^^^^^^^^^^^^
    ...<2 lines>...
        copy_threshold=copy_threshold,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "_zmq.py", line 690, in zmq.backend.cython._zmq.Socket.__init__
zmq.error.ZMQError: Too many open files
--- Logging error ---

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working🌀 solversTopic: solvers

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions