Parallel run wrapper with paramiko #Parallel

00README.md

Usage:

printf {0..100} | ./ssh_wrapper.py echo --cpus 2 --gpus 1 --email -N 2
printf myscript.py" --foo "{0..100} | ./ssh_wrapper.py python --cpus 2 --gpus 1 --email -N 2

hosts.tsv

ssh_wrapper.py

#!/bin/python3
import argparse
import os
import os.path as osp
import sys
from multiprocessing import Pool, current_process
from os.path import expanduser
import random
from subprocess import check_output
import time
import warnings

import paramiko


def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('exec', metavar='executable', type=str)
    parser.add_argument('--name', type=str, default=None, help="Default: None")
    parser.add_argument('--cpus', type=int, default=1)
    parser.add_argument('--gpus', type=int, default=0)
    parser.add_argument('--mem', type=float, default=1, help="Default: 1 (GB)")
    parser.add_argument(
        "-N", "--jobs_per_host", type=int, default=1, help="Default: 1")
    parser.add_argument(
        "--sleep",
        type=int,
        default=0,
        help="Sleep certain seconds to avoid congestion. Default: 0")
    parser.add_argument(
        "--email", action='store_true', help="Enable email notification.")

    # for SSH connection
    parser.add_argument(
        "--pw",
        metavar='password',
        type=str,
        default=None,
        help="Password. If not set, read it from ~/.parallel_run_pw ")
    parser.add_argument(
        "--host_file",
        type=str,
        default="~/hosts.tsv",
        help="Host description file.")

    return parser


def get_hosts(args):
    hosts = []
    with open(expanduser(args.host_file)) as fin:
        fin.readline()
        for line in fin:
            fields = line.strip().split()

            mem, cpus, gpus, alive = list(map(float, fields[1:]))

            if (alive and cpus > args.cpus and mem > args.mem
                    and (args.gpus <= gpus)):
                hosts.append(fields[0])

    return hosts


def runCommand(cmd, pw):

    pid = current_process().pid
    host = hosts[pid % len(hosts)]

    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        with warnings.catch_warnings():
            # NOTE: To ignore CryptographyDeprecationWarning;
            # remove it in later version of paramiko (2.4.2)
            warnings.simplefilter("ignore")
            client.connect(host, password=pw)
    except Exception:
        print("***** failed to connect host ", host, " *****", sep='')
        return False

    if args.sleep > 0:
        time.sleep(random.random() * args.sleep * len(hosts))

    print("Running on {}: {}".format(host, cmd))
    # execute the main commmand
    stdin, stdout, stderr = client.exec_command(
        "cd {}; {}".format(os.getcwd(), cmd), get_pty=True)
    for line in stdout:
        print("[{}]: {}".format(host, line), end="")
    client.close()
    return True


if __name__ == "__main__":
    # generate parameter combination
    args = get_parser().parse_args()
    if args.pw is None:
        with open(expanduser("~/.parallel_run_pw"), 'r') as fin:
            args.pw = fin.readline().strip()

    # get hosts
    hosts = get_hosts(args)
    hosts = hosts * args.jobs_per_host

    # read commmand fields
    exec_path = check_output(["which", args.exec])
    args.exec = osp.expanduser(exec_path.decode("utf")).strip()

    cmds = []
    for line in sys.stdin:
        cmds.append(args.exec + " " + line.strip())
    n_jobs = len(cmds)

    if args.email:
        text = """---Arguments---
{}

---Hosts---
{}

---Commands---
{}

---Job Name---
{}
""".format(" ".join(sys.argv[1:]), hosts, "\n".join(cmds), args.name)
        send_email(subject="Experiment Notification", text=text + "\nStart.")

    with Pool(len(hosts)) as p:
        p.starmap(runCommand, zip(cmds, [args.pw] * n_jobs))

    if args.email:
        send_email(
            subject="Experiment Notification", text=text + "\nFinished.")