Test deployment on Theta (#32)

* Switch to more compact message, drop calculator

* Add Parsl configurations

* Make sure ASE settings get propagated to workers

* Use semaphore for initial energy

jitterbug/cli.py

@@ -64,13 +64,14 @@ def main(args: Optional[list[str]] = None):
     if args.parsl_config is None:
         config = Config(run_dir=str(run_dir / 'parsl-logs'), executors=[HighThroughputExecutor(max_workers=1)])
         num_workers = 1
+        ase_options = {}
         logger.info('Running computations locally, one-at-a-time')
     else:
         config, num_workers, ase_options = load_configuration(args.parsl_config)
         logger.info(f'Running on {num_workers} workers as defined by {args.parsl_config}')
 
     # Make the function to compute energy
-    energy_fun = partial(get_energy, method=method, basis=basis)
+    energy_fun = partial(get_energy, method=method, basis=basis, **ase_options)
     update_wrapper(energy_fun, get_energy)
 
     # Create a thinker

jitterbug/parsl.py

@@ -32,8 +32,9 @@ def get_energy(atoms: ase.Atoms, method: str, basis: Optional[str], scratch_dir:
         calc = make_calculator(method, basis, directory=tmp.name, **kwargs)
         atoms.calc = calc
         atoms.get_potential_energy()
-        return write_to_string(atoms, 'json')
+        return write_to_string(atoms, 'extxyz')
     finally:
+        atoms.calc = None  # Ensure the calculator does not get passed back
         os.chdir(start_dir)
         tmp.cleanup()
 

jitterbug/thinkers/exact.py

@@ -64,6 +64,7 @@ def submit_tasks(self):
 
         # Start with the unperturbed energy
         if self.unperturbed_energy is None:
+            self.rec.acquire(None, 1)
             self.queues.send_inputs(
                 self.atoms,
                 method='get_energy',
@@ -79,7 +80,7 @@ def submit_tasks(self):
                     continue
 
                 # Submit if not done
-                self.rec.acquire(None, 1)
+                self.rec.acquire(None, 1)  # Wait until resources are free
                 count += 1
                 atom_id, axis_id, dir_id = it.multi_index
 
@@ -134,7 +135,7 @@ def store_energy(self, result: Result):
             return
 
         calc_type = result.task_info['type']
-        atoms = read_from_string(result.value, 'json')
+        atoms = read_from_string(result.value, 'extxyz')
         energy = atoms.get_potential_energy()
 
         # Store unperturbed energy

parsl-configs/README.md

@@ -0,0 +1,3 @@
+# Parsl Configuration
+
+Parsl configurations for different resources.

parsl-configs/theta_debug.py

@@ -0,0 +1,37 @@
+"""Use the debug queue on Theta"""
+from parsl import Config, HighThroughputExecutor
+from parsl.providers import CobaltProvider
+from parsl.launchers import AprunLauncher
+from parsl.addresses import address_by_interface
+
+
+def make_config():
+    return Config(
+        retries=1,
+        executors=[
+            HighThroughputExecutor(
+                label='theta_local_htex_multinode',
+                address=address_by_interface('vlan2360'),
+                max_workers=1,
+                provider=CobaltProvider(
+                    queue='debug-flat-quad',
+                    account='CSC249ADCD08',
+                    launcher=AprunLauncher(overrides="-d 64 --cc depth"),
+                    walltime='01:00:00',
+                    nodes_per_block=2,
+                    init_blocks=1,
+                    min_blocks=1,
+                    max_blocks=1,
+                    scheduler_options='#COBALT --attrs enable_ssh=1:filesystems=home,eagle',
+                    # Command to be run before starting a worker, such as:
+                    # 'module load Anaconda; source activate parsl_env'.
+                    worker_init='''
+module load miniconda-3
+source activate /lus/eagle/projects/ExaLearn/faster-molecular-hessians/env
+which python
+realpath .
+export PSIDATADIR=/lus/eagle/projects/ExaLearn/faster-molecular-hessians/env/share/psi4''',
+                    cmd_timeout=120,
+                ),
+            )
+        ],), 8, {'num_threads': 64, 'memory': '64GB'}

parsl-configs/theta_full.py

@@ -0,0 +1,36 @@
+"""Use the backfill queue on Theta"""
+from parsl import Config, HighThroughputExecutor
+from parsl.providers import CobaltProvider
+from parsl.launchers import AprunLauncher
+from parsl.addresses import address_by_interface
+
+
+def make_config():
+    return Config(
+        retries=1,
+        executors=[
+            HighThroughputExecutor(
+                label='theta_local_htex_multinode',
+                address=address_by_interface('vlan2360'),
+                max_workers=1,
+                provider=CobaltProvider(
+                    account='CSC249ADCD08',
+                    launcher=AprunLauncher(overrides="-d 64 --cc depth"),
+                    walltime='0:30:00',
+                    nodes_per_block=128,
+                    init_blocks=0,
+                    min_blocks=0,
+                    max_blocks=8,
+                    scheduler_options='#COBALT --attrs enable_ssh=1:filesystems=home,eagle',
+                    # Command to be run before starting a worker, such as:
+                    # 'module load Anaconda; source activate parsl_env'.
+                    worker_init='''
+module load miniconda-3
+source activate /lus/eagle/projects/ExaLearn/faster-molecular-hessians/env
+which python
+realpath .
+export PSIDATADIR=/lus/eagle/projects/ExaLearn/faster-molecular-hessians/env/share/psi4''',
+                    cmd_timeout=120,
+                ),
+            )
+        ],), 128 * 8 + 100, {'num_threads': 64, 'memory': '64GB'}

tests/test_parsl.py

@@ -12,7 +12,7 @@ def test_energy(xyz_path):
 
     atoms = read(xyz_path)
     atoms_msg = get_energy(atoms, 'pm7', None)
-    new_atoms = read_from_string(atoms_msg, 'json')
+    new_atoms = read_from_string(atoms_msg, 'extxyz')
     assert 'energy' in new_atoms.calc.results
 
     assert not mopac_out.exists()