From ee7d363e3d4085c4d7fa641725fd948d773f0831 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Fri, 26 Jun 2026 10:51:56 +0800 Subject: [PATCH] Support multi-GPU burden index computation --- compute_burden_index_landmarks.py | 122 +++++++++++++++++++++++++----- 1 file changed, 105 insertions(+), 17 deletions(-) diff --git a/compute_burden_index_landmarks.py b/compute_burden_index_landmarks.py index 2cd6858..189b974 100644 --- a/compute_burden_index_landmarks.py +++ b/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,27 +421,57 @@ def main() -> None: output_path.parent.mkdir(parents=True, exist_ok=True) all_rows: list[dict[str, Any]] = [] - for row in tqdm(rows, desc="Computing BI", dynamic_ncols=True): - for matrix in matrices: - all_rows.extend( - _result_rows_for_sample( - sample_row=row, - horizons=horizons.tolist(), - A=matrix["A"], - disease_ids=matrix["disease_ids"], - category_meta=matrix["category_meta"], - burden_type=matrix["burden_type"], - formed_mode=args.formed_mode, - ctx=ctx, - run_path=run_path, + 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( + _result_rows_for_sample( + sample_row=row, + horizons=horizons.tolist(), + A=matrix["A"], + disease_ids=matrix["disease_ids"], + category_meta=matrix["category_meta"], + burden_type=matrix["burden_type"], + formed_mode=args.formed_mode, + ctx=ctx, + 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]}, "