diff --git a/compute_burden_index_landmarks.py b/compute_burden_index_landmarks.py index 965d156..3b33f98 100644 --- a/compute_burden_index_landmarks.py +++ b/compute_burden_index_landmarks.py @@ -586,6 +586,26 @@ def _project_bi_rows( return out +def _write_rows_csv(rows: list[dict[str, Any]], output_path: Path) -> int: + df = pd.DataFrame(rows) + df.to_csv(output_path, index=False) + return int(len(df)) + + +def _concat_csv_shards(shard_paths: list[Path], output_path: Path) -> None: + wrote_header = False + with output_path.open("w", encoding="utf-8", newline="") as out_f: + for shard_path in shard_paths: + with shard_path.open("r", encoding="utf-8", newline="") as in_f: + header = in_f.readline() + if not wrote_header: + out_f.write(header) + wrote_header = True + for line in in_f: + out_f.write(line) + shard_path.unlink(missing_ok=True) + + def _compute_bi_from_streamed_readouts( *, rows: list[dict[str, Any]], @@ -710,6 +730,7 @@ def _compute_bi_from_streamed_readouts( def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]: device = payload["device"] run_path = Path(payload["run_path"]) + shard_path = Path(payload["shard_path"]) print( f"[BI worker {device}] starting with {len(payload['row_specs'])} rows", flush=True, @@ -745,7 +766,20 @@ def _compute_chunk_worker(payload: dict[str, Any]) -> dict[str, Any]: f"reduce={timings['reduce_sec']:.2f}s", flush=True, ) - return {"rows": out, "readout_jobs": readout_jobs, "timings": timings} + write_start = time.perf_counter() + row_count = _write_rows_csv(out, shard_path) + timings["write_csv_sec"] = time.perf_counter() - write_start + print( + f"[BI worker {device}] wrote {row_count} rows to {shard_path} " + f"in {timings['write_csv_sec']:.2f}s", + flush=True, + ) + return { + "shard_path": str(shard_path), + "row_count": row_count, + "readout_jobs": readout_jobs, + "timings": timings, + } def _attach_union_projection( @@ -959,8 +993,9 @@ def main() -> None: ) output_path.parent.mkdir(parents=True, exist_ok=True) - all_rows: list[dict[str, Any]] = [] total_readout_jobs = 0 + output_written = False + shard_paths: list[Path] = [] total_timings = { "build_readout_sec": 0.0, "forward_sec": 0.0, @@ -1000,6 +1035,10 @@ def main() -> None: ) for key, value in timings.items(): total_timings[key] += float(value) + write_start = time.perf_counter() + _write_rows_csv(all_rows, output_path) + total_timings["write_csv_sec"] = time.perf_counter() - write_start + output_written = True else: # The main-process context is only needed to build the dataset and rows. # Workers load their own model copy on the assigned device. @@ -1008,6 +1047,11 @@ def main() -> None: { "device": device, "run_path": str(run_path), + "shard_path": str( + output_path.with_name( + f"{output_path.stem}.part{part_idx:03d}{output_path.suffix}" + ) + ), "row_specs": [_row_to_worker_spec(row) for row in chunk_rows], "horizon": horizon, "matrices": matrices, @@ -1016,7 +1060,7 @@ def main() -> None: "num_workers": int(args.num_workers), "formed_mode": args.formed_mode, } - for device, chunk_rows in row_chunks + for part_idx, (device, chunk_rows) in enumerate(row_chunks) ] with ProcessPoolExecutor( max_workers=len(payloads), @@ -1030,15 +1074,18 @@ def main() -> None: dynamic_ncols=True, ): result = future.result() - all_rows.extend(result["rows"]) + shard_paths.append(Path(result["shard_path"])) total_readout_jobs += int(result["readout_jobs"]) for key, value in result["timings"].items(): total_timings[key] += float(value) - write_start = time.perf_counter() - out_df = pd.DataFrame(all_rows) - out_df.to_csv(output_path, index=False) - total_timings["write_csv_sec"] = time.perf_counter() - write_start + write_start = time.perf_counter() + _concat_csv_shards(sorted(shard_paths), output_path) + total_timings["write_csv_sec"] += time.perf_counter() - write_start + output_written = True + + if not output_written: + raise RuntimeError("BI output was not written.") print(f"Run path: {run_path}") print(f"Eval split: {eval_split}") print(f"Horizon: {horizon:g}")