[solarsys]

Problem

When a dataset builds its event cache in a Jupyter kernel, BaseDataset calls dask's progress(). In a kernel, dask draws an ipywidgets progress bar, but ipywidgets is not a PyHealth dependency. In a notebook without it, such as a plain venv kernel in VS Code or Cursor, building any dataset fails:

ModuleNotFoundError: No module named 'ipywidgets'
Error during caching, removing incomplete file …/global_event_df.parquet

The same unguarded call is in every 2.x release, so 2.0.0 to 2.1.0 are affected. I checked the release tags: git show <tag>:pyhealth/datasets/base_dataset.py has a bare dask_progress(handle) at line 500 in v2.0.0, 524 in v2.0.1, 595 in v2.0.2 and 712 in v2.1.0. PyHealth 1.x has no base_dataset.py and doesn't use dask.

Changes

• base_dataset.py:
  • When ipywidgets isn't importable, pass notebook=False to dask's progress(), so dask falls back to its text bar. With ipywidgets installed, behaviour is unchanged.
  • In notebooks, where dask uses threaded workers, wait three progress intervals (0.3 s) after the compute before closing the cluster. In distributed 2026.8, the scheduler's progress feed (Scheduler.feed) sleeps one interval between updates and stops only when its next write finds the bar's connection closed. Closing the threaded cluster during that sleep cancels the feed, and distributed logs a spurious CancelledError traceback, even though the data was written correctly. distributed has no public way to wait for the feed, so this is a timing workaround. It only runs in notebooks: console builds, with worker processes, didn't hit the race (0 of 6 runs).
• tests/core/test_base_dataset.py: simulates a kernel without ipywidgets, as in Jupyter (threaded workers, kernel detection on). It asserts that the build succeeds and that the dask scheduler logs no errors. The test fails on master with the ModuleNotFoundError above and passes with this change.
• examples/datasets_overview.ipynb: the synthetic-samples cell called SampleDataset(samples=…). That constructor no longer exists; the committed output already showed an error, and running it today raises TypeError: … missing … 'path'. It now uses create_sample_dataset(...), and its stale error output is cleared. The new cell runs and prints Loaded 2 synthetic samples. No other cell changed.
• docs/install.rst: a note on notebooks and the optional ipywidgets, under Platform-Specific Notes, naming the affected versions (2.0.0 to 2.1.0, from the tags above).

Verification

• tools/check_pr_rules.py: all rules pass.
• The cluster-close race, with distributed 2026.8.0: without the wait, the notebook build logged the CancelledError in 8 of 8 runs, and the new test failed 3 of 3. With the wait: 0 of 8, and the test passed 5 of 5.
• Full core suite: 1,382 tests. Everything passes except the 18 tests in test_medcode_icd_mappings (medcode: ICD-9 <-> ICD-10 translation and grouper vocabularies via icd-mappings #1257), which fail locally only because icd-mappings wasn't installed in my older environment. With it installed, those tests and test_base_dataset pass (47/47).

🤖 Generated with Claude Code

[jhnwu3]

/review

[github-actions]

Previous review

AI review: findings

The PR fixes dataset builds in notebook kernels that lack ipywidgets: base_dataset.py passes notebook=False to dask's progress() when ipywidgets is not importable, and in notebooks adds a 0.3 s sleep before cluster shutdown to avoid a spurious CancelledError log. It also replaces a stale SampleDataset(...) call in examples/datasets_overview.ipynb with create_sample_dataset(...), which exists with matching arguments, and adds a docs note in docs/install.rst. No existing helper in the repo duplicates the dask progress selection, so the change is not redundant. Tests: a new notebook-path test in tests/core/test_base_dataset.py simulates a kernel without ipywidgets and checks the build succeeds with no scheduler ERROR logs; the PR says it fails on master and passes here. No blocking problems found. Non-blocking: the 0.3 s sleep is a timing workaround that no test isolates, so nothing verifies it is what prevents the logged traceback. Also non-blocking: the docs claim about "2.0.2 and earlier" is broader than the PR's evidence, which only covers 2.0.2. Overall risk is low: behavior changes only in notebooks, and with ipywidgets installed nothing changes.

_{Reviewed commit ce54889}

[github-actions]

Previous review

AI Review

suggestion: human review
decision: not available

Haiku found that this change needs a human review: The change edits BaseDataset._event_transform directly, the dataset-build path every BaseDataset subclass runs through. It alters behavior for notebook users (progress-bar selection and a post-compute sleep), so a human should confirm the change is safe for the shared build path.

_{Reviewed commit ce54889}

[github-actions]

Test Coverage

Total: 69.06% (17732 / 25678 lines)

Modified file	Coverage	Lines
pyhealth/autoimprove/__init__.py	100.00%	6 / 6
pyhealth/autoimprove/evaluate.py	92.07%	151 / 164
pyhealth/autoimprove/improver.py	95.63%	372 / 389
pyhealth/autoimprove/llm.py	90.62%	29 / 32
pyhealth/autoimprove/proposers.py	93.91%	108 / 115
pyhealth/autoimprove/space.py	95.31%	122 / 128
pyhealth/datasets/base_dataset.py	89.92%	464 / 516
pyhealth/datasets/omop.py	100.00%	18 / 18
pyhealth/metrics/calibration.py	65.00%	52 / 80
pyhealth/metrics/multiclass.py	25.00%	26 / 104
pyhealth/models/__init__.py	100.00%	52 / 52
pyhealth/models/survival_mdn.py	93.83%	213 / 227
pyhealth/processors/deep_nested_sequence_processor.py	91.37%	180 / 197
pyhealth/processors/nested_sequence_processor.py	92.94%	158 / 170
pyhealth/processors/sequence_processor.py	95.77%	68 / 71
pyhealth/processors/stagenet_processor.py	84.08%	169 / 201

Workflow run · 3ed7c98

[solarsys]

Thanks for the review. Both points are addressed in 5e2f75e. test_base_dataset.py passes (7 tests) with the CI command, and the PR rules check passes.

The 0.3 s sleep after a notebook build

• Cause. I traced it in distributed 2026.8.0:
  • Scheduler.feed sleeps one interval (100 ms) between progress updates. It stops only when its next write finds that the progress bar's connection is closed.
  • With threaded workers (processes=False in a notebook), the cluster can close during that sleep. That cancels the feed, and log_errors logs the CancelledError as an error on distributed.scheduler.
• No deterministic hook. distributed has no public way to wait for the feed to stop. Which connection runs which operation is held only in the private Server._comms. incoming_comms_active also counts the client's own long-lived connection.
• Change. I kept the wait and tied it to the interval. progress() is now given interval=_DASK_PROGRESS_INTERVAL (0.1 s, dask's default), and the wait is 3 * _DASK_PROGRESS_INTERVAL. The comment names the distributed version and this cause.
• Measured:
  • Notebook mode without the wait: the traceback appeared in 8 of 8 builds. With the wait: 0 of 8.
  • Console mode (worker processes, no wait): 0 of 6. So the wait stays notebook-only, and console builds don't pay for it.
• The test. assertNoLogs("distributed.scheduler", ...) watches the right logger, because the traceback is logged there. With the wait patched out, the test failed 3 of 3 times. With it, the test passed 5 of 5. Its comment now says this.

Docs version claim

The unguarded progress() call is in every 2.x release: 2.0.0, 2.0.1, 2.0.2 and 2.1.0. PyHealth 1.x doesn't use dask. The note in docs/install.rst now says "In PyHealth 2.0.0 to 2.1.0".

[solarsys]

/review

[github-actions]

Previous review

AI review: findings

The PR makes BaseDataset._event_transform (pyhealth/datasets/base_dataset.py), the event-cache build every dataset uses, work in notebooks without ipywidgets. When ipywidgets is not importable it passes notebook=False to dask's progress(), and in notebooks it sleeps 0.3 s after compute to avoid a distributed CancelledError at cluster close. It also replaces the stale SampleDataset(samples=...) call in examples/datasets_overview.ipynb with create_sample_dataset, and adds an install note. The new regression test simulates a kernel without ipywidgets and asserts the build succeeds with no distributed.scheduler ERROR logs; the PR reports it fails on master and passes with the change, though I could not run it. Nothing in the repo already does this fallback, so it is not redundant. Blocking problems: none. Non-blocking: the fixed time.sleep(0.3) in base_dataset.py is a timing workaround with no deterministic wait, and the regression test's no-error assertion depends on it. Also non-blocking: the notebook comment says SampleDataset is on disk while create_sample_dataset defaults to in_memory=True, and the install note's "2.0.0 to 2.1.0" range is not backed by anything in the PR. Overall risk is low; the change is small and confined to the notebook build path.

_{Reviewed commit 5e2f75e}

[github-actions]

Previous review

AI Review

suggestion: human review
decision: not available

Haiku found that this change needs a human review: The change edits BaseDataset._event_transform, the shared event-cache build that every dataset subclass relies on. It changes the dask progress call and adds a sleep on the notebook path, so it changes shared behavior rather than a self-contained component.

_{Reviewed commit 5e2f75e}

[solarsys]

Thanks for the second review. All three points are covered in 3ed7c98 and the updated PR description. test_base_dataset.py passes (7 tests) with the CI command, and the PR rules check passes.

Notebook comment (examples/datasets_overview.ipynb)

Fixed. The comment now says that create_sample_dataset fits the processors named in the schemas, returns an InMemorySampleDataset, and gives a disk-backed SampleDataset with in_memory=False. I checked both by running the cell's code: the default returns InMemorySampleDataset, and in_memory=False returns SampleDataset. Only that comment changed. No other cell or output was touched, and the notebook still validates with nbformat.

Affected versions (docs/install.rst)

The range comes from the release tags. The PR description now gives the evidence: git show <tag>:pyhealth/datasets/base_dataset.py has a bare dask_progress(handle) call in each 2.x release:

Tag	Line
v2.0.0	500
v2.0.1	524
v2.0.2	595
v2.1.0	712

PyHealth 1.x has no base_dataset.py and doesn't use dask. So "2.0.0 to 2.1.0" covers every release with the bug.

I kept the note in the install guide rather than moving it to CHANGELOG.md. The changelog is written at release time: only the release PRs edit it, and it has no "Unreleased" section. The fix will go into the next release's entry with the rest. The install note serves people on 2.0.0–2.1.0 who hit the ModuleNotFoundError today and need the workaround (pip install ipywidgets). Happy to drop the version sentence from it once a release with the fix is out.

The 0.3 s wait

It stays for now. I looked for a real synchronization point in distributed 2026.8.0 and didn't find one:

• How the feed stops. Scheduler.feed sleeps one interval between updates. It stops only when its next write fails because the bar's connection is closed, or when it sees the scheduler isn't running after waking up.
• Why that can't be awaited. Nothing tells the client when the feed has exited. The per-operation state of connections is held only in the private Server._comms.

So the wait is three progress intervals. It is passed to progress() explicitly, and it only runs in notebooks. Console builds with worker processes didn't hit the race in 6 of 6 runs.

One alternative is a logging filter on distributed.scheduler during cluster shutdown that drops this one CancelledError record. It doesn't depend on timing and costs nothing. But it hides a log record instead of avoiding the race, so I didn't make that change. It's easy to switch if you prefer it.

What would let the wait go: a distributed release in which the feed ends cleanly when the cluster closes. For example, Scheduler.feed could treat cancellation as a normal exit, or log_errors could stop logging CancelledError. The regression test is the check: it watches distributed.scheduler. With distributed 2026.8.0, it fails 3 of 3 times without the wait. On a fixed release it would pass with the wait removed.

[jhnwu3]

/review

[github-actions]

AI review: findings

The PR makes BaseDataset's event-cache build work in notebooks without ipywidgets. In pyhealth/datasets/base_dataset.py, _event_transform now passes notebook=False to dask's progress() when ipywidgets is not importable, and sleeps 0.3 s after the compute in notebooks so distributed's progress feed does not log a CancelledError when the threaded cluster closes. It also replaces the stale SampleDataset(samples=...) cell in examples/datasets_overview.ipynb with create_sample_dataset (its signature accepts dataset_name and in_memory), and adds a note to docs/install.rst. No public API changes; the change reuses litdata's in_notebook and there is no existing dask-progress helper, so nothing is redundant. Blocking problems: none found. Non-blocking (low): the 0.3 s sleep is a timing heuristic tied to distributed 2026.8 internals; it reduces the spurious log but does not guarantee its absence under load (base_dataset.py:750). Non-blocking (low): the new test's patch.dict(sys.modules, ...) restores sys.modules on exit and can drop modules first imported inside the block (tests/core/test_base_dataset.py:178). The new test mirrors the reported failure and follows the file's MockDataset pattern, but its no-ERROR-log assertion is timing-dependent, so it guards the race only probabilistically. I could not run the test suite (no execution tool), so the author's reported results are unverified. The repo is at version 2.1.0, consistent with the docs' affected-range note; I could not check the release tags. Overall risk is low.

_{Reviewed commit 3ed7c98}

[github-actions]

[low] Fixed 0.3 s sleep is a timing heuristic, not a guarantee

The wait time.sleep(3 * _DASK_PROGRESS_INTERVAL) only reduces the chance that distributed's progress feed is still sleeping when the threaded LocalCluster closes. If the scheduler is slow under load, the spurious CancelledError log can still appear, though the build still succeeds. The comment is accurate about the mechanism, and the author says distributed has no public way to wait, so this is an acceptable workaround. Consider re-checking it when distributed is upgraded, and keep the comment tied to the distributed version.

[github-actions]

[low] patch.dict(sys.modules) can drop modules imported inside the block

On exit, patch.dict(sys.modules, {...}) restores the whole dict, which removes any module first imported while the block runs (for example, lazily imported dask or distributed submodules during the LocalCluster build). Later tests may then re-import those modules as new objects, which can cause order-dependent failures. Fix by blocking only the ipywidgets entry, for example by setting sys.modules['ipywidgets'] = None in a try/finally that restores the previous value, or by patching builtins.import for that name.

[github-actions]

AI Review

suggestion: human review
decision: not available

Haiku found that this change needs a human review: The change modifies BaseDataset._event_transform in pyhealth/datasets/base_dataset.py, the build path every dataset uses. It changes how dask progress is drawn and adds a notebook-only sleep to the cache build, so the shared base-class behaviour changes for all datasets built in notebooks.

_{Reviewed commit 3ed7c98}

[jhnwu3]

lgtm