diff --git a/arc/imports.py b/arc/imports.py index 860ff6020b..9ba4ee64b6 100644 --- a/arc/imports.py +++ b/arc/imports.py @@ -12,7 +12,7 @@ find_rits_ckpt, queue_deferred_warning) from arc.settings.inputs import input_files -from arc.settings.submit import incore_commands, pipe_submit, submit_scripts +from arc.settings.submit import ase_submit, incore_commands, pipe_submit, submit_scripts logger = logging.getLogger('arc') @@ -124,6 +124,7 @@ def _report_unusable_overlay(path: str, module: str, error: ImportError, what: s local_arc_submit_path = os.path.join(local_arc_path, 'submit.py') if os.path.isfile(local_arc_submit_path): local_incore_commands, local_pipe_submit, local_submit_scripts = dict(), dict(), dict() + local_ase_submit = dict() if local_arc_path not in sys.path: sys.path.insert(1, local_arc_path) try: @@ -134,6 +135,10 @@ def _report_unusable_overlay(path: str, module: str, error: ImportError, what: s from submit import pipe_submit as local_pipe_submit except ImportError as e: _report_unusable_overlay(local_arc_submit_path, 'submit', e, 'pipe_submit') + try: + from submit import ase_submit as local_ase_submit + except ImportError as e: + _report_unusable_overlay(local_arc_submit_path, 'submit', e, 'ase_submit') try: from submit import submit_scripts as local_submit_scripts except ImportError as e: @@ -142,6 +147,8 @@ def _report_unusable_overlay(path: str, module: str, error: ImportError, what: s incore_commands.update(local_incore_commands) if local_pipe_submit: pipe_submit.update(local_pipe_submit) + if local_ase_submit: + ase_submit.update(local_ase_submit) if local_submit_scripts: submit_scripts.update(local_submit_scripts) diff --git a/arc/job/adapters/ase_adapter.py b/arc/job/adapters/ase_adapter.py index c40e0d2d88..621cb80b82 100644 --- a/arc/job/adapters/ase_adapter.py +++ b/arc/job/adapters/ase_adapter.py @@ -11,9 +11,13 @@ from arc.job.adapter import JobAdapter from arc.job.adapters.common import _initialize_adapter from arc.job.factory import register_job_adapter -from arc.imports import settings +from arc.imports import ase_submit, settings from arc.settings.settings import ARC_PYTHON, UMA_LATEST_MODEL, find_executable +servers = settings['servers'] +submit_filenames = settings['submit_filenames'] +t_max_format = settings['t_max_format'] + if TYPE_CHECKING: from arc.level import Level from arc.species.species import ARCSpecies @@ -210,6 +214,25 @@ def set_scan_torsions(self) -> None: return raise ValueError(f'Could not determine scan parameters for scan job {self.job_name}') + def determine_constraints(self) -> List[Tuple[List[int], float]]: + """ + Determine the internal coordinate constraints to apply. + + A directed rotor scan is spawned by the Scheduler as one constrained optimization per + dihedral point, but ``Scheduler.run_job()`` always passes ``constraints=None`` and hands the + adapter ``torsions`` + ``dihedrals`` instead. Without translating those into a constraint the + "scan" is an unconstrained optimization repeated at every point, and every point relaxes back + to the same minimum, giving a flat V(phi). Torsions are 0-indexed; ARC constraints are + 1-indexed (as in the xTB and Gaussian adapters). + + Returns: + List[Tuple[List[int], float]]: The constraints, as (1-indexed atom indices, value) pairs. + """ + if self.constraints or self.job_type != 'directed_scan' or not self.torsions or not self.dihedrals: + return self.constraints + return [([index + 1 for index in torsion], dihedral) + for torsion, dihedral in zip(self.torsions, self.dihedrals)] + def write_input_file(self) -> None: """ Write the input file for ase_script.py. @@ -220,7 +243,7 @@ def write_input_file(self) -> None: 'charge': self.charge, 'multiplicity': self.multiplicity, 'is_ts': self.species[0].is_ts if self.species else False, - 'constraints': self.constraints, + 'constraints': self.determine_constraints(), 'irc_direction': self.irc_direction, 'settings': self.determine_settings(), } @@ -273,23 +296,26 @@ def execute_incore(self) -> None: def execute_queue(self) -> None: """ Execute a job to the server's queue. + + ``set_files()`` wrote the files and ``JobAdapter.execute()`` uploaded them, so all that is + left here is the submission itself, through the same path every other adapter uses. """ - self.write_input_file() - self.write_submit_script() - self.set_files() - if self.server_adapter is not None: - for file_dict in self.files_to_upload: - self.server_adapter.upload_file(remote_path=file_dict['remote'], - local_path=file_dict['local']) - self.server_adapter.submit_job(self.remote_path) + self.legacy_queue_execution() def set_files(self) -> None: """ - Set files to be uploaded and downloaded. + Set files to be uploaded and downloaded. Writes the files if needed. """ # 1. Upload if self.execution_type != 'incore': - self.files_to_upload.append(self.get_file_property_dictionary(file_name='submit.sh')) + # ``JobAdapter.execute()`` calls ``upload_files()`` *before* ``execute_queue()``, and + # ``_initialize_adapter()`` calls this method while the job is being constructed, so a + # queue job's files have to be written here - as the Gaussian, Orca and xTB adapters do + # - or the upload raises "InputError: Cannot upload a non-existing file". + # An incore job is not uploaded and writes its input in ``execute_incore()``. + self.write_submit_script() + self.files_to_upload.append(self.get_file_property_dictionary(file_name=self.determine_submit_filename())) + self.write_input_file() self.files_to_upload.append(self.get_file_property_dictionary(file_name='input.yml')) self.files_to_upload.append(self.get_file_property_dictionary(file_name='ase_script.py', local=self.script_path)) @@ -308,14 +334,119 @@ def set_input_file_memory(self) -> None: """ pass + def determine_submit_config(self) -> dict: + """ + Determine the cluster submission knobs for this job, taken from the level's ``args['block']``. + + Recognized keys (all optional): + + - ``env_setup``: shell lines to run on the compute node before the ASE script, e.g. + ``module load CUDA/12.1; conda activate UMA_env``. Its case is preserved (see + ``CASE_SENSITIVE_BLOCK_ARGS`` in ``arc/level.py``), so module and env names survive intact. + - ``gpu_resource``: a scheduler GPU request, appended to the PBS ``select`` statement + (e.g. ``ngpus=1``) or used as the Slurm ``--gres`` value (e.g. ``gpu:1``). + - ``python``: the python executable **on the server**. ``self.python_executable`` is + resolved against the ARC host's conda envs and generally does not exist on a remote server. + - ``queue``: the queue to submit to, if not already set on the job or in the server settings. + + Returns: + dict: The resolved submit configuration. + """ + block = (self.args or dict()).get('block', dict()) or dict() + default_queue, _ = next(iter(servers.get(self.server, dict()).get('queues', dict()).items()), (None, None)) + return {'queue': self.queue or block.get('queue') or default_queue, + 'env_setup': block.get('env_setup', ''), + 'gpu_resource': block.get('gpu_resource', ''), + 'python': block.get('python', ''), + } + + def determine_submit_filename(self) -> str: + """ + Return the filename ARC will submit for this job. + + A queue-executed PBS/Slurm job must be written under the scheduler-specific name that + ``submit_job()`` invokes (``submit_filenames``, e.g. ``submit.sl`` for Slurm), or the + submission fails because the file it names is not on disk. Everything else uses the plain + ``submit.sh`` the bare script is written to. + + Returns: + str: The submit-script filename. + """ + cluster_soft = servers.get(self.server, dict()).get('cluster_soft', '') if self.server is not None else '' + if self.execution_type != 'incore' and cluster_soft.lower() in ('pbs', 'slurm'): + return submit_filenames[cluster_soft] + return 'submit.sh' + + def get_queue_submit_script(self, command: str, config: dict, cluster_soft: str) -> str: + """ + Compose a cluster submit script for a queue-executed ASE job. + + Formats the server-independent ``ase_submit`` template (in ``arc/settings/submit.py``, + keyed by cluster software) with this job's resources and submit config. The thread-pool + exports pin the numerical libraries (torch, NumPy) to the cores the scheduler granted, so a + shared node is not oversubscribed. + + Args: + command (str): The command running the ASE script on the compute node. + config (dict): The output of ``determine_submit_config()``. + cluster_soft (str): The lowercased cluster software name ('pbs' or 'slurm'). + + Returns: + str: The submit script content. + """ + if cluster_soft not in ase_submit: + raise NotImplementedError(f"No ASE submit template for cluster software '{cluster_soft}'. " + f"Available templates: {list(ase_submit.keys())}") + memory = int(self.submit_script_memory) if isinstance(self.submit_script_memory, (int, float)) \ + else self.submit_script_memory + time_format = next((v for k, v in t_max_format.items() if k.lower() == cluster_soft), 'hours') + pwd = self.local_path if self.server is None or str(self.server).lower() == 'local' else self.remote_path + queue, gpu_resource = config['queue'], config['gpu_resource'] + format_kwargs = {'name': self.job_server_name, 'cpus': self.cpu_cores, 'memory': memory, + 't_max': self.format_max_job_time(time_format=time_format), 'pwd': pwd, + 'env_setup': config['env_setup'], 'command': command} + if cluster_soft == 'pbs': + format_kwargs['queue_directive'] = f'#PBS -q {queue}\n' if queue else '' + format_kwargs['gpu_select'] = f':{gpu_resource}' if gpu_resource else '' + else: + format_kwargs['queue_directive'] = f'#SBATCH -p {queue}\n' if queue else '' + format_kwargs['gpu_directive'] = f'#SBATCH --gres={gpu_resource}\n' if gpu_resource else '' + return ase_submit[cluster_soft].format(**format_kwargs) + def write_submit_script(self) -> None: """ Write the submission script. + + An incore job only has to invoke the ASE script. A queue job additionally needs cluster + scheduler directives, an environment setup preamble (``conda activate uma_env`` for UMA), a + server-side python executable, and - for a GPU run - a GPU resource request; a bare + ``#!/bin/bash`` script carries none of those and lands on the queue's defaults with a python + path that only exists on the ARC host. See ``determine_submit_config()`` for the knobs. """ - remote_script_path = os.path.join(self.remote_path, 'ase_script.py') - command = f"{self.python_executable} {remote_script_path} --yml_path {self.remote_path}" - content = f"#!/bin/bash\n\n{command}\n" - with open(os.path.join(self.local_path, 'submit.sh'), 'w') as f: + config = self.determine_submit_config() + cluster_soft = servers.get(self.server, dict()).get('cluster_soft', '').lower() \ + if self.server is not None else '' + queue_job = self.execution_type != 'incore' and cluster_soft in ('pbs', 'slurm') + if queue_job and not config['python']: + logger.warning(f"Job {self.job_name} is submitted to {self.server}, but no server-side python " + f"was given in args['block']['python']; falling back to {self.python_executable}, " + f"which was resolved on this machine and may not exist there.") + python_executable = config['python'] or self.python_executable + if queue_job: + # trsh_job_queue() skips queues already in attempted_queues; record the one we submit to + # here, as JobAdapter.write_submit_script() does, so a failed submission moves on. + if config['queue'] and config['queue'] not in self.attempted_queues: + self.attempted_queues.append(config['queue']) + # The script cd's into the job directory, so address the ASE script relative to it. + command = f'{python_executable} "$JOB_DIR/ase_script.py" --yml_path "$JOB_DIR"' + content = self.get_queue_submit_script(command=command, config=config, cluster_soft=cluster_soft) + else: + # A job with no server (and hence no remote path) runs out of its local directory. + path = self.remote_path or self.local_path + remote_script_path = os.path.join(path, 'ase_script.py') + command = f"{python_executable} {remote_script_path} --yml_path {path}" + content = f"#!/bin/bash\n\n{command}\n" + with open(os.path.join(self.local_path, self.determine_submit_filename()), 'w') as f: f.write(content) def parse_results(self) -> None: diff --git a/arc/job/adapters/ase_test.py b/arc/job/adapters/ase_test.py index 2f5158d83f..83eeb9d4c5 100644 --- a/arc/job/adapters/ase_test.py +++ b/arc/job/adapters/ase_test.py @@ -16,10 +16,11 @@ from ase.calculators.emt import EMT from arc.common import ARC_TESTING_PATH, read_yaml_file, save_yaml_file -from arc.job.adapters.ase_adapter import ASEAdapter +from arc.job.adapters.ase_adapter import ASEAdapter, servers from arc.parser.parser import parse_1d_scan_coords, parse_1d_scan_energies from arc.species.species import ARCSpecies -from arc.job.adapters.scripts.ase_script import (is_linear, +from arc.job.adapters.scripts.ase_script import (apply_constraints, + is_linear, merge_scan_branches, numpy_vibrational_analysis, relaxed_torsion_scan, @@ -123,6 +124,71 @@ def test_set_files(self): self.assertTrue(any('ase_script.py' in f['local'] for f in self.job_2.files_to_upload)) self.assertTrue(any('output.yml' in f['local'] for f in self.job_2.files_to_download)) + def test_set_files_writes_the_files_of_a_queue_job(self): + """Test that constructing a queue job already writes its submit script and input file""" + # JobAdapter.execute() calls upload_files() before execute_queue(), so the files must be on + # disk once the job object exists, not when it is executed. + xyz = {'symbols': ('O', 'H', 'H'), + 'isotopes': (16, 1, 1), + 'coords': ((0.0, 0.0, 0.0), (0.0, 0.75, 0.58), (0.0, -0.75, 0.58))} + fake_server = {'test_server': {'cluster_soft': 'PBS', 'un': 'test_user'}} + # arc.job.adapter and arc.job.adapters.ase_adapter share this one dict object. + with patch.dict(servers, fake_server): + job_3 = ASEAdapter(execution_type='queue', + job_type='directed_scan', + project='test_3', + project_directory=os.path.join(self.project_directory, 'test_3'), + species=[ARCSpecies(label='H2O', xyz=xyz)], + args={'keyword': {'calculator': 'xtb'}, + 'block': {'queue': 'test_q', + 'env_setup': 'conda activate uma_env', + 'python': '/remote/python'}}, + server='test_server', + testing=True) + submit_path = os.path.join(job_3.local_path, 'submit.sh') + self.assertTrue(os.path.isfile(submit_path)) + self.assertTrue(os.path.isfile(os.path.join(job_3.local_path, 'input.yml'))) + # Every file the job says it will upload must exist, or ssh.upload_file() raises an InputError. + for file_dict in job_3.files_to_upload: + self.assertTrue(os.path.isfile(file_dict['local']), msg=f"missing {file_dict['file_name']}") + with open(submit_path, 'r') as f: + content = f.read() + self.assertIn('#PBS -q test_q', content) + self.assertIn('conda activate uma_env', content) + self.assertIn('/remote/python', content) + + def test_set_files_does_not_write_for_an_incore_job(self): + """Test that an incore job writes no submit script (it writes its input when it executes)""" + self.assertFalse(os.path.isfile(os.path.join(self.job_1.local_path, 'submit.sh'))) + self.assertTrue(all('submit.sh' not in f['local'] for f in self.job_1.files_to_upload)) + + def test_determine_constraints(self): + """Test that a directed scan derives a 1-indexed dihedral constraint from torsions/dihedrals""" + xyz = {'symbols': ('O', 'H', 'H'), + 'isotopes': (16, 1, 1), + 'coords': ((0.0, 0.0, 0.0), (0.0, 0.75, 0.58), (0.0, -0.75, 0.58))} + job = ASEAdapter(execution_type='incore', + job_type='directed_scan', + project='test_c', + project_directory=os.path.join(self.project_directory, 'test_c'), + species=[ARCSpecies(label='H2O', xyz=xyz)], + torsions=[[0, 1, 2, 3]], + dihedrals=[60.0], + args={'keyword': {'calculator': 'xtb'}}, + testing=True) + # torsions are 0-indexed; the returned constraint must be 1-indexed, as xTB/Gaussian expect. + self.assertEqual(job.determine_constraints(), [([1, 2, 3, 4], 60.0)]) + + def test_apply_constraints_converts_to_zero_indexed(self): + """Test that apply_constraints translates ARC's 1-indexed dihedral to ASE's 0-indexed FixInternals""" + from ase import Atoms + atoms = Atoms('C4', positions=[(0.0, 0.0, 0.0), (1.5, 0.0, 0.0), + (2.0, 1.4, 0.0), (3.5, 1.4, 0.0)]) + apply_constraints(atoms, [([1, 2, 3, 4], 90.0)]) + self.assertEqual(len(atoms.constraints), 1) + dihedrals = atoms.constraints[0].todict()['kwargs']['dihedrals_deg'] + self.assertEqual(dihedrals, [[90.0, [0, 1, 2, 3]]]) + def test_parse_results(self): """Test parsing dummy output YAML back into object attributes""" output_data = { diff --git a/arc/job/adapters/scripts/ase_script.py b/arc/job/adapters/scripts/ase_script.py index f4c4b5c8c8..e7393e6b41 100644 --- a/arc/job/adapters/scripts/ase_script.py +++ b/arc/job/adapters/scripts/ase_script.py @@ -119,12 +119,15 @@ def get_calculator(calc_config: dict, charge: int = 0, multiplicity: int = 1): def apply_constraints(atoms: Atoms, constraints_data: list): """ Apply internal constraints to the Atoms object. + + ARC's constraint atom indices are 1-indexed (that is what Gaussian's modredundant section and + xTB's $constrain block consume); ASE's FixInternals is 0-indexed. """ if not constraints_data: return bonds, angles, dihedrals = list(), list(), list() for constraint in constraints_data: - indices = constraint[0] + indices = [index - 1 for index in constraint[0]] if len(indices) == 2: bonds.append([constraint[1], indices]) elif len(indices) == 3: diff --git a/arc/level.py b/arc/level.py index a71d6e9dca..04e5dd2d4e 100644 --- a/arc/level.py +++ b/arc/level.py @@ -14,6 +14,11 @@ levels_ess, supported_ess = settings['levels_ess'], settings['supported_ess'] +# ``block`` args whose values are case-sensitive (filesystem paths, shell commands, queue names) and +# must not be lowercased along with the rest of a level's args. Consumed by the ASE job adapter when +# running queue jobs (see ASEAdapter.determine_submit_config). +CASE_SENSITIVE_BLOCK_ARGS = ('python', 'env_setup', 'queue', 'gpu_resource') + class Level(object): """ @@ -312,7 +317,8 @@ def lower(self): if not isinstance(new_val2, str): raise ValueError(f'All entries in the args argument must be str, int, or float types.\n' f'Got {new_val2} which is a {type(new_val2)} in {self.args}.') - args[key1.lower()][key2.lower()] = new_val2.lower() + keep_case = key1.lower() == 'block' and key2.lower() in CASE_SENSITIVE_BLOCK_ARGS + args[key1.lower()][key2.lower()] = new_val2 if keep_case else new_val2.lower() elif isinstance(val1, str): args[key1.lower()]['general'] = val1.lower() elif isinstance(val1, (list, tuple)): diff --git a/arc/level_test.py b/arc/level_test.py index 1e404be628..6a7e78a778 100644 --- a/arc/level_test.py +++ b/arc/level_test.py @@ -83,6 +83,20 @@ def test_lower(self): self.assertEqual(level.solvent, 'water') self.assertEqual(level.args, {'keyword': {'general': 'iop(99/33=1)'}, 'block': {}}) + def test_lower_preserves_case_sensitive_block_args(self): + """Test that Level.lower() keeps the case of the ASE execution block args (paths, commands, queues)""" + level = Level(method='UMA', + args={'block': {'python': '/home/User/.conda/envs/UMA_env/bin/python', + 'env_setup': 'module load CUDA/12.1; conda activate UMA_env', + 'queue': 'GPU_Long', + 'general': 'Extra Block'}}) + self.assertEqual(level.method, 'uma') + self.assertEqual(level.args['block']['python'], '/home/User/.conda/envs/UMA_env/bin/python') + self.assertEqual(level.args['block']['env_setup'], 'module load CUDA/12.1; conda activate UMA_env') + self.assertEqual(level.args['block']['queue'], 'GPU_Long') + # a non-execution block arg is still lowercased as before + self.assertEqual(level.args['block']['general'], 'extra block') + def test_equal(self): """Test identifying equal levels.""" level_1 = Level(method='b3lyp', basis='def2tzvp', auxiliary_basis='aug-def2-svp', diff --git a/arc/settings/submit.py b/arc/settings/submit.py index de86d4558a..ae84c606eb 100644 --- a/arc/settings/submit.py +++ b/arc/settings/submit.py @@ -125,6 +125,68 @@ } +# Submission scripts for queue-executed ASE (e.g. UMA/fairchem MLIP) jobs, keyed by cluster +# scheduler type. These are server-independent templates. ASEAdapter.get_queue_submit_script() +# formats them with: name, cpus, memory, t_max, pwd, env_setup, command, and the optional +# directive fields queue_directive/gpu_directive (Slurm) or queue_directive/gpu_select (PBS), +# each rendered as an empty string when the job does not request it. ``pwd`` is the job's +# submission directory (the remote path, or the local path for a 'local' server). +ase_submit = { + 'slurm': """#!/bin/bash -l + +#SBATCH -J {name} +{queue_directive}#SBATCH -N 1 +#SBATCH -n {cpus} +#SBATCH --mem-per-cpu={memory} +#SBATCH -t {t_max} +#SBATCH -o out.txt +#SBATCH -e err.txt +{gpu_directive} +. ~/.bashrc + +cd "{pwd}" +JOB_DIR="$(pwd)" # resolve the submission directory to an absolute path for the ASE script below + +export OMP_NUM_THREADS={cpus} +export MKL_NUM_THREADS={cpus} +export OPENBLAS_NUM_THREADS={cpus} + +{env_setup} + +touch initial_time + +{command} + +touch final_time +""", + 'pbs': """#!/bin/bash -l + +#PBS -N {name} +{queue_directive}#PBS -l select=1:ncpus={cpus}:mem={memory}mb{gpu_select} +#PBS -l walltime={t_max} +#PBS -o out.txt +#PBS -e err.txt + +. ~/.bashrc + +cd "{pwd}" +JOB_DIR="$(pwd)" # resolve the submission directory to an absolute path for the ASE script below + +export OMP_NUM_THREADS={cpus} +export MKL_NUM_THREADS={cpus} +export OPENBLAS_NUM_THREADS={cpus} + +{env_setup} + +touch initial_time + +{command} + +touch final_time +""", +} + + # Submission scripts stored as a dictionary with server and software as primary and secondary keys submit_scripts = { 'local': { diff --git a/docs/source/advanced.rst b/docs/source/advanced.rst index c8c1325c5d..387ec55fc1 100644 --- a/docs/source/advanced.rst +++ b/docs/source/advanced.rst @@ -150,6 +150,29 @@ For multiline blocks: MaxIter 500 end +Running ASE / UMA jobs on a cluster +----------------------------------- + +For an ASE-backed level (e.g. a UMA/fairchem MLIP), the ``block`` args also carry the knobs ARC +needs to submit the job to a compute server's queue: ``python`` (the server-side interpreter of the +calculator's environment), ``env_setup`` (shell commands run before the job, e.g. loading modules +and activating the conda env), ``queue`` (the queue/partition to submit to), and ``gpu_resource`` +(the scheduler's GPU request, e.g. ``gpu:1`` on Slurm or ``ngpus=1`` on PBS): + +.. code-block:: yaml + + level: + method: uma + args: + block: + python: /home/User/.conda/envs/UMA_env/bin/python + env_setup: module load CUDA/12.1; conda activate UMA_env + queue: GPU_Long + gpu_resource: gpu:1 + +Unlike other args, the values of these four keys keep their original case, since they are +filesystem paths, shell commands, and queue names. + Multireference Methods (MRCI) -----------------------------