graham/prop
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Speed up ERA5 batch pipeline: parallel CDS + stream downloads + hot index

Browse source 
Four independent wins from a perf audit of the era5_batch worker path:

1. Parallel CDS submits per job. fetch_single_level_month and
   fetch_pressure_level_month have no ordering dependency — each blocks
   30+ min on CDS queue/download — so running them in two linked Tasks
   halves wall time for every job in the backlog.

2. Stream GRIB2 downloads directly to disk via Req `into: File.stream!`.
   A month-tile GRIB is 50–200 MB and was previously slurped into the
   BEAM heap and then written back out to a temp file before wgrib2 ran.
   New Era5Client.submit_and_download_to_file/3 skips both copies. Saves
   100–200 MB heap per worker × 12 concurrent ≈ 1.2–2.4 GB peak under
   full parallelism. Partial files are removed on HTTP error.

3. Wgrib2.extract_grid_messages_from_file/3 lets batch decoding read the
   streamed temp file directly. The binary variant still exists for the
   single-point path and now delegates through the same helper.

4. Composite index (valid_time, lat, lon) on era5_profiles, matching the
   three-range scan used by Weather.find_nearest_era5 (contact detail +
   path profiles) and StatusLive.count_era5_done (status page EXISTS
   subquery). The unique index on (lat, lon, valid_time) couldn't serve
   it efficiently.

Also bumps the bulk insert chunk size 2k → 4k to halve insert round-trips
per month-tile (still well under Postgres' 65k-parameter cap).
This commit is contained in:
Graham McIntire 2026-04-13 16:00:18 -05:00
parent ba91774b71
commit 156b170a5c
No known key found for this signature in database
GPG key ID: F4ABF488E6029E59
4 changed files with 167 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

62
lib/microwaveprop/weather/era5_batch_client.ex
View file

@ -81,14 +81,39 @@ defmodule Microwaveprop.Weather.Era5BatchClient do
Logger.info("ERA5Batch: fetching #{year}-#{pad(month)} tile #{tile_lat},#{tile_lon} (#{days} days)")
with {:ok, single_grib} <- fetch_single_level_month(year, month, days, area),
{:ok, pressure_grib} <- fetch_pressure_level_month(year, month, days, area),
{:ok, profiles} <- build_month_profiles(tile_lat, tile_lon, single_grib, pressure_grib) do
inserted = bulk_insert_profiles(profiles)
tmp_dir = System.tmp_dir!()
job_tag = "#{year}_#{pad(month)}_#{tile_lat}_#{tile_lon}_#{System.unique_integer([:positive])}"
single_path = Path.join(tmp_dir, "era5_single_#{job_tag}.grib2")
pressure_path = Path.join(tmp_dir, "era5_pressure_#{job_tag}.grib2")
Logger.info("ERA5Batch: stored #{inserted} profiles for #{year}-#{pad(month)} tile #{tile_lat},#{tile_lon}")
# Fire both CDS submits in parallel — they have no ordering dependency and
# each spends 30+ min blocking on CDS queue/download. Running them serially
# doubles wall time for zero benefit. Task.async links to the Oban worker,
# so if the worker exits (deploy, crash) both tasks die with it. We still
# shut down any unfinished task on early error so we don't waste CDS quota.
# Both downloads stream directly to disk (Req `into: File.stream!`), so
# peak heap stays O(chunk) rather than O(month-tile GRIB size ~50-200 MB).
single_task =
Task.async(fn -> fetch_single_level_month(year, month, days, area, single_path) end)
{:ok, inserted}
pressure_task =
Task.async(fn -> fetch_pressure_level_month(year, month, days, area, pressure_path) end)
try do
with :ok <- Task.await(single_task, :infinity),
:ok <- Task.await(pressure_task, :infinity),
{:ok, profiles} <- build_month_profiles(tile_lat, tile_lon, single_path, pressure_path) do
inserted = bulk_insert_profiles(profiles)
Logger.info("ERA5Batch: stored #{inserted} profiles for #{year}-#{pad(month)} tile #{tile_lat},#{tile_lon}")
{:ok, inserted}
end
after
Task.shutdown(single_task, :brutal_kill)
Task.shutdown(pressure_task, :brutal_kill)
File.rm(single_path)
File.rm(pressure_path)
end
end
@ -123,7 +148,7 @@ defmodule Microwaveprop.Weather.Era5BatchClient do
# -- CDS requests ----------------------------------------------------------
defp fetch_single_level_month(year, month, days, area) do
defp fetch_single_level_month(year, month, days, area, dest_path) do
request = %{
"product_type" => ["reanalysis"],
"variable" => Era5Client.single_level_vars(),
@ -135,10 +160,10 @@ defmodule Microwaveprop.Weather.Era5BatchClient do
"data_format" => "grib"
}
Era5Client.submit_and_download("reanalysis-era5-single-levels", request)
Era5Client.submit_and_download_to_file("reanalysis-era5-single-levels", request, dest_path)
end
defp fetch_pressure_level_month(year, month, days, area) do
defp fetch_pressure_level_month(year, month, days, area, dest_path) do
request = %{
"product_type" => ["reanalysis"],
"variable" => Era5Client.pressure_level_vars(),
@ -151,7 +176,7 @@ defmodule Microwaveprop.Weather.Era5BatchClient do
"data_format" => "grib"
}
Era5Client.submit_and_download("reanalysis-era5-pressure-levels", request)
Era5Client.submit_and_download_to_file("reanalysis-era5-pressure-levels", request, dest_path)
end
# -- GRIB extraction -------------------------------------------------------
@ -169,12 +194,21 @@ defmodule Microwaveprop.Weather.Era5BatchClient do
}
end
defp build_month_profiles(tile_lat, tile_lon, single_grib, pressure_grib) do
defp build_month_profiles(tile_lat, tile_lon, single_path, pressure_path) do
grid_spec = grid_spec_for_tile(tile_lat, tile_lon)
with {:ok, single_msgs} <-
Wgrib2.extract_grid_messages(single_grib, ":(TMP|DPT|PRES|UGRD|VGRD|PWAT|HPBL):", grid_spec),
{:ok, pressure_msgs} <- Wgrib2.extract_grid_messages(pressure_grib, ":(TMP|SPFH|HGT):", grid_spec) do
Wgrib2.extract_grid_messages_from_file(
single_path,
":(TMP|DPT|PRES|UGRD|VGRD|PWAT|HPBL):",
grid_spec
),
{:ok, pressure_msgs} <-
Wgrib2.extract_grid_messages_from_file(
pressure_path,
":(TMP|SPFH|HGT):",
grid_spec
) do
# Index both message sets by (datetime, lat, lon) → var:level map.
points_single = index_messages(single_msgs)
points_pressure = index_messages(pressure_msgs)
@ -291,7 +325,7 @@ defmodule Microwaveprop.Weather.Era5BatchClient do
# Insert in chunks to keep each query under Postgres' parameter limit
# (65,535 params / ~15 fields ≈ 4,000 rows per statement).
@chunk_size 2_000
@chunk_size 4_000
defp bulk_insert_profiles(profiles) do
now = DateTime.truncate(DateTime.utc_now(), :second)

79
lib/microwaveprop/weather/era5_client.ex
View file

@ -127,15 +127,39 @@ defmodule Microwaveprop.Weather.Era5Client do
do_submit_and_download(dataset, request)
end
@doc """
Like `submit_and_download/2` but streams the GRIB2 response directly to
`path` on disk rather than loading the whole body (50-200 MB for month
tiles) into the BEAM heap.
"""
@spec submit_and_download_to_file(String.t(), map(), Path.t()) :: :ok | {:error, term()}
def submit_and_download_to_file(dataset, request, path) do
with {:ok, source, api_key} <- submit_and_poll(dataset, request) do
case source do
{:body, body} -> File.write(path, body)
{:url, href} -> download_href_to_file(href, api_key, path)
end
end
end
defp do_submit_and_download(dataset, request) do
with {:ok, source, api_key} <- submit_and_poll(dataset, request) do
case source do
{:body, body} -> {:ok, body}
{:url, href} -> download_file(href, api_key)
end
end
end
defp submit_and_poll(dataset, request) do
api_key = api_key()
if is_nil(api_key) do
{:error, "ERA5_CDS_API_KEY not configured"}
else
case submit_request(dataset, request, api_key) do
{:ok, job_id} -> poll_and_download(job_id, api_key)
{:error, reason} -> {:error, reason}
with {:ok, job_id} <- submit_request(dataset, request, api_key),
{:ok, source} <- poll_until_complete(job_id, api_key) do
{:ok, source, api_key}
end
end
end
@ -166,24 +190,26 @@ defmodule Microwaveprop.Weather.Era5Client do
end
end
defp poll_and_download(job_id, api_key, attempt \\ 0)
# Returns {:ok, {:url, href}} when CDS returns a JSON asset pointer (new CDS)
# or {:ok, {:body, binary}} for direct binary responses (legacy CDS).
defp poll_until_complete(job_id, api_key, attempt \\ 0)
defp poll_and_download(_job_id, _api_key, attempt) when attempt >= @max_poll_attempts do
defp poll_until_complete(_job_id, _api_key, attempt) when attempt >= @max_poll_attempts do
{:error, "ERA5: poll timeout after #{@max_poll_attempts} attempts"}
end
defp poll_and_download(job_id, api_key, attempt) do
defp poll_until_complete(job_id, api_key, attempt) do
url = "#{@cds_url}/retrieve/v1/jobs/#{job_id}"
case Req.get(url, headers: [{"PRIVATE-TOKEN", api_key}], receive_timeout: 10_000) do
{:ok, %{status: 200, body: %{"status" => "successful"}}} ->
results_url = "#{url}/results"
download_result(results_url, api_key)
fetch_results_source(results_url, api_key)
{:ok, %{status: 200, body: %{"status" => status}}}
when status in ["accepted", "running"] ->
Process.sleep(@poll_interval_ms)
poll_and_download(job_id, api_key, attempt + 1)
poll_until_complete(job_id, api_key, attempt + 1)
{:ok, %{status: 200, body: %{"status" => "failed"} = body}} ->
# CDS often omits "message" — fall back to dumping the whole body so the
@ -199,18 +225,21 @@ defmodule Microwaveprop.Weather.Era5Client do
end
end
defp download_result(url, api_key) do
defp fetch_results_source(url, api_key) do
case Req.get(url, headers: [{"PRIVATE-TOKEN", api_key}], receive_timeout: 120_000) do
{:ok, %{status: 200, body: %{"asset" => %{"value" => %{"href" => href}}}}} ->
# New CDS returns a JSON with download link
download_file(href, api_key)
{:ok, {:url, href}}
{:ok, %{status: 200, headers: headers, body: body}} ->
# Direct binary download
content_type = %Req.Response{headers: headers} |> Req.Response.get_header("content-type") |> List.first("")
content_type =
%Req.Response{headers: headers}
|> Req.Response.get_header("content-type")
|> List.first("")
if String.contains?(content_type, "grib") or is_binary(body) do
{:ok, body}
{:ok, {:body, body}}
else
{:error, "ERA5: unexpected response format"}
end
@ -231,6 +260,32 @@ defmodule Microwaveprop.Weather.Era5Client do
end
end
# Streams the GRIB2 response directly to `path`. Uses Req's `into:` with a
# `File.stream!` collectable, so chunks land on disk as they arrive and peak
# heap stays O(chunk size) instead of O(full response). Target path is
# removed on HTTP error so a failed download never leaves a partial tmp file.
defp download_href_to_file(href, api_key, path) do
result =
Req.get(href,
headers: [{"PRIVATE-TOKEN", api_key}],
receive_timeout: 300_000,
into: File.stream!(path)
)
case result do
{:ok, %{status: 200}} ->
:ok
{:ok, %{status: status}} ->
File.rm(path)
{:error, "ERA5 file download HTTP #{status}"}
{:error, reason} ->
File.rm(path)
{:error, "ERA5 file download failed: #{inspect(reason)}"}
end
end
defp build_profile(lat, lon, valid_time, single_grib, pressure_grib) do
# Decode GRIB2 data using existing infrastructure
with {:ok, single_points} <- extract_point(single_grib, lat, lon),

43
lib/microwaveprop/weather/grib2/wgrib2.ex
View file

@ -354,13 +354,44 @@ defmodule Microwaveprop.Weather.Grib2.Wgrib2 do
| {:error, term()}
def extract_grid_messages(grib_binary, match_pattern, grid_spec) do
if available?() do
extract_messages_with_wgrib2(grib_binary, match_pattern, grid_spec)
tmp_grib = Path.join(System.tmp_dir!(), "era5_#{System.unique_integer([:positive])}.grib2")
try do
File.write!(tmp_grib, grib_binary)
extract_messages_with_wgrib2(tmp_grib, match_pattern, grid_spec)
after
File.rm(tmp_grib)
end
else
{:error, :wgrib2_not_available}
end
end
defp extract_messages_with_wgrib2(grib_binary, match_pattern, grid_spec) do
@doc """
Like `extract_grid_messages/3`, but reads GRIB2 from a file on disk instead
of a binary in memory. Saves 50-200 MB of heap per call for month-tile
fetches by avoiding the intermediate `File.write!(tmp, binary)` step.
"""
@spec extract_grid_messages_from_file(Path.t(), String.t(), grid_spec()) ::
{:ok,
[
%{
datetime: DateTime.t() | nil,
var: String.t(),
level: String.t(),
values: %{{float(), float()} => float()}
}
]}
| {:error, term()}
def extract_grid_messages_from_file(grib_path, match_pattern, grid_spec) do
if available?() do
extract_messages_with_wgrib2(grib_path, match_pattern, grid_spec)
else
{:error, :wgrib2_not_available}
end
end
defp extract_messages_with_wgrib2(grib_path, match_pattern, grid_spec) do
%{
lon_start: lon_start,
lon_count: lon_count,
@ -374,13 +405,10 @@ defmodule Microwaveprop.Weather.Grib2.Wgrib2 do
lon_spec = "#{wgrib2_lon_start}:#{lon_count}:#{lon_step}"
lat_spec = "#{lat_start}:#{lat_count}:#{lat_step}"
tmp_grib = Path.join(System.tmp_dir!(), "era5_#{System.unique_integer([:positive])}.grib2")
tmp_bin = tmp_grib <> ".lola.bin"
tmp_bin = grib_path <> ".lola.#{System.unique_integer([:positive])}.bin"
try do
File.write!(tmp_grib, grib_binary)
args = [tmp_grib, "-match", match_pattern, "-lola", lon_spec, lat_spec, tmp_bin, "bin"]
args = [grib_path, "-match", match_pattern, "-lola", lon_spec, lat_spec, tmp_bin, "bin"]
case System.cmd(wgrib2_path(), args, stderr_to_stdout: true) do
{output, 0} ->
@ -401,7 +429,6 @@ defmodule Microwaveprop.Weather.Grib2.Wgrib2 do
{:error, "wgrib2 failed (exit #{exit_code}): #{String.slice(output, 0, 200)}"}
end
after
File.rm(tmp_grib)
File.rm(tmp_bin)
end
end

17
priv/repo/migrations/20260413205533_add_era5_profiles_valid_time_lat_lon_index.exs Normal file
View file

@ -0,0 +1,17 @@
defmodule Microwaveprop.Repo.Migrations.AddEra5ProfilesValidTimeLatLonIndex do
use Ecto.Migration
# Matches the hot range-scan access pattern used by both
# `Weather.find_nearest_era5/3` (contact detail + path profiles) and
# `StatusLive.count_era5_done/0` (status page EXISTS correlated subquery):
# WHERE valid_time BETWEEN ? AND ?
# AND lat BETWEEN ? AND ?
# AND lon BETWEEN ? AND ?
# `valid_time` leads because it collapses the search to a ~1-hour bucket
# first; lat/lon then filter the ~9×9 tile grid within that bucket.
# The existing unique index (lat, lon, valid_time) handles upsert dedup
# but can't serve three-range scans efficiently.
def change do
create index(:era5_profiles, [:valid_time, :lat, :lon])
end
end