from multiprocessing import Pool
import sys, subprocess
from sklearn.model_selection import ParameterSampler

N_DEVICES = 9


def spawn(params):
    args_str = ' '.join('--{} {}'.format(key, params[key]) for key in params)
    return subprocess.run([sys.executable, 'train.personal.py'] + args_str.split(), env=None)


param_grid = dict(
    vnet_choice=[1, 4],
    td_num_steps=[10, 20, 50, 100],
    v_num_updates_per_step=[2, 5, 10, 20],
    v_batch_size=[8, 16, 32, 64],
    drop_rate=[0.3, 0.5, 0.7],
    v_lr=[5e-4, 1e-3, 2e-3],
    grad_clip=[10, 100]
)

param_grid = list(ParameterSampler(param_grid, 100))
for i, params in enumerate(param_grid):
    params['gpu'] = i % N_DEVICES - 1

pool = Pool((N_DEVICES * 2))
ret_codes = pool.map(spawn, param_grid)
print(ret_codes)