Support multi-GPU burden index computation

compute_burden_index_landmarks.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import argparse
-import math
+import multiprocessing as mp
+from concurrent.futures import ProcessPoolExecutor, as_completed
 from pathlib import Path
 from typing import Any, Iterable
 
@@ -13,7 +14,6 @@ from burden_index import compute_burden_index, load_burden_context
 from evaluate_auc_v2 import (
     make_eval_indices,
     parse_float_list,
-    split_indices,
 )
 from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
 
@@ -43,6 +43,15 @@ def _parse_horizons(value: Any) -> np.ndarray:
     return horizons
 
 
+def _parse_devices(args: argparse.Namespace) -> list[str | None]:
+    if args.devices is not None and str(args.devices).strip():
+        devices = [x.strip() for x in str(args.devices).split(",") if x.strip()]
+        if not devices:
+            raise ValueError("--devices was provided but no valid devices were parsed.")
+        return devices
+    return [args.device]
+
+
 def _build_burden_matrix_from_mapping(
     mapping_csv: Path,
     *,
@@ -285,6 +294,44 @@ def _result_rows_for_sample(
     return out
 
 
+def _compute_chunk_worker(payload: dict[str, Any]) -> list[dict[str, Any]]:
+    device = payload["device"]
+    run_path = Path(payload["run_path"])
+    ctx = load_burden_context(run_path, device=device)
+    out: list[dict[str, Any]] = []
+    for row in payload["rows"]:
+        for matrix in payload["matrices"]:
+            out.extend(
+                _result_rows_for_sample(
+                    sample_row=row,
+                    horizons=payload["horizons"],
+                    A=matrix["A"],
+                    disease_ids=matrix["disease_ids"],
+                    category_meta=matrix["category_meta"],
+                    burden_type=matrix["burden_type"],
+                    formed_mode=payload["formed_mode"],
+                    ctx=ctx,
+                    run_path=run_path,
+                )
+            )
+    return out
+
+
+def _split_rows_for_devices(
+    rows: list[dict[str, Any]],
+    devices: list[str | None],
+) -> list[tuple[str | None, list[dict[str, Any]]]]:
+    if len(devices) <= 1:
+        return [(devices[0], rows)]
+    index_chunks = np.array_split(np.arange(len(rows)), len(devices))
+    chunks: list[tuple[str | None, list[dict[str, Any]]]] = []
+    for device, idx in zip(devices, index_chunks):
+        if idx.size == 0:
+            continue
+        chunks.append((device, [rows[int(i)] for i in idx.tolist()]))
+    return chunks
+
+
 def main() -> None:
     parser = argparse.ArgumentParser(
         description="Compute DeepHealth Burden Indices at landmark ages."
@@ -309,14 +356,25 @@ def main() -> None:
     parser.add_argument("--min_history_events", type=int, default=1)
     parser.add_argument("--dataset_subset_size", type=int, default=0)
     parser.add_argument("--device", type=str, default=None)
+    parser.add_argument(
+        "--devices",
+        type=str,
+        default=None,
+        help=(
+            "Comma-separated devices for data-parallel BI computation, e.g. "
+            "'cuda:0,cuda:1'. Overrides --device when provided."
+        ),
+    )
     parser.add_argument("--functional_weight_col", type=str,
                         default="hfrm_normalized_weight")
     args = parser.parse_args()
 
     run_path = Path(args.run_path)
     mapping_specs = _load_mapping_specs(args)
+    devices = _parse_devices(args)
 
-    ctx = load_burden_context(run_path, device=args.device)
+    initial_device = "cpu" if len(devices) > 1 else devices[0]
+    ctx = load_burden_context(run_path, device=initial_device)
     matrices = []
     for spec in mapping_specs:
         A, disease_ids, category_meta = _build_burden_matrix_from_mapping(
@@ -363,6 +421,8 @@ def main() -> None:
     output_path.parent.mkdir(parents=True, exist_ok=True)
 
     all_rows: list[dict[str, Any]] = []
+    row_chunks = _split_rows_for_devices(rows, devices)
+    if len(row_chunks) == 1:
         for row in tqdm(rows, desc="Computing BI", dynamic_ncols=True):
             for matrix in matrices:
                 all_rows.extend(
@@ -378,12 +438,40 @@ def main() -> None:
                         run_path=run_path,
                     )
                 )
+    else:
+        # The main-process context is only needed to build the dataset and rows.
+        # Workers load their own model copy on the assigned device.
+        del ctx
+        payloads = [
+            {
+                "device": device,
+                "run_path": str(run_path),
+                "rows": chunk_rows,
+                "horizons": horizons.tolist(),
+                "matrices": matrices,
+                "formed_mode": args.formed_mode,
+            }
+            for device, chunk_rows in row_chunks
+        ]
+        with ProcessPoolExecutor(
+            max_workers=len(payloads),
+            mp_context=mp.get_context("spawn"),
+        ) as executor:
+            futures = [executor.submit(_compute_chunk_worker, p) for p in payloads]
+            for future in tqdm(
+                as_completed(futures),
+                total=len(futures),
+                desc="Computing BI chunks",
+                dynamic_ncols=True,
+            ):
+                all_rows.extend(future.result())
 
     out_df = pd.DataFrame(all_rows)
     out_df.to_csv(output_path, index=False)
     print(f"Run path: {run_path}")
     print(f"Eval split: {eval_split}")
     print(f"Landmark rows: {len(rows)}")
+    print(f"Devices: {', '.join(str(d) for d, _ in row_chunks)}")
     for matrix in matrices:
         print(
             f"{matrix['burden_type']} dimensions: {matrix['A'].shape[0]}, "