Ability to speicfy threads in docker

This commit is contained in:
Kevin Wu
2023-03-09 16:11:06 -08:00
parent 9ca513f5e1
commit 646e1b655e

View File

@@ -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__":