pipe: carry rotor scan resolution through the task payload

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:
@@ -296,19 +297,29 @@ def try_pipe_rotor_scans_1d(self, label: str, rotor_indices: list[int]) -> set[i
         """
         Batch 1D rotor scan jobs through pipe mode.
 
+        The scan resolution carried in the task payload follows the same precedence the scheduler
+        applies to ESS scans in ``set_scan_resolution()``: the run-level ``rotor_scan_resolution``
+        input key wins, and the launching host's settings value is the fallback when it is unset.
+
+        Args:
+            label (str): The species label.
+            rotor_indices (list[int]): The rotor indices to scan.
+
         Returns:
             set[int]: Rotor indices that were piped (all or empty).
         """
         level = self.sched.scan_level
         if level is None:
             return set()
+        scan_res = self.sched.rotor_scan_resolution or rotor_scan_resolution
         submitted = self._try_pipe_job(
             run_id=f'{label}_scan_1d',
             level=level,
             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=scan_res),
             log_msg=f'Routing {len(rotor_indices)} 1D rotor scans for {label} to pipe mode',
         )
         return set(rotor_indices) if submitted else set()

arc/job/pipe/pipe_planner_test.py

@@ -11,7 +11,7 @@
 from unittest.mock import MagicMock, patch
 
 from arc.job.pipe.pipe_coordinator import PipeCoordinator
-from arc.job.pipe.pipe_planner import PipePlanner
+from arc.job.pipe.pipe_planner import PipePlanner, rotor_scan_resolution
 from arc.level import Level
 from arc.species import ARCSpecies
 
@@ -65,6 +65,9 @@ def _make_mock_sched(project_directory):
     sched.species_dict = {'H2O': spc}
     sched.output = {'H2O': {'paths': {}, 'job_types': {}}}
     sched.deduce_job_adapter = MagicMock(return_value='gaussian')
+    # A real Scheduler defaults this to None (the rotor_scan_resolution input key was not given);
+    # spell it out, since a bare MagicMock attribute is truthy and would pose as a run-level value.
+    sched.rotor_scan_resolution = None
     return sched
 
 
@@ -262,6 +265,20 @@ def test_no_pipe_when_no_scan_level(self):
         handled = self.planner.try_pipe_rotor_scans_1d('H2O', list(range(12)))
         self.assertEqual(handled, set())
 
+    def test_settings_default_scan_res_when_run_level_is_none(self):
+        """Without a rotor_scan_resolution input key, the payload carries the settings default."""
+        self.planner.try_pipe_rotor_scans_1d('H2O', list(range(12)))
+        pipe = list(self.coord.active_pipes.values())[0]
+        self.assertTrue(all(task.input_payload['scan_res'] == rotor_scan_resolution
+                            for task in pipe.tasks))
+
+    def test_run_level_scan_res_overrides_settings_default(self):
+        """The rotor_scan_resolution input key wins over the launching host's settings value."""
+        self.sched.rotor_scan_resolution = 4.0
+        self.planner.try_pipe_rotor_scans_1d('H2O', list(range(12)))
+        pipe = list(self.coord.active_pipes.values())[0]
+        self.assertTrue(all(task.input_payload['scan_res'] == 4.0 for task in pipe.tasks))
+
 
 class TestTryPipeTsg(unittest.TestCase):
     """Tests for PipePlanner.try_pipe_tsg()."""

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

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

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
 

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