Skip to content

Commit 757e61c

Browse files
committed
Automate creation of month-averaged datasets
Add launch of make_averaged_dataset worker for month-averaged datasets for var_groups="biology", "chemistry", "physics" for "nowcast-green" after completion of day-averaged dataset creation on the last day of the month. Also added test cases for the functionality.
1 parent edd3f80 commit 757e61c

File tree

2 files changed

+64
-1
lines changed

2 files changed

+64
-1
lines changed

nowcast/next_workers.py

+34-1
Original file line numberDiff line numberDiff line change
@@ -1548,7 +1548,40 @@ def after_make_averaged_dataset(msg, config, checklist):
15481548
:returns: Worker(s) to launch next
15491549
:rtype: list
15501550
"""
1551-
return []
1551+
next_workers = {
1552+
"crash": [],
1553+
"failure day biology": [],
1554+
"failure day chemistry": [],
1555+
"failure day physics": [],
1556+
"failure month biology": [],
1557+
"failure month chemistry": [],
1558+
"failure month physics": [],
1559+
"success day biology": [],
1560+
"success day chemistry": [],
1561+
"success day physics": [],
1562+
"success month biology": [],
1563+
"success month chemistry": [],
1564+
"success month physics": [],
1565+
}
1566+
if msg.type.startswith("success day"):
1567+
*_, reshapr_var_group = msg.type.split()
1568+
run_date = arrow.get(msg.payload[f"day {reshapr_var_group}"]["run date"])
1569+
if run_date.shift(days=+1).day == 1:
1570+
first_of_month = run_date.format("YYYY-MM-01")
1571+
next_workers[msg.type].append(
1572+
NextWorker(
1573+
"nowcast.workers.make_averaged_dataset",
1574+
args=[
1575+
"skookum",
1576+
"month",
1577+
reshapr_var_group,
1578+
"--run-date",
1579+
first_of_month,
1580+
],
1581+
host="localhost",
1582+
)
1583+
)
1584+
return next_workers[msg.type]
15521585

15531586

15541587
def after_archive_tarball(msg, config, checklist):

tests/test_next_workers.py

+30
Original file line numberDiff line numberDiff line change
@@ -2221,6 +2221,36 @@ def test_no_next_worker_msg_types(self, msg_type, config, checklist):
22212221
)
22222222
assert workers == []
22232223

2224+
@pytest.mark.parametrize(
2225+
"msg_type",
2226+
[
2227+
"success day biology",
2228+
"success day chemistry",
2229+
"success day physics",
2230+
],
2231+
)
2232+
def test_month_end_day_success_launch_month_average(
2233+
self, msg_type, config, checklist
2234+
):
2235+
*_, reshapr_var_group = msg_type.split()
2236+
msg = Message(
2237+
"make_averaged_dataset",
2238+
msg_type,
2239+
payload={
2240+
f"day {reshapr_var_group}": {
2241+
"run date": "2024-02-29",
2242+
"file path": "SalishSea_1d_20240301_20240301_biol_T.nc",
2243+
}
2244+
},
2245+
)
2246+
workers = next_workers.after_make_averaged_dataset(msg, config, checklist)
2247+
expected = NextWorker(
2248+
"nowcast.workers.make_averaged_dataset",
2249+
args=["skookum", "month", reshapr_var_group, "--run-date", "2024-02-01"],
2250+
host="localhost",
2251+
)
2252+
assert expected in workers
2253+
22242254

22252255
class TestAfterArchiveTarball:
22262256
"""Unit tests for the after_archive_tarball function."""

0 commit comments

Comments
 (0)