From f417a91a7466b3ce7f7db1283c74f0aeaef497db Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Wed, 1 Jul 2026 13:21:08 +0800 Subject: [PATCH] Parallelize extra-info summary reduction --- evaluate_extra_info_attribution.py | 139 ++++++++++++++++++++++++----- run_missing_evaluations.sh | 12 ++- 2 files changed, 128 insertions(+), 23 deletions(-) diff --git a/evaluate_extra_info_attribution.py b/evaluate_extra_info_attribution.py index 85e3359..cdd8fa9 100644 --- a/evaluate_extra_info_attribution.py +++ b/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() - for key_rows, values in zip(death_key_chunks, death_value_chunks): - update_death_summary( - death_summary, - key_rows=key_rows, - values=values, - ) - - 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, disease_parameter_summary = reduce_attribution_chunks( + death_key_chunks=death_key_chunks, + death_value_chunks=death_value_chunks, + disease_stat_chunks=disease_stat_chunks, + group_names=group_names, + group_labels=group_labels, + cpu_reduce_workers=cpu_reduce_workers, + ) death_summary_path = output_dir / "summary_extra_info_death_parameters.csv" disease_summary_path = output_dir / "summary_extra_info_disease_parameters.csv" diff --git a/run_missing_evaluations.sh b/run_missing_evaluations.sh index 83960d3..b6faf43 100644 --- a/run_missing_evaluations.sh +++ b/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