From a35f1117c009edb76161412abc8c15acffaf0d65 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Wed, 8 Jul 2026 16:34:28 +0800 Subject: [PATCH] Parallelize exposure cache reads --- prepare_exposure_cache.py | 213 ++++++++++++++++++++++++++++---------- 1 file changed, 159 insertions(+), 54 deletions(-) diff --git a/prepare_exposure_cache.py b/prepare_exposure_cache.py index 7d4dee7..c794075 100644 --- a/prepare_exposure_cache.py +++ b/prepare_exposure_cache.py @@ -36,7 +36,9 @@ converts back to these raw tokens before reading this cache. from __future__ import annotations import argparse +from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, wait import json +import os from pathlib import Path from typing import Iterable @@ -162,6 +164,75 @@ def _reshape_window(df: pd.DataFrame, cols: list[str], length: int, n_channels: return arr.reshape(len(df), n_channels, length).transpose(0, 2, 1) +def _quality_column(df: pd.DataFrame, name: str, n_rows: int) -> np.ndarray: + if name not in df: + return np.full(n_rows, np.nan, dtype=np.float32) + return df[name].to_numpy(dtype=np.float32, copy=True) + + +def _process_exposure_task(task: dict) -> dict: + daily_file = Path(task["daily_path"]) + monthly_file = Path(task["monthly_path"]) + wanted = task["wanted"] + daily_cols = task["daily_cols"] + monthly_cols = task["monthly_cols"] + + daily_read_cols = [ + "eid", + "onset_date", + "token", + *_safe_columns(daily_file, daily_cols), + *_safe_columns(daily_file, ["n_days_nonmissing", "n_rh_days_nonmissing"]), + ] + monthly_read_cols = [ + "eid", + "onset_date", + "token", + *_safe_columns(monthly_file, monthly_cols), + *_safe_columns(monthly_file, ["n_months_nonmissing", "n_rh_months_nonmissing"]), + ] + daily_df = _read_matching_parquet_rows(daily_file, daily_read_cols, wanted) + monthly_df = _read_matching_parquet_rows(monthly_file, monthly_read_cols, wanted) + if daily_df.empty: + return {"positions": np.empty(0, dtype=np.int64)} + + common_positions = np.intersect1d( + daily_df["position"].to_numpy(dtype=np.int64), + monthly_df["position"].to_numpy(dtype=np.int64), + ) + if len(common_positions) == 0: + return {"positions": np.empty(0, dtype=np.int64)} + + daily_df = daily_df.set_index("position").loc[common_positions].reset_index() + monthly_df = monthly_df.set_index("position").loc[common_positions].reset_index() + n_match = len(common_positions) + quality = np.stack( + [ + _quality_column(daily_df, "n_days_nonmissing", n_match), + _quality_column(daily_df, "n_rh_days_nonmissing", n_match), + _quality_column(monthly_df, "n_months_nonmissing", n_match), + _quality_column(monthly_df, "n_rh_months_nonmissing", n_match), + ], + axis=1, + ) + return { + "positions": common_positions.astype(np.int64), + "daily": _reshape_window( + daily_df, + daily_cols, + DAILY_LENGTH, + len(DAILY_CHANNELS), + ), + "monthly": _reshape_window( + monthly_df, + monthly_cols, + MONTHLY_LENGTH, + len(MONTHLY_CHANNELS), + ), + "quality": quality, + } + + def _load_summary( exposure_dir: Path, summary_file: str, @@ -257,6 +328,8 @@ def build_exposure_cache( data_prefix: str = "ukb", labels_file: str | Path = "labels.csv", summary_file: str = "summary.csv", + workers: int = 1, + max_in_flight: int = 0, overwrite: bool = False, show_progress: bool = True, ) -> int: @@ -359,71 +432,83 @@ def build_exposure_cache( } write_offset = 0 - iterator = tqdm( - summary.itertuples(index=False), - total=len(summary), - desc="Writing eid-sequence exposure cache", - unit="file", - disable=not show_progress, - ) - for row in iterator: - daily_file = Path(row.daily_path) - monthly_file = Path(row.monthly_path) + tasks: list[dict] = [] + for row in summary.itertuples(index=False): token = int(row.raw_token) wanted = wanted_by_token.get(token) if wanted is None or wanted.empty: continue - - daily_read_cols = [ - "eid", - "onset_date", - "token", - *_safe_columns(daily_file, daily_cols), - *_safe_columns(daily_file, ["n_days_nonmissing", "n_rh_days_nonmissing"]), - ] - monthly_read_cols = [ - "eid", - "onset_date", - "token", - *_safe_columns(monthly_file, monthly_cols), - *_safe_columns(monthly_file, ["n_months_nonmissing", "n_rh_months_nonmissing"]), - ] - daily_df = _read_matching_parquet_rows(daily_file, daily_read_cols, wanted) - monthly_df = _read_matching_parquet_rows(monthly_file, monthly_read_cols, wanted) - if daily_df.empty: - continue - - common_positions = np.intersect1d( - daily_df["position"].to_numpy(dtype=np.int64), - monthly_df["position"].to_numpy(dtype=np.int64), + tasks.append( + { + "daily_path": str(row.daily_path), + "monthly_path": str(row.monthly_path), + "wanted": wanted, + "daily_cols": daily_cols, + "monthly_cols": monthly_cols, + } ) - if len(common_positions) == 0: - continue - daily_df = daily_df.set_index("position").loc[common_positions].reset_index() - monthly_df = monthly_df.set_index("position").loc[common_positions].reset_index() - positions = common_positions.astype(np.int64) + workers = max(1, int(workers)) + max_in_flight = int(max_in_flight) + + def write_result(result: dict) -> None: + nonlocal write_offset + positions = result["positions"] + if len(positions) == 0: + return n_match = len(positions) end_offset = write_offset + n_match - daily_mm[write_offset:end_offset] = _reshape_window( - daily_df, - daily_cols, - DAILY_LENGTH, - len(DAILY_CHANNELS), - ) - monthly_mm[write_offset:end_offset] = _reshape_window( - monthly_df, - monthly_cols, - MONTHLY_LENGTH, - len(MONTHLY_CHANNELS), - ) - quality_mm[write_offset:end_offset, 0] = daily_df.get("n_days_nonmissing", np.nan) - quality_mm[write_offset:end_offset, 1] = daily_df.get("n_rh_days_nonmissing", np.nan) - quality_mm[write_offset:end_offset, 2] = monthly_df.get("n_months_nonmissing", np.nan) - quality_mm[write_offset:end_offset, 3] = monthly_df.get("n_rh_months_nonmissing", np.nan) + daily_mm[write_offset:end_offset] = result["daily"] + monthly_mm[write_offset:end_offset] = result["monthly"] + quality_mm[write_offset:end_offset] = result["quality"] row_index_mm[positions] = np.arange(write_offset, end_offset, dtype=np.int64) write_offset = end_offset + if workers == 1: + iterator = tqdm( + map(_process_exposure_task, tasks), + total=len(tasks), + desc="Reading exposure parquet and writing cache", + unit="file", + disable=not show_progress, + ) + for result in iterator: + write_result(result) + else: + with ProcessPoolExecutor(max_workers=workers) as executor: + task_iter = iter(tasks) + iterator = tqdm( + total=len(tasks), + desc=f"Reading exposure parquet ({workers} workers)", + unit="file", + disable=not show_progress, + ) + if max_in_flight <= 0: + in_flight = { + executor.submit(_process_exposure_task, task) + for task in task_iter + } + while in_flight: + done, in_flight = wait(in_flight, return_when=FIRST_COMPLETED) + for future in done: + write_result(future.result()) + iterator.update(1) + else: + in_flight = { + executor.submit(_process_exposure_task, task) + for task in [next(task_iter, None) for _ in range(max_in_flight)] + if task is not None + } + while in_flight: + done, in_flight = wait(in_flight, return_when=FIRST_COMPLETED) + for future in done: + write_result(future.result()) + iterator.update(1) + task = next(task_iter, None) + if task is not None: + in_flight.add(executor.submit(_process_exposure_task, task)) + iterator.close() + row_index_mm.flush() daily_mm.flush() monthly_mm.flush() @@ -460,6 +545,24 @@ def main() -> None: parser.add_argument("--data-prefix", default="ukb") parser.add_argument("--labels-file", default="labels.csv") parser.add_argument("--summary-file", default="summary.csv") + parser.add_argument( + "--workers", + type=int, + default=max(1, os.cpu_count() or 1), + help=( + "Worker processes for parquet reading and window extraction. " + "The main process remains the only writer to the output memmaps." + ), + ) + parser.add_argument( + "--max-in-flight", + type=int, + default=0, + help=( + "Maximum submitted parquet tasks waiting/running at once. " + "Use 0 to submit all tasks, which is the default for high-memory servers." + ), + ) parser.add_argument( "--no-progress", action="store_true", @@ -474,6 +577,8 @@ def main() -> None: data_prefix=args.data_prefix, labels_file=args.labels_file, summary_file=args.summary_file, + workers=args.workers, + max_in_flight=args.max_in_flight, overwrite=args.overwrite, show_progress=not args.no_progress, )