diff --git a/arc/job/pipe/pipe_planner.py b/arc/job/pipe/pipe_planner.py index 660ed5af8e..c39df31cf4 100644 --- a/arc/job/pipe/pipe_planner.py +++ b/arc/job/pipe/pipe_planner.py @@ -52,6 +52,7 @@ logger = get_logger() pipe_settings = settings['pipe_settings'] +rotor_scan_resolution = settings['rotor_scan_resolution'] class PipePlanner: @@ -308,7 +309,8 @@ def try_pipe_rotor_scans_1d(self, label: str, rotor_indices: list[int]) -> set[i job_type='scan', build_tasks_fn=lambda adapter: build_rotor_scan_1d_tasks( self.sched.species_dict[label], label, rotor_indices, - self._level_dict(level), adapter, self._memory_mb), + self._level_dict(level), adapter, self._memory_mb, + scan_res=rotor_scan_resolution), log_msg=f'Routing {len(rotor_indices)} 1D rotor scans for {label} to pipe mode', ) return set(rotor_indices) if submitted else set() diff --git a/arc/job/pipe/pipe_run.py b/arc/job/pipe/pipe_run.py index ba4bc34245..c4f344d2a3 100644 --- a/arc/job/pipe/pipe_run.py +++ b/arc/job/pipe/pipe_run.py @@ -1049,8 +1049,15 @@ def build_ts_opt_tasks(species, label: str, xyzs: list[dict], def build_rotor_scan_1d_tasks(species, label: str, rotor_indices: list[int], level_dict: dict, job_adapter: str, - memory_mb: int) -> list[TaskSpec]: - """Build TaskSpec objects for 1D rotor scan tasks.""" + memory_mb: int, scan_res: float | None = None) -> list[TaskSpec]: + """ + Build TaskSpec objects for 1D rotor scan tasks. + + The scan resolution is captured here, at staging time on the ARC host, and carried in each + task's payload so the worker runs the resolution ARC intended rather than re-reading whatever + ``rotor_scan_resolution`` the worker node's settings happen to hold. When ``scan_res`` is + ``None`` the payload omits it and the worker falls back to its own settings default. + """ cores = default_job_settings.get('job_cpu_cores', 8) species_dict_payload = species.as_dict() tasks = [] @@ -1059,6 +1066,13 @@ def build_rotor_scan_1d_tasks(species, label: str, rotor_indices: list[int], torsions = rotor['torsion'] if isinstance(torsions[0], int): torsions = [torsions] + input_payload = { + 'species_dicts': [species_dict_payload], + 'torsions': torsions, + 'rotor_index': ri, + } + if scan_res is not None: + input_payload['scan_res'] = scan_res tasks.append(TaskSpec( task_id=f'{label}_scan_r{ri}', task_family='rotor_scan_1d', @@ -1069,11 +1083,7 @@ def build_rotor_scan_1d_tasks(species, label: str, rotor_indices: list[int], level=level_dict, required_cores=cores, required_memory_mb=memory_mb, - input_payload={ - 'species_dicts': [species_dict_payload], - 'torsions': torsions, - 'rotor_index': ri, - }, + input_payload=input_payload, ingestion_metadata={'rotor_index': ri}, )) return tasks diff --git a/arc/job/pipe/pipe_run_test.py b/arc/job/pipe/pipe_run_test.py index 0556cf204b..e929125a1e 100644 --- a/arc/job/pipe/pipe_run_test.py +++ b/arc/job/pipe/pipe_run_test.py @@ -17,7 +17,8 @@ from arc.job.adapters.mockter import MockAdapter from arc.job.pipe.pipe_state import (TaskState, TaskStateRecord, PipeRunState, TaskSpec, get_task_attempt_dir, read_task_state, update_task_state) -from arc.job.pipe.pipe_run import PipeRun, local_cpu_budget, local_worker_limit, worker_cpu_cores +from arc.job.pipe.pipe_run import (PipeRun, build_rotor_scan_1d_tasks, local_cpu_budget, + local_worker_limit, worker_cpu_cores) import arc.parser.parser as parser from arc.common import ARC_TESTING_PATH from arc.level import Level @@ -783,5 +784,31 @@ def test_ingestion_of_an_unmatched_conformer_index_is_a_no_op(self): self.assertEqual(sorted(tsg.index for tsg in self.ts_species.ts_guesses), [3, 7]) +class TestBuildRotorScan1dTasks(unittest.TestCase): + """Unit tests for build_rotor_scan_1d_tasks: one task per rotor and scan_res propagation.""" + + def setUp(self): + self.spc = ARCSpecies(label='propane', smiles='CCC') + # Two rotors, set explicitly to keep the test hermetic (no conformer generation). + self.spc.rotors_dict = {0: {'torsion': [3, 0, 1, 2], 'pivots': [1, 2]}, + 1: {'torsion': [1, 2, 3, 8], 'pivots': [2, 3]}} + self.level_dict = {'method': 'uma-s-1p2'} + + def test_one_task_per_rotor(self): + tasks = build_rotor_scan_1d_tasks(self.spc, 'propane', [0, 1], self.level_dict, 'ase', 4096) + self.assertEqual([t.task_id for t in tasks], ['propane_scan_r0', 'propane_scan_r1']) + self.assertTrue(all(t.task_family == 'rotor_scan_1d' and t.engine == 'ase' for t in tasks)) + self.assertEqual([t.input_payload['rotor_index'] for t in tasks], [0, 1]) + + def test_scan_res_carried_into_payload_when_given(self): + tasks = build_rotor_scan_1d_tasks(self.spc, 'propane', [0, 1], self.level_dict, 'ase', 4096, + scan_res=8.0) + self.assertTrue(all(t.input_payload['scan_res'] == 8.0 for t in tasks)) + + def test_scan_res_omitted_when_none(self): + tasks = build_rotor_scan_1d_tasks(self.spc, 'propane', [0], self.level_dict, 'ase', 4096) + self.assertNotIn('scan_res', tasks[0].input_payload) + + if __name__ == '__main__': unittest.main(testRunner=unittest.TextTestRunner(verbosity=2)) diff --git a/arc/scripts/pipe_worker.py b/arc/scripts/pipe_worker.py index 8b16a6338d..99a79341d6 100644 --- a/arc/scripts/pipe_worker.py +++ b/arc/scripts/pipe_worker.py @@ -274,10 +274,15 @@ def _get_family_extra_kwargs(spec: TaskSpec) -> dict: elif spec.task_family == 'rotor_scan_1d': torsions = payload.get('torsions') rotor_index = payload.get('rotor_index') + scan_res = payload.get('scan_res') if torsions is not None: kwargs['torsions'] = torsions if rotor_index is not None: kwargs['rotor_index'] = rotor_index + if scan_res is not None: + # Carry the staging-time resolution into the reconstructed job the way trsh does, so the + # ASE scan runs ARC's intended resolution and not the worker node's settings default. + kwargs['args'] = {'trsh': {'scan_res': scan_res}} return kwargs diff --git a/arc/scripts/pipe_worker_test.py b/arc/scripts/pipe_worker_test.py index 4ead0144d4..d70735534d 100644 --- a/arc/scripts/pipe_worker_test.py +++ b/arc/scripts/pipe_worker_test.py @@ -22,7 +22,8 @@ read_task_state, update_task_state, ) -from arc.scripts.pipe_worker import claim_task, run_task, main, logger as worker_logger +from arc.scripts.pipe_worker import (claim_task, run_task, main, logger as worker_logger, + _get_family_extra_kwargs) from arc.species import ARCSpecies @@ -351,5 +352,33 @@ def test_no_duplicate_log_handlers(self): self.assertLessEqual(len(worker_logger.handlers), 2) +class TestFamilyExtraKwargs(unittest.TestCase): + """Unit tests for _get_family_extra_kwargs, which maps a task payload to adapter kwargs.""" + + def _rotor_spec(self, payload): + spc = ARCSpecies(label='ethane', smiles='CC') + return TaskSpec( + task_id='ethane_scan_r0', task_family='rotor_scan_1d', owner_type='species', + owner_key='ethane', input_fingerprint='fp', engine='ase', + level={'method': 'uma-s-1p2'}, required_cores=1, required_memory_mb=512, + input_payload={'species_dicts': [spc.as_dict()], **payload}, + ingestion_metadata={'rotor_index': 0}) + + def test_rotor_scan_forwards_torsions_and_rotor_index(self): + kwargs = _get_family_extra_kwargs(self._rotor_spec({'torsions': [[2, 0, 1, 5]], 'rotor_index': 0})) + self.assertEqual(kwargs['torsions'], [[2, 0, 1, 5]]) + self.assertEqual(kwargs['rotor_index'], 0) + self.assertNotIn('args', kwargs) + + def test_rotor_scan_forwards_scan_res_as_trsh_args(self): + kwargs = _get_family_extra_kwargs(self._rotor_spec( + {'torsions': [[2, 0, 1, 5]], 'rotor_index': 0, 'scan_res': 8.0})) + self.assertEqual(kwargs['args'], {'trsh': {'scan_res': 8.0}}) + + def test_rotor_scan_omits_args_when_scan_res_absent(self): + kwargs = _get_family_extra_kwargs(self._rotor_spec({'torsions': [[2, 0, 1, 5]], 'rotor_index': 0})) + self.assertNotIn('args', kwargs) + + if __name__ == '__main__': unittest.main(testRunner=unittest.TextTestRunner(verbosity=2))