Build datasets in notebooks without ipywidgets; fix datasets overview example

docs/install.rst

@@ -117,6 +117,20 @@ the shell before running:
 
     ulimit -n 8192
 
+**Jupyter notebooks (including VS Code and Cursor)**
+
+When a dataset builds its event cache inside a notebook, PyHealth shows a Dask
+progress bar. With ``ipywidgets`` installed it is an interactive widget; without
+it, PyHealth falls back to a plain text bar. ``ipywidgets`` is optional:
+
+.. code-block:: bash
+
+    pip install ipywidgets
+
+In PyHealth 2.0.2 and earlier, building a dataset in a notebook without
+``ipywidgets`` failed with ``ModuleNotFoundError: No module named 'ipywidgets'``.
+Installing ``ipywidgets`` works around it on those versions.
+
 **Other Platforms**
 
 PyHealth should work without additional configuration on:

examples/datasets_overview.ipynb

@@ -532,51 +532,26 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Label label vocab: {0: 0, 1: 1}\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Processing samples:   0%|          | 0/2 [00:00<?, ?it/s]\n"
-     ]
-    },
-    {
-     "ename": "TypeError",
-     "evalue": "unhashable type: 'list'",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
-      "\u001b[31mTypeError\u001b[39m                                 Traceback (most recent call last)",
-      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[60]\u001b[39m\u001b[32m, line 9\u001b[39m\n\u001b[32m      3\u001b[39m \u001b[38;5;66;03m# Create synthetic samples\u001b[39;00m\n\u001b[32m      4\u001b[39m samples = [\n\u001b[32m      5\u001b[39m     {\u001b[33m\"\u001b[39m\u001b[33mpatient_id\u001b[39m\u001b[33m\"\u001b[39m: \u001b[33m\"\u001b[39m\u001b[33m1\u001b[39m\u001b[33m\"\u001b[39m, \u001b[33m\"\u001b[39m\u001b[33mconditions\u001b[39m\u001b[33m\"\u001b[39m: [[\u001b[33m\"\u001b[39m\u001b[33mC001\u001b[39m\u001b[33m\"\u001b[39m, \u001b[33m\"\u001b[39m\u001b[33mC002\u001b[39m\u001b[33m\"\u001b[39m]], \u001b[33m\"\u001b[39m\u001b[33mlabel\u001b[39m\u001b[33m\"\u001b[39m: \u001b[32m0\u001b[39m},\n\u001b[32m      6\u001b[39m     {\u001b[33m\"\u001b[39m\u001b[33mpatient_id\u001b[39m\u001b[33m\"\u001b[39m: \u001b[33m\"\u001b[39m\u001b[33m2\u001b[39m\u001b[33m\"\u001b[39m, \u001b[33m\"\u001b[39m\u001b[33mconditions\u001b[39m\u001b[33m\"\u001b[39m: [[\u001b[33m\"\u001b[39m\u001b[33mC003\u001b[39m\u001b[33m\"\u001b[39m]], \u001b[33m\"\u001b[39m\u001b[33mlabel\u001b[39m\u001b[33m\"\u001b[39m: \u001b[32m1\u001b[39m}\n\u001b[32m      7\u001b[39m ]\n\u001b[32m----> \u001b[39m\u001b[32m9\u001b[39m dataset = \u001b[43mSampleDataset\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m     10\u001b[39m \u001b[43m    \u001b[49m\u001b[43msamples\u001b[49m\u001b[43m=\u001b[49m\u001b[43msamples\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     11\u001b[39m \u001b[43m    \u001b[49m\u001b[43minput_schema\u001b[49m\u001b[43m=\u001b[49m\u001b[43m{\u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mconditions\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43msequence\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m}\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     12\u001b[39m \u001b[43m    \u001b[49m\u001b[43moutput_schema\u001b[49m\u001b[43m=\u001b[49m\u001b[43m{\u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mlabel\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mbinary\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m}\u001b[49m\n\u001b[32m     13\u001b[39m \u001b[43m)\u001b[49m\n\u001b[32m     14\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mLoaded \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mlen\u001b[39m(dataset)\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m synthetic samples.\u001b[39m\u001b[33m\"\u001b[39m)\n",
-      "\u001b[36mFile \u001b[39m\u001b[32m~/PyHealth/pyhealth/datasets/sample_dataset.py:88\u001b[39m, in \u001b[36mSampleDataset.__init__\u001b[39m\u001b[34m(self, samples, input_schema, output_schema, dataset_name, task_name, input_processors, output_processors)\u001b[39m\n\u001b[32m     85\u001b[39m         \u001b[38;5;28mself\u001b[39m.record_to_index[record_id].append(i)\n\u001b[32m     87\u001b[39m \u001b[38;5;28mself\u001b[39m.validate()\n\u001b[32m---> \u001b[39m\u001b[32m88\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mbuild\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
-      "\u001b[36mFile \u001b[39m\u001b[32m~/PyHealth/pyhealth/datasets/sample_dataset.py:150\u001b[39m, in \u001b[36mSampleDataset.build\u001b[39m\u001b[34m(self)\u001b[39m\n\u001b[32m    148\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m k, v \u001b[38;5;129;01min\u001b[39;00m sample.items():\n\u001b[32m    149\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m.input_processors:\n\u001b[32m--> \u001b[39m\u001b[32m150\u001b[39m         sample[k] = \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43minput_processors\u001b[49m\u001b[43m[\u001b[49m\u001b[43mk\u001b[49m\u001b[43m]\u001b[49m\u001b[43m.\u001b[49m\u001b[43mprocess\u001b[49m\u001b[43m(\u001b[49m\u001b[43mv\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    151\u001b[39m     \u001b[38;5;28;01melif\u001b[39;00m k \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m.output_processors:\n\u001b[32m    152\u001b[39m         sample[k] = \u001b[38;5;28mself\u001b[39m.output_processors[k].process(v)\n",
-      "\u001b[36mFile \u001b[39m\u001b[32m~/PyHealth/pyhealth/processors/sequence_processor.py:38\u001b[39m, in \u001b[36mSequenceProcessor.process\u001b[39m\u001b[34m(self, value)\u001b[39m\n\u001b[32m     36\u001b[39m     indices.append(\u001b[38;5;28mself\u001b[39m.code_vocab[\u001b[33m\"\u001b[39m\u001b[33m<unk>\u001b[39m\u001b[33m\"\u001b[39m])\n\u001b[32m     37\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m---> \u001b[39m\u001b[32m38\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[43mtoken\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mnot\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mcode_vocab\u001b[49m:\n\u001b[32m     39\u001b[39m         \u001b[38;5;28mself\u001b[39m.code_vocab[token] = \u001b[38;5;28mself\u001b[39m._next_index\n\u001b[32m     40\u001b[39m         \u001b[38;5;28mself\u001b[39m._next_index += \u001b[32m1\u001b[39m\n",
-      "\u001b[31mTypeError\u001b[39m: unhashable type: 'list'"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "from pyhealth.datasets import SampleDataset\n",
+    "from pyhealth.datasets import create_sample_dataset\n",
     "\n",
     "# Create synthetic samples\n",
     "samples = [\n",
     "    {\"patient_id\": \"1\", \"conditions\": [\"C001\", \"C002\"], \"label\": 0},\n",
     "    {\"patient_id\": \"2\", \"conditions\": [\"C003\"], \"label\": 1}\n",
     "]\n",
     "\n",
-    "dataset = SampleDataset(\n",
+    "# SampleDataset is stored on disk; create_sample_dataset builds one from\n",
+    "# in-memory samples (fitting the processors named in the schemas).\n",
+    "dataset = create_sample_dataset(\n",
     "    samples=samples,\n",
     "    input_schema={\"conditions\": \"sequence\"},\n",
-    "    output_schema={\"label\": \"binary\"}\n",
+    "    output_schema={\"label\": \"binary\"},\n",
+    "    dataset_name=\"synthetic_demo\",\n",
     ")\n",
-    "print(f\"Loaded {len(dataset)} synthetic samples.\")"
+    "print(f\"Loaded {len(dataset)} synthetic samples.\")\n",
+    "print(dataset[0])"
    ]
   },
   {

pyhealth/datasets/base_dataset.py

@@ -14,6 +14,7 @@
 import multiprocessing
 import multiprocessing.queues
 import shutil
+import time
 
 from filelock import FileLock
 import litdata
@@ -150,6 +151,24 @@ def _litdata_merge(cache_dir: Path) -> None:
     )
 
 
+def _dask_progress_notebook() -> bool | None:
+    """
+    Chooses the ``notebook`` argument for dask's ``progress()``.
+
+    In a Jupyter kernel dask draws an ipywidgets progress bar, but ipywidgets is
+    not a PyHealth dependency. Without it, fall back to dask's text bar.
+
+    Returns:
+        bool | None: None to let dask detect the environment, or False to force
+            the text progress bar when ipywidgets is not installed.
+    """
+    try:
+        import ipywidgets  # noqa: F401
+    except ImportError:
+        return False
+    return None
+
+
 class _ProgressContext:
     def __init__(
         self, queue: multiprocessing.queues.Queue | None, total: int, **kwargs
@@ -595,9 +614,14 @@ def _event_transform(self, output_dir: Path) -> None:
                         compute=False,
                     )
                     handle = client.compute(collection)
-                    dask_progress(handle)
+                    dask_progress(handle, notebook=_dask_progress_notebook())
                     handle.result()  # type: ignore
                     compute_ok = True  # Data is fully written to disk
+                    if in_notebook():
+                        # With threaded workers, closing the cluster while the
+                        # scheduler's 100ms progress feed is asleep logs a spurious
+                        # CancelledError traceback. Let the feed see the bar close.
+                        time.sleep(0.3)
         except TimeoutError:
             if compute_ok:
                 # Cluster shutdown timed out after successful compute — data is intact

tests/core/test_base_dataset.py

@@ -1,4 +1,5 @@
 import os
+import sys
 import tempfile
 import unittest
 from unittest.mock import patch
@@ -160,6 +161,32 @@ def build(i: int) -> None:
             leaked, 4, f"{leaked} file descriptors leaked across 4 dataset builds"
         )
 
+    def test_event_df_builds_in_notebook_without_ipywidgets(self):
+        # In a Jupyter kernel dask's progress() draws an ipywidgets bar, but
+        # ipywidgets is not a PyHealth dependency. Simulate a kernel without it;
+        # in_notebook() also switches dask to threaded workers, as in Jupyter.
+        with tempfile.TemporaryDirectory() as cache_root, patch(
+            "pyhealth.datasets.base_dataset.platformdirs.user_cache_dir",
+            return_value=cache_root,
+        ), patch(
+            "pyhealth.datasets.base_dataset.in_notebook", return_value=True
+        ), patch(
+            "distributed.diagnostics.progressbar.is_kernel", return_value=True
+        ), patch.dict(sys.modules, {"ipywidgets": None}), self.assertNoLogs(
+            "distributed.scheduler", level="ERROR"
+        ):
+            dataset = MockDataset(
+                data=self._single_row_data(),
+                root="/data/root_no_ipywidgets",
+                tables=["table_a"],
+                dataset_name="NoIpywidgetsDataset",
+                dev=False,
+            )
+            _ = dataset.global_event_df
+            self.assertTrue(
+                (dataset.cache_dir / "global_event_df.parquet").exists()
+            )
+
     def test_empty_string_handling(self):
         import os
         from dataclasses import dataclass