Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion arc/job/pipe/pipe_planner.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@
logger = get_logger()

pipe_settings = settings['pipe_settings']
rotor_scan_resolution = settings['rotor_scan_resolution']


class PipePlanner:
Expand Down Expand Up @@ -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()
24 changes: 17 additions & 7 deletions arc/job/pipe/pipe_run.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = []
Expand All @@ -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',
Expand All @@ -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
29 changes: 28 additions & 1 deletion arc/job/pipe/pipe_run_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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))
5 changes: 5 additions & 0 deletions arc/scripts/pipe_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
31 changes: 30 additions & 1 deletion arc/scripts/pipe_worker_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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))
Loading