diff --git a/nwp/jobs.py b/nwp/jobs.py index c585889..3e52189 100644 --- a/nwp/jobs.py +++ b/nwp/jobs.py @@ -84,19 +84,29 @@ class NWPConsumerDagDefinition: """A class to define the NWPConsumerDagDefinition.""" def __init__( - self, area: str, source: str, storage_path: str | None = None + self, source: str, storage_path: str | None = None, env_overrides: dict[str, str] | None = None ) -> "NWPConsumerDagDefinition": """Create a NWPConsumerDagDefinition.""" - self.area = area self.source = source + area = env_overrides.get("ECMWF_AREA", "no-area") self.storage_path = \ storage_path or \ - f'/mnt/storage_b/data/ocf/solar_pv_nowcasting/nowcasting_dataset_pipeline/NWP/ECMWF/{self.area}' + f'/mnt/storage_b/data/ocf/solar_pv_nowcasting/nowcasting_dataset_pipeline/NWP/ECMWF/{area}' + self.env_overrides = env_overrides nwp_consumer_jobs: dict[str, NWPConsumerDagDefinition] = { - "uk": NWPConsumerDagDefinition(area="uk", source="ecmwf-mars"), - "india": NWPConsumerDagDefinition(area="nw-india", source="ecmwf-mars"), - "malta": NWPConsumerDagDefinition(area="malta", source="ecmwf-mars") + "uk": NWPConsumerDagDefinition( + source="ecmwf-mars", + env_overrides={"ECMWF_AREA": "uk"} + ), + "india": NWPConsumerDagDefinition( + source="ecmwf-mars", + env_overrides={"ECMWF_AREA": "nw-india", "ECMWF_HOURS": "84"} + ), + "malta": NWPConsumerDagDefinition( + source="ecmwf-mars", + env_overrides={"ECMWF_AREA": "malta"} + ), } @@ -110,7 +120,7 @@ def partitioned_config_func(partition_key: str) -> dict[str, Any]: date_from=time_window.start.strftime("%Y-%m-%d"), date_to=time_window.start.strftime("%Y-%m-%d"), source=dagdef.source, - env_overrides={"ECMWF_AREA": dagdef.area}, + env_overrides=dagdef.env_overrides, zarr_dir=f"{dagdef.storage_path}/zarr", raw_dir=f"{dagdef.storage_path}/raw", ) diff --git a/pyproject.toml b/pyproject.toml index 2dbdfa9..009426b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ dependencies = [ "dagster-docker == 0.20.11", "huggingface-hub == 0.16.4", "numpy >= 1.23.0", - "nwp-consumer == 0.1.21", + "nwp-consumer == 0.1.23", "ocf-blosc2 == 0.0.3", "pathlib == 1.0.1", "requests >= 2.28.0",