diff --git a/pyproject.toml b/pyproject.toml index d088330f1..2293614e4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,7 +18,7 @@ build-backend = "setuptools.build_meta" [project] name = "swell" -version = "20260603" +version = "20260728" authors = [{name = "NASA Global Modeling and Assimilation Office"}] description = "Workflow suites, tasks and configuration for coupled data assimilation" requires-python = ">=3.11" diff --git a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/r2d2.py b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/r2d2.py index 03011d946..22ffef2e3 100644 --- a/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/r2d2.py +++ b/src/swell/configuration/jedi/interfaces/geos_atmosphere/model/r2d2.py @@ -31,6 +31,7 @@ def r2d2(template_dict: Mapping) -> Mapping: 'store': { 'fc': [ {'file_type': 'bkg', + 'filename': f'{cycle_dir}/bkg.%Y%m%dT%H%M%SZ.nc4', 'r2d2_model': 'geos'} ] } diff --git a/src/swell/configuration/jedi/interfaces/geos_cf/model/r2d2.py b/src/swell/configuration/jedi/interfaces/geos_cf/model/r2d2.py index fed69a755..4f21b8fc1 100644 --- a/src/swell/configuration/jedi/interfaces/geos_cf/model/r2d2.py +++ b/src/swell/configuration/jedi/interfaces/geos_cf/model/r2d2.py @@ -63,6 +63,7 @@ def r2d2(template_dict: Mapping) -> Mapping: 'store': { 'fc': [ {'file_type': 'bkg', + 'filename': f'{cycle_dir}/bkg.%Y%m%dT%H%M%SZ.nc4', 'r2d2_model': 'geos_cf'} ] } diff --git a/src/swell/configuration/jedi/interfaces/geos_marine/model/r2d2.py b/src/swell/configuration/jedi/interfaces/geos_marine/model/r2d2.py index 680ee3ecc..53b104454 100644 --- a/src/swell/configuration/jedi/interfaces/geos_marine/model/r2d2.py +++ b/src/swell/configuration/jedi/interfaces/geos_marine/model/r2d2.py @@ -12,18 +12,20 @@ def r2d2(template_dict: Mapping) -> Mapping: + analysis_time_iso = template_dict['analysis_time_iso'] cycle_dir = template_dict['cycle_dir'] - local_background_time = template_dict['local_background_time'] experiment_id = template_dict['experiment_id'] - analysis_time_iso = template_dict['analysis_time_iso'] + local_background_time = template_dict['local_background_time'] fc_list = [{'file_type': 'MOM.res', 'r2d2_model': 'mom6', + 'file_extension': 'tar.gz', 'filename': f'{cycle_dir}/MOM6.res.{local_background_time}.nc'}] if 'cice6' in template_dict['marine_models']: fc_list.append({'file_type': 'cice.res', 'r2d2_model': 'cice6', + 'file_extension': 'tar.gz', 'filename': f'{cycle_dir}/cice.res.{local_background_time}.nc'}) an_list = [{'file_type': 'ocn.incr', diff --git a/src/swell/suites/3dfgat_marine_cycle/flow.cylc b/src/swell/suites/3dfgat_marine_cycle/flow.cylc index ee00e36d0..a36d54a43 100644 --- a/src/swell/suites/3dfgat_marine_cycle/flow.cylc +++ b/src/swell/suites/3dfgat_marine_cycle/flow.cylc @@ -106,14 +106,11 @@ SaveRestart-{{model_component}} => MoveDaRestart-{{model_component}} => CleanCycle-{{model_component}} {% if not skip_r2d2 %} - # Save analysis output - # RunJediFgatExecutable-{{model_component}} => SaveAnalysis-{{model_component}} + # Compress forecast on compute node, then store in R2D2 on login node + LinkCoupledGeosOutput-{{model_component}} => RunCompressForecast-{{model_component}} => SaveForecast-{{model_component}} => MoveDaRestart-{{model_component}} RunJediFgatExecutable-{{model_component}} => SaveObsDiags-{{model_component}} => CleanCycle-{{model_component}} {% endif %} - # Save model output - # MoveBackground-{{model_component}} => StoreBackground-{{model_component}} - # Clean up large files EvaObservations-{{model_component}} & EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} => CleanCycle-{{model_component}} @@ -246,6 +243,18 @@ --{{key}} = {{value}} {%- endfor %} + [[RunCompressForecast-{{model_component}}]] + script = "swell task RunCompressForecast $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunCompressForecast"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunCompressForecast"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + + [[SaveForecast-{{model_component}}]] + script = "swell task SaveForecast $config -d $datetime -m {{model_component}}" + [[SaveRestart-{{model_component}}]] script = "swell task SaveRestart $config -d $datetime -m {{model_component}}" diff --git a/src/swell/suites/3dvar_marine/suite_config.py b/src/swell/suites/3dvar_marine/suite_config.py index c9db952fb..f6809d91e 100644 --- a/src/swell/suites/3dvar_marine/suite_config.py +++ b/src/swell/suites/3dvar_marine/suite_config.py @@ -37,7 +37,6 @@ class SuiteConfig(QuestionContainer, Enum): qd.horizontal_resolution("72x36"), qd.vertical_resolution("50"), qd.total_processors(6), - qd.obs_experiment("s2s_v1"), qd.observations([ "adt_cryosat2n", "adt_jason3", @@ -54,6 +53,7 @@ class SuiteConfig(QuestionContainer, Enum): "temp_profile_xbt" ]), qd.background_time_offset("PT18H"), + qd.background_experiment("pr790_test"), qd.clean_patterns(['*.nc4', '*.txt']), ] ) @@ -68,3 +68,31 @@ class SuiteConfig(QuestionContainer, Enum): ) # -------------------------------------------------------------------------------------------------- + + _3dvar_marine_tier2 = QuestionList( + list_name="3dvar_marine_tier2", + questions=[ + _3dvar_marine, + qd.start_cycle_point("2023-01-02T12:00:00Z"), + qd.final_cycle_point("2023-01-03T12:00:00Z"), + ], + geos_marine=[ + qd.cycle_times(['T12']), + qd.marine_models(['mom6', 'cice6']), + qd.analysis_variables([ + "sea_water_salinity", + "sea_water_potential_temperature", + "sea_surface_height_above_geoid", + "sea_water_cell_thickness", + "sea_ice_area_fraction", + "sea_ice_thickness", + "sea_ice_snow_thickness" + ]), + qd.window_length("P1D"), + qd.horizontal_resolution("1440x1080"), + qd.vertical_resolution("75"), + qd.total_processors(720), + ] + ) + + # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/suites/3dvar_marine_cycle/flow.cylc b/src/swell/suites/3dvar_marine_cycle/flow.cylc index 4c7e10aab..d7619156a 100644 --- a/src/swell/suites/3dvar_marine_cycle/flow.cylc +++ b/src/swell/suites/3dvar_marine_cycle/flow.cylc @@ -105,14 +105,11 @@ SaveRestart-{{model_component}} => MoveDaRestart-{{model_component}} => CleanCycle-{{model_component}} {% if not skip_r2d2 %} - # Save analysis output - # RunJediVariationalExecutable-{{model_component}} => SaveAnalysis-{{model_component}} + # Compress forecast on compute node, then store in R2D2 on login node + LinkCoupledGeosOutput-{{model_component}} => RunCompressForecast-{{model_component}} => SaveForecast-{{model_component}} => MoveDaRestart-{{model_component}} RunJediVariationalExecutable-{{model_component}} => SaveObsDiags-{{model_component}} => CleanCycle-{{model_component}} {% endif %} - # Save model output - # MoveBackground-{{model_component}} => StoreBackground-{{model_component}} - # Clean up large files EvaObservations-{{model_component}} & EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} => CleanCycle-{{model_component}} @@ -217,6 +214,15 @@ {% endif %} + [[RunCompressForecast-{{model_component}}]] + script = "swell task RunCompressForecast $config -d $datetime -m {{model_component}}" + platform = {{platform}} + execution time limit = {{scheduling["RunCompressForecast"]["execution_time_limit"]}} + [[[directives]]] + {%- for key, value in scheduling["RunCompressForecast"]["directives"][model_component].items() %} + --{{key}} = {{value}} + {%- endfor %} + [[RenderJediObservations-{{model_component}}]] script = "swell task RenderJediObservations $config -d $datetime -m {{model_component}}" @@ -245,6 +251,9 @@ --{{key}} = {{value}} {%- endfor %} + [[SaveForecast-{{model_component}}]] + script = "swell task SaveForecast $config -d $datetime -m {{model_component}}" + [[SaveRestart-{{model_component}}]] script = "swell task SaveRestart $config -d $datetime -m {{model_component}}" diff --git a/src/swell/tasks/geos_marine/get_background_geos_marine.py b/src/swell/tasks/geos_marine/get_background_geos_marine.py new file mode 100644 index 000000000..b657cc529 --- /dev/null +++ b/src/swell/tasks/geos_marine/get_background_geos_marine.py @@ -0,0 +1,111 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + + +# -------------------------------------------------------------------------------------------------- + +import os +import tarfile + +import r2d2 + +from swell.tasks.base.task_base import taskBase +from swell.utilities.r2d2 import load_r2d2_credentials + +# -------------------------------------------------------------------------------------------------- + + +def _marine_archive_path(cycle_dir: str, model_name: str, local_background_time: str) -> str: + """Return the canonical path for a marine background tar.gz archive. + + Args: + cycle_dir: Directory where the cycle artifacts are written. + model_name: Name of the marine model, such as "mom6" or "cice6". + local_background_time: Background time string used in the archive name. + + Returns: + The full path to the expected archive file. + """ + return os.path.join(cycle_dir, f"{model_name}.{local_background_time}.tar.gz") + + +# -------------------------------------------------------------------------------------------------- + + +class GetBackground(taskBase): + """Fetch marine forecast archives from R2D2 and unpack them into the cycle directory. + + This task retrieves the pre-compressed marine background archives produced + by the marine forecast save task and extracts them into the current cycle + directory for downstream use. + """ + + def execute(self) -> None: + """Retrieve the archived marine forecast outputs for the current cycle. + + The method loads the R2D2 credentials, resolves the active background + time, fetches the archive for each configured marine model, and unpacks + it into the cycle directory. + """ + + load_r2d2_credentials( + self.logger, + self.platform(), + r2d2_server=self.config.r2d2_server(default=None), + ) + + r2d2_datastore = self.config.r2d2_datastore(default=None) + marine_models = self.config.marine_models(None) or [] + window_type = self.config.window_type() + window_length = self.config.window_length() + horizontal_resolution = self.config.horizontal_resolution() + background_experiment = self.config.background_experiment() + + local_background_time, local_background_time_dto = \ + self.da_window_params.local_background_time(window_length, window_type, dto=True) + + os.makedirs(self.cycle_dir(), 0o755, exist_ok=True) + + marine_model_configs = [('mom6', 'MOM.res')] + if 'cice6' in marine_models: + marine_model_configs.append(('cice6', 'cice.res')) + + for model_name, file_type in marine_model_configs: + archive_path = _marine_archive_path( + self.cycle_dir(), model_name, local_background_time) + + self.logger.info( + f'Fetching marine background archive {os.path.basename(archive_path)} ' + f'({file_type}) for {model_name}' + ) + + fetch_kwargs = dict( + item='forecast', + target_file=archive_path, + model=model_name, + experiment=background_experiment, + file_extension='tar.gz', + resolution=horizontal_resolution, + step='PT00', + date=local_background_time_dto.strftime('%Y-%m-%d %H:%M:%S'), + file_type=file_type, + ) + if r2d2_datastore: + fetch_kwargs['data_store'] = r2d2_datastore + + r2d2.fetch(**fetch_kwargs) + + if not os.path.exists(archive_path): + self.logger.abort(f'Marine archive not fetched: {archive_path}') + + self.logger.info(f'Unpacking {os.path.basename(archive_path)} into {self.cycle_dir()}') + with tarfile.open(archive_path, 'r:gz') as tar: + tar.extractall(path=self.cycle_dir()) + + os.remove(archive_path) + self.logger.info(f'Unpacked {os.path.basename(archive_path)} into {self.cycle_dir()}') + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/tasks/geos_marine/run_compress_forecast_geos_marine.py b/src/swell/tasks/geos_marine/run_compress_forecast_geos_marine.py new file mode 100644 index 000000000..350fc486d --- /dev/null +++ b/src/swell/tasks/geos_marine/run_compress_forecast_geos_marine.py @@ -0,0 +1,175 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + + +# -------------------------------------------------------------------------------------------------- + + +from datetime import datetime as dt +import os +import tarfile + +from swell.tasks.base.task_base import taskBase +from swell.utilities.compress import compress_file + +# -------------------------------------------------------------------------------------------------- + +"""Create compressed marine forecast archives for GEOS marine components.""" + + +def _marine_archive_path(cycle_dir: str, model_name: str, local_background_time: str) -> str: + """Return the canonical path for a marine background tar.gz archive. + + Args: + cycle_dir: Directory where the cycle artifacts are written. + model_name: Name of the marine model, such as "mom6" or "cice6". + local_background_time: Background time string used in the archive name. + + Returns: + The full path to the expected archive file. + """ + return os.path.join(cycle_dir, f"{model_name}.{local_background_time}.tar.gz") + +# -------------------------------------------------------------------------------------------------- + + +class RunCompressForecast(taskBase): + """Compress marine state files into a tar.gz archive for later storage. + + The task runs on a compute node so that pigz parallelism is available. + """ + + # ---------------------------------------------------------------------------------------------- + + def execute(self) -> None: + """Gather marine states for the active cycle and create per-model archives. + + The method resolves the background window parameters, prepares the + interface rendering data, and compresses each requested marine model's + state files into a tar.gz archive. + """ + + marine_models = self.config.marine_models(None) or [] + window_type = self.config.window_type() + window_length = self.config.window_length() + window_begin_iso = self.da_window_params.window_begin_iso(window_length) + horizontal_resolution = self.config.horizontal_resolution() + + is_4d = window_type == '4D' or 'fgat' in self.suite_name() + + if is_4d: + background_frequency = self.config.background_frequency() + + (local_background_time, local_background_time_dto) = \ + self.da_window_params.local_background_time(window_length, window_type, dto=True) + + analysis_time_iso = self.da_window_params.analysis_time_iso() + local_background_time_iso = self.da_window_params.local_background_time_iso( + window_length, window_type) + + self.jedi_rendering.add_key('local_background_time', local_background_time) + self.jedi_rendering.add_key('local_background_time_iso', local_background_time_iso) + self.jedi_rendering.add_key('marine_models', marine_models) + self.jedi_rendering.add_key('horizontal_resolution', horizontal_resolution) + self.jedi_rendering.add_key('analysis_time_iso', analysis_time_iso) + + if is_4d: + self.jedi_rendering.add_key('background_frequency', background_frequency) + + r2d2_dict = self.jedi_rendering.render_interface_model('r2d2') + + marine_model_configs = [('mom6', 'ocn_filename', 'MOM.res')] + if 'cice6' in marine_models: + marine_model_configs.append(('cice6', 'ice_filename', 'cice.res')) + + for model_name, filename_key, _ in marine_model_configs: + self._compress_marine( + model_name=model_name, + filename_key=filename_key, + is_4d=is_4d, + background_frequency=background_frequency if is_4d else None, + window_length=window_length, + window_begin_iso=window_begin_iso, + marine_models=marine_models, + r2d2_dict=r2d2_dict, + local_background_time=local_background_time, + local_background_time_dto=local_background_time_dto, + ) + + # ---------------------------------------------------------------------------------------------- + + def _compress_marine(self, + model_name: str, + filename_key: str, + is_4d: bool, + background_frequency, + window_length: str, + window_begin_iso: str, + marine_models: list, + r2d2_dict: dict, + local_background_time: str, + local_background_time_dto: dt, + ) -> None: + """Create and compress the archive for a single marine model. + + Args: + model_name: Marine model identifier. + filename_key: Configuration key used to locate state files. + is_4d: Whether the current window is a 4D or FGAT-style window. + background_frequency: Frequency used to discover background states. + window_length: DA window length string. + window_begin_iso: ISO-formatted start of the DA window. + marine_models: List of configured marine models. + r2d2_dict: Rendered interface model data for R2D2 file discovery. + local_background_time: Local background time string. + local_background_time_dto: Parsed local background datetime. + """ + + # Gather file paths + files_to_archive = [] + + for fc in r2d2_dict['store']['fc']: + if fc.get('r2d2_model') == model_name: + source_file = local_background_time_dto.strftime(fc['filename']) + files_to_archive.append(source_file) + + if is_4d: + states = self.geos.states_generator( + background_frequency, window_length, + window_begin_iso, self.get_model(), marine_models) + for state in states: + files_to_archive.append(os.path.join(self.cycle_dir(), state[filename_key])) + + for f in files_to_archive: + if not os.path.exists(f): + if os.path.islink(f): + self.logger.abort(f"Marine state is a broken symbolic link: {f}") + else: + self.logger.abort(f"Required marine state file does not exist: {f}") + + archive_gz = _marine_archive_path(self.cycle_dir(), model_name, local_background_time) + archive_tar = archive_gz[:-3] # strip '.gz' + + self.logger.info(f"Archiving {len(files_to_archive)} marine state file(s) into " + f"{archive_gz}") + self.logger.debug(f"Files to archive: {files_to_archive}") + + try: + with tarfile.open(archive_tar, 'w', dereference=True) as tar: + for f in files_to_archive: + tar.add(f, arcname=os.path.basename(f)) + except Exception as e: + self.logger.abort(f"Failed to create tar archive for marine states: {e}") + + try: + compress_file(archive_tar, algorithm='pigz') + except Exception as e: + self.logger.abort(f"Failed to compress marine archive with pigz: {e}") + + os.remove(archive_tar) + + +# ---------------------------------------------------------------------------------------------- diff --git a/src/swell/tasks/geos_marine/save_forecast_geos_marine.py b/src/swell/tasks/geos_marine/save_forecast_geos_marine.py new file mode 100644 index 000000000..ce12d8c1e --- /dev/null +++ b/src/swell/tasks/geos_marine/save_forecast_geos_marine.py @@ -0,0 +1,113 @@ +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + + +# -------------------------------------------------------------------------------------------------- + + +import os +from r2d2 import store + +from swell.tasks.base.task_base import taskBase +from swell.utilities.datetime_util import datetime_formats +from swell.utilities.r2d2 import load_r2d2_credentials + + +# -------------------------------------------------------------------------------------------------- + +"""Store pre-compressed marine forecast archives in R2D2.""" + + +def _marine_archive_path(cycle_dir: str, model_name: str, local_background_time: str) -> str: + """Return the canonical path for a marine background tar.gz archive. + + Args: + cycle_dir: Directory where the cycle artifacts are written. + model_name: Name of the marine model, such as "mom6" or "cice6". + local_background_time: Background time string used in the archive name. + + Returns: + The full path to the expected archive file. + """ + return os.path.join(cycle_dir, f"{model_name}.{local_background_time}.tar.gz") + +# -------------------------------------------------------------------------------------------------- + + +class SaveForecast(taskBase): + """Store the pre-compressed marine archive in R2D2. + + The task runs on a login node, where internet access is available for + R2D2. The archive must have been produced by RunCompressForecast in the + same cycle. + """ + + # ---------------------------------------------------------------------------------------------- + + def execute(self) -> None: + """Resolve the target background time and store the generated archives. + + The method loads the R2D2 credentials, determines the active window + parameters, and uploads each marine model archive that was produced in + the current cycle. + + The archive files are stored at background time, which changes according to the configured + window type (and/or suite type) and length. + """ + + self.marine_models = self.config.marine_models(None) or [] + window_type = self.config.window_type() + self.window_length = self.config.window_length() + self.horizontal_resolution = self.config.horizontal_resolution() + + load_r2d2_credentials(self.logger, self.platform()) + + self.local_background_time, self.local_background_time_dto = \ + self.da_window_params.local_background_time(self.window_length, window_type, dto=True) + + marine_model_configs = [('mom6', 'MOM.res')] + if 'cice6' in self.marine_models: + marine_model_configs.append(('cice6', 'cice.res')) + + for model_name, file_type in marine_model_configs: + self._store_forecast(model_name, file_type) + + # ---------------------------------------------------------------------------------------------- + + def _store_forecast(self, model_name: str, file_type: str) -> None: + """Upload a single marine archive to R2D2. + + Args: + model_name: Marine model identifier. + file_type: File type label used when storing the archive. + """ + + archive_path = _marine_archive_path( + self.cycle_dir(), model_name, self.local_background_time) + + if not os.path.exists(archive_path): + self.logger.abort( + f"Marine archive not found (did RunCompressForecast complete?): {archive_path}") + + self.logger.info( + f'Storing {os.path.basename(archive_path)} ({file_type}) ' + f'step=PT00 ' + f'at {self.local_background_time_dto.strftime(datetime_formats["iso_format"])}') + + store( + item='forecast', + model=model_name, + experiment=self.config.r2d2_experiment_id(), + resolution=self.horizontal_resolution, + date=self.local_background_time_dto.strftime('%Y-%m-%d %H:%M:%S'), + source_file=archive_path, + file_type=file_type, + file_extension='tar.gz', + step='PT00', + store_as_symlink=False, + ) + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/tasks/task_questions.py b/src/swell/tasks/task_questions.py index 7bb0cab89..a8d3436d3 100644 --- a/src/swell/tasks/task_questions.py +++ b/src/swell/tasks/task_questions.py @@ -153,6 +153,16 @@ class TaskQuestions(QuestionContainer, Enum): # -------------------------------------------------------------------------------------------------- + CleanEda = QuestionList( + list_name="CleanEda", + questions=[ + run_jedi_executable, + qd.ensemble_num_members(), + ] + ) + + # -------------------------------------------------------------------------------------------------- + CloneGeos = QuestionList( list_name="CloneGeos", questions=[ @@ -840,6 +850,18 @@ class TaskQuestions(QuestionContainer, Enum): # -------------------------------------------------------------------------------------------------- + RunCompressForecast = QuestionList( + list_name="RunCompressForecast", + questions=[ + window_questions, + qd.r2d2_experiment_id(), + qd.background_frequency(), + qd.horizontal_resolution(), + ] + ) + + # -------------------------------------------------------------------------------------------------- + RunJediEdaExecutable = QuestionList( list_name="RunJediEdaExecutable", questions=[ @@ -854,11 +876,25 @@ class TaskQuestions(QuestionContainer, Enum): # -------------------------------------------------------------------------------------------------- - CleanEda = QuestionList( - list_name="CleanEda", + SaveForecast = QuestionList( + list_name="SaveForecast", questions=[ - run_jedi_executable, - qd.ensemble_num_members(), + window_questions, + qd.r2d2_experiment_id(), + qd.background_frequency(), + qd.horizontal_resolution(), + ] + ) + + # -------------------------------------------------------------------------------------------------- + + SaveForecastCf = QuestionList( + list_name="SaveForecastCf", + questions=[ + qd.forecast_length(), + qd.forecast_output_frequency(), + qd.horizontal_resolution(), + qd.window_length(), ] ) @@ -910,18 +946,6 @@ class TaskQuestions(QuestionContainer, Enum): # -------------------------------------------------------------------------------------------------- - SaveForecastCf = QuestionList( - list_name="SaveForecastCf", - questions=[ - qd.forecast_length(), - qd.forecast_output_frequency(), - qd.horizontal_resolution(), - qd.window_length(), - ] - ) - - # -------------------------------------------------------------------------------------------------- - SaveBackground = QuestionList( list_name="SaveBackground", questions=[ diff --git a/src/swell/utilities/compress.py b/src/swell/utilities/compress.py new file mode 100644 index 000000000..fda7de22b --- /dev/null +++ b/src/swell/utilities/compress.py @@ -0,0 +1,108 @@ + +# (C) Copyright 2021- United States Government as represented by the Administrator of the +# National Aeronautics and Space Administration. All Rights Reserved. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. + + +# -------------------------------------------------------------------------------------------------- + +import gzip +import os +import shutil +import subprocess + +# -------------------------------------------------------------------------------------------------- + + +def compress_file(source_file: str, + algorithm: str = 'gzip', + level: int = 6, + num_threads: int = 8) -> str: + + if not os.path.isfile(source_file): + raise FileNotFoundError(f"Source file not found: {source_file}") + + compressed_path = source_file + '.gz' + + if algorithm == 'gzip': + _compress_gzip(source_file, compressed_path, level) + elif algorithm == 'pigz': + _compress_pigz(source_file, level, num_threads) + + return compressed_path + + +# -------------------------------------------------------------------------------------------------- + + +def _compress_gzip(source_file: str, compressed_path: str, level: int) -> None: + """Compress using Python stdlib gzip (single-threaded, streaming).""" + + with open(source_file, 'rb') as f_in: + with gzip.open(compressed_path, 'wb', compresslevel=level) as f_out: + shutil.copyfileobj(f_in, f_out) + + +# -------------------------------------------------------------------------------------------------- + + +def _compress_pigz(source_file: str, level: int, num_threads: int) -> None: + """Compress using the external ``pigz`` binary (parallel gzip). + + ``pigz -k`` keeps the original file and writes ``source_file.gz`` + alongside it — exactly the same convention as :func:`compress_file`. + """ + + pigz_bin = shutil.which('pigz') + if pigz_bin is None: + raise FileNotFoundError( + "pigz binary not found on PATH. " + "Install pigz (e.g. 'module load pigz' on NCCS Discover) " + "or set compress_algorithm to 'gzip'." + ) + + # -k : keep the original file (do not delete it) + # -p : number of threads + # -N : compression level (1–9) + # -f : Force overwrite, compress .gz, links, and to terminal + subprocess.run( + [pigz_bin, f'-{level}', '-k', '-f', f'-p{num_threads}', source_file], + check=True, + ) + + +# -------------------------------------------------------------------------------------------------- + + +def decompress_file(compressed_file: str, target_file: str) -> None: + if not os.path.isfile(compressed_file): + raise FileNotFoundError(f"Compressed file not found: {compressed_file}") + + with gzip.open(compressed_file, 'rb') as f_in: + with open(target_file, 'wb') as f_out: + shutil.copyfileobj(f_in, f_out) + + os.remove(compressed_file) + + +# -------------------------------------------------------------------------------------------------- + + +def decompress_if_needed(file_path: str) -> str: + if not file_path.endswith('.gz'): + return file_path + + decompressed_path = file_path[:-3] # strip the trailing '.gz' + decompress_file(file_path, decompressed_path) + return decompressed_path + + +# -------------------------------------------------------------------------------------------------- + + +def compressed_extension(original_ext: str) -> str: + return original_ext + '.gz' + +# -------------------------------------------------------------------------------------------------- diff --git a/src/swell/utilities/slurm.py b/src/swell/utilities/slurm.py index 26622af72..0d3a0ccea 100644 --- a/src/swell/utilities/slurm.py +++ b/src/swell/utilities/slurm.py @@ -79,6 +79,7 @@ def prepare_scheduling_dict( 'EvaComparisonObservations', 'EvaTimeseries', 'GenerateBClimatology', + 'RunCompressForecast', 'RunGeos', 'RunJediEdaExecutable', 'RunJediEnsembleMeanVariance',