diff --git a/compute_burden_index_landmarks.py b/compute_burden_index_landmarks.py index 0ea1a31..49055f4 100644 --- a/compute_burden_index_landmarks.py +++ b/compute_burden_index_landmarks.py @@ -228,6 +228,50 @@ def _eligible_landmark_rows( return rows +def _row_to_worker_spec(row: dict[str, Any]) -> dict[str, Any]: + return { + "patient_id": int(row["patient_id"]), + "dataset_index": int(row["dataset_index"]), + "landmark_age": float(row["landmark_age"]), + "t_query": float(row["t_query"]), + "followup_end_time": float(row["followup_end_time"]), + } + + +def _materialize_worker_rows( + dataset: Any, + row_specs: list[dict[str, Any]], +) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for spec in row_specs: + sample = dataset.samples[int(spec["dataset_index"])] + seq_event = np.asarray(sample["event_seq"], dtype=np.int64) + seq_time = np.asarray(sample["time_seq"], dtype=np.float32) + tgt_event = np.asarray(sample["target_event_seq"], dtype=np.int64) + tgt_time = np.asarray(sample["target_time_seq"], dtype=np.float32) + full_event = np.concatenate([seq_event, tgt_event[-1:]]) + full_time = np.concatenate([seq_time, tgt_time[-1:]]) + t_query = np.float32(float(spec["t_query"])) + prefix_mask = full_time <= t_query + rows.append( + { + "patient_id": int(spec["patient_id"]), + "dataset_index": int(spec["dataset_index"]), + "sex": int(sample["sex"]), + "landmark_age": np.float32(float(spec["landmark_age"])), + "t_query": t_query, + "followup_end_time": np.float32(float(spec["followup_end_time"])), + "event_seq": full_event[prefix_mask].astype(np.int64, copy=False), + "time_seq": full_time[prefix_mask].astype(np.float32, copy=False), + "other_type": np.asarray(sample["other_type"], dtype=np.int64), + "other_value": np.asarray(sample["other_value"], dtype=np.float32), + "other_value_kind": np.asarray(sample["other_value_kind"], dtype=np.int64), + "other_time": np.asarray(sample["other_time"], dtype=np.float32), + } + ) + return rows + + def _config_split_indices( n: int, cfg: dict[str, Any], @@ -620,9 +664,25 @@ def _compute_bi_from_readout_table( def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]: device = payload["device"] run_path = Path(payload["run_path"]) + print( + f"[BI worker {device}] starting with {len(payload['row_specs'])} rows", + flush=True, + ) + load_start = time.perf_counter() ctx = load_burden_context(run_path, device=device) + print( + f"[BI worker {device}] context loaded in {time.perf_counter() - load_start:.2f}s", + flush=True, + ) + materialize_start = time.perf_counter() + rows = _materialize_worker_rows(ctx.dataset, payload["row_specs"]) + print( + f"[BI worker {device}] rows materialized in " + f"{time.perf_counter() - materialize_start:.2f}s", + flush=True, + ) out, readout_jobs, timings = _compute_bi_from_readout_table( - rows=payload["rows"], + rows=rows, horizon=payload["horizon"], matrices=payload["matrices"], union_disease_ids=payload["union_disease_ids"], @@ -630,6 +690,13 @@ def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]: readout_batch_size=int(payload["readout_batch_size"]), ctx=ctx, ) + print( + f"[BI worker {device}] done: readout_jobs={readout_jobs}, " + f"build={timings['build_readout_sec']:.2f}s, " + f"forward={timings['forward_sec']:.2f}s, " + f"reduce={timings['reduce_sec']:.2f}s", + flush=True, + ) return {"rows": out, "readout_jobs": readout_jobs, "timings": timings} @@ -774,6 +841,13 @@ def main() -> None: "'cuda:0,cuda:1'. Overrides --device when provided." ), ) + parser.add_argument( + "--mp_start_method", + type=str, + default="fork", + choices=["fork", "forkserver"], + help="Multiprocessing start method for Linux multi-GPU runs.", + ) parser.add_argument("--functional_weight_col", type=str, default="hfrm_normalized_weight") args = parser.parse_args() @@ -845,6 +919,19 @@ def main() -> None: formed_mode=args.formed_mode, horizon=horizon, ) + for device, chunk_rows in row_chunks: + estimated_jobs = sum( + _estimate_readout_jobs_for_row( + row, + formed_mode=args.formed_mode, + horizon=horizon, + ) + for row in chunk_rows + ) + print( + f"Assigned {len(chunk_rows)} rows / ~{estimated_jobs} readout jobs to {device}", + flush=True, + ) if len(row_chunks) == 1: all_rows, total_readout_jobs, timings = _compute_bi_from_readout_table( rows=rows, @@ -865,7 +952,7 @@ def main() -> None: { "device": device, "run_path": str(run_path), - "rows": chunk_rows, + "row_specs": [_row_to_worker_spec(row) for row in chunk_rows], "horizon": horizon, "matrices": matrices, "union_disease_ids": union_disease_ids, @@ -876,7 +963,7 @@ def main() -> None: ] with ProcessPoolExecutor( max_workers=len(payloads), - mp_context=mp.get_context("spawn"), + mp_context=mp.get_context(args.mp_start_method), ) as executor: futures = [executor.submit(_compute_chunk_worker, p) for p in payloads] for future in tqdm( @@ -901,6 +988,7 @@ def main() -> None: print(f"Landmark rows: {len(rows)}") print(f"Readout jobs: {total_readout_jobs}") print(f"Readout batch size per worker: {int(args.readout_batch_size)}") + print(f"Multiprocessing start method: {args.mp_start_method}") print( "Timing seconds: " f"build_readout={total_timings['build_readout_sec']:.2f}, "