Test script for tracking down problems on cbio cluster when using multiple GPUs in parallel

openmm-gpu-mpi-test.py

import socket
import mpi4py.MPI
import gzip
import simtk.openmm as openmm
import simtk.openmm.app as app
import simtk.unit as unit

comm = mpi4py.MPI.COMM_WORLD
rank = comm.rank
size = comm.size
gpuid = rank % 4
hostname = socket.gethostname()

timestep = 2.0 * unit.femtoseconds
temperature = 300.0 * unit.kelvin
collision_rate = 20.0 / unit.picoseconds
pH = 8.0
forcefields_to_use = ['amber99sbildn.xml', 'amber99_obc.xml']
forcefield = app.ForceField(*forcefields_to_use)

model_filename = '/cbio/jclab/home/parton/tmp/debugging/AAPK2_HUMAN_D0_4CFE_C/model.pdb.gz'
with gzip.open(model_filename) as model_file:
    pdb = app.PDBFile(model_file)
modeller = app.Modeller(pdb.topology, pdb.positions)
modeller.addHydrogens(forcefield, pH=pH)
topology = modeller.getTopology()
positions = modeller.getPositions()

system = forcefield.createSystem(topology, nonbondedMethod=app.NoCutoff, constraints=app.HBonds)

platform = openmm.Platform.getPlatformByName('CUDA')
platform.setPropertyDefaultValue('CudaDeviceIndex', str(gpuid))
# platform = openmm.Platform.getPlatformByName('OpenCL')
# platform.setPropertyDefaultValue('OpenCLDeviceIndex', str(gpuid))

integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)

print 'rank %d host %r gpuid %d default CudaDeviceIndex %r about to build context...' % (rank, hostname, gpuid, platform.getPropertyDefaultValue('CudaDeviceIndex'))

# try:
context = openmm.Context(system, integrator, platform)
print 'rank %d host %r gpuid %d CudaDeviceIndex %r built context: %r' % (rank, hostname, gpuid, platform.getPropertyValue(context, 'CudaDeviceIndex'), context)
# finally:
#     del context
#     print 'rank %d host %r gpuid %d deleted context' % (rank, hostname, gpuid)