Coverage for src/CSET/cset_workflow/app/fetch_fcst/bin/fetch_data.py: 78%
117 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-04 08:32 +0000
1# © Crown copyright, Met Office (2022-2025) and CSET contributors.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
15"""Retrieve the files from the filesystem for the current cycle point."""
17import abc
18import ast
19import glob
20import itertools
21import logging
22import os
23import ssl
24import sys
25import urllib.parse
26import urllib.request
27from concurrent.futures import ThreadPoolExecutor
28from datetime import datetime, timedelta
29from pathlib import Path
30from typing import Literal, Self
32import isodate
34logging.basicConfig(
35 level=os.getenv("LOGLEVEL", "INFO"),
36 format="%(asctime)s %(levelname)s %(message)s",
37 stream=sys.stdout,
38)
39logger = logging.getLogger(__name__)
42class FileRetrieverABC(abc.ABC):
43 """Abstract base class for retrieving files from a data source.
45 The `get_file` method must be defined. Optionally the __enter__ and __exit__
46 methods maybe be overridden to add setup or cleanup code.
48 The class is designed to be used as a context manager, so that resources can
49 be cleaned up after the retrieval is complete. All the files of a model are
50 retrieved within a single context manager block, within which the `get_file`
51 method is called for each file path.
52 """
54 def __enter__(self) -> Self:
55 """Initialise the file retriever."""
56 logger.debug("Initialising FileRetriever.")
57 return self
59 def __exit__(self, exc_type, exc_value, traceback):
60 """Clean up the file retriever."""
61 logger.debug("Tearing down FileRetriever.")
63 @abc.abstractmethod
64 def get_file(self, file_path: str, output_dir: str) -> bool: # pragma: no cover
65 """Save a file from the data source to the output directory.
67 Not all of the given paths will exist, so FileNotFoundErrors should be
68 logged, but not raised.
70 Implementations should be thread safe, as the method is called from
71 multiple threads.
73 Parameters
74 ----------
75 file_path: str
76 Path of the file to copy on the data source. It may contain patterns
77 like globs, which will be expanded in a system specific manner.
78 output_dir: str
79 Path to filesystem directory into which the file should be copied.
81 Returns
82 -------
83 bool:
84 True if files were transferred, otherwise False.
85 """
86 raise NotImplementedError
89class FilesystemFileRetriever(FileRetrieverABC):
90 """Retrieve files from the filesystem."""
92 def get_file(self, file_path: str, output_dir: str) -> bool:
93 """Save a file from the filesystem to the output directory.
95 Parameters
96 ----------
97 file_path: str
98 Path of the file to copy on the filesystem. It may contain patterns
99 like globs, which will be expanded in a system specific manner.
100 output_dir: str
101 Path to filesystem directory into which the file should be copied.
103 Returns
104 -------
105 bool:
106 True if files were transferred, otherwise False.
107 """
108 file_paths = glob.glob(os.path.expanduser(file_path))
109 logger.debug("Copying files:\n%s", "\n".join(file_paths))
110 if not file_paths:
111 logger.warning("file_path does not match any files: %s", file_path)
112 any_files_copied = False
113 for f in file_paths:
114 file = Path(f).absolute()
115 try:
116 # Save to a filename derived from the full path, to
117 # differentiate similarly named files from different
118 # directories.
119 # `}` replaces `/` as it can be in file names.
120 os.symlink(file, f"{output_dir}/{'}'.join(file.parts)}")
121 any_files_copied = True
122 except OSError as err:
123 logger.warning("Failed to copy %s, error: %s", file, err)
124 return any_files_copied
127class HTTPFileRetriever(FileRetrieverABC):
128 """Retrieve files via HTTP."""
130 def get_file(self, file_path: str, output_dir: str) -> bool:
131 """Save a file from a HTTP address to the output directory.
133 Parameters
134 ----------
135 file_path: str
136 Path of the file to copy on MASS. It may contain patterns like
137 globs, which will be expanded in a system specific manner.
138 output_dir: str
139 Path to filesystem directory into which the file should be copied.
141 Returns
142 -------
143 bool:
144 True if files were transferred, otherwise False.
145 """
146 ctx = ssl.create_default_context()
147 # Needed to enable compatibility with malformed iBoss TLS certificates.
148 ctx.verify_flags &= ~ssl.VERIFY_X509_STRICT
149 save_path = (
150 f"{output_dir.removesuffix('/')}/"
151 + urllib.parse.urlparse(file_path).path.split("/")[-1]
152 )
153 any_files_copied = False
154 try:
155 with urllib.request.urlopen(file_path, timeout=30, context=ctx) as response:
156 with open(save_path, "wb") as fp:
157 # Read in 1 MiB chunks so data needn't fit in memory.
158 while data := response.read(1024 * 1024):
159 fp.write(data)
160 any_files_copied = True
161 except OSError as err:
162 logger.warning("Failed to retrieve %s, error: %s", file_path, err)
163 return any_files_copied
166def _get_needed_environment_variables() -> dict:
167 """Load the needed variables from the environment."""
168 variables = {
169 "raw_path": os.environ["DATA_PATH"],
170 "date_type": os.environ["DATE_TYPE"],
171 "data_time": datetime.fromisoformat(os.environ["CYLC_TASK_CYCLE_POINT"]),
172 "forecast_length": isodate.parse_duration(os.environ["ANALYSIS_LENGTH"]),
173 "forecast_offset": isodate.parse_duration(os.environ["ANALYSIS_OFFSET"]),
174 "model_identifier": os.environ["MODEL_IDENTIFIER"],
175 "rose_datac": os.environ["ROSE_DATAC"],
176 }
177 try:
178 variables["data_period"] = isodate.parse_duration(os.environ["DATA_PERIOD"])
179 except KeyError:
180 # Data period is not needed for initiation time.
181 if variables["date_type"] != "initiation":
182 raise
183 variables["data_period"] = None
184 logger.debug("Environment variables loaded: %s", variables)
185 return variables
188def _get_needed_environment_variables_obs() -> dict:
189 """Load the needed variables from the environment."""
190 variables = {
191 "subtype": ast.literal_eval(os.environ.get("OBS_SUBTYPE")),
192 "data_time": datetime.fromisoformat(os.environ["CYLC_TASK_CYCLE_POINT"]),
193 "forecast_length": isodate.parse_duration(os.environ["ANALYSIS_LENGTH"]),
194 "obs_fields": ast.literal_eval(os.environ["SURFACE_SYNOP_FIELDS"]),
195 "model_identifier": "OBS",
196 "wmo_nmbrs": ast.literal_eval(os.environ.get("WMO_BLOCK_STTN_NMBRS"))
197 if len(os.environ.get("WMO_BLOCK_STTN_NMBRS")) > 0
198 else None,
199 "subarea_extent": ast.literal_eval(os.environ.get("SUBAREA_EXTENT"))
200 if len(os.environ.get("SUBAREA_EXTENT")) > 0
201 else None,
202 "obs_interval": isodate.parse_duration(os.environ["SURFACE_SYNOP_INTERVAL"]),
203 "obs_offset": isodate.parse_duration(os.environ["SURFACE_SYNOP_OFFSET"]),
204 "rose_datac": os.environ["ROSE_DATAC"],
205 }
206 logger.debug("Environment variables loaded: %s", variables)
207 return variables
210def _template_file_path(
211 raw_path: str,
212 date_type: Literal["validity", "initiation"],
213 data_time: datetime,
214 forecast_length: timedelta,
215 forecast_offset: timedelta,
216 data_period: timedelta,
217) -> list[str]:
218 """Fill time placeholders to generate a file path to fetch."""
219 placeholder_times: list[datetime] = []
220 lead_times: list[timedelta] = []
221 match date_type:
222 case "validity":
223 date = data_time
224 while date < data_time + forecast_length:
225 placeholder_times.append(date)
226 date += data_period
227 case "initiation":
228 placeholder_times.append(data_time)
229 lead_time = forecast_offset
230 while lead_time < forecast_length:
231 lead_times.append(lead_time)
232 lead_time += data_period
233 case _:
234 raise ValueError(f"Invalid date type: {date_type}")
236 paths: set[str] = set()
237 for placeholder_time in placeholder_times:
238 # Expand out all other format strings.
239 path = placeholder_time.strftime(os.path.expandvars(raw_path))
240 if lead_times:
241 # Expand out lead time format strings, %N.
242 for lead_time in lead_times:
243 # BUG: Will not respect escaped % signs, e.g: %%N.
244 paths.add(
245 path.replace("%N", f"{int(lead_time.total_seconds()) // 3600:03d}")
246 )
247 else:
248 paths.add(path)
249 return sorted(paths)
252def fetch_data(file_retriever: FileRetrieverABC):
253 """Fetch the data for a model.
255 The following environment variables need to be set:
256 * ANALYSIS_OFFSET
257 * ANALYSIS_LENGTH
258 * CYLC_TASK_CYCLE_POINT
259 * DATA_PATH
260 * DATA_PERIOD
261 * DATE_TYPE
262 * MODEL_IDENTIFIER
263 * ROSE_DATAC
265 Parameters
266 ----------
267 file_retriever: FileRetriever
268 FileRetriever implementation to use.
270 Raises
271 ------
272 FileNotFound:
273 If no files are found for the model, across all tried paths.
274 """
275 v = _get_needed_environment_variables()
277 # Prepare output directory.
278 cycle_data_dir = f"{v['rose_datac']}/data/{v['model_identifier']}"
279 os.makedirs(cycle_data_dir, exist_ok=True)
280 logger.debug("Output directory: %s", cycle_data_dir)
282 # Get file paths.
283 paths = _template_file_path(
284 v["raw_path"],
285 v["date_type"],
286 v["data_time"],
287 v["forecast_length"],
288 v["forecast_offset"],
289 v["data_period"],
290 )
291 logger.info("Retrieving paths:\n%s", "\n".join(paths))
293 # Use file retriever to transfer data with multiple threads.
294 with file_retriever() as retriever, ThreadPoolExecutor() as executor:
295 files_found = executor.map(
296 retriever.get_file, paths, itertools.repeat(cycle_data_dir)
297 )
298 # Exhaust the iterator with list so all futures get resolved before we
299 # exit the with block, ensuring all files are retrieved.
300 any_files_found = any(list(files_found))
301 if not any_files_found:
302 raise FileNotFoundError("No files found for model!")
305def fetch_obs(obs_retriever: FileRetrieverABC):
306 """Fetch the observations corresponding to a model run.
308 The following environment variables need to be set:
309 * ANALYSIS_OFFSET
310 * ANALYSIS_LENGTH
311 * CYLC_TASK_CYCLE_POINT
312 * DATA_PATH
313 * DATA_PERIOD
314 * DATE_TYPE
315 * MODEL_IDENTIFIER
316 * ROSE_DATAC
318 Parameters
319 ----------
320 obs_retriever: ObsRetriever
321 ObsRetriever implementation to use. Defaults to FilesystemFileRetriever.
323 Raises
324 ------
325 FileNotFound:
326 If no observations are available.
327 """
328 v = _get_needed_environment_variables_obs()
330 # Prepare output directory.
331 cycle_obs_dir = f"{v['rose_datac']}/data/OBS"
332 os.makedirs(cycle_obs_dir, exist_ok=True)
333 logger.debug("Output directory: %s", cycle_obs_dir)
335 # Loop over requested obs subtypes.
336 for subtype in v["subtype"]:
337 obs_base_path = (
338 subtype
339 + "_"
340 + "%Y%m%dT%H%MZ_dt_"
341 + str(int(v["forecast_length"].total_seconds() // 3600)).zfill(3)
342 + ".nc"
343 )
344 paths = _template_file_path(
345 obs_base_path,
346 "initiation",
347 v["data_time"],
348 v["forecast_length"],
349 timedelta(seconds=0),
350 v["obs_interval"],
351 )
352 logger.info("Retrieving paths:\n%s", "\n".join(paths))
354 # Use obs retriever to transfer data with multiple threads.
355 # We shouldn't need to iterate as we do for the forecast data
356 # because these files will be smaller.
357 try:
358 obs_retriever.get_file(
359 paths[0],
360 subtype,
361 v["obs_fields"],
362 v["data_time"],
363 v["obs_offset"],
364 v["forecast_length"],
365 v["obs_interval"],
366 cycle_obs_dir,
367 wmo_nmbrs=v["wmo_nmbrs"],
368 subarea_extent=v["subarea_extent"],
369 )
370 except Exception as exc:
371 raise ValueError("No observations available.") from exc