diff --git a/dpgen2/superop/prep_run_calypso.py b/dpgen2/superop/prep_run_calypso.py index daa48143..998adec8 100644 --- a/dpgen2/superop/prep_run_calypso.py +++ b/dpgen2/superop/prep_run_calypso.py @@ -290,7 +290,7 @@ def _prep_run_caly( key="%s--run-caly-model-devi-{{item}}" % (prep_run_caly_steps.inputs.parameters["block_id"],), executor=run_executor, - **prep_config, + **run_config, ) prep_run_caly_steps.add(run_caly_model_devi) diff --git a/tests/test_prep_run_caly.py b/tests/test_prep_run_caly.py index e949410d..f303e0f5 100644 --- a/tests/test_prep_run_caly.py +++ b/tests/test_prep_run_caly.py @@ -4,9 +4,15 @@ import shutil import time import unittest +from collections import ( + defaultdict, +) from pathlib import ( Path, ) +from types import ( + SimpleNamespace, +) from typing import ( List, Set, @@ -36,6 +42,10 @@ OPIOSign, PythonOPTemplate, ) +from mock import ( + Mock, + patch, +) from dpgen2.constants import ( calypso_check_opt_file, @@ -89,6 +99,7 @@ ) from dpgen2.superop.prep_run_calypso import ( PrepRunCaly, + _prep_run_caly, ) from dpgen2.utils.step_config import normalize as normalize_step_dict @@ -101,6 +112,99 @@ ) +class TestPrepRunCalyConfiguration(unittest.TestCase): + def test_model_deviation_step_uses_run_config(self): + """Route every step-level control from the intended phase config.""" + step_config_keys = ( + "continue_on_failed", + "continue_on_num_success", + "continue_on_success_ratio", + "parallelism", + ) + + def make_mapping(): + return defaultdict(Mock) + + prep_config = normalize_step_dict( + { + "continue_on_failed": False, + "continue_on_num_success": 1, + "continue_on_success_ratio": 0.1, + "parallelism": 3, + } + ) + run_config = normalize_step_dict( + { + "continue_on_failed": True, + "continue_on_num_success": 9, + "continue_on_success_ratio": 0.9, + "parallelism": 7, + } + ) + + for expl_mode in ("default", "merge"): + with self.subTest(expl_mode=expl_mode): + step_calls = [] + + def make_step(name, *args, **kwargs): + step_calls.append((name, kwargs)) + return SimpleNamespace( + outputs=SimpleNamespace( + parameters=make_mapping(), + artifacts=make_mapping(), + ) + ) + + prep_run_steps = SimpleNamespace( + inputs=SimpleNamespace( + parameters=make_mapping(), + artifacts=make_mapping(), + ), + outputs=SimpleNamespace(artifacts=make_mapping()), + add=Mock(), + ) + + with ( + patch( + "dpgen2.superop.prep_run_calypso.Step", + side_effect=make_step, + ), + patch("dpgen2.superop.prep_run_calypso.PythonOPTemplate"), + patch("dpgen2.superop.prep_run_calypso.Slices"), + patch("dpgen2.superop.prep_run_calypso.argo_range"), + patch( + "dpgen2.superop.prep_run_calypso.init_executor", + side_effect=lambda value: value, + ), + ): + _prep_run_caly( + prep_run_steps, + defaultdict(str), + Mock(), + Mock(), + Mock(), + Mock(), + expl_mode=expl_mode, + prep_config=prep_config, + run_config=run_config, + ) + + calls_by_name = dict(step_calls) + expected_configs = { + "prep-caly-input": prep_config, + "caly-evo-step": ( + prep_config if expl_mode == "default" else run_config + ), + "run-caly-model-devi": run_config, + } + for step_name, expected_config in expected_configs.items(): + actual_kwargs = calls_by_name[step_name] + self.assertEqual( + {key: actual_kwargs[key] for key in step_config_keys}, + {key: expected_config[key] for key in step_config_keys}, + ) + + def make_task_group_list(njobs): tgrp = BaseExplorationTaskGroup() for ii in range(njobs):