| 36 | return (' '.join(envs)) |
| 37 | |
| 38 | def submit(args): |
| 39 | assert args.host_file is not None |
| 40 | with open(args.host_file) as f: |
| 41 | tmp = f.readlines() |
| 42 | assert len(tmp) > 0 |
| 43 | hosts=[] |
| 44 | for h in tmp: |
| 45 | if len(h.strip()) > 0: |
| 46 | # parse addresses of the form ip:port |
| 47 | h = h.strip() |
| 48 | |
| 49 | # parse mpi host file form ip slots=?? |
| 50 | # this is to create an unified api for mpi and ssh |
| 51 | i = h.find("slots=") |
| 52 | if i != -1: |
| 53 | h = h[:i].strip() |
| 54 | |
| 55 | i = h.find(":") |
| 56 | p = "22" |
| 57 | if i != -1: |
| 58 | p = h[i+1:] |
| 59 | h = h[:i] |
| 60 | # hosts now contain the pair ip, port |
| 61 | hosts.append((h, p)) |
| 62 | |
| 63 | def ssh_submit(nworker, nserver, pass_envs): |
| 64 | """ |
| 65 | customized submit script |
| 66 | """ |
| 67 | # thread func to run the job |
| 68 | def run(prog): |
| 69 | subprocess.check_call(prog, shell = True) |
| 70 | |
| 71 | # sync programs if necessary |
| 72 | local_dir = os.getcwd()+'/' |
| 73 | working_dir = local_dir |
| 74 | if args.sync_dst_dir is not None and args.sync_dst_dir != 'None': |
| 75 | working_dir = args.sync_dst_dir |
| 76 | pool = Pool(processes=len(hosts)) |
| 77 | for h in hosts: |
| 78 | pool.apply_async(sync_dir, args=(local_dir, h, working_dir)) |
| 79 | pool.close() |
| 80 | pool.join() |
| 81 | |
| 82 | |
| 83 | # launch jobs |
| 84 | for i in range(nworker + nserver): |
| 85 | pass_envs['DMLC_ROLE'] = 'server' if i < nserver else 'worker' |
| 86 | (node, port) = hosts[i % len(hosts)] |
| 87 | pass_envs['DMLC_NODE_HOST'] = node |
| 88 | prog = get_env(pass_envs) + ' cd ' + working_dir + '; ' + (' '.join(args.command)) |
| 89 | prog = 'ssh -o StrictHostKeyChecking=no ' + node + ' -p ' + port + ' \'' + prog + '\'' |
| 90 | thread = Thread(target = run, args=(prog,)) |
| 91 | thread.setDaemon(True) |
| 92 | thread.start() |
| 93 | |
| 94 | return ssh_submit |
| 95 | |