Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Refactory dynamic chunking to plugin #62

Merged
merged 11 commits into from
Nov 15, 2023
Merged

Refactory dynamic chunking to plugin #62

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
2503d74
Refactory dynamic chunking to plugin
jbusecke Nov 15, 2023
12cc429
Update requirements.txt
jbusecke Nov 15, 2023
a5b6881
Update recipe.py
jbusecke Nov 15, 2023
0604b59
Update recipe.py
jbusecke Nov 15, 2023
6e9b65d
fix indent and add warnings
jbusecke Nov 15, 2023
69348fd
Fix output type for chunking func
jbusecke Nov 15, 2023
4c4ad14
reduce submitted iids
jbusecke Nov 15, 2023
a81423f
Update recipe.py
jbusecke Nov 15, 2023
907a771
Improved error handling
jbusecke Nov 15, 2023
5481c85
rebase on recipes 0.10.4
jbusecke Nov 15, 2023
77b033e
Update recipe.py
jbusecke Nov 15, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
148 changes: 96 additions & 52 deletions feedstock/recipe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import apache_beam as beam
from dataclasses import dataclass
from datetime import datetime
from typing import List
from typing import List, Dict
from pangeo_forge_esgf import get_urls_from_esgf, setup_logging
from pangeo_forge_esgf.parsing import parse_instance_ids
from pangeo_forge_recipes.patterns import pattern_from_file_sequence
Expand All @@ -14,6 +14,7 @@
import asyncio
import xarray as xr
import zarr
import warnings

# setup_logging('DEBUG')
setup_logging('INFO')
Expand Down Expand Up @@ -276,52 +277,54 @@ def expand(self, pcoll: beam.PCollection) -> beam.PCollection:
)

iids_raw = [
# from https://github.com/Timh37/CMIP6cex/issues/2
'CMIP6.*.*.*.historical.*.day.psl.*.*',
'CMIP6.*.*.*.historical.*.day.sfcWind.*.*',
'CMIP6.*.*.*.historical.*.day.pr.*.*',
'CMIP6.*.*.*.ssp245.*.day.psl.*.*',
'CMIP6.*.*.*.ssp245.*.day.sfcWind.*.*',
'CMIP6.*.*.*.ssp245.*.day.pr.*.*',
'CMIP6.*.*.*.ssp585.*.day.psl.*.*',
'CMIP6.*.*.*.ssp585.*.day.sfcWind.*.*',
'CMIP6.*.*.*.ssp585.*.day.pr.*.*',
# from https://github.com/pangeo-forge/cmip6-feedstock/issues/22
'CMIP6.*.*.*.historical.*.Omon.zmeso.*.*'
'CMIP6.*.*.*.ssp126.*.Omon.zmeso.*.*',
'CMIP6.*.*.*.ssp245.*.Omon.zmeso.*.*',
'CMIP6.*.*.*.ssp585.*.Omon.zmeso.*.*',
# PMIP velocities
'CMIP6.PMIP.MIROC.MIROC-ES2L.lgm.r1i1p1f2.Omon.uo.gn.v20191002',
'CMIP6.PMIP.AWI.AWI-ESM-1-1-LR.lgm.r1i1p1f1.Odec.uo.gn.v20200212',
'CMIP6.PMIP.AWI.AWI-ESM-1-1-LR.lgm.r1i1p1f1.Omon.uo.gn.v20200212',
'CMIP6.PMIP.MIROC.MIROC-ES2L.lgm.r1i1p1f2.Omon.uo.gr1.v20200911',
'CMIP6.PMIP.MPI-M.MPI-ESM1-2-LR.lgm.r1i1p1f1.Omon.uo.gn.v20200909',
'CMIP6.PMIP.AWI.AWI-ESM-1-1-LR.lgm.r1i1p1f1.Omon.vo.gn.v20200212',
'CMIP6.PMIP.MIROC.MIROC-ES2L.lgm.r1i1p1f2.Omon.vo.gn.v20191002',
'CMIP6.PMIP.AWI.AWI-ESM-1-1-LR.lgm.r1i1p1f1.Odec.vo.gn.v20200212',
'CMIP6.PMIP.MIROC.MIROC-ES2L.lgm.r1i1p1f2.Omon.vo.gr1.v20200911',
'CMIP6.PMIP.MPI-M.MPI-ESM1-2-LR.lgm.r1i1p1f1.Omon.vo.gn.v20190710',
# from https://github.com/leap-stc/cmip6-leap-feedstock/issues/41
'CMIP6.*.*.*.historical.*.SImon.sifb.*.*',
'CMIP6.*.*.*.ssp126.*.SImon.sifb.*.*',
'CMIP6.*.*.*.ssp245.*.SImon.sifb.*.*',
'CMIP6.*.*.*.ssp585.*.SImon.sifb.*.*',
'CMIP6.*.*.*.historical.*.SImon.siitdthick.*.*',
'CMIP6.*.*.*.ssp126.*.SImon.siitdthick.*.*',
'CMIP6.*.*.*.ssp245.*.SImon.siitdthick.*.*',
'CMIP6.*.*.*.ssp585.*.SImon.siitdthick.*.*',
# for CMIP6 pco2 testbed
"CMIP6.*.*.*.historical.*.Omon.spco2.*.*",
"CMIP6.*.*.*.historical.*.Omon.tos.*.*",
"CMIP6.*.*.*.historical.*.Omon.sos.*.*",
"CMIP6.*.*.*.historical.*.Omon.chl.*.*",
"CMIP6.*.*.*.historical.*.Omon.mlotst.*.*",
"CMIP6.*.*.*.ssp245.*.Omon.spco2.*.*",
"CMIP6.*.*.*.ssp245.*.Omon.tos.*.*",
"CMIP6.*.*.*.ssp245.*.Omon.sos.*.*",
"CMIP6.*.*.*.ssp245.*.Omon.chl.*.*",
"CMIP6.*.*.*.ssp245.*.Omon.mlotst.*.*",
# to test the latest PR
"CMIP6.*.*.*.ssp585.*.Omon.o2.*.*",
# # from https://github.com/Timh37/CMIP6cex/issues/2
# 'CMIP6.*.*.*.historical.*.day.psl.*.*',
# 'CMIP6.*.*.*.historical.*.day.sfcWind.*.*',
# 'CMIP6.*.*.*.historical.*.day.pr.*.*',
# 'CMIP6.*.*.*.ssp245.*.day.psl.*.*',
# 'CMIP6.*.*.*.ssp245.*.day.sfcWind.*.*',
# 'CMIP6.*.*.*.ssp245.*.day.pr.*.*',
# 'CMIP6.*.*.*.ssp585.*.day.psl.*.*',
# 'CMIP6.*.*.*.ssp585.*.day.sfcWind.*.*',
# 'CMIP6.*.*.*.ssp585.*.day.pr.*.*',
# # from https://github.com/pangeo-forge/cmip6-feedstock/issues/22
# 'CMIP6.*.*.*.historical.*.Omon.zmeso.*.*'
# 'CMIP6.*.*.*.ssp126.*.Omon.zmeso.*.*',
# 'CMIP6.*.*.*.ssp245.*.Omon.zmeso.*.*',
# 'CMIP6.*.*.*.ssp585.*.Omon.zmeso.*.*',
# # PMIP velocities
# 'CMIP6.PMIP.MIROC.MIROC-ES2L.lgm.r1i1p1f2.Omon.uo.gn.v20191002',
# 'CMIP6.PMIP.AWI.AWI-ESM-1-1-LR.lgm.r1i1p1f1.Odec.uo.gn.v20200212',
# 'CMIP6.PMIP.AWI.AWI-ESM-1-1-LR.lgm.r1i1p1f1.Omon.uo.gn.v20200212',
# 'CMIP6.PMIP.MIROC.MIROC-ES2L.lgm.r1i1p1f2.Omon.uo.gr1.v20200911',
# 'CMIP6.PMIP.MPI-M.MPI-ESM1-2-LR.lgm.r1i1p1f1.Omon.uo.gn.v20200909',
# 'CMIP6.PMIP.AWI.AWI-ESM-1-1-LR.lgm.r1i1p1f1.Omon.vo.gn.v20200212',
# 'CMIP6.PMIP.MIROC.MIROC-ES2L.lgm.r1i1p1f2.Omon.vo.gn.v20191002',
# 'CMIP6.PMIP.AWI.AWI-ESM-1-1-LR.lgm.r1i1p1f1.Odec.vo.gn.v20200212',
# 'CMIP6.PMIP.MIROC.MIROC-ES2L.lgm.r1i1p1f2.Omon.vo.gr1.v20200911',
# 'CMIP6.PMIP.MPI-M.MPI-ESM1-2-LR.lgm.r1i1p1f1.Omon.vo.gn.v20190710',
# # from https://github.com/leap-stc/cmip6-leap-feedstock/issues/41
# 'CMIP6.*.*.*.historical.*.SImon.sifb.*.*',
# 'CMIP6.*.*.*.ssp126.*.SImon.sifb.*.*',
# 'CMIP6.*.*.*.ssp245.*.SImon.sifb.*.*',
# 'CMIP6.*.*.*.ssp585.*.SImon.sifb.*.*',
# 'CMIP6.*.*.*.historical.*.SImon.siitdthick.*.*',
# 'CMIP6.*.*.*.ssp126.*.SImon.siitdthick.*.*',
# 'CMIP6.*.*.*.ssp245.*.SImon.siitdthick.*.*',
# 'CMIP6.*.*.*.ssp585.*.SImon.siitdthick.*.*',
# # for CMIP6 pco2 testbed
# "CMIP6.*.*.*.historical.*.Omon.spco2.*.*",
# "CMIP6.*.*.*.historical.*.Omon.tos.*.*",
# "CMIP6.*.*.*.historical.*.Omon.sos.*.*",
# "CMIP6.*.*.*.historical.*.Omon.chl.*.*",
# "CMIP6.*.*.*.historical.*.Omon.mlotst.*.*",
# "CMIP6.*.*.*.ssp245.*.Omon.spco2.*.*",
# "CMIP6.*.*.*.ssp245.*.Omon.tos.*.*",
# "CMIP6.*.*.*.ssp245.*.Omon.sos.*.*",
# "CMIP6.*.*.*.ssp245.*.Omon.chl.*.*",
# "CMIP6.*.*.*.ssp245.*.Omon.mlotst.*.*",
]

def parse_wildcards(iids:List[str]) -> List[str]:
Expand Down Expand Up @@ -394,8 +397,52 @@ def parse_wildcards(iids:List[str]) -> List[str]:
# Print the actual urls
print(url_dict)

## Dynamic Chunking Wrapper
from dynamic_chunks.algorithms import even_divisor_algo, iterative_ratio_increase_algo, NoMatchingChunks


def dynamic_chunking_func(ds: xr.Dataset) -> Dict[str, int]:

target_chunk_size='150MB'
target_chunks_aspect_ratio = {'time': 1}
size_tolerance=0.5

try:
target_chunks = even_divisor_algo(
ds,
target_chunk_size,
target_chunks_aspect_ratio,
size_tolerance,
)

except NoMatchingChunks:
warnings.warn(
"Primary algorithm using even divisors along each dimension failed "
f"with {e}. Trying secondary algorithm."
)
try:
target_chunks = iterative_ratio_increase_algo(
ds,
target_chunk_size,
target_chunks_aspect_ratio,
size_tolerance,
)
except NoMatchingChunks:
raise ValueError(
(
"Could not find any chunk combinations satisfying "
"the size constraint with either algorithm."
)
)
# If something fails
except Exception as e:
raise e
except Exception as e:
raise e

return target_chunks

## Create the recipes
target_chunks_aspect_ratio = {'time': 1}
recipes = {}

for iid, urls in url_dict.items():
Expand All @@ -412,10 +459,7 @@ def parse_wildcards(iids:List[str]) -> List[str]:
| StoreToZarr(
store_name=f"{iid}.zarr",
combine_dims=pattern.combine_dim_keys,
target_chunk_size='150MB',
target_chunks_aspect_ratio = target_chunks_aspect_ratio,
size_tolerance=0.5,
allow_fallback_algo=True,
dynamic_chunking_fn=dynamic_chunking_func,
)
| "Logging to non-QC table" >> LogToBigQuery(iid=iid, table_id=table_id_nonqc)
| TestDataset(iid=iid)
Expand Down
7 changes: 3 additions & 4 deletions feedstock/requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# git+https://github.com/jbusecke/pangeo-forge-recipes.git@dynamic_chunks_2
# git+https://github.com/pangeo-forge/pangeo-forge-recipes.git@main#egg=pangeo_forge_recipes #testing first without dyn chunking
git+https://github.com/jbusecke/pangeo-forge-recipes.git@dynamic_chunks_2#egg=pangeo_forge_recipes
git+https://github.com/jbusecke/pangeo-forge-esgf.git@beam-refactor#egg=pangeo_forge_esgf
git+https://github.com/jbusecke/pangeo-forge-esgf.git@beam-refactor#egg=pangeo_forge_esgf
pangeo-forge-recipes==0.10.4
dynamic-chunks==0.0.2