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

Parallelize lake tile uploads (#17)
All checks were successful
secrets-guard / encrypted (push) Successful in 7s
Details
Build + push backend image (Forgejo registry) / build-push (push) Successful in 1m2s
Details
Deploy backend to LAN dev server / build (push) Successful in 1m12s
Details
Deploy backend to LAN dev server / deploy (push) Successful in 21s
Details

Browse source
This commit is contained in:
emi 2026-07-23 21:32:58 +00:00
parent d9537798c6
commit 21eea3a4c2
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. as seed_era5.py.
""" """
import argparse import argparse
import concurrent.futures
import datetime import datetime
import io import io
import os 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: 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 = [] rows = []
for (i, j), pdf in tile_df.group_by(["lat_idx", "lon_idx"]): for (i, j), pdf in tile_df.group_by(["lat_idx", "lon_idx"]):
pdf = pdf.sort("date") 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) lat, lon = era5lake.to_coords(i, j)
rows.append({"lat_idx": i, "lon_idx": j, "lat": lat, "lon": lon, rows.append({"lat_idx": i, "lon_idx": j, "lat": lat, "lon": lon,
"tile": f"{ti}_{tj}", "rows": pdf.height, "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.year().alias("year"),
pl.col("date").dt.month().alias("month")) pl.col("date").dt.month().alias("month"))
for (year, month), mdf in parts.group_by(["year", "month"]): for (year, month), mdf in parts.group_by(["year", "month"]):
sink.put(era5lake.daily_part_key(ti, tj, year, month), puts.append((era5lake.daily_part_key(ti, tj, year, month),
mdf.drop(["year", "month"]).sort(["lat_idx", "lon_idx", "date"])) 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) return pl.DataFrame(rows)