Add option to skip storing files in R2D2

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
     # --------------------------

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]
 
     # ----------------------------------------------------------------------------------------------
 

src/swell/suites/3dfgat_atmos/flow.cylc

@@ -89,8 +89,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}} &

src/swell/suites/3dfgat_marine_cycle/flow.cylc

@@ -102,15 +102,17 @@
             # Move restart to next cycle and then erase current forecast folder
             SaveRestart-{{model_component}} => MoveDaRestart-{{model_component}}
 
+            {% if not save_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 %}
         """

src/swell/suites/3dvar_atmos/flow.cylc

@@ -86,8 +86,10 @@
             # EvaIncrement
             RunJediVariationalExecutable-{{model_component}} => EvaIncrement-{{model_component}}
 
+	    {% if not save_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}} &

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}} &

src/swell/suites/3dvar_marine_cycle/flow.cylc

@@ -101,9 +101,11 @@
             # Move restart to next cycle and then erase current forecast folder
             SaveRestart-{{model_component}} => MoveDaRestart-{{model_component}}
 
+            {% if 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}}

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 %}

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 %}

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(),
         ]
     )
 

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)
 
 
 # --------------------------------------------------------------------------------------------------

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"