[WIP] Distributed computing with celery

README.md

@@ -30,3 +30,29 @@ The `openmmtools.testsystems` module contains a large suite of test systems---in
 If differences in energies in excess of `ENERGY_TOLERANCE` (default: 0.06 kcal/mol) are detected, these systems will be serialized to XML for further debugging.
 
 This is installed onto the command line when the repository is installed.
+
+## Distributed computing with OpenMM
+
+You can use `openmmtools.distributed` to run various OpenMM tasks in a distributed manner using [`celery`](http://www.celeryproject.org/).
+
+To run, first start one or more workers:
+```bash
+$ celery -A openmmtools.distributed worker -l info
+```
+Then 
+```python
+from simtk import openmm, unit
+from openmmtools import testsystems, distributed
+
+# Set up a test system
+testsystem = testsystems.AlanineDipeptideVacuum()
+[system, positions] = [testsystem.system, testsystem.positions]
+
+# Run distributed simulation
+from celery import group, chain, chord
+for iteration in range(10):
+    c = chain(distributed.tasks.propagate.s(system, positions) | distributed.tasks.compute_energy.s(system=system))
+    c().get()
+```
+
+See [First Steps with Celery](http://docs.celeryproject.org/en/latest/getting-started/first-steps-with-celery.html) for information on getting started with celery.

examples/README.md

@@ -2,3 +2,4 @@ Examples
 --------
 
 * `integrator-benchmarks/` - Timing benchmarks of various integrators with standard systems.
+* `distributed/` - distributed computing example

examples/distributed/README.md

@@ -0,0 +1,22 @@
+# Distributed computing example with `celery`
+
+See [First Steps with Celery](http://docs.celeryproject.org/en/latest/getting-started/first-steps-with-celery.html).
+
+First, set your environment variables:
+```
+RABBITMQ_SERVER (defaults to 'localhost')
+RABBITMQ_USERNAME (defaults to 'user')
+RABBITMQ_PASSWORD (defaults to 'password')
+RABBITMQ_PORT (defaults to '5762')
+RABBITMQ_VHOST (defaults to 'celery')
+```
+
+To run, first start one or more workers:
+```bash
+$ celery -A openmmtools.distributed worker -l info
+```
+
+Then run the application:
+```bash
+$ python distributed-example.py
+```

examples/distributed/distributed-example.py

@@ -0,0 +1,21 @@
+from simtk import openmm, unit
+from openmmtools import testsystems, distributed
+
+testsystem = testsystems.AlanineDipeptideVacuum()
+[system, positions] = [testsystem.system, testsystem.positions]
+
+from celery import group, chain, chord
+from openmmtools.distributed.tasks import propagate, compute_energy, mix_replicas
+print('Start group...')
+
+# Run 10 iterations of 10 parallel propagations, mixing replicas after each propagation.
+# Efficiency could be improved by constructing the DAG for many iterations up front and then submitting it to celery all at once.
+niterations = 10
+nreplicas = 10
+for iteration in range(niterations):
+    print('iteration %5d / %5d' % (iteration, niterations))
+    propagate_replicas_stage = group( chain(propagate.s(system, positions) | compute_energy.s(system=system)) for replica_index in range(nreplicas) )
+    mix_replicas_stage = mix_replicas.s()
+    c = ( propagate_replicas_stage | mix_replicas_stage )
+    c().get()
+print('Done.')

openmmtools/distributed/celery.py

@@ -0,0 +1,28 @@
+from __future__ import absolute_import, unicode_literals
+from celery import Celery
+import os
+
+server = os.getenv('RABBITMQ_SERVER', 'localhost')
+username = os.getenv('RABBITMQ_USERNAME', 'user')
+password = os.getenv('RABBITMQ_PASSWORD', 'password')
+port = os.getenv('RABBITMQ_PORT', '5672')
+vhost = os.getenv('RABBITMQ_VHOST', 'celery')
+
+broker = 'amqp://%(username)s:%(password)s@%(server)s:%(port)s/%(vhost)s' % vars()
+backend = 'amqp://%(username)s:%(password)s@%(server)s:%(port)s/%(vhost)s' % vars()
+
+app = Celery('openmmtools.distributed',
+             broker=broker,
+             backend=backend,
+             include=['openmmtools.distributed.tasks'])
+
+# Optional configuration, see the application user guide.
+app.conf.update(
+    result_expires=3600,
+    task_serializer = 'pickle',
+    result_serializer = 'pickle',
+    accept_content = {'pickle'},
+)
+
+if __name__ == '__main__':
+    app.start()

openmmtools/distributed/tasks.py

@@ -0,0 +1,35 @@
+from __future__ import absolute_import, unicode_literals
+from openmmtools.distributed.celery import app
+from simtk import openmm, unit
+import numpy as np
+
+@app.task
+def propagate(system, positions):
+    print('Propagate start')
+    temperature = 300 * unit.kelvin
+    collision_rate = 5.0 / unit.picoseconds
+    timestep = 2.0 * unit.femtoseconds
+    nsteps = 500
+    integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
+    context = openmm.Context(system, integrator)
+    context.setPositions(positions)
+    positions = context.getState(getPositions=True).getPositions(asNumpy=True)
+    integrator.step(nsteps)
+    del context, integrator
+    print('Propagate end')
+    return positions
+
+@app.task
+def compute_energy(positions, system=None):
+    timestep = 2.0 * unit.femtoseconds
+    integrator = openmm.VerletIntegrator(timestep)
+    print('Calling openmm.Context with %s, %s' % (str(system), str(integrator)))
+    context = openmm.Context(system, integrator)
+    context.setPositions(positions)
+    potential = context.getState(getEnergy=True).getPotentialEnergy()
+    del context, integrator
+    return potential
+
+@app.task
+def mix_replicas(energies):
+    print(energies)

setup.py

@@ -168,7 +168,7 @@ def check_dependencies():
     url='https://github.com/choderalab/openmmtools',
     platforms=['Linux', 'Mac OS-X', 'Unix', 'Windows'],
     classifiers=CLASSIFIERS.splitlines(),
-    packages=['openmmtools', 'openmmtools.tests', 'openmmtools.scripts'],
+    packages=['openmmtools', 'openmmtools.tests', 'openmmtools.scripts', 'openmmtools.distributed'],
     package_dir={'openmmtools': 'openmmtools'},
     package_data={'openmmtools': find_package_data('openmmtools/data', 'openmmtools')},
     install_requires=['numpy', 'scipy', 'nose'],