benchmark.py 25.1 KB
Newer Older
1
from __future__ import print_function
2
3
4
import openmm.app as app
import openmm as mm
import openmm.unit as unit
bdenhollander's avatar
bdenhollander committed
5
from datetime import datetime, timezone
6
import argparse
7
import os
8
9
10

def cpuinfo():
    """Return CPU info"""
11
    import platform, subprocess, re
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
    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
37
        The test result to append to the 'benchmarks' block
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
    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

    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']")
98

99
def timeIntegration(context, steps, initialSteps):
100
    """Integrate a Context for a specified number of steps, then return how many seconds it took."""
101
    context.getIntegrator().step(initialSteps) # Make sure everything is fully initialized
Peter Eastman's avatar
Peter Eastman committed
102
    context.getState(energy=True)
103
104
    start = datetime.now()
    context.getIntegrator().step(steps)
Peter Eastman's avatar
Peter Eastman committed
105
    context.getState(energy=True)
106
    end = datetime.now()
peastman's avatar
peastman committed
107
    elapsed = end-start
108
109
    return elapsed.seconds + elapsed.microseconds*1e-6

110
111
112
113
114
115
116
117
118
119
120
121
122
123
def downloadAmberSuite():
    """Download and extract Amber benchmark to Amber20_Benchmark_Suite/ in current directory."""
    dirname = 'Amber20_Benchmark_Suite'
    url = 'https://ambermd.org/Amber20_Benchmark_Suite.tar.gz'
    if not os.path.exists(dirname):
        import urllib.request
        print('Downloading', url)
        filename, headers = urllib.request.urlretrieve(url, filename='Amber20_Benchmark_Suite.tar.gz')
        import tarfile
        print('Extracting', filename)
        tarfh = tarfile.open(filename, 'r:gz')
        tarfh.extractall(path=dirname)
    return dirname

124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
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

    """
153
    explicit = (testName not in ('gbsa', 'amoebagk'))
154
    amoeba = (testName in ('amoebagk', 'amoebapme'))
peastman's avatar
peastman committed
155
    apoa1 = testName.startswith('apoa1')
156
    amber = (testName.startswith('amber'))
157
    hydrogenMass = None
158

159
160
161
162
163
164
165
    # Create dictionary of test parameters
    test_parameters = dict()

    # Store the test name
    test_parameters['test'] = testName

    # Create the System.
166
167
    if amoeba:
        constraints = None
168
169
        test_parameters['epsilon'] = epsilon
        epsilon = float(epsilon)
170
171
172
173
174
        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
175
            system = ff.createSystem(pdb.topology, nonbondedMethod=app.PME, nonbondedCutoff=cutoff, vdwCutoff=vdwCutoff, constraints=constraints, ewaldErrorTolerance=0.00075, mutualInducedTargetEpsilon=epsilon, polarization=polarization)
176
177
178
        else:
            ff = app.ForceField('amoeba2009.xml', 'amoeba2009_gk.xml')
            pdb = app.PDBFile('5dfr_minimized.pdb')
179
            system = ff.createSystem(pdb.topology, nonbondedMethod=app.NoCutoff, constraints=constraints, mutualInducedTargetEpsilon=epsilon, polarization=polarization)
180
181
182
        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)
183
        positions = pdb.positions
184
185
        test_parameters['constraints'] = 'None'
        test_parameters['hydrogen_mass'] = '1'
186
187
    elif amber:
        dirname = downloadAmberSuite()
188
        names = {'amber20-dhfr':'JAC', 'amber20-cellulose':'Cellulose', 'amber20-stmv':'STMV'}
189
190
191
192
193
194
        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
        method = app.PME
        constraints = app.HBonds
195
        cutoff = pme_cutoff
196
        system = prmtop.createSystem(nonbondedMethod=method, nonbondedCutoff=cutoff, constraints=constraints)
197
198
199
200
201
        if inpcrd.boxVectors is not None:
            system.setDefaultPeriodicBoxVectors(*inpcrd.boxVectors)
        test_parameters['cutoff'] = cutoff
        test_parameters['constraints'] = 'HBonds'
        test_parameters['hydrogen_mass'] = '1.5'
202
    else:
peastman's avatar
peastman committed
203
204
205
206
207
        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
208
                cutoff = pme_cutoff
peastman's avatar
peastman committed
209
210
            elif testName == 'apoa1ljpme':
                method = app.LJPME
211
                cutoff = pme_cutoff
peastman's avatar
peastman committed
212
            else:
213
                # Reaction field uses hard-coded cutoff
peastman's avatar
peastman committed
214
                method = app.CutoffPeriodic
215
                cutoff = 1.0 # nanometers ; JDC: Shouldn't this be larger for reaction field?
peastman's avatar
peastman committed
216
        elif explicit:
217
218
219
220
            ff = app.ForceField('amber99sb.xml', 'tip3p.xml')
            pdb = app.PDBFile('5dfr_solv-cube_equil.pdb')
            if testName == 'pme':
                method = app.PME
221
                cutoff = pme_cutoff
222
            else:
223
                # Reaction field uses hard-coded cutoff
224
                method = app.CutoffPeriodic
225
                cutoff = 1.0 # nanometers; JDC: Shouldn't this be larger for reaction field?
226
227
228
229
        else:
            ff = app.ForceField('amber99sb.xml', 'amber99_obc.xml')
            pdb = app.PDBFile('5dfr_minimized.pdb')
            method = app.CutoffNonPeriodic
230
231
232
233
234
235
236
237
            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.
238
            hydrogenMass = 4*unit.amu
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
            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
    """
    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)])
323
        else:
324
325
326
327
328
329
330
331
332
            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':
Peter Eastman's avatar
Peter Eastman committed
333
            dt = 0.004*unit.picoseconds
334
335
336
337
            if options.ensemble == 'NVE':
                integ = mm.VerletIntegrator(dt)
            else:
                integ = mm.LangevinMiddleIntegrator(temperature, friction, dt)
338
339
340
341
342
343
344
345
346
347
        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)
348
    properties = {}
349
    initialSteps = 5
Peter Eastman's avatar
Peter Eastman committed
350
    platform = mm.Platform.getPlatform(options.platform)
351
    if options.device is not None and 'DeviceIndex' in platform.getPropertyNames():
352
        properties['DeviceIndex'] = options.device
353
354
        if ',' in options.device or ' ' in options.device:
            initialSteps = 250
bdenhollander's avatar
bdenhollander committed
355
356
    if options.disable_pme_stream:
        properties['DisablePmeStream'] = 'true'
357
358
    if options.opencl_platform is not None and 'OpenCLPlatformIndex' in platform.getPropertyNames():
        properties['OpenCLPlatformIndex'] = options.opencl_platform
359
    if (options.precision is not None) and ('Precision' in platform.getPropertyNames()):
360
        properties['Precision'] = options.precision
361

362
363
364
365
366
    # Add barostat if requested
    if options.ensemble == 'NPT':
        system.addForce(mm.MonteCarloBarostat(1*unit.bar, temperature, 100))

    # Create the Context
367
368
369
370
371
    integ.setConstraintTolerance(1e-5)
    if len(properties) > 0:
        context = mm.Context(system, integ, platform, properties)
    else:
        context = mm.Context(system, integ, platform)
372
373
374
375
376
377
378

    # 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
379
380
    context.setPositions(positions)
    if amber:
381
382
        mm.LocalEnergyMinimizer.minimize(context, 100*unit.kilojoules_per_mole/unit.nanometer)
    context.setVelocitiesToTemperature(temperature)
383

384
385
386
387
388
    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)
Peter Eastman's avatar
Peter Eastman committed
389
        state = context.getState(positions=True, velocities=True, energy=True, forces=True, parameters=True)
390
391

    # Time integration, ensuring we trigger kernel compilation before we start timing
392
393
    steps = 20
    while True:
394
395
        elapsed_time = timeIntegration(context, steps, initialSteps)
        if elapsed_time >= 0.5*options.seconds:
396
            break
397
398
        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.
399
        else:
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
            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)
449
450
451

# Parse the command line options.

452
453
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)]
454
TESTS = ('gbsa', 'rf', 'pme', 'apoa1rf', 'apoa1pme', 'apoa1ljpme', 'amoebagk', 'amoebapme', 'amber20-dhfr', 'amber20-cellulose', 'amber20-stmv')
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
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]')
479
parser.add_argument('--seconds', default=60, dest='seconds', type=float, help='target simulation length in seconds [default: 60]')
480
parser.add_argument('--polarization', default='mutual', dest='polarization', choices=POLARIZATION_MODES, help=f'the polarization method for AMOEBA: {POLARIZATION_MODES} [default: mutual]')
481
parser.add_argument('--mutual-epsilon', default=1e-5, dest='epsilon', type=float, help='mutual induced epsilon for AMOEBA [default: 1e-5]')
482
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]')
bdenhollander's avatar
bdenhollander committed
483
484
parser.add_argument('--disable-pme-stream', default=False, action='store_true', dest='disable_pme_stream', help='disable use of a separate GPU stream for PME')
parser.add_argument('--device', default=None, dest='device', help='device index for CUDA, HIP, or OpenCL')
485
parser.add_argument('--opencl-platform', default=None, dest='opencl_platform', help='platform index for OpenCL')
bdenhollander's avatar
bdenhollander committed
486
parser.add_argument('--precision', default='single', dest='precision', help=f'precision modes for CUDA, HIP, or OpenCL: {PRECISIONS} [default: single]')
487
488
489
490
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')
491
492
args = parser.parse_args()
if args.platform is None:
493
494
    parser.error('No platform specified')

495
496
497
498
# Collect system information
system_info = dict()
import socket, platform
system_info['hostname'] = socket.gethostname()
bdenhollander's avatar
bdenhollander committed
499
system_info['timestamp'] = datetime.now(timezone.utc).isoformat()
500
501
502
503
504
505
506
507
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:
bdenhollander's avatar
bdenhollander committed
508
    import shutil
509
    import subprocess
bdenhollander's avatar
bdenhollander committed
510
511
512
513
    if shutil.which('nvidia-smi') is not None:
        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(', ')
514
515
516
517
518
519
520
521
522
523
524
525
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
    if os.path.exists(args.outfile):
        os.unlink(args.outfile)
    # Write system info
    appendTestResult(args.outfile, system_info=system_info)
526

527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
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):
563
        try:
564
            runOneTest(test, args)
565
566
567
568
        except OpenMMException as e:
            if args.verbose:
                print(e)
            pass
569
else:
570
    raise ValueError(f"style {args.style} unknown; must be one of ['simple', 'table']")