import multiprocessing
import subprocess
import os

def init(queue):
    global gpuid
    gpuid = queue.get()

gpus_list = list(range(8))
num_gpus = len(gpus_list)
manager = multiprocessing.Manager()
gpuQueue = manager.Queue()
[gpuQueue.put(i) for i in gpus_list]

filename = 'evaluate_experiment.py'

def f(snap_dir, env):
    global gpuid
    env = env.copy()
    env['CUDA_VISIBLE_DEVICES'] = str(gpuid)
    print(gpuid, snap_dir)
    command = ['python', filename, '--snap_dir', snap_dir]
    command = list(map(str, command))
    subprocess.run(command, env=env)
    return ' '.join(command)

snap_dirs = os.listdir('/homes/agolinsk/code/sylvester-flows-ortho/snapshots')
env = os.environ.copy()

with multiprocessing.Pool(num_gpus, init, (gpuQueue,)) as p:
    res = [p.apply_async(f, (snap_dir, env)) for snap_dir in snap_dirs]
    print([r.get() for r in res])

print("Done")