Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "ab98fc9adf42354394848375c7e495efa375239f"
Unverified Commit 4c2f3323 authored by Peter Eastman's avatar Peter Eastman Committed by GitHub
Browse files

Overhaul benchmarking script to add lots of new features (#3887) 

* Overhaul benchmarking script to add lots of new features

* Added another example
parent bf1cae4f
Hide whitespace changes
Inline Side-by-side
Showing with 441 additions and 90 deletions
examples/benchmark.py
View file @ 4c2f3323
......@@ -3,8 +3,98 @@ import openmm.app as app
import openmm as mm
import openmm.unit as unit
from datetime import datetime
import os
from argparse import ArgumentParser
import argparse
def cpuinfo():
"""Return CPU info"""
import os, platform, subprocess, re
if platform.system() == "Windows":
return platform.processor()
elif platform.system() == "Darwin":
os.environ['PATH'] = os.environ['PATH'] + os.pathsep + '/usr/sbin'
command ="sysctl -n machdep.cpu.brand_string"
return subprocess.check_output(command, shell=True, text=True).strip()
elif platform.system() == "Linux":
command = "cat /proc/cpuinfo"
all_info = subprocess.check_output(command, shell=True, text=True).strip()
for line in all_info.split("\n"):
if "model name" in line:
return re.sub( ".*model name.*: ", "", line,1)
return ""
def appendTestResult(filename=None, test_result=None, system_info=None):
"""Append a test result to a JSON or YAML file.
TODO: The efficiency of this could be improved.
Parameters
----------
filename : str, optional, default=None
The filename to append a result to, ending either in .yaml or .json
If None, no result is written
test_result : dict, optional, default=None
The test result to append to the 'benchmarks' blcok
system_info : dict, optional, default=None
System info to append to the 'system' block
"""
# Do nothing if filename is None
if filename is None:
return
import os
all_results = { 'benchmarks' : list() }
if system_info is not None:
all_results['system'] = system_info
if filename.endswith('.yaml'):
# Append test result to a YAML file, creating if file does not exist.
import yaml
if os.path.exists(filename):
with open(filename, 'rt') as infile:
all_results = yaml.safe_load(infile)
if test_result is not None:
all_results['benchmarks'].append(test_result)
with open(filename, 'wt') as outfile:
yaml.dump(all_results, outfile, sort_keys=False)
elif filename.endswith('.json'):
# Append test result to a JSON file, creating if file does not exist.
import json
if os.path.exists(filename):
with open(filename, 'rt') as infile:
all_results = json.load(infile)
if test_result is not None:
all_results['benchmarks'].append(test_result)
with open(filename, 'wt') as outfile:
json.dump(all_results, outfile, sort_keys=False, indent=4)
else:
raise ValueError('--output filename must end with .json or .yaml')
def printTestResult(test_result, options):
"""Render a test result
Parameters
----------
test_result : dict
The test result
options :
Options structure
"""
if options.style == 'simple':
for (key, value) in test_result.items():
print(f'{key}: {value}')
print('')
elif options.style == 'table':
print('%-18s%-12s%-14s%-15s%-10g%-11s%-11s%-g' %
(test_result['test'],
test_result['precision'],
test_result['constraints'],
test_result['hydrogen_mass'],
test_result['timestep_in_fs'],
test_result['ensemble'],
test_result['platform'],
test_result['ns_per_day']))
else:
raise ValueError(f"style '{style}' must be one of ['simple', 'table']")
def timeIntegration(context, steps, initialSteps):
"""Integrate a Context for a specified number of steps, then return how many seconds it took."""
......@@ -31,183 +121,444 @@ def downloadAmberSuite():
tarfh.extractall(path=dirname)
return dirname
def runOneTest(testName, options):
"""Perform a single benchmarking simulation."""
import functools
@functools.lru_cache(maxsize=None)
def retrieveTestSystem(testName, pme_cutoff=0.9, bond_constraints='hbonds', polarization='mutual', epsilon=1e-5):
"""Retrieve a benchmark system
Parameters
----------
testName : str
The name of the test
pme_cutoff : float
PME cutoff, in nm
bond_constraints : str, optional, default='hbonds'
'hbonds' : use app.HBonds and set H mass to 1.5*amu
'allbonds' : use app.AllBonds and set H mass to 4*amu
polarization : str, optional, default='mutual'
Polarization scheme for Amoeba
epsilon : str or float, optional, default=1e-5
mutualInducedTargetEpsilon for Amoeba
Returns
-------
system : openmm.System
The test system object
positions : openmm.unit.Quantity with shape (natoms,3)
The initial positions
test_parameters : dict of str : str
Special test parameters to report in test results
"""
explicit = (testName not in ('gbsa', 'amoebagk'))
amoeba = (testName in ('amoebagk', 'amoebapme'))
apoa1 = testName.startswith('apoa1')
amber = (testName.startswith('amber'))
hydrogenMass = None
print()
if amoeba:
print('Test: %s (epsilon=%g)' % (testName, options.epsilon))
elif testName == 'pme':
print('Test: pme (cutoff=%g)' % options.cutoff)
else:
print('Test: %s' % testName)
print('Ensemble: %s' % options.ensemble)
platform = mm.Platform.getPlatformByName(options.platform)
# Create the System.
temperature = 300*unit.kelvin
if explicit:
friction = 1*(1/unit.picoseconds)
else:
friction = 91*(1/unit.picoseconds)
# Create dictionary of test parameters
test_parameters = dict()
# Store the test name
test_parameters['test'] = testName
# Create the System.
if amoeba:
constraints = None
epsilon = float(options.epsilon)
test_parameters['epsilon'] = epsilon
epsilon = float(epsilon)
if explicit:
ff = app.ForceField('amoeba2009.xml')
pdb = app.PDBFile('5dfr_solv-cube_equil.pdb')
cutoff = 0.7*unit.nanometers
vdwCutoff = 0.9*unit.nanometers
system = ff.createSystem(pdb.topology, nonbondedMethod=app.PME, nonbondedCutoff=cutoff, vdwCutoff=vdwCutoff, constraints=constraints, ewaldErrorTolerance=0.00075, mutualInducedTargetEpsilon=epsilon, polarization=options.polarization)
system = ff.createSystem(pdb.topology, nonbondedMethod=app.PME, nonbondedCutoff=cutoff, vdwCutoff=vdwCutoff, constraints=constraints, ewaldErrorTolerance=0.00075, mutualInducedTargetEpsilon=epsilon, polarization=polarization)
else:
ff = app.ForceField('amoeba2009.xml', 'amoeba2009_gk.xml')
pdb = app.PDBFile('5dfr_minimized.pdb')
system = ff.createSystem(pdb.topology, nonbondedMethod=app.NoCutoff, constraints=constraints, mutualInducedTargetEpsilon=epsilon, polarization=options.polarization)
system = ff.createSystem(pdb.topology, nonbondedMethod=app.NoCutoff, constraints=constraints, mutualInducedTargetEpsilon=epsilon, polarization=polarization)
for f in system.getForces():
if isinstance(f, mm.AmoebaMultipoleForce) or isinstance(f, mm.AmoebaVdwForce) or isinstance(f, mm.AmoebaGeneralizedKirkwoodForce) or isinstance(f, mm.AmoebaWcaDispersionForce):
f.setForceGroup(1)
dt = 0.002*unit.picoseconds
if options.ensemble == 'NVE':
integ = mm.MTSIntegrator(dt, [(0,2), (1,1)])
else:
integ = mm.MTSLangevinIntegrator(temperature, friction, dt, [(0,2), (1,1)])
positions = pdb.positions
test_parameters['constraints'] = 'None'
test_parameters['hydrogen_mass'] = '1'
elif amber:
dirname = downloadAmberSuite()
names = {'amber20-dhfr':'JAC', 'amber20-factorix':'FactorIX', 'amber20-cellulose':'Cellulose', 'amber20-stmv':'STMV'}
names = {'amber20-dhfr':'JAC', 'amber20-factorix':'FactorIX', 'amber20-cellulose':'Cellulose', 'amber20-stmv':'STMV'}
fileName = names[testName]
prmtop = app.AmberPrmtopFile(os.path.join(dirname, f'PME/Topologies/{fileName}.prmtop'))
inpcrd = app.AmberInpcrdFile(os.path.join(dirname, f'PME/Coordinates/{fileName}.inpcrd'))
positions = inpcrd.positions
dt = 0.004*unit.picoseconds
method = app.PME
cutoff = options.cutoff
constraints = app.HBonds
hydrogenMass = 1.5*unit.amu
cutoff = pme_cutoff
system = prmtop.createSystem(nonbondedMethod=method, nonbondedCutoff=cutoff, constraints=constraints)
if options.ensemble == 'NVE':
integ = mm.VerletIntegrator(dt)
else:
integ = mm.LangevinMiddleIntegrator(temperature, friction, dt)
if inpcrd.boxVectors is not None:
system.setDefaultPeriodicBoxVectors(*inpcrd.boxVectors)
test_parameters['cutoff'] = cutoff
test_parameters['constraints'] = 'HBonds'
test_parameters['hydrogen_mass'] = '1.5'
else:
if apoa1:
ff = app.ForceField('amber14/protein.ff14SB.xml', 'amber14/lipid17.xml', 'amber14/tip3p.xml')
pdb = app.PDBFile('apoa1.pdb')
if testName == 'apoa1pme':
method = app.PME
cutoff = options.cutoff
cutoff = pme_cutoff
elif testName == 'apoa1ljpme':
method = app.LJPME
cutoff = options.cutoff
cutoff = pme_cutoff
else:
# Reaction field uses hard-coded cutoff
method = app.CutoffPeriodic
cutoff = 1*unit.nanometers
cutoff = 1.0 # nanometers ; JDC: Shouldn't this be larger for reaction field?
elif explicit:
ff = app.ForceField('amber99sb.xml', 'tip3p.xml')
pdb = app.PDBFile('5dfr_solv-cube_equil.pdb')
if testName == 'pme':
method = app.PME
cutoff = options.cutoff
cutoff = pme_cutoff
else:
# Reaction field uses hard-coded cutoff
method = app.CutoffPeriodic
cutoff = 1*unit.nanometers
cutoff = 1.0 # nanometers; JDC: Shouldn't this be larger for reaction field?
else:
ff = app.ForceField('amber99sb.xml', 'amber99_obc.xml')
pdb = app.PDBFile('5dfr_minimized.pdb')
method = app.CutoffNonPeriodic
cutoff = 2*unit.nanometers
if options.heavy:
dt = 0.005*unit.picoseconds
constraints = app.AllBonds
cutoff = 2.0 # nanometers
if bond_constraints == 'hbonds':
constraints = app.HBonds
hydrogenMass = 1.5*unit.amu
test_parameters['constraints'] = 'HBonds'
test_parameters['hydrogen_mass'] = '1.5'
elif bond_constraints == 'allbonds':
constraints = app.AllBonds # CAUTION: Constraining all bonds will perturb the effective dihedral marginal distributions.
hydrogenMass = 4*unit.amu
if options.ensemble == 'NVE':
integ = mm.VerletIntegrator(dt)
else:
integ = mm.LangevinIntegrator(temperature, friction, dt)
test_parameters['constraints'] = 'AllBonds'
test_parameters['hydrogen_mass'] = '4'
else:
raise ValueError(f"bond_constraints must be one of 'hbonds', 'allbonds': found {bond_constraints}")
test_parameters['cutoff'] = cutoff
positions = pdb.positions
system = ff.createSystem(pdb.topology, nonbondedMethod=method, nonbondedCutoff=cutoff, constraints=constraints, hydrogenMass=hydrogenMass)
return system, positions, test_parameters
def serializeTest(directory=None, system=None, integrator=None, state=None, coredata=None, metadata=None):
"""Serialize test system for benchmarking on Folding@home.
Parameters
----------
directory : str
Directory to write serialized test XML files
system : openmm.System
The System to serialize
integrator : openmm.Integrator
The Integrator to serialize
state : openmm.State
The State to serialize
coredata : dict
Core parameters to serialize to XML
metadata : dict
Metadata to dump to YAML
"""
import os
os.makedirs(directory, exist_ok=True)
import openmm
def serialize(obj, filename):
if filename.endswith('.bz2'):
import bz2
with bz2.open(filename, 'wt') as outfile:
outfile.write(openmm.XmlSerializer.serialize(obj))
elif filename.endswith('.gz'):
import gzip
with gzip.open(filename, 'wt') as outfile:
outfile.write(openmm.XmlSerializer.serialize(obj))
else:
with open(filename, 'wt') as outfile:
outfile.write(openmm.XmlSerializer.serialize(obj))
serialize(system, os.path.join(directory, 'system.xml.bz2'))
serialize(integrator, os.path.join(directory, 'integrator.xml.bz2'))
serialize(state, os.path.join(directory, 'state.xml.bz2'))
with open(os.path.join(directory, 'core.xml'), 'wt') as outfile:
outfile.write('<config>\n')
for key, value in coredata.items():
outfile.write(f' <{key} v="{value}"/>\n')
outfile.write('</config>\n')
import yaml
with open(os.path.join(directory, 'metadata.yaml'), 'wt') as outfile:
outfile.write(yaml.dump(metadata))
def runOneTest(testName, options):
"""Perform a single benchmarking simulation."""
system, positions, test_parameters = retrieveTestSystem(testName, pme_cutoff=options.pme_cutoff, bond_constraints=options.bond_constraints, polarization=options.polarization, epsilon=options.epsilon)
# Create a copy of the basic test_parameters dict (which may be cached) to report the test results
test_result = test_parameters.copy()
test_result['ensemble'] = options.ensemble
test_result['precision'] = options.precision
explicit = (testName not in ('gbsa', 'amoebagk'))
amoeba = (testName in ('amoebagk', 'amoebapme'))
apoa1 = testName.startswith('apoa1')
amber = (testName.startswith('amber'))
# Create the integrator
temperature = 300*unit.kelvin
if explicit:
friction = 1*(1/unit.picoseconds)
else:
friction = 91*(1/unit.picoseconds)
if amoeba:
dt = 0.002*unit.picoseconds
if options.ensemble == 'NVE':
integ = mm.MTSIntegrator(dt, [(0,2), (1,1)])
else:
integ = mm.MTSLangevinIntegrator(temperature, friction, dt, [(0,2), (1,1)])
elif amber:
dt = 0.004*unit.picoseconds
if options.ensemble == 'NVE':
integ = mm.VerletIntegrator(dt)
else:
integ = mm.LangevinMiddleIntegrator(temperature, friction, dt)
else:
if options.bond_constraints == 'hbonds':
dt = 0.004*unit.picoseconds
constraints = app.HBonds
hydrogenMass = 1.5*unit.amu
if options.ensemble == 'NVE':
integ = mm.VerletIntegrator(dt)
else:
integ = mm.LangevinMiddleIntegrator(temperature, friction, dt)
positions = pdb.positions
system = ff.createSystem(pdb.topology, nonbondedMethod=method, nonbondedCutoff=cutoff, constraints=constraints, hydrogenMass=hydrogenMass)
if options.ensemble == 'NPT':
system.addForce(mm.MonteCarloBarostat(1*unit.bar, temperature, 100))
print('Step Size: %g fs' % dt.value_in_unit(unit.femtoseconds))
elif options.bond_constraints == 'allbonds':
dt = 0.005*unit.picoseconds
if options.ensemble == 'NVE':
integ = mm.VerletIntegrator(dt)
else:
integ = mm.LangevinIntegrator(temperature, friction, dt)
else:
raise ValueError(f'Unknown bond_constraints {options.bond_constraints}')
test_result['timestep_in_fs'] = dt.value_in_unit(unit.femtoseconds)
properties = {}
initialSteps = 5
if options.device is not None and platform.getName() in ('CUDA', 'OpenCL'):
platform = mm.Platform.getPlatformByName(options.platform)
if options.device is not None and 'DeviceIndex' in platform.getPropertyNames():
properties['DeviceIndex'] = options.device
if ',' in options.device or ' ' in options.device:
initialSteps = 250
if options.precision is not None and platform.getName() in ('CUDA', 'OpenCL'):
if (options.precision is not None) and ('Precision' in platform.getPropertyNames()):
properties['Precision'] = options.precision
# Run the simulation.
# Add barostat if requested
if options.ensemble == 'NPT':
system.addForce(mm.MonteCarloBarostat(1*unit.bar, temperature, 100))
# Create the Context
integ.setConstraintTolerance(1e-5)
if len(properties) > 0:
context = mm.Context(system, integ, platform, properties)
else:
context = mm.Context(system, integ, platform)
# Store information about the Platform used by the Context
platform = context.getPlatform()
test_result['platform'] = platform.getName()
test_result['platform_properties'] = { property_name : platform.getPropertyValue(context, property_name) for property_name in platform.getPropertyNames() }
# Prepare the simulation
context.setPositions(positions)
if amber:
if inpcrd.boxVectors is not None:
context.setPeriodicBoxVectors(*inpcrd.boxVectors)
mm.LocalEnergyMinimizer.minimize(context, 100*unit.kilojoules_per_mole/unit.nanometer)
context.setVelocitiesToTemperature(temperature)
if options.serialize:
# Apply constraints on positions and velocities if needing to serialize for Folding@home
tol = 1.0e-8
context.applyConstraints(tol)
context.applyVelocityConstraints(tol)
state = context.getState(getPositions=True, getVelocities=True, getEnergy=True, getForces=True, getParameters=True)
# Time integration, ensuring we trigger kernel compilation before we start timing
steps = 20
while True:
time = timeIntegration(context, steps, initialSteps)
if time >= 0.5*options.seconds:
elapsed_time = timeIntegration(context, steps, initialSteps)
if elapsed_time >= 0.5*options.seconds:
break
if time < 0.5:
steps = int(steps*1.0/time) # Integrate enough steps to get a reasonable estimate for how many we'll need.
if elapsed_time < 0.5:
steps = int(steps*1.0/elapsed_time) # Integrate enough steps to get a reasonable estimate for how many we'll need.
else:
steps = int(steps*options.seconds/time)
print('Integrated %d steps in %g seconds' % (steps, time))
print('%g ns/day' % (dt*steps*86400/time).value_in_unit(unit.nanoseconds))
steps = int(steps*options.seconds/elapsed_time)
test_result['steps'] = steps
test_result['elapsed_time'] = elapsed_time
time_per_step = elapsed_time * unit.seconds / steps
ns_per_day = (integ.getStepSize() / time_per_step) / (unit.nanoseconds/unit.day)
test_result['ns_per_day'] = ns_per_day
# Serialize XML files for Folding@home benchmark if requested
if options.serialize:
wu_duration = 5*unit.minutes
ncheckpoints_per_wu = 2
nsteps_per_wu = int(wu_duration / time_per_step)
coredata = dict()
coredata['checkpointFreq'] = int(nsteps_per_wu / ncheckpoints_per_wu)
coredata['numSteps'] = ncheckpoints_per_wu * coredata['checkpointFreq']
coredata['xtcFreq'] = coredata['numSteps']
coredata['precision'] = options.precision
coredata['xtcAtoms'] = 'solute'
coredata['disableCheckpointStateTests'] = 1 # disable checkpoint state tests
coredata['DisablePmeStream'] = 0 # don't disable separate PME streams
coredata['forceTolerance'] = 200000 # for some reason, we have to make this large (TODO: Investigate why)
coredata['energyTolerance'] = 200000 # for some reason, we have to make this large (TODO: Investigate why)
forces = { force.getName() : force for force in system.getForces() }
NONBONDED_METHODS = { 0 : 'NoCutoff', 1 : 'CutoffNonPeriodic', 2 : 'CutoffPeriodic', 3 : 'Ewald', 4 : 'PME', 5 : 'LJPME' }
if 'NonbondedForce' in forces:
nonbonded_method = NONBONDED_METHODS[forces['NonbondedForce'].getNonbondedMethod()]
else:
nonbonded_method = 'unknown'
metadata = {
'name' : testName,
'description' : f'OpenMM Benchmark Suite : {testName}',
'precision' : options.precision,
'num_atoms' : context.getSystem().getNumParticles(),
'nonbonded_method' : f'{nonbonded_method}',
'timestep' : integ.getStepSize() / unit.picoseconds,
'integrator' : f'{integ.__class__.__name__}',
}
directory = os.path.join(options.serialize, f'{testName}-{options.precision}')
serializeTest(directory=directory, system=context.getSystem(), integrator=context.getIntegrator(), state=state, coredata=coredata, metadata=metadata)
# Clean up
del context, integ
printTestResult(test_result, options)
appendTestResult(test_result=test_result, filename=options.outfile)
# Parse the command line options.
parser = ArgumentParser()
platformNames = [mm.Platform.getPlatform(i).getName() for i in range(mm.Platform.getNumPlatforms())]
parser.add_argument('--platform', dest='platform', choices=platformNames, help='name of the platform to benchmark')
parser.add_argument('--test', dest='test', choices=('gbsa', 'rf', 'pme', 'apoa1rf', 'apoa1pme', 'apoa1ljpme', 'amoebagk', 'amoebapme', 'amber20-dhfr', 'amber20-factorix', 'amber20-cellulose', 'amber20-stmv'),
help='the test to perform: gbsa, rf, pme, apoa1rf, apoa1pme, apoa1ljpme, amoebagk, amoebapme, amber20-dhfr, amber20-factorix, amber20-cellulose, amber20-stmv [default: all except amber-*]')
parser.add_argument('--ensemble', default='NVT', dest='ensemble', choices=('NPT', 'NVE', 'NVT'), help='the thermodynamic ensemble to simulate [default: NVT]')
parser.add_argument('--pme-cutoff', default=0.9, dest='cutoff', type=float, help='direct space cutoff for PME in nm [default: 0.9]')
platform_speeds = { mm.Platform.getPlatform(i).getName() : mm.Platform.getPlatform(i).getSpeed() for i in range(mm.Platform.getNumPlatforms()) }
PLATFORMS = [platform for platform, speed in sorted(platform_speeds.items(), key=lambda item: item[1], reverse=True)]
TESTS = ('gbsa', 'rf', 'pme', 'apoa1rf', 'apoa1pme', 'apoa1ljpme', 'amoebagk', 'amoebapme', 'amber20-dhfr', 'amber20-factorix', 'amber20-cellulose', 'amber20-stmv')
ENSEMBLES = ('NVE', 'NVT', 'NPT')
BOND_CONSTRAINTS = ('hbonds', 'allbonds')
PRECISIONS = ('single', 'mixed', 'double')
POLARIZATION_MODES = ('direct', 'extrapolated', 'mutual')
STYLES = ('simple', 'table')
parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,
description="Run one or more benchmarks of OpenMM",
epilog="""
Example: run the full suite of benchmarks for the CUDA platform, printing the results as a table
python benchmark.py --platform=CUDA --style=table
Example: run the apoa1pme benchmark for the CPU platform with a reduced cutoff distance
python benchmark.py --platform=CPU --test=apoa1pme --pme-cutoff=0.8
Example: run the full suite in mixed precision mode, saving the results to a YAML file
python benchmark.py --platform=CUDA --precision=mixed --outfile=benchmark.yaml""")
parser.add_argument('--platform', dest='platform', choices=PLATFORMS, help='name of the platform to benchmark')
parser.add_argument('--test', default=','.join(TESTS), dest='test', help=f'the test to perform, or comma-separated list: {TESTS} [default: all]')
parser.add_argument('--ensemble', default='NVT', dest='ensemble', help=f'the thermodynamic ensemble to simulate: {ENSEMBLES} [default: NVT]')
parser.add_argument('--pme-cutoff', default=0.9, dest='pme_cutoff', type=float, help='direct space cutoff for PME in nm [default: 0.9]')
parser.add_argument('--seconds', default=60, dest='seconds', type=float, help='target simulation length in seconds [default: 60]')
parser.add_argument('--polarization', default='mutual', dest='polarization', choices=('direct', 'extrapolated', 'mutual'), help='the polarization method for AMOEBA: direct, extrapolated, or mutual [default: mutual]')
parser.add_argument('--polarization', default='mutual', dest='polarization', choices=POLARIZATION_MODES, help=f'the polarization method for AMOEBA: {POLARIZATION_MODES} [default: mutual]')
parser.add_argument('--mutual-epsilon', default=1e-5, dest='epsilon', type=float, help='mutual induced epsilon for AMOEBA [default: 1e-5]')
parser.add_argument('--heavy-hydrogens', action='store_true', default=False, dest='heavy', help='repartition mass to allow a larger time step')
parser.add_argument('--bond-constraints', default='hbonds', dest='bond_constraints', help=f'hbonds: constrain bonds to hydrogen, use 1.5*amu H mass; allbonds: constrain all bonds, use 4*amu H mass, and use larger timestep. This option is ignored for AMOEBA: {BOND_CONSTRAINTS} [default: hbonds]')
parser.add_argument('--device', default=None, dest='device', help='device index for CUDA or OpenCL')
parser.add_argument('--precision', default='single', dest='precision', choices=('single', 'mixed', 'double'), help='precision mode for CUDA or OpenCL: single, mixed, or double [default: single]')
parser.add_argument('--precision', default='single', dest='precision', help=f'precision modes for CUDA or OpenCL: {PRECISIONS} [default: single]')
parser.add_argument('--style', default='simple', dest='style', choices=STYLES, help=f'output style: {STYLES} [default: simple]')
parser.add_argument('--outfile', default=None, dest='outfile', help='output filename for benchmark logging (must end with .yaml or .json)')
parser.add_argument('--serialize', default=None, dest='serialize', help='if specified, output serialized test systems for Folding@home or other uses')
parser.add_argument('--verbose', default=False, action='store_true', dest='verbose', help='if specified, print verbose output')
args = parser.parse_args()
if args.platform is None:
parser.error('No platform specified')
print('Platform:', args.platform)
if args.platform in ('CUDA', 'OpenCL'):
print('Precision:', args.precision)
if args.device is not None:
print('Device:', args.device)
# Run the simulations.
# Collect system information
system_info = dict()
import socket, platform
system_info['hostname'] = socket.gethostname()
system_info['timestamp'] = datetime.utcnow().isoformat()
system_info['openmm_version'] = mm.version.version
system_info['cpuinfo'] = cpuinfo()
system_info['cpuarch'] = platform.processor()
system_info['system'] = platform.system()
# TODO: Capture information on how many CPU threads will be used
# Attempt to get GPU info
try:
import subprocess
cmd = 'nvidia-smi --query-gpu=driver_version,gpu_name --format=csv,noheader'
output = subprocess.check_output(cmd, shell=True, text=True)
system_info['nvidia_driver'], system_info['gpu'] = output.strip().split(', ')
except Exception as e:
pass
for key, value in system_info.items():
print(f'{key}: {value}')
if args.outfile is not None:
# Remove output file
import os
if os.path.exists(args.outfile):
os.unlink(args.outfile)
# Write system info
appendTestResult(args.outfile, system_info=system_info)
if args.test is None:
for test in ('gbsa', 'rf', 'pme', 'apoa1rf', 'apoa1pme', 'apoa1ljpme', 'amoebagk', 'amoebapme'):
tests = args.test.split(',')
if not set(tests).issubset(TESTS):
parser.error(f'Available tests: {TESTS}')
precisions = args.precision.split(',')
if args.platform == 'Reference':
precisions = ['double']
if args.platform == 'CPU':
precisions = ['mixed']
if not set(precisions).issubset(PRECISIONS):
parser.error(f'Available precisions: {PRECISIONS}')
ensembles = args.ensemble.split(',')
if not set(ensembles).issubset(ENSEMBLES):
parser.error(f'Available ensembles: {ENSEMBLES}')
bond_constraints = args.bond_constraints.split(',')
if not set(bond_constraints).issubset(BOND_CONSTRAINTS):
parser.error(f'Available bond constraints: {BOND_CONSTRAINTS}')
# Combinatorially run all requested benchmarks, ignoring combinations that cannot be run
from openmm import OpenMMException
from itertools import product
if args.style == 'simple':
for (test, args.bond_constraints, args.ensemble, args.precision) in product(tests, bond_constraints, ensembles, precisions):
try:
runOneTest(test, args)
except OpenMMException as e:
if args.verbose:
print(e)
pass
elif args.style == 'table':
print()
print('Test Precision Constraints H mass (amu) dt (fs) Ensemble Platform ns/day')
for (test, args.bond_constraints, args.ensemble, args.precision) in product(tests, bond_constraints, ensembles, precisions):
try:
runOneTest(test, args)
except Exception as ex:
print('Test failed: %s' % ex.message)
except OpenMMException as e:
if args.verbose:
print(e)
pass
else:
runOneTest(args.test, args)
raise ValueError(f"style {args.style} unkown; must be one of ['simple', 'table']")
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment