Use DataLoader for burden readouts

compute_burden_index_landmarks.py

@@ -10,6 +10,8 @@ from typing import Any, Iterable
 import numpy as np
 import pandas as pd
 import torch
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader, IterableDataset, get_worker_info
 from tqdm.auto import tqdm
 
 from burden_index import (
@@ -294,85 +296,31 @@ def _config_split_indices(
     return make_eval_indices(_Sized(), args, cfg)
 
 
-def _iter_readout_batches(
+class ReadoutJobIterableDataset(IterableDataset):
-    n: int,
+    def __init__(
-    batch_size: int,
+        self,
-) -> Iterable[slice]:
-    batch_size = max(1, int(batch_size))
-    for start in range(0, n, batch_size):
-        yield slice(start, min(start + batch_size, n))
-
-
-@torch.inference_mode()
-def _query_hidden_jobs(
-    *,
-    ctx: Any,
-    jobs: list[tuple[dict[str, Any], float]],
-) -> torch.Tensor:
-    if not jobs:
-        return torch.empty(0, ctx.model.n_embd, device=ctx.device)
-
-    batch_size = len(jobs)
-    max_event_len = max(int(np.asarray(row["event_seq"]).size) for row, _ in jobs)
-    max_other_len = max(int(np.asarray(row["other_type"]).size) for row, _ in jobs)
-
-    event = np.full((batch_size, max_event_len), PAD_IDX, dtype=np.int64)
-    time = np.zeros((batch_size, max_event_len), dtype=np.float32)
-    other_type = np.zeros((batch_size, max_other_len), dtype=np.int64)
-    other_value = np.zeros((batch_size, max_other_len), dtype=np.float32)
-    other_value_kind = np.zeros((batch_size, max_other_len), dtype=np.int64)
-    other_time = np.zeros((batch_size, max_other_len), dtype=np.float32)
-    sex = np.zeros(batch_size, dtype=np.int64)
-    query_times = np.zeros(batch_size, dtype=np.float32)
-
-    for i, (row, query_time) in enumerate(jobs):
-        event_seq = np.asarray(row["event_seq"], dtype=np.int64)
-        time_seq = np.asarray(row["time_seq"], dtype=np.float32)
-        other_type_seq = np.asarray(row["other_type"], dtype=np.int64)
-        other_value_seq = np.asarray(row["other_value"], dtype=np.float32)
-        other_value_kind_seq = np.asarray(row["other_value_kind"], dtype=np.int64)
-        other_time_seq = np.asarray(row["other_time"], dtype=np.float32)
-
-        event[i, : event_seq.size] = event_seq
-        time[i, : time_seq.size] = time_seq
-        other_type[i, : other_type_seq.size] = other_type_seq
-        other_value[i, : other_value_seq.size] = other_value_seq
-        other_value_kind[i, : other_value_kind_seq.size] = other_value_kind_seq
-        other_time[i, : other_time_seq.size] = other_time_seq
-        sex[i] = int(row["sex"])
-        query_times[i] = np.float32(query_time)
-
-    event_t = torch.from_numpy(event).long().to(ctx.device)
-    return ctx.model(
-        event_seq=event_t,
-        time_seq=torch.from_numpy(time).float().to(ctx.device),
-        sex=torch.from_numpy(sex).long().to(ctx.device),
-        padding_mask=event_t > PAD_IDX,
-        t_query=torch.from_numpy(query_times).float().to(ctx.device),
-        other_type=torch.from_numpy(other_type).long().to(ctx.device),
-        other_value=torch.from_numpy(other_value).float().to(ctx.device),
-        other_value_kind=torch.from_numpy(other_value_kind).long().to(ctx.device),
-        other_time=torch.from_numpy(other_time).float().to(ctx.device),
-        target_mode="all_future",
-    )
-
-
-def _build_readout_table(
         *,
         rows: list[dict[str, Any]],
         formed_mode: str,
         horizon: float,
-) -> dict[str, Any]:
+    ) -> None:
-    jobs: list[tuple[dict[str, Any], float]] = []
+        super().__init__()
-    row_indices: list[int] = []
+        self.rows = rows
-    kinds: list[str] = []
+        self.formed_mode = str(formed_mode)
-    deltas: list[float] = []
+        self.horizon = float(horizon)
+        if self.formed_mode not in {"observed", "model_weighted"}:
+            raise ValueError(f"Unknown formed_mode={self.formed_mode!r}")
 
-    if formed_mode not in {"observed", "model_weighted"}:
+    def __iter__(self) -> Iterable[dict[str, torch.Tensor]]:
-        raise ValueError(f"Unknown formed_mode={formed_mode!r}")
+        worker = get_worker_info()
+        if worker is None:
+            start, step = 0, 1
+        else:
+            start, step = int(worker.id), int(worker.num_workers)
 
-    for row_idx, row in enumerate(rows):
+        for row_idx in range(start, len(self.rows), step):
-        if formed_mode == "model_weighted":
+            row = self.rows[row_idx]
+            if self.formed_mode == "model_weighted":
                 grid = _build_readout_grid(
                     event_seq=row["event_seq"],
                     time_seq=row["time_seq"],
@@ -387,24 +335,96 @@ def _build_readout_table(
                     row_deltas = np.maximum(end_times - grid, 0.0).astype(np.float32)
                     valid = row_deltas > 0
                     for query_time, delta in zip(grid[valid].tolist(), row_deltas[valid].tolist()):
-                    jobs.append((row, float(query_time)))
+                        yield _make_readout_job(row, row_idx, "formed", query_time, delta)
-                    row_indices.append(row_idx)
+            if self.horizon > 0:
-                    kinds.append("formed")
+                yield _make_readout_job(
-                    deltas.append(float(delta))
+                    row,
-        if horizon > 0:
+                    row_idx,
-            jobs.append((row, float(row["t_query"])))
+                    "future",
-            row_indices.append(row_idx)
+                    float(row["t_query"]),
-            kinds.append("future")
+                    self.horizon,
-            deltas.append(float(horizon))
+                )
 
+
+def _make_readout_job(
+    row: dict[str, Any],
+    row_idx: int,
+    kind: str,
+    query_time: float,
+    delta: float,
+) -> dict[str, torch.Tensor]:
     return {
-        "jobs": jobs,
+        "event_seq": torch.from_numpy(np.asarray(row["event_seq"], dtype=np.int64)).long(),
-        "row_indices": np.asarray(row_indices, dtype=np.int64),
+        "time_seq": torch.from_numpy(np.asarray(row["time_seq"], dtype=np.float32)).float(),
-        "kinds": np.asarray(kinds, dtype=object),
+        "sex": torch.tensor(int(row["sex"]), dtype=torch.long),
-        "deltas": np.asarray(deltas, dtype=np.float32),
+        "other_type": torch.from_numpy(np.asarray(row["other_type"], dtype=np.int64)).long(),
+        "other_value": torch.from_numpy(np.asarray(row["other_value"], dtype=np.float32)).float(),
+        "other_value_kind": torch.from_numpy(
+            np.asarray(row["other_value_kind"], dtype=np.int64)
+        ).long(),
+        "other_time": torch.from_numpy(np.asarray(row["other_time"], dtype=np.float32)).float(),
+        "query_time": torch.tensor(float(query_time), dtype=torch.float32),
+        "delta": torch.tensor(float(delta), dtype=torch.float32),
+        "row_idx": torch.tensor(int(row_idx), dtype=torch.long),
+        "kind": torch.tensor(0 if kind == "formed" else 1, dtype=torch.long),
     }
 
 
+def _collate_readout_jobs(batch: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
+    event_seq = pad_sequence(
+        [x["event_seq"] for x in batch], batch_first=True, padding_value=PAD_IDX
+    )
+    time_seq = pad_sequence(
+        [x["time_seq"] for x in batch], batch_first=True, padding_value=0.0
+    )
+    other_type = pad_sequence(
+        [x["other_type"] for x in batch], batch_first=True, padding_value=0
+    )
+    other_value = pad_sequence(
+        [x["other_value"] for x in batch], batch_first=True, padding_value=0.0
+    )
+    other_value_kind = pad_sequence(
+        [x["other_value_kind"] for x in batch], batch_first=True, padding_value=0
+    )
+    other_time = pad_sequence(
+        [x["other_time"] for x in batch], batch_first=True, padding_value=0.0
+    )
+    return {
+        "event_seq": event_seq,
+        "time_seq": time_seq,
+        "padding_mask": event_seq > PAD_IDX,
+        "sex": torch.stack([x["sex"] for x in batch]),
+        "other_type": other_type,
+        "other_value": other_value,
+        "other_value_kind": other_value_kind,
+        "other_time": other_time,
+        "query_time": torch.stack([x["query_time"] for x in batch]),
+        "delta": torch.stack([x["delta"] for x in batch]),
+        "row_idx": torch.stack([x["row_idx"] for x in batch]),
+        "kind": torch.stack([x["kind"] for x in batch]),
+    }
+
+
+@torch.inference_mode()
+def _query_hidden_readout_batch(
+    *,
+    ctx: Any,
+    batch: dict[str, torch.Tensor],
+) -> torch.Tensor:
+    event = batch["event_seq"].long().to(ctx.device, non_blocking=True)
+    return ctx.model(
+        event_seq=event,
+        time_seq=batch["time_seq"].float().to(ctx.device, non_blocking=True),
+        sex=batch["sex"].long().to(ctx.device, non_blocking=True),
+        padding_mask=event > PAD_IDX,
+        t_query=batch["query_time"].float().to(ctx.device, non_blocking=True),
+        other_type=batch["other_type"].long().to(ctx.device, non_blocking=True),
+        other_value=batch["other_value"].float().to(ctx.device, non_blocking=True),
+        other_value_kind=batch["other_value_kind"].long().to(ctx.device, non_blocking=True),
+        other_time=batch["other_time"].float().to(ctx.device, non_blocking=True),
+        target_mode="all_future",
+    )
+
 @torch.inference_mode()
 def _probabilities_from_hidden_torch(
     *,
@@ -442,30 +462,21 @@ def _probabilities_from_hidden_torch(
 
     return -torch.expm1(-rate * exposure)
 
-
 @torch.inference_mode()
-def _readout_probabilities(
+def _readout_probabilities_from_batch(
     *,
     ctx: Any,
-    readout_table: dict[str, Any],
+    batch: dict[str, torch.Tensor],
     union_disease_ids: np.ndarray,
-    readout_batch_size: int,
 ) -> torch.Tensor:
-    jobs = readout_table["jobs"]
+    hidden = _query_hidden_readout_batch(ctx=ctx, batch=batch)
-    if not jobs:
+    deltas = batch["delta"].detach().cpu().numpy().astype(np.float32, copy=False)
-        return torch.empty((0, union_disease_ids.size), dtype=torch.float32, device=ctx.device)
+    return _probabilities_from_hidden_torch(
-
-    out = torch.empty((len(jobs), union_disease_ids.size), dtype=torch.float32, device=ctx.device)
-    deltas = np.asarray(readout_table["deltas"], dtype=np.float32)
-    for slc in _iter_readout_batches(len(jobs), readout_batch_size):
-        hidden = _query_hidden_jobs(ctx=ctx, jobs=jobs[slc])
-        out[slc] = _probabilities_from_hidden_torch(
         ctx=ctx,
         hidden=hidden,
         disease_ids=union_disease_ids,
-            deltas=deltas[slc],
+        deltas=deltas,
-        ).to(dtype=out.dtype)
+    ).to(dtype=torch.float32)
-    return out
 
 
 def _observed_formed_for_rows(
@@ -484,66 +495,22 @@ def _observed_formed_for_rows(
     return formed
 
 
-def _reduce_readout_table_to_bi_rows(
+def _apply_readout_batch_to_accumulators(
     *,
-    rows: list[dict[str, Any]],
+    batch: dict[str, torch.Tensor],
-    horizon: float,
-    matrices: list[dict[str, Any]],
-    union_disease_ids: np.ndarray,
-    formed_mode: str,
-    readout_table: dict[str, Any],
     readout_prob: torch.Tensor,
+    survival_by_row: torch.Tensor | None,
+    future_prob_by_row: torch.Tensor,
     ctx: Any,
-) -> list[dict[str, Any]]:
+) -> None:
-    if formed_mode == "observed":
+    if readout_prob.numel() == 0:
-        formed_by_row = torch.as_tensor(
+        return
-            _observed_formed_for_rows(
-                rows=rows,
-                union_disease_ids=union_disease_ids,
-            ),
-            dtype=readout_prob.dtype,
-            device=ctx.device,
-        )
-    elif formed_mode == "model_weighted":
-        formed_by_row = torch.zeros(
-            (len(rows), union_disease_ids.size),
-            dtype=readout_prob.dtype,
-            device=ctx.device,
-        )
-    else:
-        raise ValueError(f"Unknown formed_mode={formed_mode!r}")
 
-    future_prob_by_row = torch.zeros(
+    row_indices = batch["row_idx"].long().to(ctx.device, non_blocking=True)
-        (len(rows), union_disease_ids.size),
+    kinds = batch["kind"].long().to(ctx.device, non_blocking=True)
-        dtype=readout_prob.dtype,
+    kind_is_formed = kinds == 0
-        device=ctx.device,
+    kind_is_future = kinds == 1
-    )
-    row_indices = torch.as_tensor(
-        np.asarray(readout_table["row_indices"], dtype=np.int64),
-        dtype=torch.long,
-        device=ctx.device,
-    )
-    kinds = np.asarray(readout_table["kinds"], dtype=object)
-    if formed_mode == "model_weighted":
-        survival_by_row = torch.ones(
-            (len(rows), union_disease_ids.size),
-            dtype=readout_prob.dtype,
-            device=ctx.device,
-        )
-    else:
-        survival_by_row = None
 
-    if readout_prob.numel() > 0:
-        kind_is_formed = torch.as_tensor(
-            np.asarray(kinds == "formed", dtype=np.bool_),
-            dtype=torch.bool,
-            device=ctx.device,
-        )
-        kind_is_future = torch.as_tensor(
-            np.asarray(kinds == "future", dtype=np.bool_),
-            dtype=torch.bool,
-            device=ctx.device,
-        )
     if survival_by_row is not None and bool(kind_is_formed.any().item()):
         formed_rows = row_indices[kind_is_formed]
         formed_survival = 1.0 - readout_prob[kind_is_formed].clamp(0.0, 1.0)
@@ -564,14 +531,22 @@ def _reduce_readout_table_to_bi_rows(
         future_rows = row_indices[kind_is_future]
         future_prob_by_row[future_rows] = readout_prob[kind_is_future]
 
-    if survival_by_row is not None:
-        formed_by_row = 1.0 - survival_by_row
 
+def _project_bi_rows(
+    *,
+    rows: list[dict[str, Any]],
+    horizon: float,
+    matrices: list[dict[str, Any]],
+    formed_mode: str,
+    formed_by_row: torch.Tensor,
+    future_prob_by_row: torch.Tensor,
+    ctx: Any,
+) -> list[dict[str, Any]]:
     disease_future_by_row = (1.0 - formed_by_row) * future_prob_by_row
     disease_total_by_row = formed_by_row + disease_future_by_row
     projected: list[dict[str, Any]] = []
     for matrix in matrices:
-        A = torch.as_tensor(matrix["A_union"], dtype=readout_prob.dtype, device=ctx.device)
+        A = torch.as_tensor(matrix["A_union"], dtype=formed_by_row.dtype, device=ctx.device)
         projected.append(
             {
                 "matrix": matrix,
@@ -611,7 +586,7 @@ def _reduce_readout_table_to_bi_rows(
     return out
 
 
-def _compute_bi_from_readout_table(
+def _compute_bi_from_streamed_readouts(
     *,
     rows: list[dict[str, Any]],
     horizon: float,
@@ -619,46 +594,117 @@ def _compute_bi_from_readout_table(
     union_disease_ids: np.ndarray,
     formed_mode: str,
     readout_batch_size: int,
+    num_workers: int,
     ctx: Any,
+    log_prefix: str | None = None,
 ) -> tuple[list[dict[str, Any]], int, dict[str, float]]:
     horizon = float(horizon)
     if horizon < 0:
         raise ValueError(f"horizon must be non-negative, got {horizon}")
-    t0 = time.perf_counter()
+
-    readout_table = _build_readout_table(
+    dtype = torch.float32
+    if formed_mode == "observed":
+        formed_by_row = torch.as_tensor(
+            _observed_formed_for_rows(
+                rows=rows,
+                union_disease_ids=union_disease_ids,
+            ),
+            dtype=dtype,
+            device=ctx.device,
+        )
+        survival_by_row = None
+    elif formed_mode == "model_weighted":
+        survival_by_row = torch.ones(
+            (len(rows), union_disease_ids.size),
+            dtype=dtype,
+            device=ctx.device,
+        )
+        formed_by_row = None
+    else:
+        raise ValueError(f"Unknown formed_mode={formed_mode!r}")
+    future_prob_by_row = torch.zeros(
+        (len(rows), union_disease_ids.size),
+        dtype=dtype,
+        device=ctx.device,
+    )
+
+    readout_jobs = 0
+    n_batches = 0
+    build_readout_sec = 0.0
+    forward_sec = 0.0
+    reduce_sec = 0.0
+    t_loader0 = time.perf_counter()
+    readout_dataset = ReadoutJobIterableDataset(
         rows=rows,
         formed_mode=formed_mode,
         horizon=horizon,
     )
-    t1 = time.perf_counter()
+    readout_loader = DataLoader(
-    readout_prob = _readout_probabilities(
+        readout_dataset,
+        batch_size=max(1, int(readout_batch_size)),
+        collate_fn=_collate_readout_jobs,
+        num_workers=max(0, int(num_workers)),
+        pin_memory=ctx.device.type == "cuda",
+        persistent_workers=int(num_workers) > 0,
+        prefetch_factor=2 if int(num_workers) > 0 else None,
+    )
+    build_readout_sec += time.perf_counter() - t_loader0
+
+    for batch in readout_loader:
+        t_forward0 = time.perf_counter()
+        readout_prob = _readout_probabilities_from_batch(
             ctx=ctx,
-        readout_table=readout_table,
+            batch=batch,
             union_disease_ids=union_disease_ids,
-        readout_batch_size=readout_batch_size,
         )
         if ctx.device.type == "cuda":
             torch.cuda.synchronize(ctx.device)
-    t2 = time.perf_counter()
+        forward_sec += time.perf_counter() - t_forward0
-    rows_out = _reduce_readout_table_to_bi_rows(
+
+        t_reduce0 = time.perf_counter()
+        _apply_readout_batch_to_accumulators(
+            batch=batch,
+            readout_prob=readout_prob,
+            survival_by_row=survival_by_row,
+            future_prob_by_row=future_prob_by_row,
+            ctx=ctx,
+        )
+        if ctx.device.type == "cuda":
+            torch.cuda.synchronize(ctx.device)
+        reduce_sec += time.perf_counter() - t_reduce0
+
+        n_batches += 1
+        readout_jobs += int(batch["row_idx"].numel())
+        if log_prefix and (n_batches == 1 or n_batches % 50 == 0):
+            print(
+                f"{log_prefix} processed {readout_jobs} readout jobs "
+                f"in {n_batches} batches",
+                flush=True,
+            )
+
+    if survival_by_row is not None:
+        formed_by_row = 1.0 - survival_by_row
+    assert formed_by_row is not None
+
+    t_project0 = time.perf_counter()
+    rows_out = _project_bi_rows(
         rows=rows,
         horizon=horizon,
         matrices=matrices,
-        union_disease_ids=union_disease_ids,
         formed_mode=formed_mode,
-        readout_table=readout_table,
+        formed_by_row=formed_by_row,
-        readout_prob=readout_prob,
+        future_prob_by_row=future_prob_by_row,
         ctx=ctx,
     )
     if ctx.device.type == "cuda":
         torch.cuda.synchronize(ctx.device)
-    t3 = time.perf_counter()
+    reduce_sec += time.perf_counter() - t_project0
     timings = {
-        "build_readout_sec": t1 - t0,
+        "build_readout_sec": build_readout_sec,
-        "forward_sec": t2 - t1,
+        "forward_sec": forward_sec,
-        "reduce_sec": t3 - t2,
+        "reduce_sec": reduce_sec,
     }
-    return rows_out, len(readout_table["jobs"]), timings
+    return rows_out, readout_jobs, timings
 
 
 def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]:
@@ -681,14 +727,16 @@ def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]:
         f"{time.perf_counter() - materialize_start:.2f}s",
         flush=True,
     )
-    out, readout_jobs, timings = _compute_bi_from_readout_table(
+    out, readout_jobs, timings = _compute_bi_from_streamed_readouts(
         rows=rows,
         horizon=payload["horizon"],
         matrices=payload["matrices"],
         union_disease_ids=payload["union_disease_ids"],
         formed_mode=payload["formed_mode"],
         readout_batch_size=int(payload["readout_batch_size"]),
+        num_workers=int(payload["num_workers"]),
         ctx=ctx,
+        log_prefix=f"[BI worker {device}]",
     )
     print(
         f"[BI worker {device}] done: readout_jobs={readout_jobs}, "
@@ -831,6 +879,12 @@ def main() -> None:
             "Increase this to improve GPU utilization if memory allows."
         ),
     )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=4,
+        help="DataLoader workers per GPU process for readout job generation.",
+    )
     parser.add_argument("--device", type=str, default=None)
     parser.add_argument(
         "--devices",
@@ -933,14 +987,16 @@ def main() -> None:
             flush=True,
         )
     if len(row_chunks) == 1:
-        all_rows, total_readout_jobs, timings = _compute_bi_from_readout_table(
+        all_rows, total_readout_jobs, timings = _compute_bi_from_streamed_readouts(
             rows=rows,
             horizon=horizon,
             matrices=matrices,
             union_disease_ids=union_disease_ids,
             formed_mode=args.formed_mode,
             readout_batch_size=int(args.readout_batch_size),
+            num_workers=int(args.num_workers),
             ctx=ctx,
+            log_prefix="[BI main]",
         )
         for key, value in timings.items():
             total_timings[key] += float(value)
@@ -957,6 +1013,7 @@ def main() -> None:
                 "matrices": matrices,
                 "union_disease_ids": union_disease_ids,
                 "readout_batch_size": int(args.readout_batch_size),
+                "num_workers": int(args.num_workers),
                 "formed_mode": args.formed_mode,
             }
             for device, chunk_rows in row_chunks
@@ -988,6 +1045,7 @@ def main() -> None:
     print(f"Landmark rows: {len(rows)}")
     print(f"Readout jobs: {total_readout_jobs}")
     print(f"Readout batch size per worker: {int(args.readout_batch_size)}")
+    print(f"DataLoader workers per GPU process: {int(args.num_workers)}")
     print(f"Multiprocessing start method: {args.mp_start_method}")
     print(
         "Timing seconds: "