diff --git a/VERSION.txt b/VERSION.txt index ee1372d..7179039 100644 --- a/VERSION.txt +++ b/VERSION.txt @@ -1 +1 @@ -0.2.2 +0.2.3 diff --git a/weather_sp/setup.py b/weather_sp/setup.py index 6241643..ddba566 100644 --- a/weather_sp/setup.py +++ b/weather_sp/setup.py @@ -44,7 +44,7 @@ packages=find_packages(), author='Anthromets', author_email='anthromets-ecmwf@google.com', - version='0.3.10', + version='0.3.11', url='https://weather-tools.readthedocs.io/en/latest/weather_sp/', description='A tool to split weather data files into per-variable files.', install_requires=beam_gcp_requirements + base_requirements, diff --git a/weather_sp/splitter_pipeline/file_splitters.py b/weather_sp/splitter_pipeline/file_splitters.py index a6a57db..8370b15 100644 --- a/weather_sp/splitter_pipeline/file_splitters.py +++ b/weather_sp/splitter_pipeline/file_splitters.py @@ -162,17 +162,12 @@ def split_data(self) -> None: if not self.output_info.split_dims(): raise ValueError('No splitting specified in template.') - if self.should_skip(): - metrics.Metrics.counter('file_splitters', 'skipped').inc() - self.logger.info('Skipping %s, file already split.', - repr(self.input_path)) - return - # Here, we keep a map of open file objects (`outputs`). We need these since # each output grib file (named `key`) will include multiple `grb` messages # each. By writing data to the cache of open file objects, we can keep a # minimal amount of data in memory at a time. outputs = dict() + skipped_keys = set() with self._open_grib_locally() as grbs: self.logger.info('Splitting & uploading %r...', self.input_path) try: @@ -188,10 +183,20 @@ def split_data(self) -> None: 'Variable not found in grib: %s', dim) key = self.output_info.formatted_output_path(splits) + if key in skipped_keys: + del grb + continue + # Append the current grib message to a set number of output files. # If the target shard doesn't exist, create it. if key not in outputs: + if self.should_skip_file(key): + self.logger.info('Skipping %s, file already split.', repr(key)) + skipped_keys.add(key) + del grb + continue outputs[key] = FileSystems.create(key) + outputs[key].write(grb.tostring()) outputs[key].flush() @@ -202,6 +207,12 @@ def split_data(self) -> None: finally: for out in outputs.values(): out.close() + + if not outputs: + metrics.Metrics.counter('file_splitters', 'skipped').inc() + self.logger.info('Skipping %s, file already split.', + repr(self.input_path)) + else: self.logger.info('Split %s into %d files', self.input_path, len(outputs)) @@ -276,6 +287,7 @@ def split_data(self) -> None: splits = dict(zip(split_dims, line.split(' '))) output_path = self.output_info.formatted_output_path(splits) if self.should_skip_file(output_path): + self.logger.info('Skipping %s, file already split.', repr(output_path)) skipped_paths.append(output_path) continue output_paths.append(output_path) @@ -285,6 +297,7 @@ def split_data(self) -> None: repr(self.input_path), ', '.join(skipped_paths)) return + output_paths_set = set(output_paths) with tempfile.TemporaryDirectory() as tmpdir: self.logger.info('Performing split.') dest = os.path.join(tmpdir, flat_output_template) @@ -296,9 +309,17 @@ def split_data(self) -> None: subprocess.run([grib_copy_cmd, local_file.name, dest], check=True) - files = os.listdir(tmpdir) - num_files = len(files) - for f in files: + # grib_copy generates a file for every message, including ones + # that were already split. Drop those so copy_dir won't re-upload + # them (issue #538), then flatten the DELIMITER-encoded names back + # into their nested directory structure for the recursive copy. + num_files = 0 + for f in os.listdir(tmpdir): + dest_file_path = f'{prefix}{f.replace(delimiter, "/")}' + if dest_file_path not in output_paths_set: + os.remove(os.path.join(tmpdir, f)) + continue + num_files += 1 if delimiter in f: new_rel_path = f.replace(delimiter, '/').lstrip('/') new_path = os.path.join(tmpdir, new_rel_path) @@ -322,11 +343,6 @@ def split_data(self) -> None: raise ValueError('No splitting specified in template.') if any(dim in self._UNSUPPORTED_DIMENSIONS for dim in self.output_info.split_dims()): raise ValueError('Unsupported split dimension (lat, lng).') - if self.should_skip(): - metrics.Metrics.counter('file_splitters', 'skipped').inc() - self.logger.info('Skipping %s, file already split.', - repr(self.input_path)) - return with self._open_dataset_locally() as dataset: if any(split not in dataset.dims and split not in ('variable') for split in self.output_info.split_dims()): @@ -344,13 +360,20 @@ def split_data(self) -> None: iterlists.append(dataset[dim]) combinations = itertools.product(*iterlists) self.logger.info('Splitting & uploading %r...', self.input_path) + files_written = 0 for comb in combinations: selected = comb[0] for da in comb[1:]: for dim in da.coords: selected = selected.sel({dim: getattr(da, dim)}) - self._write_dataset(selected, filtered_split_dims) - self.logger.info('Finished splitting & uploading %r.', self.input_path) + if self._write_dataset(selected, filtered_split_dims): + files_written += 1 + if files_written == 0: + metrics.Metrics.counter('file_splitters', 'skipped').inc() + self.logger.info('Skipping %s, file already split.', + repr(self.input_path)) + else: + self.logger.info('Finished splitting & uploading %r.', self.input_path) @contextmanager def _open_dataset_locally(self) -> t.Iterator[xr.Dataset]: @@ -359,14 +382,19 @@ def _open_dataset_locally(self) -> t.Iterator[xr.Dataset]: yield ds ds.close() - def _write_dataset(self, dataset: xr.Dataset, split_dims: t.List[str]) -> None: - """Write destination NetCDF file in NETCDF4 format.""" + def _write_dataset(self, dataset: xr.Dataset, split_dims: t.List[str]) -> bool: + """Write destination NetCDF file in NETCDF4 format. Returns True if the file was written.""" # Here, we need to write the file locally, since only the scipy engine supports file objects or # returning bytes. Further, the scipy engine does not support NETCDF4 (which is HDF5 compliant). # Storing data in HDF5 is advantageous since it allows opening NetCDF files with buffered readers. + output_path = self._get_output_for_dataset(dataset, split_dims) + if self.should_skip_file(output_path): + self.logger.info('Skipping %s, file already split.', repr(output_path)) + return False with tempfile.NamedTemporaryFile() as tmp: dataset.to_netcdf(path=tmp.name, engine='netcdf4', format='NETCDF4') - copy(tmp.name, self._get_output_for_dataset(dataset, split_dims)) + copy(tmp.name, output_path) + return True def _get_output_for_dataset(self, dataset: xr.Dataset, split_dims: t.List[str]) -> str: splits = {'variable': list(dataset.data_vars.keys())[0]} diff --git a/weather_sp/splitter_pipeline/file_splitters_test.py b/weather_sp/splitter_pipeline/file_splitters_test.py index e9614e1..61840f5 100644 --- a/weather_sp/splitter_pipeline/file_splitters_test.py +++ b/weather_sp/splitter_pipeline/file_splitters_test.py @@ -191,6 +191,36 @@ def test_skips_existing_split(self, data_dir, grib_splitter): assert os.path.exists(f'{data_dir}/split_files/') assert splitter.should_skip() + def test_splits_only_missing_files(self, data_dir, grib_splitter): + input_path = f'{data_dir}/era5_sample.grib' + split_dir = f'{data_dir}/split_files/' + splitter = grib_splitter( + input_path, + OutFileInfo( + f'{data_dir}/split_files/era5_sample', + formatting='_{typeOfLevel}_{shortName}', + ending='.grib', + template_folders=[]) + ) + splitter.split_data() + assert os.path.exists(split_dir) + + original_mtimes = { + f: os.path.getmtime(os.path.join(split_dir, f)) + for f in os.listdir(split_dir) + } + + missing_file = 'era5_sample_isobaricInhPa_z.grib' + os.remove(os.path.join(split_dir, missing_file)) + + splitter.split_data() + + assert os.path.exists(os.path.join(split_dir, missing_file)) + for fname, orig_mtime in original_mtimes.items(): + if fname == missing_file: + continue + assert os.path.getmtime(os.path.join(split_dir, fname)) == orig_mtime + @patch('weather_sp.splitter_pipeline.file_splitters.FileSplitter.should_skip_file') def test_skips_existing_split_with_filter(self, mock_should_skip_file, data_dir): input_path = f'{data_dir}/era5_sample.grib' @@ -304,6 +334,34 @@ def test_skips_existing_split(self, data_dir): assert os.path.exists(f'{data_dir}/split_files/') assert splitter.should_skip() + def test_splits_only_missing_files(self, data_dir): + input_path = f'{data_dir}/era5_sample.nc' + split_dir = f'{data_dir}/split_files/' + splitter = NetCdfSplitter( + input_path, + OutFileInfo(f'{data_dir}/split_files/era5_sample', + formatting='_{time}_{variable}', + ending='.nc', + template_folders=[])) + splitter.split_data() + assert os.path.exists(split_dir) + + original_mtimes = { + f: os.path.getmtime(os.path.join(split_dir, f)) + for f in os.listdir(split_dir) + } + + missing_file = 'era5_sample_2015-01-15T00:00_z.nc' + os.remove(os.path.join(split_dir, missing_file)) + + splitter.split_data() + + assert os.path.exists(os.path.join(split_dir, missing_file)) + for fname, orig_mtime in original_mtimes.items(): + if fname == missing_file: + continue + assert os.path.getmtime(os.path.join(split_dir, fname)) == orig_mtime + def test_does_not_skip__if_forced(self, data_dir): input_path = f'{data_dir}/era5_sample.nc' output_base = f'{data_dir}/split_files/era5_sample'