mirror of
https://github.com/microsoft/foldingdiff.git
synced 2026-06-04 13:30:33 +08:00
Ability to speicfy threads in docker
This commit is contained in:
@@ -16,10 +16,11 @@ def build_parser():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("input_file", help="Input file to run GROMACS on")
|
||||
parser.add_argument("output_dir", help="Output dir to write output files to")
|
||||
parser.add_argument("--gpu", type=int, default=0, help="GPU to use")
|
||||
return parser
|
||||
|
||||
|
||||
def run_gromacs_in_docker(fname: str, out_dir: str):
|
||||
def run_gromacs_in_docker(fname: str, out_dir: str, gpu: int = 0):
|
||||
"""
|
||||
Run gromacs in docker
|
||||
"""
|
||||
@@ -32,7 +33,8 @@ def run_gromacs_in_docker(fname: str, out_dir: str):
|
||||
# Copy the file into the directory
|
||||
shutil.copy(fname, tmpdir)
|
||||
# Build and run the command
|
||||
cmd = f"nvidia-docker run -it --rm -v {tmpdir}:/host_pwd --workdir /host_pwd wukevin:gromacs-latest {os.path.basename(fname)}"
|
||||
# https://github.com/NVIDIA/nvidia-docker/wiki/Frequently-Asked-Questions#i-have-multiple-gpu-devices-how-can-i-isolate-them-between-my-containers
|
||||
cmd = f"nvidia-docker run -it --rm -e NVIDIA_VISIBLE_DEVICES={gpu} -v {tmpdir}:/host_pwd --workdir /host_pwd wukevin:gromacs-latest {os.path.basename(fname)}"
|
||||
subprocess.call(cmd, shell=True)
|
||||
|
||||
for fname in os.listdir(tmpdir):
|
||||
@@ -44,7 +46,7 @@ def main():
|
||||
args = build_parser().parse_args()
|
||||
if not os.path.isdir(args.output_dir):
|
||||
os.makedirs(args.output_dir)
|
||||
run_gromacs_in_docker(args.input_file, args.output_dir)
|
||||
run_gromacs_in_docker(args.input_file, args.output_dir, gpu=args.gpu)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user