From 83956aad9a34178e2463c6adf9a8b44108c6d1cf Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:38:22 +0000 Subject: [PATCH 1/2] Initial plan From e0f31ef3b7b63345a71edae5cf39ae11fa53b12d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:40:53 +0000 Subject: [PATCH 2/2] refactor: simplify shared JEDI bias staging Co-authored-by: ADCollard <40322596+ADCollard@users.noreply.github.com> --- .../unittests/test_jedi_bias_staging.py | 138 ++++++++++++++++++ ush/python/pygfs/jedi/jedi.py | 33 ++--- 2 files changed, 151 insertions(+), 20 deletions(-) create mode 100644 dev/ci/scripts/unittests/test_jedi_bias_staging.py diff --git a/dev/ci/scripts/unittests/test_jedi_bias_staging.py b/dev/ci/scripts/unittests/test_jedi_bias_staging.py new file mode 100644 index 00000000000..d4461179122 --- /dev/null +++ b/dev/ci/scripts/unittests/test_jedi_bias_staging.py @@ -0,0 +1,138 @@ +import importlib.util +import logging +import shutil +import sys +import types +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[4] +JEDI_PATH = REPO_ROOT / 'ush/python/pygfs/jedi/jedi.py' + + +def load_jedi_module(monkeypatch): + wxflow = types.ModuleType('wxflow') + + class AttrDict(dict): + def __getattr__(self, name): + try: + return self[name] + except KeyError as err: + raise AttributeError(name) from err + + def __setattr__(self, name, value): + self[name] = value + + def deepcopy(self): + return AttrDict(self) + + class WorkflowKeyError(KeyError): + pass + + def logit(_logger): + def decorator(func): + return func + return decorator + + wxflow.AttrDict = AttrDict + wxflow.FileHandler = object + wxflow.Task = object + wxflow.Executable = object + wxflow.WorkflowException = Exception + wxflow.WorkflowKeyError = WorkflowKeyError + wxflow.WorkflowTypeError = TypeError + wxflow.chdir = lambda *args, **kwargs: None + wxflow.parse_j2yaml = lambda *args, **kwargs: {} + wxflow.save_as_yaml = lambda *args, **kwargs: None + wxflow.logit = logit + + jcb = types.ModuleType('jcb') + jcb.render = lambda *args, **kwargs: None + + monkeypatch.setitem(sys.modules, 'wxflow', wxflow) + monkeypatch.setitem(sys.modules, 'jcb', jcb) + + spec = importlib.util.spec_from_file_location('test_jedi_module', JEDI_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class RecordingFileHandler: + calls = [] + + def __init__(self, config): + self.config = config + RecordingFileHandler.calls.append(config) + + def sync(self): + for directory in self.config.get('mkdir', []): + Path(directory).mkdir(parents=True, exist_ok=True) + for source, destination in self.config.get('copy_opt', []): + Path(destination).mkdir(parents=True, exist_ok=True) + shutil.copy(source, destination) + + +def test_stage_obsbiasin_stages_shared_varbc_tarball_once(tmp_path, monkeypatch): + module = load_jedi_module(monkeypatch) + RecordingFileHandler.calls = [] + monkeypatch.setattr(module, 'FileHandler', RecordingFileHandler) + + extracted = [] + monkeypatch.setattr(module.Jedi, 'extract_tar', staticmethod(lambda tar_file: extracted.append(tar_file))) + + comin = tmp_path / 'comin' + obsbiasin = tmp_path / 'obsbiasin' + obsbiasout = tmp_path / 'obsbiasout' + comin.mkdir() + (comin / 'gdas.varbc_params.tar').write_text('stub tarball', encoding='ascii') + + jedi = object.__new__(module.Jedi) + jedi.component = 'gdas' + jedi.jcb_config = { + 'observations': ['obs_a', 'obs_b', 'obs_a'], + 'gdas_obsbiasin_path': str(obsbiasin), + 'gdas_obsbiasout_path': str(obsbiasout), + 'gdas_obsbiasin_prefix': 'gdas.', + } + + jedi.stage_obsbiasin(str(comin)) + + assert RecordingFileHandler.calls == [{ + 'mkdir': [str(obsbiasin), str(obsbiasout)], + 'copy_opt': [[str(comin / 'gdas.varbc_params.tar'), str(obsbiasin)]], + }] + assert extracted == [str(obsbiasin / 'gdas.varbc_params.tar')] + + +def test_stage_obsbiasin_warns_and_skips_missing_shared_tarball(tmp_path, monkeypatch, caplog): + module = load_jedi_module(monkeypatch) + RecordingFileHandler.calls = [] + monkeypatch.setattr(module, 'FileHandler', RecordingFileHandler) + + extracted = [] + monkeypatch.setattr(module.Jedi, 'extract_tar', staticmethod(lambda tar_file: extracted.append(tar_file))) + + comin = tmp_path / 'comin' + obsbiasin = tmp_path / 'obsbiasin' + obsbiasout = tmp_path / 'obsbiasout' + comin.mkdir() + + jedi = object.__new__(module.Jedi) + jedi.component = 'gdas' + jedi.jcb_config = { + 'observations': ['obs_a', 'obs_a'], + 'gdas_obsbiasin_path': str(obsbiasin), + 'gdas_obsbiasout_path': str(obsbiasout), + 'gdas_obsbiasin_prefix': 'gdas.', + } + + with caplog.at_level(logging.WARNING): + jedi.stage_obsbiasin(str(comin)) + + assert RecordingFileHandler.calls == [{ + 'mkdir': [str(obsbiasin), str(obsbiasout)], + 'copy_opt': [], + }] + assert extracted == [] + assert f"Bias correction file {comin / 'gdas.varbc_params.tar'} does not exist and will be skipped" in caplog.text diff --git a/ush/python/pygfs/jedi/jedi.py b/ush/python/pygfs/jedi/jedi.py index 137ce43f4ac..00121c8af4f 100644 --- a/ush/python/pygfs/jedi/jedi.py +++ b/ush/python/pygfs/jedi/jedi.py @@ -14,7 +14,7 @@ logit) logger = getLogger(__name__.split('.')[-1]) - +VARBC_PARAMS_TAR = 'varbc_params.tar' required_jedi_keys = ['jedi_app_name', 'rundir', 'exe_src', 'mpi_cmd', 'jcb_base_yaml'] optional_jedi_keys = ['jedi_args', 'jcb_algo', 'jcb_algo_yaml', 'obs_list_yaml', 'bias_files_yaml', 'app_test_yaml'] @@ -406,32 +406,25 @@ def stage_obsbiasin(self, comin) -> None: fh_dict['mkdir'].append(self.jcb_config[f'{self.component}_obsbiasin_path']) fh_dict['mkdir'].append(self.jcb_config[f'{self.component}_obsbiasout_path']) - # Copy files - files_already_copied = [] bias_dest = self.jcb_config[f'{self.component}_obsbiasin_path'] - for observation_from_jcb in self.jcb_config['observations']: - if observation_from_jcb in self.jcb_config.bias_files_dict and observation_from_jcb not in files_already_copied: - bias_src = os.path.join(comin, self.jcb_config[f'{self.component}_obsbiasin_prefix'] + self.jcb_config.bias_files_dict[observation_from_jcb]) - + bias_tarball = self.jcb_config[f'{self.component}_obsbiasin_prefix'] + VARBC_PARAMS_TAR + if self.jcb_config['observations']: + bias_src = os.path.join(comin, bias_tarball) + if os.path.exists(bias_src): fh_dict['copy_opt'].append([bias_src, bias_dest]) - - # Don't copy same file multiple times - files_already_copied.append(observation_from_jcb) + else: + logger.warning(f"Bias correction file {bias_src} does not exist and will be skipped") # Execute FileHandler sync FileHandler(fh_dict).sync() # Untar bias corrections - bias_file_list = [] - for ob in self.jcb_config['observations']: - if ob in self.jcb_config.bias_files_dict and not self.jcb_config.bias_files_dict[ob] in bias_file_list: - bias_file_list.append(self.jcb_config.bias_files_dict[ob]) - bias_file_path = os.path.join(self.jcb_config[f"{self.component}_obsbiasin_path"], - self.jcb_config[f"{self.component}_obsbiasin_prefix"] + self.jcb_config.bias_files_dict[ob]) - if os.path.exists(bias_file_path): - Jedi.extract_tar(bias_file_path) - else: - logger.warning(f"Bias correction file {bias_file_path} does not exist and will be skipped") + if fh_dict['copy_opt']: + bias_file_path = os.path.join(bias_dest, bias_tarball) + if os.path.exists(bias_file_path): + Jedi.extract_tar(bias_file_path) + else: + logger.warning(f"Bias correction file {bias_file_path} does not exist and will be skipped") @logit(logger) def save_obsbiasout(self, comout: str, archive_name: str) -> None: