Skip to content

Commit 7e959b9

Browse files
authored
Fix: Defer TableWidget import to prevent ZMQ port conflicts (#2309)
This PR addresses a flaky ZMQError: Address already in use that occasionally occurred during parallel notebook test execution. **Problem:** The bigframes.display module eagerly imported anywidget and traitlets at module load time (`bigframes/display/__init__.py`). This meant that when multiple Jupyter kernels were spun up simultaneously by nox for parallel testing, they would all try to initialize traitlets.HasTraits objects with sync=True properties. This led to race conditions and ZMQ port conflicts, causing notebook tests (including those that did not directly use anywidget like `streaming_dataframe.ipynb`) to fail. Log is [here](https://fusion2.corp.google.com/invocations/72088900-0196-4441-944b-ad68e491a8f8/targets/bigframes%2Fpresubmit%2Fnotebook/log). **Solution:** The TableWidget class import in `bigframes/display/__init__.py` has been refactored to use Python's `__getattr__` for lazy loading. This ensures that anywidget and traitlets are only imported and their associated kernel communication channels are initialized when display.TableWidget is actually accessed by the code. This prevents premature initialization and eliminates the port collision race condition during parallel test startup. Fixes #<465768150> 🦕
1 parent 719b278 commit 7e959b9

File tree

1 file changed

+28
-6
lines changed

1 file changed

+28
-6
lines changed

bigframes/display/__init__.py

Lines changed: 28 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,33 @@
1616

1717
from __future__ import annotations
1818

19-
try:
20-
import anywidget # noqa
19+
from typing import Any
2120

22-
from bigframes.display.anywidget import TableWidget
2321

24-
__all__ = ["TableWidget"]
25-
except Exception:
26-
pass
22+
def __getattr__(name: str) -> Any:
23+
"""Lazily import TableWidget to avoid ZMQ port conflicts.
24+
25+
anywidget and traitlets eagerly initialize kernel communication channels on
26+
import. This can lead to race conditions and ZMQ port conflicts when
27+
multiple Jupyter kernels are started in parallel, such as during notebook
28+
tests. By using __getattr__, we defer the import of TableWidget until it is
29+
explicitly accessed, preventing premature initialization and avoiding port
30+
collisions.
31+
"""
32+
if name == "TableWidget":
33+
try:
34+
import anywidget # noqa
35+
36+
from bigframes.display.anywidget import TableWidget
37+
38+
return TableWidget
39+
except Exception:
40+
raise AttributeError(
41+
f"module '{__name__}' has no attribute '{name}'. "
42+
"TableWidget requires anywidget and traitlets to be installed. "
43+
"Please `pip install anywidget traitlets` or `pip install 'bigframes[anywidget]'`."
44+
)
45+
raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
46+
47+
48+
__all__ = ["TableWidget"]

0 commit comments

Comments
 (0)