From c38b42e0f12be3cf4d5912adf9fe1dc96978696b Mon Sep 17 00:00:00 2001 From: Emi Griffith Date: Thu, 23 Jul 2026 14:30:57 -0700 Subject: [PATCH] Parallelize lake tile uploads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A tile is ~1.2k small parquet objects; sequential PUTs to the bucket are round-trip-bound (measured 2.5 obj/s, ~8 min per tile — ~78 h for the city stage). A 16-thread pool per tile turns that into well under a minute; boto3 clients are thread-safe and failures surface via future.result(). --- backend/gen_era5_lake.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/backend/gen_era5_lake.py b/backend/gen_era5_lake.py index f058674..c7be634 100644 --- a/backend/gen_era5_lake.py +++ b/backend/gen_era5_lake.py @@ -20,6 +20,7 @@ Seed-only deps (icechunk/xarray/zarr) stay out of the app image, same policy as seed_era5.py. """ import argparse +import concurrent.futures import datetime import io import os @@ -126,11 +127,15 @@ def _tile_frame(ds, ti: int, tj: int, end: str) -> "pl.DataFrame | None": def _write_tile(sink: _Sink, ti: int, tj: int, tile_df: pl.DataFrame) -> pl.DataFrame: - """Write one tile's point files + hive partitions; return its manifest rows.""" + """Write one tile's point files + hive partitions; return its manifest rows. + Uploads run on a thread pool: a tile is ~1.2k small objects and sequential + PUTs are pure round-trip latency (measured 2.5 obj/s -> 8 min/tile; the + pool brings that to well under a minute). boto3 clients are thread-safe.""" + puts = [] rows = [] for (i, j), pdf in tile_df.group_by(["lat_idx", "lon_idx"]): pdf = pdf.sort("date") - sink.put(era5lake.point_key(i, j), pdf.drop(["lat_idx", "lon_idx"])) + puts.append((era5lake.point_key(i, j), pdf.drop(["lat_idx", "lon_idx"]))) lat, lon = era5lake.to_coords(i, j) rows.append({"lat_idx": i, "lon_idx": j, "lat": lat, "lon": lon, "tile": f"{ti}_{tj}", "rows": pdf.height, @@ -140,8 +145,11 @@ def _write_tile(sink: _Sink, ti: int, tj: int, tile_df: pl.DataFrame) -> pl.Data pl.col("date").dt.year().alias("year"), pl.col("date").dt.month().alias("month")) for (year, month), mdf in parts.group_by(["year", "month"]): - sink.put(era5lake.daily_part_key(ti, tj, year, month), - mdf.drop(["year", "month"]).sort(["lat_idx", "lon_idx", "date"])) + puts.append((era5lake.daily_part_key(ti, tj, year, month), + mdf.drop(["year", "month"]).sort(["lat_idx", "lon_idx", "date"]))) + with concurrent.futures.ThreadPoolExecutor(max_workers=16) as pool: + for f in [pool.submit(sink.put, k, df) for k, df in puts]: + f.result() # surface the first failure instead of hiding it return pl.DataFrame(rows)