From 4f48070cd2f6feabae24a095b4fbc1a4311704b0 Mon Sep 17 00:00:00 2001 From: Michael Anstett Date: Tue, 24 Feb 2026 16:40:26 -0500 Subject: [PATCH 01/16] Add experiment ID for R2D2 --- .../prepare_config_and_suite.py | 51 ++++++++++++++++++- src/swell/suites/suite_questions.py | 3 +- src/swell/tasks/save_obs_diags.py | 2 +- src/swell/tasks/save_restart.py | 4 +- src/swell/utilities/question_defaults.py | 9 ++++ 5 files changed, 63 insertions(+), 6 deletions(-) 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 7ba60dd75..5d0542fae 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 @@ -1,5 +1,4 @@ -# (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. @@ -13,6 +12,7 @@ from ruamel.yaml import YAML from collections.abc import Mapping from typing import Union, Tuple, Optional +import random from swell.swell_path import get_swell_path from swell.deployment.prepare_config_and_suite.question_and_answer_cli import GetAnswerCli @@ -23,6 +23,7 @@ from swell.utilities.dictionary import update_dict from swell.tasks.task_questions import TaskQuestions as task_questions from swell.suites.all_suites import AllSuites +from swell.utilities.r2d2 import load_r2d2_credentials # -------------------------------------------------------------------------------------------------- @@ -326,6 +327,7 @@ def override_with_defaults(self) -> None: # Look for defer_to_code in the model_ind dictionary # -------------------------------------------------- for key, val in self.question_dictionary_model_ind.items(): + if key == 'model_components': if val['default_value'] == 'defer_to_code': val['default_value'] = self.possible_model_components @@ -335,6 +337,13 @@ def override_with_defaults(self) -> None: if key == 'experiment_id' and val['default_value'] == 'defer_to_code': val['default_value'] = f'swell-{self.suite}' + if key == 'r2d2_experiment_id' and val['default_value'] == 'defer_to_code': + swell_id = self.question_dictionary_model_ind['experiment_id']['default_value'] + if swell_id == 'defer_to_code': + swell_id = f'swell-{self.suite}' + self.question_dictionary_model_ind['experiment_id']['default_value'] = swell_id + val['default_value'] = self.create_r2d2_id(swell_id) + # ---------------------------------------------------------------------------------------------- def override_with_external(self) -> None: @@ -689,4 +698,42 @@ def get_dynamic_tasks(self, question_list: list) -> list: return tasks + # ---------------------------------------------------------------------------------------------- + + def random_hex_id(self, swell_id: str, length: int = 8): + return f"{swell_id}-{random.randrange(16**length):0{length}x}" + + # ---------------------------------------------------------------------------------------------- + + def create_r2d2_id(self, swell_id: str) -> str: + + # Load credentials to allow search + load_r2d2_credentials(self.logger, self.platform) + + import r2d2 + + self.logger.info('Generating Experiment ID for R2D2') + + # Only try this 10 times + for i in range(10): + temp_id = self.random_hex_id(swell_id, length=8) + try: + r2d2.get(item='experiment', name=temp_id) + except Exception as e: + + if '400 Client Error' in str(e): + user = r2d2.get_client_user() + host = r2d2.get_client_host() + compiler = r2d2.get_client_compiler() + + r2d2.register(item='experiment', + name=temp_id, + user=user, + compute_host=f'{host}-{compiler}', + lifetime='debug') + + return temp_id + + raise Exception('Could not find a valid experiment_id for R2D2') + # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/suites/suite_questions.py b/src/swell/suites/suite_questions.py index 247bee869..a4da99431 100644 --- a/src/swell/suites/suite_questions.py +++ b/src/swell/suites/suite_questions.py @@ -40,7 +40,8 @@ class SuiteQuestions(QuestionContainer, Enum): qd.start_cycle_point(), qd.final_cycle_point(), qd.model_components(), - qd.runahead_limit() + qd.runahead_limit(), + qd.r2d2_experiment_id() ] ) diff --git a/src/swell/tasks/save_obs_diags.py b/src/swell/tasks/save_obs_diags.py index d03c6ba0c..0cd489434 100644 --- a/src/swell/tasks/save_obs_diags.py +++ b/src/swell/tasks/save_obs_diags.py @@ -81,7 +81,7 @@ def execute(self) -> None: try: r2d2.store( item='feedback', - experiment=self.experiment_id(), + experiment=self.config.r2d2_experiment_id(), observation_type=name, file_extension=obs_path_file.split('.')[-1], window_length='PT6H', diff --git a/src/swell/tasks/save_restart.py b/src/swell/tasks/save_restart.py index c251cebe1..91445adb6 100644 --- a/src/swell/tasks/save_restart.py +++ b/src/swell/tasks/save_restart.py @@ -68,7 +68,7 @@ def execute(self): step=window_length, resolution=self.config.horizontal_resolution(), type='fc', - experiment=self.experiment_id()) + experiment=self.config.r2d2_experiment_id()) # Loop over an for an in r2d2_dict['store']['an']: @@ -79,7 +79,7 @@ def execute(self): fc_date_rendering='analysis', resolution=self.config.horizontal_resolution(), type='an', - experiment=self.experiment_id()) + experiment=self.config.r2d2_experiment_id()) # Oceanstats needs special handling from the forecast folder. It is produced at the end of # the forecast and could be saved as a good metric. We are replicating the same structure as diff --git a/src/swell/utilities/question_defaults.py b/src/swell/utilities/question_defaults.py index bf11de750..a63faa2c6 100644 --- a/src/swell/utilities/question_defaults.py +++ b/src/swell/utilities/question_defaults.py @@ -152,6 +152,15 @@ class parser_options(SuiteQuestion): # -------------------------------------------------------------------------------------------------- + @dataclass + class r2d2_experiment_id(SuiteQuestion): + default_value: str = "defer_to_code" + question_name: str = "r2d2_experiment_id" + prompt: str = "What experiment_id should r2d2 reference for experiment?" + widget_type: WType = WType.STRING + + # -------------------------------------------------------------------------------------------------- + @dataclass class runahead_limit(SuiteQuestion): default_value: str = "P4" From d32c3b9133c1bdf7fdbcc6f2760ca027fddb1b20 Mon Sep 17 00:00:00 2001 From: Michael Anstett Date: Tue, 24 Feb 2026 16:56:30 -0500 Subject: [PATCH 02/16] Change location that registry happens in --- src/swell/deployment/create_experiment.py | 19 +++++++++++++++++++ .../prepare_config_and_suite.py | 10 ---------- 2 files changed, 19 insertions(+), 10 deletions(-) diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index 1c2105e9f..c565c019a 100644 --- a/src/swell/deployment/create_experiment.py +++ b/src/swell/deployment/create_experiment.py @@ -237,6 +237,25 @@ def create_experiment_directory( with open(os.path.join(exp_suite_path, 'experiment.yaml'), 'w') as file: file.write(experiment_dict_str) + # Register the experiment in R2D2 + # ------------------------------- + + import r2d2 + from swell.utilities.r2d2 import load_r2d2_credentials + + r2d2_id = experiment_dict['r2d2_experiment_id'] + + load_r2d2_credentials(logger, platform) + user = r2d2.get_client_user() + host = r2d2.get_client_host() + compiler = r2d2.get_client_compiler() + + r2d2.register(item='experiment', + name=r2d2_id, + user=user, + compute_host=f'{host}-{compiler}', + lifetime='debug') + # At this point we need to write the complete suite file with all templates resolved. Call the # function to build the scheduling dictionary, combine with the experiment dictionary, # resolve the templates and write the suite file to the experiment suite directory. 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 5d0542fae..db5b0e870 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 @@ -722,16 +722,6 @@ def create_r2d2_id(self, swell_id: str) -> str: except Exception as e: if '400 Client Error' in str(e): - user = r2d2.get_client_user() - host = r2d2.get_client_host() - compiler = r2d2.get_client_compiler() - - r2d2.register(item='experiment', - name=temp_id, - user=user, - compute_host=f'{host}-{compiler}', - lifetime='debug') - return temp_id raise Exception('Could not find a valid experiment_id for R2D2') From 28be5ba804576df852637e76c600412fa2e45f26 Mon Sep 17 00:00:00 2001 From: Michael Anstett Date: Tue, 24 Feb 2026 17:03:25 -0500 Subject: [PATCH 03/16] add key for experiment lifetime --- src/swell/deployment/create_experiment.py | 31 ++++++++++++----------- src/swell/suites/suite_questions.py | 3 ++- src/swell/utilities/question_defaults.py | 10 ++++++++ 3 files changed, 28 insertions(+), 16 deletions(-) diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index c565c019a..14dede765 100644 --- a/src/swell/deployment/create_experiment.py +++ b/src/swell/deployment/create_experiment.py @@ -239,22 +239,23 @@ def create_experiment_directory( # Register the experiment in R2D2 # ------------------------------- - - import r2d2 - from swell.utilities.r2d2 import load_r2d2_credentials + if 'r2d2_experiment_id' in experiment_dict: + import r2d2 + from swell.utilities.r2d2 import load_r2d2_credentials - r2d2_id = experiment_dict['r2d2_experiment_id'] - - load_r2d2_credentials(logger, platform) - user = r2d2.get_client_user() - host = r2d2.get_client_host() - compiler = r2d2.get_client_compiler() - - r2d2.register(item='experiment', - name=r2d2_id, - user=user, - compute_host=f'{host}-{compiler}', - lifetime='debug') + r2d2_id = experiment_dict['r2d2_experiment_id'] + r2d2_lifetime = experiment_dict['r2d2_experiment_lifetime'] + + load_r2d2_credentials(logger, platform) + user = r2d2.get_client_user() + host = r2d2.get_client_host() + compiler = r2d2.get_client_compiler() + + r2d2.register(item='experiment', + name=r2d2_id, + user=user, + compute_host=f'{host}-{compiler}', + lifetime=r2d2_lifetime) # At this point we need to write the complete suite file with all templates resolved. Call the # function to build the scheduling dictionary, combine with the experiment dictionary, diff --git a/src/swell/suites/suite_questions.py b/src/swell/suites/suite_questions.py index a4da99431..2073e69b9 100644 --- a/src/swell/suites/suite_questions.py +++ b/src/swell/suites/suite_questions.py @@ -41,7 +41,8 @@ class SuiteQuestions(QuestionContainer, Enum): qd.final_cycle_point(), qd.model_components(), qd.runahead_limit(), - qd.r2d2_experiment_id() + qd.r2d2_experiment_id(), + qd.r2d2_experiment_lifetime() ] ) diff --git a/src/swell/utilities/question_defaults.py b/src/swell/utilities/question_defaults.py index a63faa2c6..521220ae5 100644 --- a/src/swell/utilities/question_defaults.py +++ b/src/swell/utilities/question_defaults.py @@ -161,6 +161,16 @@ class r2d2_experiment_id(SuiteQuestion): # -------------------------------------------------------------------------------------------------- + @dataclass + class r2d2_experiment_lifetime(SuiteQuestion): + default_value: str = "debug" + question_name: str = "r2d2_experiment_lifetime" + options: list = mutable_field(['debug', 'science', 'publication', 'release']) + prompt: str = "What lifetime should the experiment have in R2D2?" + widget_type: WType = WType.STRING + + # -------------------------------------------------------------------------------------------------- + @dataclass class runahead_limit(SuiteQuestion): default_value: str = "P4" From 3d571cf7bd5dc41ae07496d9f1ad4cc09659292e Mon Sep 17 00:00:00 2001 From: Michael Anstett Date: Tue, 24 Feb 2026 17:13:58 -0500 Subject: [PATCH 04/16] pycodestyle --- src/swell/deployment/create_experiment.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index 14dede765..286d22939 100644 --- a/src/swell/deployment/create_experiment.py +++ b/src/swell/deployment/create_experiment.py @@ -242,7 +242,7 @@ def create_experiment_directory( if 'r2d2_experiment_id' in experiment_dict: import r2d2 from swell.utilities.r2d2 import load_r2d2_credentials - + r2d2_id = experiment_dict['r2d2_experiment_id'] r2d2_lifetime = experiment_dict['r2d2_experiment_lifetime'] From 1e65b35ff797cb4b6c2ff9804f435dc88c597087 Mon Sep 17 00:00:00 2001 From: Michael Anstett Date: Wed, 25 Feb 2026 10:22:00 -0500 Subject: [PATCH 05/16] Refactor --- src/swell/deployment/create_experiment.py | 45 +++++++++-------- .../prepare_config_and_suite.py | 33 +------------ src/swell/utilities/r2d2.py | 48 ++++++++++++++++++- 3 files changed, 73 insertions(+), 53 deletions(-) diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index 286d22939..b73ee23ca 100644 --- a/src/swell/deployment/create_experiment.py +++ b/src/swell/deployment/create_experiment.py @@ -163,6 +163,31 @@ def prepare_config( logger.abort(f'SLURM file contains invalid keys: {slurm_invalid_keys}') experiment_dict = {**experiment_dict, **slurm_dict} + # Register the experiment in R2D2 + # ------------------------------- + if 'r2d2_experiment_id' in experiment_dict: + + import r2d2 + from swell.utilities.r2d2 import load_r2d2_credentials, unique_r2d2_id + + r2d2_id = experiment_dict['r2d2_experiment_id'] + + unique_id = unique_r2d2_id(r2d2_id, platform) + experiment_dict['r2d2_experiment_id'] = unique_id + + r2d2_lifetime = experiment_dict['r2d2_experiment_lifetime'] + + load_r2d2_credentials(logger, platform) + user = r2d2.get_client_user() + host = r2d2.get_client_host() + compiler = r2d2.get_client_compiler() + + r2d2.register(item='experiment', + name=unique_id, + user=user, + compute_host=f'{host}-{compiler}', + lifetime=r2d2_lifetime) + # Expand all environment vars in the dictionary # --------------------------------------------- output = io.StringIO() @@ -237,26 +262,6 @@ def create_experiment_directory( with open(os.path.join(exp_suite_path, 'experiment.yaml'), 'w') as file: file.write(experiment_dict_str) - # Register the experiment in R2D2 - # ------------------------------- - if 'r2d2_experiment_id' in experiment_dict: - import r2d2 - from swell.utilities.r2d2 import load_r2d2_credentials - - r2d2_id = experiment_dict['r2d2_experiment_id'] - r2d2_lifetime = experiment_dict['r2d2_experiment_lifetime'] - - load_r2d2_credentials(logger, platform) - user = r2d2.get_client_user() - host = r2d2.get_client_host() - compiler = r2d2.get_client_compiler() - - r2d2.register(item='experiment', - name=r2d2_id, - user=user, - compute_host=f'{host}-{compiler}', - lifetime=r2d2_lifetime) - # At this point we need to write the complete suite file with all templates resolved. Call the # function to build the scheduling dictionary, combine with the experiment dictionary, # resolve the templates and write the suite file to the experiment suite directory. 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 db5b0e870..e66f2110e 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,6 @@ from ruamel.yaml import YAML from collections.abc import Mapping from typing import Union, Tuple, Optional -import random from swell.swell_path import get_swell_path from swell.deployment.prepare_config_and_suite.question_and_answer_cli import GetAnswerCli @@ -23,8 +22,6 @@ from swell.utilities.dictionary import update_dict from swell.tasks.task_questions import TaskQuestions as task_questions from swell.suites.all_suites import AllSuites -from swell.utilities.r2d2 import load_r2d2_credentials - # -------------------------------------------------------------------------------------------------- @@ -342,7 +339,7 @@ def override_with_defaults(self) -> None: if swell_id == 'defer_to_code': swell_id = f'swell-{self.suite}' self.question_dictionary_model_ind['experiment_id']['default_value'] = swell_id - val['default_value'] = self.create_r2d2_id(swell_id) + val['default_value'] = swell_id # ---------------------------------------------------------------------------------------------- @@ -698,32 +695,4 @@ def get_dynamic_tasks(self, question_list: list) -> list: return tasks - # ---------------------------------------------------------------------------------------------- - - def random_hex_id(self, swell_id: str, length: int = 8): - return f"{swell_id}-{random.randrange(16**length):0{length}x}" - - # ---------------------------------------------------------------------------------------------- - - def create_r2d2_id(self, swell_id: str) -> str: - - # Load credentials to allow search - load_r2d2_credentials(self.logger, self.platform) - - import r2d2 - - self.logger.info('Generating Experiment ID for R2D2') - - # Only try this 10 times - for i in range(10): - temp_id = self.random_hex_id(swell_id, length=8) - try: - r2d2.get(item='experiment', name=temp_id) - except Exception as e: - - if '400 Client Error' in str(e): - return temp_id - - raise Exception('Could not find a valid experiment_id for R2D2') - # -------------------------------------------------------------------------------------------------- diff --git a/src/swell/utilities/r2d2.py b/src/swell/utilities/r2d2.py index 5810f4b7f..963632a0d 100644 --- a/src/swell/utilities/r2d2.py +++ b/src/swell/utilities/r2d2.py @@ -9,10 +9,11 @@ import os from ruamel.yaml import YAML +import random from swell.swell_path import get_swell_path from swell.utilities.jinja2 import template_string_jinja2 -from swell.utilities.logger import Logger +from swell.utilities.logger import get_logger, Logger # -------------------------------------------------------------------------------------------------- @@ -59,6 +60,8 @@ def create_r2d2_config( with open(r2d2_config_file, 'w') as f: f.write(r2d2_config_file_template_str) +# -------------------------------------------------------------------------------------------------- + def _get_platform_r2d2_config(logger: Logger, platform: str = None) -> tuple: if not platform: @@ -95,6 +98,8 @@ def _get_platform_r2d2_config(logger: Logger, platform: str = None) -> tuple: logger.warning(f"Unknown platform '{platform}', cannot determine R2D2 host/compiler") return None, None +# -------------------------------------------------------------------------------------------------- + def load_r2d2_credentials( logger: Logger, @@ -161,5 +166,46 @@ def load_r2d2_credentials( logger.info("R2D2 v3 credentials loaded successfully") +# ---------------------------------------------------------------------------------------------- + + +def random_hex_id(swell_id: str, length: int = 8): + return f"{swell_id}-{random.randrange(16**length):0{length}x}" + +# ---------------------------------------------------------------------------------------------- + + +def experiment_exists(r2d2_id: str): + import r2d2 + + try: + r2d2.get(item='experiment', name=r2d2_id) + except Exception as e: + if '400 Client Error' in str(e): + return False + + return True # ---------------------------------------------------------------------------------------------- + + +def unique_r2d2_id(swell_id: str, platform: str) -> str: + logger = get_logger('CreateR2D2ID') + + # Load credentials to allow search + load_r2d2_credentials(logger, platform) + + # Just use the ID if it doesn't exist + if not experiment_exists(swell_id): + return swell_id + + # If not, append an unused hex id + # Only try this 10 times + for i in range(10): + temp_id = random_hex_id(swell_id, length=8) + if not experiment_exists(temp_id): + return temp_id + + raise Exception('Could not find a valid experiment_id for R2D2') + +# -------------------------------------------------------------------------------------------------- From f69dbe55feeebfc4a6e5eaa08732a50cc4b09e4a Mon Sep 17 00:00:00 2001 From: ftgoktas Date: Wed, 25 Feb 2026 12:44:51 -0500 Subject: [PATCH 06/16] Load r2d2 modules during experiment creation (#701). --- src/swell/deployment/create_experiment.py | 7 ++-- src/swell/swell.py | 6 ++++ src/swell/utilities/r2d2.py | 43 ++++++++++++++++++++++- 3 files changed, 53 insertions(+), 3 deletions(-) diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index b73ee23ca..bf8196017 100644 --- a/src/swell/deployment/create_experiment.py +++ b/src/swell/deployment/create_experiment.py @@ -167,8 +167,12 @@ def prepare_config( # ------------------------------- if 'r2d2_experiment_id' in experiment_dict: + from swell.utilities.r2d2 import load_r2d2_credentials, load_r2d2_module, unique_r2d2_id + + load_r2d2_credentials(logger, platform) + load_r2d2_module(logger, platform) + import r2d2 - from swell.utilities.r2d2 import load_r2d2_credentials, unique_r2d2_id r2d2_id = experiment_dict['r2d2_experiment_id'] @@ -177,7 +181,6 @@ def prepare_config( r2d2_lifetime = experiment_dict['r2d2_experiment_lifetime'] - load_r2d2_credentials(logger, platform) user = r2d2.get_client_user() host = r2d2.get_client_host() compiler = r2d2.get_client_compiler() diff --git a/src/swell/swell.py b/src/swell/swell.py index 944a43779..923ab779b 100644 --- a/src/swell/swell.py +++ b/src/swell/swell.py @@ -112,6 +112,12 @@ def create( suite (str): Name of the suite you wish to run. \n """ + # Load R2D2 credentials + from swell.utilities.r2d2 import load_r2d2_credentials + from swell.utilities.logger import get_logger + logger = get_logger("Swell Create") + load_r2d2_credentials(logger, platform) + # Create the experiment directory create_experiment_directory(suite, input_method, platform, override, advanced, slurm) diff --git a/src/swell/utilities/r2d2.py b/src/swell/utilities/r2d2.py index 963632a0d..b3d60e558 100644 --- a/src/swell/utilities/r2d2.py +++ b/src/swell/utilities/r2d2.py @@ -16,7 +16,48 @@ from swell.utilities.logger import get_logger, Logger # -------------------------------------------------------------------------------------------------- - +import subprocess + +# Platform-specific R2D2 module config +_R2D2_MODULE_CONFIG = { + 'nccs_discover_sles15': { + 'module_path': '/discover/nobackup/projects/gmao/advda/JediOpt/modulefiles/core', + 'module_name': 'r2d2-client/112025', + }, + 'nccs_discover_cascade': { + 'module_path': '/discover/nobackup/projects/gmao/advda/JediOpt/modulefiles/core', + 'module_name': 'r2d2-client/112025', + }, +} + +def load_r2d2_module(logger: Logger, platform: str) -> None: + """Load R2D2 module via bash, capture env, apply to current process.""" + if platform not in _R2D2_MODULE_CONFIG: + return + config = _R2D2_MODULE_CONFIG[platform] + cmd = ( + f'source /usr/share/lmod/lmod/init/bash && ' + f'module use -a {config["module_path"]} && ' + f'module load {config["module_name"]} && env' + ) + try: + result = subprocess.run(['bash', '-c', cmd], capture_output=True, text=True, timeout=30) + if result.returncode != 0: + logger.warning(f'Failed to load R2D2 module: {result.stderr}') + return + for line in result.stdout.strip().split('\n'): + if '=' in line: + key, _, value = line.partition('=') + os.environ[key] = value + # PYTHONPATH needs to be added to sys.path for import to work + if key == 'PYTHONPATH': + import sys + for p in value.split(':'): + if p and p not in sys.path: + sys.path.insert(0, p) + logger.info(f'Loaded R2D2 module: {config["module_name"]}') + except Exception as e: + logger.warning(f'Could not load R2D2 module: {e}') def create_r2d2_config( logger: Logger, From 7041a44a3d729913f5aab597ba6dd840b25f9365 Mon Sep 17 00:00:00 2001 From: Michael Anstett Date: Wed, 25 Feb 2026 15:33:21 -0500 Subject: [PATCH 07/16] few touches --- src/swell/swell.py | 5 ----- src/swell/utilities/r2d2.py | 8 +++++++- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/src/swell/swell.py b/src/swell/swell.py index 923ab779b..94675da52 100644 --- a/src/swell/swell.py +++ b/src/swell/swell.py @@ -112,11 +112,6 @@ def create( suite (str): Name of the suite you wish to run. \n """ - # Load R2D2 credentials - from swell.utilities.r2d2 import load_r2d2_credentials - from swell.utilities.logger import get_logger - logger = get_logger("Swell Create") - load_r2d2_credentials(logger, platform) # Create the experiment directory create_experiment_directory(suite, input_method, platform, override, advanced, slurm) diff --git a/src/swell/utilities/r2d2.py b/src/swell/utilities/r2d2.py index b3d60e558..bc4f08c2d 100644 --- a/src/swell/utilities/r2d2.py +++ b/src/swell/utilities/r2d2.py @@ -10,13 +10,13 @@ import os from ruamel.yaml import YAML import random +import subprocess from swell.swell_path import get_swell_path from swell.utilities.jinja2 import template_string_jinja2 from swell.utilities.logger import get_logger, Logger # -------------------------------------------------------------------------------------------------- -import subprocess # Platform-specific R2D2 module config _R2D2_MODULE_CONFIG = { @@ -30,6 +30,9 @@ }, } +# -------------------------------------------------------------------------------------------------- + + def load_r2d2_module(logger: Logger, platform: str) -> None: """Load R2D2 module via bash, capture env, apply to current process.""" if platform not in _R2D2_MODULE_CONFIG: @@ -59,6 +62,9 @@ def load_r2d2_module(logger: Logger, platform: str) -> None: except Exception as e: logger.warning(f'Could not load R2D2 module: {e}') +# ---------------------------------------------------------------------------------------------- + + def create_r2d2_config( logger: Logger, platform: str, From 07714eab8bed721f83c78e1e9b055dab4152fe80 Mon Sep 17 00:00:00 2001 From: Michael Anstett Date: Thu, 26 Feb 2026 13:02:22 -0500 Subject: [PATCH 08/16] fix mistake --- .../prepare_config_and_suite/prepare_config_and_suite.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 e66f2110e..0b230f845 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 @@ -1,4 +1,5 @@ - +# (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 07add942f54cd29d1a4146606706d2acb663bc3c Mon Sep 17 00:00:00 2001 From: Michael Anstett Date: Thu, 26 Feb 2026 13:10:34 -0500 Subject: [PATCH 09/16] Hard-code debug --- src/swell/deployment/create_experiment.py | 4 +--- src/swell/suites/suite_questions.py | 1 - src/swell/utilities/question_defaults.py | 10 ---------- 3 files changed, 1 insertion(+), 14 deletions(-) diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index bf8196017..dc57f9dad 100644 --- a/src/swell/deployment/create_experiment.py +++ b/src/swell/deployment/create_experiment.py @@ -179,8 +179,6 @@ def prepare_config( unique_id = unique_r2d2_id(r2d2_id, platform) experiment_dict['r2d2_experiment_id'] = unique_id - r2d2_lifetime = experiment_dict['r2d2_experiment_lifetime'] - user = r2d2.get_client_user() host = r2d2.get_client_host() compiler = r2d2.get_client_compiler() @@ -189,7 +187,7 @@ def prepare_config( name=unique_id, user=user, compute_host=f'{host}-{compiler}', - lifetime=r2d2_lifetime) + lifetime='debug') # Expand all environment vars in the dictionary # --------------------------------------------- diff --git a/src/swell/suites/suite_questions.py b/src/swell/suites/suite_questions.py index 2073e69b9..cc2824591 100644 --- a/src/swell/suites/suite_questions.py +++ b/src/swell/suites/suite_questions.py @@ -42,7 +42,6 @@ class SuiteQuestions(QuestionContainer, Enum): qd.model_components(), qd.runahead_limit(), qd.r2d2_experiment_id(), - qd.r2d2_experiment_lifetime() ] ) diff --git a/src/swell/utilities/question_defaults.py b/src/swell/utilities/question_defaults.py index 521220ae5..a63faa2c6 100644 --- a/src/swell/utilities/question_defaults.py +++ b/src/swell/utilities/question_defaults.py @@ -161,16 +161,6 @@ class r2d2_experiment_id(SuiteQuestion): # -------------------------------------------------------------------------------------------------- - @dataclass - class r2d2_experiment_lifetime(SuiteQuestion): - default_value: str = "debug" - question_name: str = "r2d2_experiment_lifetime" - options: list = mutable_field(['debug', 'science', 'publication', 'release']) - prompt: str = "What lifetime should the experiment have in R2D2?" - widget_type: WType = WType.STRING - - # -------------------------------------------------------------------------------------------------- - @dataclass class runahead_limit(SuiteQuestion): default_value: str = "P4" From d3abccac6da703287588b7a9e15bcaf4ff6ad92a Mon Sep 17 00:00:00 2001 From: Michael Anstett Date: Thu, 26 Feb 2026 17:28:30 -0500 Subject: [PATCH 10/16] Add option to skip R2D2 --- src/swell/deployment/create_experiment.py | 36 +++++++++++-- .../prepare_config_and_suite.py | 50 +++++++------------ src/swell/suites/3dfgat_atmos/flow.cylc | 6 ++- src/swell/suites/3dfgat_cycle/flow.cylc | 6 ++- src/swell/suites/3dvar/flow.cylc | 8 +-- src/swell/suites/3dvar_atmos/flow.cylc | 7 +-- src/swell/suites/3dvar_cycle/flow.cylc | 7 +-- src/swell/suites/hofx/flow.cylc | 7 +-- src/swell/suites/hofx_cf/flow.cylc | 9 ++-- src/swell/suites/suite_questions.py | 1 + src/swell/swell.py | 8 ++- src/swell/utilities/question_defaults.py | 9 ++++ 12 files changed, 94 insertions(+), 60 deletions(-) diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index dc57f9dad..23e883d14 100644 --- a/src/swell/deployment/create_experiment.py +++ b/src/swell/deployment/create_experiment.py @@ -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,9 +87,9 @@ def prepare_config( suite_config: str, method: str, platform: str, - override: Union[dict, str, None], + override: dict, advanced: bool, - slurm: str + slurm: str, ) -> str: # Create a logger @@ -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 0b230f845..571a69484 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 @@ -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 @@ -346,40 +346,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 372a6faa2..73e4e89b0 100644 --- a/src/swell/suites/3dfgat_atmos/flow.cylc +++ b/src/swell/suites/3dfgat_atmos/flow.cylc @@ -89,12 +89,14 @@ # 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 - EvaObservations-{{model_component}} & SaveObsDiags-{{model_component}} => - CleanCycle-{{model_component}} + EvaObservations-{{model_component}} => CleanCycle-{{model_component}} {% endif %} {% endfor %} diff --git a/src/swell/suites/3dfgat_cycle/flow.cylc b/src/swell/suites/3dfgat_cycle/flow.cylc index db6b339ba..86600410f 100644 --- a/src/swell/suites/3dfgat_cycle/flow.cylc +++ b/src/swell/suites/3dfgat_cycle/flow.cylc @@ -104,9 +104,11 @@ # Move restart to next cycle # 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}} @@ -116,7 +118,7 @@ MoveDaRestart-{{model_component}} => RemoveForecastDir # 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/flow.cylc b/src/swell/suites/3dvar/flow.cylc index 00e18ee8c..0bf10070c 100644 --- a/src/swell/suites/3dvar/flow.cylc +++ b/src/swell/suites/3dvar/flow.cylc @@ -71,13 +71,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 - EvaObservations-{{model_component}} & SaveObsDiags-{{model_component}} => - CleanCycle-{{model_component}} + EvaObservations-{{model_component}} => CleanCycle-{{model_component}} {% endif %} {% endfor %} diff --git a/src/swell/suites/3dvar_atmos/flow.cylc b/src/swell/suites/3dvar_atmos/flow.cylc index fab25dcfe..48c84f0fd 100644 --- a/src/swell/suites/3dvar_atmos/flow.cylc +++ b/src/swell/suites/3dvar_atmos/flow.cylc @@ -86,12 +86,13 @@ # 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 - EvaObservations-{{model_component}} & SaveObsDiags-{{model_component}} => - CleanCycle-{{model_component}} + EvaObservations-{{model_component}} => CleanCycle-{{model_component}} {% endif %} {% endfor %} diff --git a/src/swell/suites/3dvar_cycle/flow.cylc b/src/swell/suites/3dvar_cycle/flow.cylc index 8db2144fd..9aadcb786 100644 --- a/src/swell/suites/3dvar_cycle/flow.cylc +++ b/src/swell/suites/3dvar_cycle/flow.cylc @@ -102,9 +102,11 @@ # Move restart to next cycle # 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}} @@ -115,8 +117,7 @@ # Clean up large files # EvaObservations-{{model_component}} & EvaJediLog-{{model_component}} & SaveObsDiags-{{model_component}} & RemoveForecastDir => - EvaObservations-{{model_component}} & EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} & SaveObsDiags-{{model_component}} => - CleanCycle-{{model_component}} + EvaObservations-{{model_component}} & EvaJediLog-{{model_component}} & EvaIncrement-{{model_component}} => CleanCycle-{{model_component}} {% endfor %} """ {% 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 94675da52..9b6b6d12c 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 a63faa2c6..ffcf0fa46 100644 --- a/src/swell/utilities/question_defaults.py +++ b/src/swell/utilities/question_defaults.py @@ -184,6 +184,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" From 08e8674314c39e9bfb94f4a6c0f1d4708a0e8b9b Mon Sep 17 00:00:00 2001 From: Michael Anstett Date: Tue, 10 Mar 2026 13:54:55 -0400 Subject: [PATCH 11/16] Fixes --- src/swell/deployment/create_experiment.py | 5 +++-- .../prepare_config_and_suite/prepare_config_and_suite.py | 2 +- src/swell/suites/3dvar_marine/flow.cylc | 2 +- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index e60515505..29050c7e6 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 \ @@ -178,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 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 c260c0bde..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 diff --git a/src/swell/suites/3dvar_marine/flow.cylc b/src/swell/suites/3dvar_marine/flow.cylc index ef036e0be..7afdb11af 100644 --- a/src/swell/suites/3dvar_marine/flow.cylc +++ b/src/swell/suites/3dvar_marine/flow.cylc @@ -71,7 +71,7 @@ # EvaIncrement RunJediVariationalExecutable-{{model_component}} => EvaIncrement-{{model_component}} - + {% if not skip_r2d2 %} # Save observations RunJediVariationalExecutable-{{model_component}} => SaveObsDiags-{{model_component}} From e42e3b9297d225419f022bbb137d99cbde01fc4f Mon Sep 17 00:00:00 2001 From: Michael Anstett Date: Tue, 10 Mar 2026 14:37:35 -0400 Subject: [PATCH 12/16] remove comma --- src/swell/deployment/create_experiment.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/swell/deployment/create_experiment.py b/src/swell/deployment/create_experiment.py index 29050c7e6..0b7fa81e3 100644 --- a/src/swell/deployment/create_experiment.py +++ b/src/swell/deployment/create_experiment.py @@ -89,7 +89,7 @@ def prepare_config( platform: str, override: dict, advanced: bool, - slurm: str, + slurm: str ) -> str: # Create a logger From a97c229b7cca9fe5d7b719a371330485b9ac9816 Mon Sep 17 00:00:00 2001 From: Michael Anstett Date: Tue, 10 Mar 2026 14:47:05 -0400 Subject: [PATCH 13/16] code_test fix --- src/swell/utilities/r2d2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/swell/utilities/r2d2.py b/src/swell/utilities/r2d2.py index b735dbc0c..2d96a82de 100644 --- a/src/swell/utilities/r2d2.py +++ b/src/swell/utilities/r2d2.py @@ -14,7 +14,7 @@ from swell.swell_path import get_swell_path from swell.utilities.jinja2 import template_string_jinja2 -from swell.utilities.logger import get_logger, Logger +from swell.utilities.logger import Logger # -------------------------------------------------------------------------------------------------- From b07544334d4c5400fec5ba3ff62867ca95cfa175 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 18 Mar 2026 18:55:34 -0400 Subject: [PATCH 14/16] swell.py is refactored to commands.clone/create.py etc. The outcome did not speedup swell CLI, but simplified the source swell.py --- src/swell/commands/clone.py | 36 +++ src/swell/commands/create.py | 54 +++++ src/swell/commands/help_strings.py | 53 +++++ src/swell/commands/launch.py | 24 ++ src/swell/commands/t1test.py | 23 ++ src/swell/commands/t2test.py | 24 ++ src/swell/commands/task.py | 62 ++++++ src/swell/commands/test.py | 18 ++ src/swell/commands/utility.py | 18 ++ src/swell/swell.py | 339 +++++------------------------ 10 files changed, 372 insertions(+), 279 deletions(-) create mode 100644 src/swell/commands/clone.py create mode 100644 src/swell/commands/create.py create mode 100644 src/swell/commands/help_strings.py create mode 100644 src/swell/commands/launch.py create mode 100644 src/swell/commands/t1test.py create mode 100644 src/swell/commands/t2test.py create mode 100644 src/swell/commands/task.py create mode 100644 src/swell/commands/test.py create mode 100644 src/swell/commands/utility.py diff --git a/src/swell/commands/clone.py b/src/swell/commands/clone.py new file mode 100644 index 000000000..bad9f23f9 --- /dev/null +++ b/src/swell/commands/clone.py @@ -0,0 +1,36 @@ +import click +from swell.commands.help_strings import input_method_help, platform_help, advanced_help + + +@click.command() +@click.argument('configuration') +@click.argument('experiment_id') +@click.option('-m', '--input_method', 'input_method', default='defaults', + type=click.Choice(['defaults', 'cli']), help=input_method_help) +@click.option('-p', '--platform', 'platform', default=None, help=platform_help()) +@click.option('-a', '--advanced', 'advanced', default=False, help=advanced_help) +def clone( + configuration: str, + experiment_id: str, + input_method: str, + platform: str, + advanced: bool +) -> None: + """ + Clone an existing experiment + + This command creates an experiment directory based on the provided experiment configuration. + + Arguments: \n + configuration (str): Path to a YAML containing the experiment configuration you wish to + clone from. \n + + """ + + from swell.deployment.create_experiment import clone_config, create_experiment_directory + # Create experiment configuration by cloning from existing experiment + experiment_dict_str = clone_config(configuration, experiment_id, input_method, platform, + advanced) + + # Create the experiment directory + create_experiment_directory(experiment_dict_str) diff --git a/src/swell/commands/create.py b/src/swell/commands/create.py new file mode 100644 index 000000000..568e31d4c --- /dev/null +++ b/src/swell/commands/create.py @@ -0,0 +1,54 @@ +import os +import click +from typing import Union +from swell.swell_path import get_swell_path +from swell.deployment.platforms.platforms import get_platforms +from swell.commands.help_strings import (input_method_help, platform_help, + override_help, advanced_help, slurm_help, skip_r2d2_help) + + +class LazySuiteChoice(click.ParamType): + name = "suite" + + def convert(self, value, param, ctx): + suites_dir = os.path.join(get_swell_path(), 'suites') + suite_names = [ + name for name in os.listdir(suites_dir) + if os.path.isdir(os.path.join(suites_dir, name)) + and not name.startswith('__') + ] + if value not in suite_names: + self.fail(f"{value} is not a valid suite", param, ctx) + return value + + +@click.command() +@click.argument('suite', type=LazySuiteChoice()) +@click.option('-m', '--input_method', 'input_method', default='defaults', + type=click.Choice(['defaults', 'cli']), help=input_method_help) +@click.option('-p', '--platform', 'platform', default='nccs_discover_sles15', + type=click.Choice(get_platforms()), help=platform_help()) +@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, + skip_r2d2: bool +) -> None: + """ + Create a new experiment + + This command creates an experiment directory based on the provided suite name and options. + + Arguments: \n + suite (str): Name of the suite you wish to run. \n + + """ + from swell.deployment.create_experiment import create_experiment_directory + create_experiment_directory(suite, input_method, platform, override, advanced, slurm, skip_r2d2) diff --git a/src/swell/commands/help_strings.py b/src/swell/commands/help_strings.py new file mode 100644 index 000000000..e945e4b41 --- /dev/null +++ b/src/swell/commands/help_strings.py @@ -0,0 +1,53 @@ +# Help strings for optional arguments + +input_method_help = 'Method by which to create the YAML configuration file. If choosing ' + \ + 'defaults the setting for the default suite test will be used. If using ' + \ + 'CLI you will be led through the questions to configure the experiment.' + + +def platform_help(): + from swell.deployment.platforms.platforms import get_platforms + return ( + "If using defaults for input_method, this option is used to determine which " + "platform to use for platform specific defaults. Options are " + + str(get_platforms()) + ) + + +override_help = 'After generating the config file, parameters inside can be overridden ' + \ + 'using values from the override config file.' + + +advanced_help = 'Show configuration questions which are otherwise not shown to the user.' + + +no_detach_help = 'Tells the workflow manager not to detach. That is to say run the entire ' + \ + 'run the entire workflow in the foreground and pass back a return code.' + + +def log_path_help(): + return ( + 'Directory to receive workflow manager logging output (instead of ' + '$HOME/cylc-run/' + ) + + +datetime_help = 'Datetime to use for task execution. Format is yyyy-mm-ddThh:mm:ss. Note that ' + \ + 'non-numeric characters will be stripped from the string. Minutes and seconds ' + \ + 'are optional.' + + +model_help = 'Data assimilation system. I.e. the model being initialized by data assimilation.' + + +ensemble_help = 'When handling ensemble workflows using a parallel strategy, ' + \ + 'specify which packet of ensemble members to consider.' + + +slurm_help = """ +Customize SLURM directives, globally (e.g., account name), for specific tasks, +or for task-model combinations. +""" + + +skip_r2d2_help = """Skip registering this experiment and storing products in R2D2.""" diff --git a/src/swell/commands/launch.py b/src/swell/commands/launch.py new file mode 100644 index 000000000..4462f3e70 --- /dev/null +++ b/src/swell/commands/launch.py @@ -0,0 +1,24 @@ +import click +from swell.commands.help_strings import no_detach_help, log_path_help + + +@click.command() +@click.argument('suite_path') +@click.option('-b', '--no-detach', 'no_detach', is_flag=True, default=False, help=no_detach_help) +@click.option('-l', '--log_path', 'log_path', default=None, help=log_path_help()) +def launch( + suite_path: str, + no_detach: bool, + log_path: str +) -> None: + """ + Launch an experiment with the cylc workflow manager + + This command launches an experiment using the provided suite path and options. + + Arguments: \n + suite_path (str): Path to where the flow.cylc and associated suite files are located. \n + + """ + from swell.deployment.launch_experiment import launch_experiment + launch_experiment(suite_path, no_detach, log_path) diff --git a/src/swell/commands/t1test.py b/src/swell/commands/t1test.py new file mode 100644 index 000000000..996f51811 --- /dev/null +++ b/src/swell/commands/t1test.py @@ -0,0 +1,23 @@ +import click +from swell.deployment.platforms.platforms import get_platforms +from swell.commands.help_strings import platform_help +from typing import Optional, Literal + + +@click.command() +@click.option('-p', '--platform', 'platform', type=click.Choice(get_platforms()), + default="nccs_discover_sles15", help=platform_help()) +@click.argument('suite', type=click.Choice(("hofx", "3dvar_marine", "3dvar_atmos", + "localensembleda", "3dvar_cycle"))) +def t1test( + suite: Literal["hofx", "3dvar_marine", "3dvar_atmos", "localensembleda", "3dvar_cycle"], + platform: Optional[str] = "nccs_discover_sles15" +) -> None: + """ + Run a particular swell suite from the tier 1 tests. + + Arguments: + suite (str): Name of the suite to run (e.g., 3dvar_marine, 3dvar_atmos, localensembleda) + """ + from swell.test.suite_tests.suite_tests import run_suite, TestSuite + run_suite(suite, platform, TestSuite.TIER1) diff --git a/src/swell/commands/t2test.py b/src/swell/commands/t2test.py new file mode 100644 index 000000000..ddcedb27a --- /dev/null +++ b/src/swell/commands/t2test.py @@ -0,0 +1,24 @@ +import click +from swell.deployment.platforms.platforms import get_platforms +from swell.commands.help_strings import platform_help +from typing import Optional, Literal + + +@click.command() +@click.option('-p', '--platform', 'platform', type=click.Choice(get_platforms()), + default="nccs_discover_sles15", help=platform_help()) +@click.argument('suite', type=click.Choice(("hofx", "3dvar_marine", "ufo_testing", + "convert_ncdiags", "3dfgat_atmos", "build_jedi"))) +def t2test( + suite: Literal["hofx", "3dvar_marine", "ufo_testing", + "convert_ncdiags", "3dfgat_atmos", "build_jedi"], + platform: Optional[str] = "nccs_discover_sles15" +) -> None: + """ + Run a particular swell suite from the tier 2 tests. + + Arguments: + suite (str): Name of the suite to run (e.g., hofx, 3dvar_marine, ufo_testing) + """ + from swell.test.suite_tests.suite_tests import run_suite, TestSuite + run_suite(suite, platform, TestSuite.TIER2) diff --git a/src/swell/commands/task.py b/src/swell/commands/task.py new file mode 100644 index 000000000..24599c5c7 --- /dev/null +++ b/src/swell/commands/task.py @@ -0,0 +1,62 @@ +import os +import glob +import click +from typing import Union, Optional +from swell.swell_path import get_swell_path +from swell.commands.help_strings import datetime_help, model_help, ensemble_help + + +def snake_case_to_camel_case(name: str) -> str: + return ''.join(word.capitalize() for word in name.split('_')) + + +class TaskChoice(click.ParamType): + """Click parameter type for dynamically listing tasks.""" + + name = "task" + + def convert(self, value, param, ctx): + tasks = self.get_tasks() # lazy evaluation + if value not in tasks: + self.fail(f"{value} is not a valid task", param, ctx) + return value + + @staticmethod + def get_tasks() -> list[str]: + """Return list of task names in CamelCase.""" + tasks_dir = os.path.join(get_swell_path(), "tasks", "*.py") + return sorted( + snake_case_to_camel_case(os.path.splitext(os.path.basename(f))[0]) + for f in glob.glob(tasks_dir) + if "__" not in os.path.basename(f) + ) + + +TASK = TaskChoice() + + +@click.command() +@click.argument('task', type=TASK) +@click.argument('config') +@click.option('-d', '--datetime', 'datetime', default=None, help=datetime_help) +@click.option('-m', '--model', 'model', default=None, help=model_help) +@click.option('-p', '--ensemblePacket', 'ensemblePacket', default=None, help=ensemble_help) +def task( + task: str, + config: str, + datetime: Optional[str], + model: Optional[str], + ensemblePacket: Optional[str] +) -> None: + """ + Run a workflow task + + This command executes a task using the provided task name, configuration file and options. + + Arguments:\n + task (str): Name of the task to execute.\n + config (str): Path to the configuration file for the task.\n + + """ + from swell.tasks.base.task_base import task_wrapper + task_wrapper(task, config, datetime, model, ensemblePacket) diff --git a/src/swell/commands/test.py b/src/swell/commands/test.py new file mode 100644 index 000000000..bff3ee7b2 --- /dev/null +++ b/src/swell/commands/test.py @@ -0,0 +1,18 @@ +import click +from swell.test.test_driver import valid_tests + + +@click.command() +@click.argument('test', type=click.Choice(valid_tests)) +def test(test: str) -> None: + """ + Run one of the test suites + + This command performs the test specified by the test argument. + + Arguments:\n + test (str): Name of the test to execute. + + """ + from swell.test.test_driver import test_wrapper + test_wrapper(test) diff --git a/src/swell/commands/utility.py b/src/swell/commands/utility.py new file mode 100644 index 000000000..0d729298f --- /dev/null +++ b/src/swell/commands/utility.py @@ -0,0 +1,18 @@ +import click +from swell.utilities.scripts.utility_driver import get_utilities + + +@click.command() +@click.argument('utility', type=click.Choice(get_utilities())) +def utility(utility: str) -> None: + """ + Run a utility script + + This command performs a utility operation specified by the utility argument. + + Arguments:\n + utility (str): Name of the utility operation to perform.\n + + """ + from swell.utilities.scripts.utility_driver import utility_wrapper + utility_wrapper(utility) diff --git a/src/swell/swell.py b/src/swell/swell.py index 89f55a3ad..59c90a1fc 100644 --- a/src/swell/swell.py +++ b/src/swell/swell.py @@ -1,32 +1,61 @@ -# (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 click -from typing import Union, Optional, Literal - -from swell.deployment.platforms.platforms import get_platforms -from swell.deployment.create_experiment import clone_config, create_experiment_directory -from swell.deployment.launch_experiment import launch_experiment -from swell.tasks.base.task_base import task_wrapper, get_tasks -from swell.test.test_driver import test_wrapper, valid_tests -from swell.test.suite_tests.suite_tests import run_suite, TestSuite -from swell.suites.all_suites import AllSuites -from swell.utilities.welcome_message import write_welcome_message -from swell.utilities.scripts.utility_driver import get_utilities, utility_wrapper - +import importlib +import sys -# -------------------------------------------------------------------------------------------------- +# Import the package to access __version__ +import swell +__version__ = swell.__version__ +from swell.utilities.welcome_message import write_welcome_message -@click.group() -def swell_driver() -> None: +COMMANDS = { + "clone": "swell.commands.clone", + "create": "swell.commands.create", + "launch": "swell.commands.launch", + "t1test": "swell.commands.t1test", + "t2test": "swell.commands.t2test", + "task": "swell.commands.task", + "test": "swell.commands.test", + "utility": "swell.commands.utility", +} + + +class SwellCLI(click.MultiCommand): + def list_commands(self, ctx): + """Returns sorted list of command names for the Commands: section""" + return sorted(COMMANDS.keys()) + + def get_command(self, ctx, name): + """Lazy-load the requested command module and return its Click command object""" + if name not in COMMANDS: + raise click.UsageError(f"Unknown command: {name}") + + module = importlib.import_module(COMMANDS[name]) + + # The command function is usually named the same as the subcommand + cmd_func_name = name.replace("-", "_") # handles possible future hyphenated names + if hasattr(module, cmd_func_name): + return getattr(module, cmd_func_name) + + # Fallback: look for any click.Command in the module + for attr_name in dir(module): + attr = getattr(module, attr_name) + if isinstance(attr, click.Command): + return attr + + raise click.UsageError( + f"Command '{name}' not found in {COMMANDS[name]}. " + f"Expected a Click command named '{cmd_func_name}' or similar." + ) + + +@click.group(cls=SwellCLI) +@click.version_option( + version=__version__, + prog_name="swell", + message="%(prog)s version %(version)s", +) +def swell_driver(): """ Welcome to swell! @@ -40,266 +69,18 @@ def swell_driver() -> None: followed by swell launch - """ pass -# -------------------------------------------------------------------------------------------------- - -# Help strings for optional arguments - -input_method_help = 'Method by which to create the YAML configuration file. If choosing ' + \ - 'defaults the setting for the default suite test will be used. If using ' + \ - 'CLI you will be led through the questions to configure the experiment.' - -platform_help = 'If using defaults for input_method, this option is used to determine which ' + \ - 'platform to use for platform specific defaults. Options are ' + \ - str(get_platforms()) - -override_help = 'After generating the config file, parameters inside can be overridden ' + \ - 'using values from the override config file.' - -advanced_help = 'Show configuration questions which are otherwise not shown to the user.' - -no_detach_help = 'Tells the workflow manager not to detach. That is to say run the entire ' + \ - 'run the entire workflow in the foreground and pass back a return code.' - -log_path_help = 'Directory to receive workflow manager logging output (instead of ' + \ - '$HOME/cylc-run/)' - -datetime_help = 'Datetime to use for task execution. Format is yyyy-mm-ddThh:mm:ss. Note that ' + \ - 'non-numeric characters will be stripped from the string. Minutes and seconds ' + \ - 'are optional.' - -model_help = 'Data assimilation system. I.e. the model being initialized by data assimilation.' - -ensemble_help = 'When handling ensemble workflows using a parallel strategy, ' + \ - 'specify which packet of ensemble members to consider.' - -slurm_help = """ -Customize SLURM directives, globally (e.g., account name), for specific tasks, -or for task-model combinations. -""" - -skip_r2d2_help = """Skip registering this experiment and storing products in R2D2.""" - - -# -------------------------------------------------------------------------------------------------- - - -@swell_driver.command() -@click.argument('suite', type=click.Choice(AllSuites.config_names())) -@click.option('-m', '--input_method', 'input_method', default='defaults', - type=click.Choice(['defaults', 'cli']), help=input_method_help) -@click.option('-p', '--platform', 'platform', default='nccs_discover_sles15', - type=click.Choice(get_platforms()), help=platform_help) -@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, - skip_r2d2: bool -) -> None: - """ - Create a new experiment - - This command creates an experiment directory based on the provided suite name and options. - - Arguments: \n - suite (str): Name of the suite you wish to run. \n - - """ - - # Create the experiment directory - create_experiment_directory(suite, input_method, platform, override, advanced, slurm, skip_r2d2) - - -# -------------------------------------------------------------------------------------------------- - - -@swell_driver.command() -@click.argument('configuration') -@click.argument('experiment_id') -@click.option('-m', '--input_method', 'input_method', default='defaults', - type=click.Choice(['defaults', 'cli']), help=input_method_help) -@click.option('-p', '--platform', 'platform', default=None, help=platform_help) -@click.option('-a', '--advanced', 'advanced', default=False, help=advanced_help) -def clone( - configuration: str, - experiment_id: str, - input_method: str, - platform: str, - advanced: bool -) -> None: - """ - Clone an existing experiment - - This command creates an experiment directory based on the provided experiment configuration. - - Arguments: \n - configuration (str): Path to a YAML containing the experiment configuration you wish to - clone from. \n - - """ - # Create experiment configuration by cloning from existing experiment - experiment_dict_str = clone_config(configuration, experiment_id, input_method, platform, - advanced) - - # Create the experiment directory - create_experiment_directory(experiment_dict_str) - - -# -------------------------------------------------------------------------------------------------- - - -@swell_driver.command() -@click.argument('suite_path') -@click.option('-b', '--no-detach', 'no_detach', is_flag=True, default=False, help=no_detach_help) -@click.option('-l', '--log_path', 'log_path', default=None, help=log_path_help) -def launch( - suite_path: str, - no_detach: bool, - log_path: str -) -> None: - """ - Launch an experiment with the cylc workflow manager - - This command launches an experiment using the provided suite path and options. - - Arguments: \n - suite_path (str): Path to where the flow.cylc and associated suite files are located. \n - - """ - launch_experiment(suite_path, no_detach, log_path) - - -# -------------------------------------------------------------------------------------------------- - - -@swell_driver.command() -@click.argument('task', type=click.Choice(get_tasks())) -@click.argument('config') -@click.option('-d', '--datetime', 'datetime', default=None, help=datetime_help) -@click.option('-m', '--model', 'model', default=None, help=model_help) -@click.option('-p', '--ensemblePacket', 'ensemblePacket', default=None, help=ensemble_help) -def task( - task: str, - config: str, - datetime: Optional[str], - model: Optional[str], - ensemblePacket: Optional[str] -) -> None: - """ - Run a workflow task - - This command executes a task using the provided task name, configuration file and options. - - Arguments:\n - task (str): Name of the task to execute.\n - config (str): Path to the configuration file for the task.\n - - """ - task_wrapper(task, config, datetime, model, ensemblePacket) - - -# -------------------------------------------------------------------------------------------------- - - -@swell_driver.command() -@click.argument('utility', type=click.Choice(get_utilities())) -def utility(utility: str) -> None: - """ - Run a utility script - - This command performs a utility operation specified by the utility argument. - - Arguments:\n - utility (str): Name of the utility operation to perform.\n - - """ - utility_wrapper(utility) - - -# -------------------------------------------------------------------------------------------------- - - -@swell_driver.command() -@click.argument('test', type=click.Choice(valid_tests)) -def test(test: str) -> None: - """ - Run one of the test suites - - This command performs the test specified by the test argument. - - Arguments:\n - test (str): Name of the test to execute. - - """ - test_wrapper(test) - - -# -------------------------------------------------------------------------------------------------- - - -@swell_driver.command() -@click.option('-p', '--platform', 'platform', type=click.Choice(get_platforms()), - default="nccs_discover_sles15", help=platform_help) -@click.argument('suite', type=click.Choice(("hofx", "3dvar_marine", "3dvar_atmos", - "localensembleda", "3dvar_cycle"))) -def t1test( - suite: Literal["hofx", "3dvar_marine", "3dvar_atmos", "localensembleda", "3dvar_cycle"], - platform: Optional[str] = "nccs_discover_sles15" -) -> None: - """ - Run a particular swell suite from the tier 1 tests. - - Arguments: - suite (str): Name of the suite to run (e.g., 3dvar_marine, 3dvar_atmos, localensembleda) - """ - run_suite(suite, platform, TestSuite.TIER1) - - -# -------------------------------------------------------------------------------------------------- - - -@swell_driver.command() -@click.option('-p', '--platform', 'platform', type=click.Choice(get_platforms()), - default="nccs_discover_sles15", help=platform_help) -@click.argument('suite', type=click.Choice(("hofx", "3dvar_marine", "ufo_testing", - "convert_ncdiags", "3dfgat_atmos", "build_jedi"))) -def t2test( - suite: Literal["hofx", "3dvar_marine", "ufo_testing", - "convert_ncdiags", "3dfgat_atmos", "build_jedi"], - platform: Optional[str] = "nccs_discover_sles15" -) -> None: - """ - Run a particular swell suite from the tier 2 tests. - - Arguments: - suite (str): Name of the suite to run (e.g., hofx, 3dvar_marine, ufo_testing) - """ - run_suite(suite, platform, TestSuite.TIER2) - - -# -------------------------------------------------------------------------------------------------- - - def main() -> None: - """ - Main Function + if len(sys.argv) > 1: + first_arg = sys.argv[1] + if first_arg not in ("--help", "-h", "--version", "-V") and "--help" not in sys.argv: + write_welcome_message() - This function is the entry point for swell. It writes a welcome message and - sets up the driver group. - """ - write_welcome_message() swell_driver() -# -------------------------------------------------------------------------------------------------- +if __name__ == "__main__": + main() From 35634550d937660f09390c30c983640eda5a0f78 Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Wed, 18 Mar 2026 20:00:26 -0400 Subject: [PATCH 15/16] sovle code style issue --- src/swell/swell.py | 47 +++++++++------------------------------------- 1 file changed, 9 insertions(+), 38 deletions(-) diff --git a/src/swell/swell.py b/src/swell/swell.py index 59c90a1fc..c5b545c9f 100644 --- a/src/swell/swell.py +++ b/src/swell/swell.py @@ -2,12 +2,6 @@ import importlib import sys -# Import the package to access __version__ -import swell -__version__ = swell.__version__ - -from swell.utilities.welcome_message import write_welcome_message - COMMANDS = { "clone": "swell.commands.clone", "create": "swell.commands.create", @@ -22,46 +16,24 @@ class SwellCLI(click.MultiCommand): def list_commands(self, ctx): - """Returns sorted list of command names for the Commands: section""" return sorted(COMMANDS.keys()) def get_command(self, ctx, name): - """Lazy-load the requested command module and return its Click command object""" if name not in COMMANDS: - raise click.UsageError(f"Unknown command: {name}") - + return None + # Move imports inside the getter so they only trigger when the command is called module = importlib.import_module(COMMANDS[name]) - - # The command function is usually named the same as the subcommand - cmd_func_name = name.replace("-", "_") # handles possible future hyphenated names - if hasattr(module, cmd_func_name): - return getattr(module, cmd_func_name) - - # Fallback: look for any click.Command in the module - for attr_name in dir(module): - attr = getattr(module, attr_name) - if isinstance(attr, click.Command): - return attr - - raise click.UsageError( - f"Command '{name}' not found in {COMMANDS[name]}. " - f"Expected a Click command named '{cmd_func_name}' or similar." - ) + cmd_func_name = name.replace("-", "_") + return getattr(module, cmd_func_name, None) @click.group(cls=SwellCLI) -@click.version_option( - version=__version__, - prog_name="swell", - message="%(prog)s version %(version)s", -) +# Fetch version lazily to avoid importing the whole package on every call +@click.version_option(package_name='swell') def swell_driver(): """ Welcome to swell! - This is the top level driver for swell. It serves as a container for various commands - related to experiment creation, launching, tasks, and utilities. - The normal process for creating and running an experiment is to issue: swell create @@ -74,10 +46,9 @@ def swell_driver(): def main() -> None: - if len(sys.argv) > 1: - first_arg = sys.argv[1] - if first_arg not in ("--help", "-h", "--version", "-V") and "--help" not in sys.argv: - write_welcome_message() + if len(sys.argv) > 1 and sys.argv[1] not in ("--help", "-h", "--version", "-V"): + from swell.utilities.welcome_message import write_welcome_message + write_welcome_message() swell_driver() From 2ce9fe28095e3cf7604b8abc014a8ef9f8b72aec Mon Sep 17 00:00:00 2001 From: Yonggang Yu Date: Mon, 23 Mar 2026 14:51:30 -0400 Subject: [PATCH 16/16] fixed one code_tests error --- src/swell/commands/task.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/swell/commands/task.py b/src/swell/commands/task.py index 24599c5c7..523751dec 100644 --- a/src/swell/commands/task.py +++ b/src/swell/commands/task.py @@ -1,7 +1,7 @@ import os import glob import click -from typing import Union, Optional +from typing import Optional from swell.swell_path import get_swell_path from swell.commands.help_strings import datetime_help, model_help, ensemble_help