From 646e1b655ea5ba010236a3fc4d1e7a6360671467 Mon Sep 17 00:00:00 2001 From: Kevin Wu Date: Thu, 9 Mar 2023 16:11:06 -0800 Subject: [PATCH] Ability to speicfy threads in docker --- scripts/gromacs/gromacs_docker.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/scripts/gromacs/gromacs_docker.py b/scripts/gromacs/gromacs_docker.py index 67f7d71..bafb829 100644 --- a/scripts/gromacs/gromacs_docker.py +++ b/scripts/gromacs/gromacs_docker.py @@ -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__":