Повторный wandb.init() в параллельных проходах wandb ⇐ Python

[Anonymous]

Я написал некоторый код, пытаясь распараллелить свои проверки wandb, поскольку сходимость модели, с которой я работаю, занимает много времени, и мне нужно пройти множество подпроцессов. По сути, у меня сейчас нет такой роскоши, как время. Вот обобщенный фрагмент моего кода:

Код: Выделить всё

def run_pipeline(args):
# Stuff happens here

# Wandb init
group = "within_session" if session_config["within_session"] else "across_session"
run = wandb.init(name=f"{sessions[i]}_{group}_decoder_run", group=group, config=sweep_config, reinit=True)

# Model training

return results

def run_pipeline_wrapper(args):
# Stuff happens here
run_pipeline(args)

return None

if __name__ == "__main__":
total_runs = 30
agents = 5
runs_per_agent = total_runs // agents

sweep_config = {'method': 'random'}
parameters_dict = {
# Lota of parameters to sweep
}
sweep_config['parameters'] = parameters_dict

# Create a sweep id that stores sweep ids
sweep_id_json_path = 'sweep_id.json'
if not os.path.exists(sweep_id_json_path):
with open(sweep_id_json_path, 'w') as f:
json.dump({}, f)
sweep_id_json = json.load(open(sweep_id_json_path, 'r'))

# Sessions_list = number of unique data that I need to run my sweeps
for i in range(len(sessions_list)):

# Preparing a partial method to pass
run_pipeline_with_args = partial(run_pipeline_wrapper, args)

# I cache the existing sweep_ids in a json file to help in attaching sweep ids if I rerun the code again
if f"{sessions_list[i]}_{is_within}" not in sweep_id_json:
sweep_id = wandb.sweep(sweep_config, project=f"HPC_model_{sess}_session_{data}_{data_type}")
else:
sweep_id = wandb.sweep(sweep_config, project=f"HPC_model_{sess}_session_{data}_{data_type}"
, prior_runs=sweep_id_json[f"{sessions_list[i]}_{is_within}"])

# This is the parallelization logic, where I parallelize the sweeps
with concurrent.futures.ThreadPoolExecutor(max_workers=agents) as executor:
futures = [
executor.submit(wandb.agent, sweep_id, run_pipeline_with_args, count=runs_per_agent)
for _ in range(agents)
]

concurrent.futures.wait(futures)

Когда я запускаю этот код, он зависает в wandb.init(), и этот процесс в конечном итоге завершается из-за тайм-аута. Я не думаю, что это проблема увеличения таймаута wandb. Как мне это исправить? Как вы думаете, это может быть проблемой из-за моей логики распараллеливания? Если да, то как вы, разработчики, распараллеливаете проверки wandb в коде?



Подробнее здесь: https://stackoverflow.com/questions/792 ... ndb-sweeps