Multi gpu

[BK33333]

Hi,
I was given some advice to use
the TaskMapper and WorkerMapper to run multiple GPUs but unfortunately I still have not managed to get this to work.
Does any one have a snippet of code using the TaskMapper or WorkerMapper?
thanks

[salotz]

You should provide some code for what you have tried so far and we can help from there.

[BK33333]

thanks for getting back to me: So far this is what I am using.

import os
import multiprocessing

def worker(gpu_id):
# Set CUDA_VISIBLE_DEVICES to the specified GPU ID for this worker
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)

# Initialize your CUDA context here or run your task that uses CUDA

List of GPUs you want to use

gpus = [0, 1]

Create a process for each GPU

processes = []
for gpu_id in gpus:
p = multiprocessing.Process(target=worker, args=(gpu_id,))
p.start()
processes.append(p)

Wait for all processes to complete

for p in processes:
p.join()

print("lets do this")
import sys

from copy import copy, deepcopy

import openmm as omm
import openmm.app as app
import openmm.unit as unit

from openmm.app import *
from openmm import *
from openmm.unit import *
from sys import stdout
import numpy as np

from copy import copy

#####now wepy import
from wepy.resampling.resamplers.resampler import NoResampler
from wepy.runners.openmm import OpenMMRunner, OpenMMState, gen_walker_state
from wepy.walker import Walker
from wepy.sim_manager import Manager
from wepy.reporter.hdf5 import WepyHDF5Reporter

import mdtraj
from wepy.util.mdtraj import mdtraj_to_json_topology

from geomm.rmsd import calc_rmsd
from geomm.superimpose import superimpose

from scipy.spatial.distance import euclidean
from wepy.resampling.distances.distance import Distance
from wepy.resampling.resamplers.wexplore import WExploreResampler
from wepy.reporter.dashboard import DashboardReporter
from wepy.reporter.dashboard import DashboardReporter
from wepy.reporter.receptor.dashboard import RebindingBCDashboardSection, ReceptorBCDashboardSection
from wepy.reporter.wexplore.dashboard import WExploreDashboardSection
from wepy.reporter.revo.dashboard import REVODashboardSection
from wepy.reporter.openmm import OpenMMRunnerDashboardSection

####################

OUTPUT_FIELDS = (
"box_vectors",
"box_volume",
# "forces",
"kinetic_energy",
"positions",
"potential_energy",
"time",
# "velocities",
## weights included automatically
)

JUSTPOS = {
'getPositions': True,
'getVelocities': False,
'getForces': False,
'getEnergy': False,
'getParameters': False,
'getParameterDerivatives': False,
'enforcePeriodicBox': True
}

ALMOST_ALL_TRUE = {
'getPositions': True,
'getVelocities': True,
'getForces': True,
'getEnergy': True,
'getParameters': True,
'getParameterDerivatives': True,
'enforcePeriodicBox': True ###switch for implicit
}

###############

Input Files

import os
cwd = os.getcwd().split("/")[-1]
print(cwd)

filestring = cwd
OUT = filestring
print(filestring, OUT + "...")

prmtop = app.AmberPrmtopFile(filestring+'.prmtop')
inpcrd = app.AmberInpcrdFile(filestring+'.inpcrd')

mdtraj_topology = mdtraj.Topology.from_openmm(prmtop.topology)
json_str_top = mdtraj_to_json_topology(mdtraj_topology)

nonbondedMethod = PME
nonbondedCutoff = 1*nanometers
ewaldErrorTolerance = 0.0005
constraints = None
rigidWater = False

temperature = 300kelvin
friction = 1/picosecond
pressure = 1atmospheres
barostatInterval = 25

system = prmtop.createSystem(nonbondedMethod=nonbondedMethod, nonbondedCutoff=nonbondedCutoff,
constraints=constraints, ewaldErrorTolerance=ewaldErrorTolerance)
system.addForce(MonteCarloBarostat(pressure, temperature, barostatInterval))
integrator = omm.LangevinIntegrator(300.0 * kelvin, 1/picosecond, 0.001*picoseconds)

#################
################

mini simulation just to equilibrate positions

presim = app.Simulation(prmtop.topology, system, integrator, platform=omm.Platform_getPlatformByName("CUDA"))
presim.context.setPositions(inpcrd.positions)
if inpcrd.boxVectors is not None:
presim.context.setPeriodicBoxVectors(*inpcrd.boxVectors)
presim.minimizeEnergy()

#anneal
print("annealing")
for i in range(20):
print(i)
integrator.setTemperature(6*(100-i*5)*kelvin)
presim.step(1000)

prestate = presim.context.getState(**ALMOST_ALL_TRUE)
#################

init_state = OpenMMState(prestate)

runner = OpenMMRunner(system, prmtop.topology, integrator,enforce_box=True, platform='CUDA')
##########################

#################

resampler = NoResampler()

############################
from wepy.work_mapper.task_mapper import TaskMapper

Initialize the TaskMapper with the desired number of workers (GPUs)

num_workers = 2 # Adjust this number based on your available GPUs
task_mapper = TaskMapper(num_workers=num_workers)

############################

now run it

n_cycles = 50
n_steps = 100 # 1000
steps = [n_steps for i in range(n_cycles)]

n_walkers = 5

init_weight = 1.0/n_walkers
init_walkers = [Walker(copy(init_state), init_weight) for i in range(n_walkers)]

h5_reporter = WepyHDF5Reporter(
save_fields=OUTPUT_FIELDS,
file_path=filestring+".h5", mode='w',
topology=json_str_top,
resampler=resampler
)
STEP_SIZE = 0.001*picoseconds
dashboard_reporter = DashboardReporter(
file_path=filestring+".wepy.dash.org",
mode='w+',
runner_dash=OpenMMRunnerDashboardSection(runner, STEP_SIZE),
bc_dash=None
)

sim_manager = Manager(init_walkers,
runner=runner,
resampler=resampler,
work_mapper=task_mapper,
reporters=[h5_reporter,dashboard_reporter])

final_walkers, _ = sim_manager.run_simulation(n_cycles, steps)

#################

[salotz]

You need to format the code so it is readable in this thread, see the documentation on Markdown or use the editor for code samples.

[salotz]

You should also give an errors or describe the incorrect behavior you are observing when you run this.