Jinemi/thermograph
4
0
Fork 0
Code Issues 2 Pull requests 9 Projects Releases Packages Wiki Activity Actions

Parallelize lake tile uploads
All checks were successful
PR build (required check) / changes (pull_request) Successful in 10s
Details
secrets-guard / encrypted (pull_request) Successful in 8s
Details
PR build (required check) / build-frontend (pull_request) Has been skipped
Details
PR build (required check) / validate-observability (pull_request) Has been skipped
Details
PR build (required check) / build-backend (pull_request) Successful in 56s
Details
PR build (required check) / gate (pull_request) Successful in 2s
Details

Browse source 
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().
This commit is contained in:
Emi Griffith 2026-07-23 14:30:57 -07:00
parent d9537798c6
commit c38b42e0f1
1 changed files with 12 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
backend/gen_era5_lake.py
View file

@ -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)