Skip to content

Commit

Permalink
Automate creation of month-averaged datasets
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
douglatornell committed Mar 4, 2024
1 parent edd3f80 commit 757e61c
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 1 deletion.
35 changes: 34 additions & 1 deletion nowcast/next_workers.py
Original file line number Diff line number Diff line change
Expand Up @@ -1548,7 +1548,40 @@ def after_make_averaged_dataset(msg, config, checklist):
:returns: Worker(s) to launch next
:rtype: list
"""
return []
next_workers = {
"crash": [],
"failure day biology": [],
"failure day chemistry": [],
"failure day physics": [],
"failure month biology": [],
"failure month chemistry": [],
"failure month physics": [],
"success day biology": [],
"success day chemistry": [],
"success day physics": [],
"success month biology": [],
"success month chemistry": [],
"success month physics": [],
}
if msg.type.startswith("success day"):
*_, reshapr_var_group = msg.type.split()
run_date = arrow.get(msg.payload[f"day {reshapr_var_group}"]["run date"])
if run_date.shift(days=+1).day == 1:
first_of_month = run_date.format("YYYY-MM-01")
next_workers[msg.type].append(
NextWorker(
"nowcast.workers.make_averaged_dataset",
args=[
"skookum",
"month",
reshapr_var_group,
"--run-date",
first_of_month,
],
host="localhost",
)
)
return next_workers[msg.type]


def after_archive_tarball(msg, config, checklist):
Expand Down
30 changes: 30 additions & 0 deletions tests/test_next_workers.py
Original file line number Diff line number Diff line change
Expand Up @@ -2221,6 +2221,36 @@ def test_no_next_worker_msg_types(self, msg_type, config, checklist):
)
assert workers == []

@pytest.mark.parametrize(
"msg_type",
[
"success day biology",
"success day chemistry",
"success day physics",
],
)
def test_month_end_day_success_launch_month_average(
self, msg_type, config, checklist
):
*_, reshapr_var_group = msg_type.split()
msg = Message(
"make_averaged_dataset",
msg_type,
payload={
f"day {reshapr_var_group}": {
"run date": "2024-02-29",
"file path": "SalishSea_1d_20240301_20240301_biol_T.nc",
}
},
)
workers = next_workers.after_make_averaged_dataset(msg, config, checklist)
expected = NextWorker(
"nowcast.workers.make_averaged_dataset",
args=["skookum", "month", reshapr_var_group, "--run-date", "2024-02-01"],
host="localhost",
)
assert expected in workers


class TestAfterArchiveTarball:
"""Unit tests for the after_archive_tarball function."""
Expand Down

0 comments on commit 757e61c

Please sign in to comment.