Write burden CSV shards in workers

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
+        _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}")