diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index 2ac7dd029..0b7fa81e3 100644 --- a/src/swell/deployment/create_experiment.py +++ b/src/swell/deployment/create_experiment.py @@ -15,7 +15,7 @@ import shutil import sys from ruamel.yaml import YAML -from typing import Union, Optional +from typing import Optional from swell.suites.all_suites import AllSuites from swell.deployment.prepare_config_and_suite.prepare_config_and_suite import \ @@ -31,6 +31,19 @@ # -------------------------------------------------------------------------------------------------- +def read_override_file(override_path: str | None) -> dict: + + yaml = YAML(typ='safe') + + if override_path is None: + return {} + else: + with open(override_path, 'r') as f: + return yaml.load(f) + +# -------------------------------------------------------------------------------------------------- + + def clone_config( configuration: str, experiment_id: str, @@ -74,7 +87,7 @@ def prepare_config( suite_config: str, method: str, platform: str, - override: Union[dict, str, None], + override: dict, advanced: bool, slurm: str ) -> str: @@ -165,7 +178,8 @@ def prepare_config( # Register the experiment in R2D2 # ------------------------------- - if 'r2d2_experiment_id' in experiment_dict: + if 'r2d2_experiment_id' in experiment_dict and 'skip_r2d2' in experiment_dict \ + and not experiment_dict['skip_r2d2']: from swell.utilities.r2d2 import load_r2d2_credentials, load_r2d2_module, unique_r2d2_id @@ -221,7 +235,8 @@ def create_experiment_directory( platform: str, override: str, advanced: bool, - slurm: Optional[str] + slurm: str | None, + skip_r2d2: bool ) -> None: # Get the base name of the suite @@ -232,10 +247,21 @@ def create_experiment_directory( # --------------- logger = get_logger('SwellCreateExperiment') + # Read override file + # ------------------ + override_dict = read_override_file(override) + + # Specify whether to skip registering and storing in R2D2 + # ------------------------------------------------------- + if skip_r2d2: + + # Only override this if it is true, otherwise let the suite decide + override_dict['skip_r2d2'] = skip_r2d2 + # Call the experiment config and suite generation # ------------------------------------------------ experiment_dict_str = prepare_config(suite, suite_config, method, platform, - override, advanced, slurm) + override_dict, advanced, slurm) # Load the string using yaml # -------------------------- diff --git a/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py b/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py index 7cafb1033..69987bf5f 100644 --- a/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py +++ b/src/swell/deployment/prepare_config_and_suite/prepare_config_and_suite.py @@ -12,7 +12,7 @@ import os from ruamel.yaml import YAML from collections.abc import Mapping -from typing import Union, Tuple, Optional +from typing import Tuple, Optional from swell.swell_path import get_swell_path from swell.deployment.prepare_config_and_suite.question_and_answer_cli import GetAnswerCli @@ -58,7 +58,7 @@ def __init__( suite_config: str, platform: str, config_client: str, - override: Union[str, dict, None] + override: dict ) -> None: # Store local copy of the inputs @@ -353,40 +353,24 @@ def override_with_defaults(self) -> None: def override_with_external(self) -> None: - # Append with any user provide overrides - if self.override is not None: + # In this case the user is sending in a dictionary that looks like the experiment + # dictionary that they will ultimately be looking at. This means the dictionary does + # not contain default_value or options and the override cannot be performed. - # Create an override dictionary - override_dict = {} - - if isinstance(self.override, Mapping): - override_dict.update_dict(override_dict, self.override) + # Iterate over the model_ind dictionary and override + # -------------------------------------------------- + for key, val in self.question_dictionary_model_ind.items(): + if key in self.override: + val['default_value'] = self.override[key] - elif isinstance(self.override, str): - yaml = YAML(typ='safe') - with open(self.override, 'r') as ymlfile: - override_dict = update_dict(override_dict, yaml.load(ymlfile)) - else: - self.logger.abort(f'Override must be a dictionary or a path to a yaml file.') - - # In this case the user is sending in a dictionary that looks like the experiment - # dictionary that they will ultimately be looking at. This means the dictionary does - # not contain default_value or options and the override cannot be performed. - - # Iterate over the model_ind dictionary and override - # -------------------------------------------------- - for key, val in self.question_dictionary_model_ind.items(): - if key in override_dict: - val['default_value'] = override_dict[key] - - # Iterate over the model_dep dictionary and override - # -------------------------------------------------- - if self.suite_needs_model_components and 'models' in override_dict.keys(): - for model, model_dict in self.question_dictionary_model_dep.items(): - for key, val in model_dict.items(): - if model in override_dict['models']: - if key in override_dict['models'][model]: - val['default_value'] = override_dict['models'][model][key] + # Iterate over the model_dep dictionary and override + # -------------------------------------------------- + if self.suite_needs_model_components and 'models' in self.override.keys(): + for model, model_dict in self.question_dictionary_model_dep.items(): + for key, val in model_dict.items(): + if model in self.override['models']: + if key in self.override['models'][model]: + val['default_value'] = self.override['models'][model][key] # ---------------------------------------------------------------------------------------------- diff --git a/src/swell/suites/3dfgat_atmos/flow.cylc b/src/swell/suites/3dfgat_atmos/flow.cylc index b4d238308..878637314 100644 --- a/src/swell/suites/3dfgat_atmos/flow.cylc +++ b/src/swell/suites/3dfgat_atmos/flow.cylc @@ -89,12 +89,15 @@ # EvaIncrement RunJediVariationalExecutable-{{model_component}} => EvaIncrement-{{model_component}} + {% if not skip_r2d2 %} # Save observations RunJediVariationalExecutable-{{model_component}} => SaveObsDiags-{{model_component}} + SaveObsDiags-{{model_component}} => CleanCycle-{{model_component}} + {% endif %} # Clean up large files EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} & - EvaObservations-{{model_component}} & SaveObsDiags-{{model_component}} => + EvaObservations-{{model_component}} => CleanCycle-{{model_component}} {% endif %} diff --git a/src/swell/suites/3dfgat_marine_cycle/flow.cylc b/src/swell/suites/3dfgat_marine_cycle/flow.cylc index fa1e956f8..854ba71bb 100644 --- a/src/swell/suites/3dfgat_marine_cycle/flow.cylc +++ b/src/swell/suites/3dfgat_marine_cycle/flow.cylc @@ -100,17 +100,19 @@ {% endif %} # Move restart to next cycle and then erase current forecast folder - SaveRestart-{{model_component}} => MoveDaRestart-{{model_component}} + SaveRestart-{{model_component}} => MoveDaRestart-{{model_component}} => CleanCycle-{{model_component}} + {% if not skip_r2d2 %} # Save analysis output # RunJediFgatExecutable-{{model_component}} => SaveAnalysis-{{model_component}} - RunJediFgatExecutable-{{model_component}} => SaveObsDiags-{{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}} & SaveObsDiags-{{model_component}} => + EvaObservations-{{model_component}} & EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} => CleanCycle-{{model_component}} {% endfor %} """ diff --git a/src/swell/suites/3dvar_atmos/flow.cylc b/src/swell/suites/3dvar_atmos/flow.cylc index e0861d645..b71bba42a 100644 --- a/src/swell/suites/3dvar_atmos/flow.cylc +++ b/src/swell/suites/3dvar_atmos/flow.cylc @@ -86,12 +86,14 @@ # EvaIncrement RunJediVariationalExecutable-{{model_component}} => EvaIncrement-{{model_component}} + {% if not skip_r2d2 %} # Save observations - RunJediVariationalExecutable-{{model_component}} => SaveObsDiags-{{model_component}} + RunJediVariationalExecutable-{{model_component}} => SaveObsDiags-{{model_component}} => CleanCycle-{{model_component}} + {% endif %} # Clean up large files EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} & - EvaObservations-{{model_component}} & SaveObsDiags-{{model_component}} => + EvaObservations-{{model_component}} => CleanCycle-{{model_component}} {% endif %} diff --git a/src/swell/suites/3dvar_marine/flow.cylc b/src/swell/suites/3dvar_marine/flow.cylc index 9b1d063fd..7afdb11af 100644 --- a/src/swell/suites/3dvar_marine/flow.cylc +++ b/src/swell/suites/3dvar_marine/flow.cylc @@ -72,8 +72,11 @@ # EvaIncrement RunJediVariationalExecutable-{{model_component}} => EvaIncrement-{{model_component}} + {% if not skip_r2d2 %} # Save observations RunJediVariationalExecutable-{{model_component}} => SaveObsDiags-{{model_component}} + SaveObsDiags-{{model_component}} => CleanCycle-{{model_component}} + {% endif %} # Clean up large files EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} & diff --git a/src/swell/suites/3dvar_marine_cycle/flow.cylc b/src/swell/suites/3dvar_marine_cycle/flow.cylc index d16910113..8b76edf85 100644 --- a/src/swell/suites/3dvar_marine_cycle/flow.cylc +++ b/src/swell/suites/3dvar_marine_cycle/flow.cylc @@ -99,17 +99,19 @@ {% endif %} # Move restart to next cycle and then erase current forecast folder - SaveRestart-{{model_component}} => MoveDaRestart-{{model_component}} + SaveRestart-{{model_component}} => MoveDaRestart-{{model_component}} => CleanCycle-{{model_component}} + {% if not skip_r2d2 %} # Save analysis output # RunJediVariationalExecutable-{{model_component}} => SaveAnalysis-{{model_component}} - RunJediVariationalExecutable-{{model_component}} => SaveObsDiags-{{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}} & SaveObsDiags-{{model_component}} => + EvaObservations-{{model_component}} & EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} => CleanCycle-{{model_component}} {% endfor %} """ diff --git a/src/swell/suites/hofx/flow.cylc b/src/swell/suites/hofx/flow.cylc index 666399b9d..b465c6fe1 100644 --- a/src/swell/suites/hofx/flow.cylc +++ b/src/swell/suites/hofx/flow.cylc @@ -72,12 +72,13 @@ # EvaObservations RunJediHofxExecutable-{{model_component}} => EvaObservations-{{model_component}} + {% if not skip_r2d2 %} # Save observations - RunJediHofxExecutable-{{model_component}} => SaveObsDiags-{{model_component}} + RunJediHofxExecutable-{{model_component}} => SaveObsDiags-{{model_component}} => CleanCycle-{{model_component}} + {% endif %} # Clean up large files - EvaObservations-{{model_component}} & SaveObsDiags-{{model_component}} => - CleanCycle-{{model_component}} + EvaObservations-{{model_component}} => CleanCycle-{{model_component}} {% endif %} {% endfor %} diff --git a/src/swell/suites/hofx_cf/flow.cylc b/src/swell/suites/hofx_cf/flow.cylc index 562482ecf..2129dafd0 100644 --- a/src/swell/suites/hofx_cf/flow.cylc +++ b/src/swell/suites/hofx_cf/flow.cylc @@ -65,12 +65,13 @@ # EvaObservations RunJediHofxExecutable-{{model_component}} => EvaObservations-{{model_component}} - # Save feedback - RunJediHofxExecutable-{{model_component}} => SaveObsDiags-{{model_component}} + {% if not skip_r2d2 %} + # Save observations + RunJediHofxExecutable-{{model_component}} => SaveObsDiags-{{model_component}} => CleanCycle-{{model_component}} + {% endif %} # Clean up large files - EvaObservations-{{model_component}} & SaveObsDiags-{{model_component}} => - CleanCycle-{{model_component}} + EvaObservations-{{model_component}} => CleanCycle-{{model_component}} {% endif %} {% endfor %} diff --git a/src/swell/suites/suite_questions.py b/src/swell/suites/suite_questions.py index cc2824591..ee08b6d08 100644 --- a/src/swell/suites/suite_questions.py +++ b/src/swell/suites/suite_questions.py @@ -42,6 +42,7 @@ class SuiteQuestions(QuestionContainer, Enum): qd.model_components(), qd.runahead_limit(), qd.r2d2_experiment_id(), + qd.skip_r2d2(), ] ) diff --git a/src/swell/swell.py b/src/swell/swell.py index 3812b4770..89f55a3ad 100644 --- a/src/swell/swell.py +++ b/src/swell/swell.py @@ -82,6 +82,8 @@ def swell_driver() -> None: or for task-model combinations. """ +skip_r2d2_help = """Skip registering this experiment and storing products in R2D2.""" + # -------------------------------------------------------------------------------------------------- @@ -95,13 +97,15 @@ def swell_driver() -> None: @click.option('-o', '--override', 'override', default=None, help=override_help) @click.option('-a', '--advanced', 'advanced', default=False, help=advanced_help) @click.option('-s', '--slurm', 'slurm', default=None, help=slurm_help) +@click.option('-k', '--skip-r2d2', 'skip_r2d2', is_flag=True, default=False, help=skip_r2d2_help) def create( suite: str, input_method: str, platform: str, override: Union[dict, str, None], advanced: bool, - slurm: str + slurm: str, + skip_r2d2: bool ) -> None: """ Create a new experiment @@ -114,7 +118,7 @@ def create( """ # Create the experiment directory - create_experiment_directory(suite, input_method, platform, override, advanced, slurm) + create_experiment_directory(suite, input_method, platform, override, advanced, slurm, skip_r2d2) # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/utilities/question_defaults.py b/src/swell/utilities/question_defaults.py index 987f64322..fe59001b9 100644 --- a/src/swell/utilities/question_defaults.py +++ b/src/swell/utilities/question_defaults.py @@ -185,6 +185,15 @@ class skip_ensemble_hofx(SuiteQuestion): # -------------------------------------------------------------------------------------------------- + @dataclass + class skip_r2d2(SuiteQuestion): + default_value: bool = False + question_name: str = "skip_r2d2" + prompt: str = "Skip registering and storing results of this experiment in R2D2?" + widget_type: WType = WType.BOOLEAN + + # -------------------------------------------------------------------------------------------------- + @dataclass class start_cycle_point(SuiteQuestion): default_value: str = "2023-10-10T00:00:00Z"