Parallelize extra-info summary reduction

evaluate_extra_info_attribution.py

@@ -1,4 +1,4 @@
-"""Evaluate extra-info attribution to death parameters and future disease risks.
+"""Evaluate extra-info attribution to death and disease distribution parameters.
 
 For each landmark query, this script scans selected extra-info types that are
 available at or before the query age. For each such type it re-runs the model
@@ -15,6 +15,7 @@ from __future__ import annotations
 import argparse
 import json
 import re
+from concurrent.futures import ProcessPoolExecutor, as_completed
 from pathlib import Path
 from typing import Any, Sequence
 
@@ -462,6 +463,98 @@ def update_disease_parameter_summary_from_group_stats(
                 acc[f"sumsq__{column}"] += float(sumsq[int(row_idx), int(group_idx), int(col_idx)])
 
 
+def merge_summary_dict(
+    dst: dict[tuple[Any, ...], dict[str, float]],
+    src: dict[tuple[Any, ...], dict[str, float]],
+) -> None:
+    for key, src_acc in src.items():
+        dst_acc = dst.setdefault(key, {name: 0.0 for name in src_acc})
+        for name, value in src_acc.items():
+            dst_acc[name] = dst_acc.get(name, 0.0) + float(value)
+
+
+def reduce_attribution_chunk_bundle(
+    payload: tuple[
+        list[tuple[pd.DataFrame, np.ndarray]],
+        list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]],
+        list[str],
+        list[str],
+    ],
+) -> tuple[dict[tuple[Any, ...], dict[str, float]], dict[tuple[Any, ...], dict[str, float]]]:
+    death_items, disease_items, group_names, group_labels = payload
+    death_summary: dict[tuple[Any, ...], dict[str, float]] = {}
+    disease_summary: dict[tuple[Any, ...], dict[str, float]] = {}
+
+    for key_rows, values in death_items:
+        update_death_summary(
+            death_summary,
+            key_rows=key_rows,
+            values=values,
+        )
+
+    for key_rows, sums, sumsq, counts in disease_items:
+        update_disease_parameter_summary_from_group_stats(
+            disease_summary,
+            key_rows=key_rows,
+            group_names=group_names,
+            group_labels=group_labels,
+            sums=sums,
+            sumsq=sumsq,
+            counts=counts,
+        )
+
+    return death_summary, disease_summary
+
+
+def reduce_attribution_chunks(
+    *,
+    death_key_chunks: list[pd.DataFrame],
+    death_value_chunks: list[np.ndarray],
+    disease_stat_chunks: list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]],
+    group_names: list[str],
+    group_labels: list[str],
+    cpu_reduce_workers: int,
+) -> tuple[dict[tuple[Any, ...], dict[str, float]], dict[tuple[Any, ...], dict[str, float]]]:
+    n_chunks = max(len(death_key_chunks), len(disease_stat_chunks))
+    if n_chunks == 0:
+        return {}, {}
+
+    worker_count = max(1, min(int(cpu_reduce_workers), n_chunks))
+    if worker_count == 1:
+        return reduce_attribution_chunk_bundle(
+            (
+                list(zip(death_key_chunks, death_value_chunks)),
+                disease_stat_chunks,
+                group_names,
+                group_labels,
+            )
+        )
+
+    bundles = []
+    for worker_idx in range(worker_count):
+        start = worker_idx * n_chunks // worker_count
+        stop = (worker_idx + 1) * n_chunks // worker_count
+        if start >= stop:
+            continue
+        death_items = [
+            (death_key_chunks[i], death_value_chunks[i])
+            for i in range(start, min(stop, len(death_key_chunks)))
+        ]
+        disease_items = disease_stat_chunks[start:min(stop, len(disease_stat_chunks))]
+        bundles.append((death_items, disease_items, group_names, group_labels))
+
+    merged_death: dict[tuple[Any, ...], dict[str, float]] = {}
+    merged_disease: dict[tuple[Any, ...], dict[str, float]] = {}
+    with ProcessPoolExecutor(max_workers=len(bundles)) as executor:
+        futures = [executor.submit(reduce_attribution_chunk_bundle, bundle) for bundle in bundles]
+        for future in tqdm(as_completed(futures), total=len(futures), desc="CPU summary reduction", dynamic_ncols=True):
+            death_part, disease_part = future.result()
+            merge_summary_dict(merged_death, death_part)
+            merge_summary_dict(merged_disease, disease_part)
+
+    return merged_death, merged_disease
+
+
 def write_death_summary_csv(
     path: Path,
     summary: dict[tuple[Any, ...], dict[str, float]],
@@ -532,7 +625,7 @@ def write_disease_parameter_summary_csv(
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
-        description="Compute extra-info ablation attribution for death parameters and future disease risks."
+        description="Compute extra-info ablation attribution for death and disease distribution parameters."
     )
     parser.add_argument("--run_path", type=str, required=True)
     parser.add_argument(
@@ -563,6 +656,12 @@ def parse_args() -> argparse.Namespace:
         help="Forward batch size for expanded extra-info ablation queries.",
     )
     parser.add_argument("--num_workers", type=int, default=None)
+    parser.add_argument(
+        "--cpu_reduce_workers",
+        type=int,
+        default=None,
+        help="Worker processes for CPU-side summary reduction. Defaults to --num_workers.",
+    )
     parser.add_argument("--device", type=str, default=None)
     return parser.parse_args()
 
@@ -685,6 +784,13 @@ def main() -> None:
         raise ValueError("attribution_batch_size must be positive")
 
     num_workers = int(cfg_get(args, cfg, "num_workers", 4))
+    cpu_reduce_workers = int(
+        args.cpu_reduce_workers
+        if args.cpu_reduce_workers is not None
+        else max(1, num_workers)
+    )
+    if cpu_reduce_workers <= 0:
+        raise ValueError("--cpu_reduce_workers must be positive")
     loader = DataLoader(
         IndexedLandmarkDataset(landmark_dataset),
         batch_size=batch_size,
@@ -713,10 +819,9 @@ def main() -> None:
     print(f"Extra-info types: {selected_extra_info_types}")
     print(f"Landmark rows: {len(landmark_dataset)}")
     print(f"Attribution batch size: {attribution_batch_size}")
+    print(f"CPU reduce workers: {cpu_reduce_workers}")
     print(f"Output directory: {output_dir}")
 
-    death_summary: dict[tuple[Any, ...], dict[str, float]] = {}
-    disease_parameter_summary: dict[tuple[Any, ...], dict[str, float]] = {}
     death_key_chunks: list[pd.DataFrame] = []
     death_value_chunks: list[np.ndarray] = []
     disease_stat_chunks: list[tuple[pd.DataFrame, np.ndarray, np.ndarray, np.ndarray]] = []
@@ -799,24 +904,14 @@ def main() -> None:
             )
             disease_stat_chunks.append((key_table, sums, sumsq, counts))
 
-    death_summary.clear()
+    death_summary, disease_parameter_summary = reduce_attribution_chunks(
-    for key_rows, values in zip(death_key_chunks, death_value_chunks):
+        death_key_chunks=death_key_chunks,
-        update_death_summary(
+        death_value_chunks=death_value_chunks,
-            death_summary,
+        disease_stat_chunks=disease_stat_chunks,
-            key_rows=key_rows,
+        group_names=group_names,
-            values=values,
+        group_labels=group_labels,
-        )
+        cpu_reduce_workers=cpu_reduce_workers,
-
+    )
-    for key_rows, sums, sumsq, counts in disease_stat_chunks:
-        update_disease_parameter_summary_from_group_stats(
-            disease_parameter_summary,
-            key_rows=key_rows,
-            group_names=group_names,
-            group_labels=group_labels,
-            sums=sums,
-            sumsq=sumsq,
-            counts=counts,
-        )
 
     death_summary_path = output_dir / "summary_extra_info_death_parameters.csv"
     disease_summary_path = output_dir / "summary_extra_info_disease_parameters.csv"

run_missing_evaluations.sh

@@ -10,6 +10,7 @@ PYTHON_BIN="${PYTHON_BIN:-python}"
 DEVICE="${DEVICE:-cuda}"
 EVAL_SPLIT="${EVAL_SPLIT:-test}"
 NUM_WORKERS="${NUM_WORKERS:-4}"
+CPU_REDUCE_WORKERS="${CPU_REDUCE_WORKERS:-}"
 NUM_WORKERS_AUC="${NUM_WORKERS_AUC:-}"
 BATCH_SIZE="${BATCH_SIZE:-}"
 DATASET_SUBSET_SIZE="${DATASET_SUBSET_SIZE:-}"
@@ -41,6 +42,12 @@ auc_args() {
   fi
 }
 
+cpu_reduce_args() {
+  if [[ -n "${CPU_REDUCE_WORKERS}" ]]; then
+    printf '%s\n' --cpu_reduce_workers "${CPU_REDUCE_WORKERS}"
+  fi
+}
+
 has_completed_dir() {
   local dir="$1"
   shift
@@ -127,6 +134,9 @@ for run_path in runs/*; do
   auc_extra=()
   while IFS= read -r arg; do auc_extra+=("${arg}"); done < <(auc_args)
 
+  cpu_reduce_extra=()
+  while IFS= read -r arg; do cpu_reduce_extra+=("${arg}"); done < <(cpu_reduce_args)
+
   run_dir_result_if_missing \
     "evaluate_auc.py" \
     "${run_path}" \
@@ -153,7 +163,7 @@ for run_path in runs/*; do
         "${run_path}/extra_info_attribution_${EVAL_SPLIT}" \
         "manifest.json" \
         "summary_extra_info_disease_parameters.csv" \
-        "${PYTHON_BIN}" evaluate_extra_info_attribution.py "${common[@]}"
+        "${PYTHON_BIN}" evaluate_extra_info_attribution.py "${common[@]}" "${cpu_reduce_extra[@]}"
     else
       echo "    skip evaluate_extra_info_attribution.py: run has no extra-info types"
     fi