MCPcopy Create free account
hub / github.com/PaddlePaddle/Paddle / _set_trainer_env

Function _set_trainer_env

python/paddle/distributed/spawn.py:350–371  ·  view source on GitHub ↗
(env_dict, backend)

Source from the content-addressed store, hash-verified

348
349
350def _set_trainer_env(env_dict, backend):
351 # NOTE(chenweihang): [ Why need set FLAGS_selected_gpus or FLAGS_selected_xpus here? ]
352 # When the child process starts, it will inherit the configuration of the
353 # main process and set the FLAGS once, but the environment variable has
354 # not been set at this time, which leads to the FLAGS_selected_gpus or FLAGS_selected_xpus
355 # is keep same with mainprocess(usually empty), so manually update the flags here
356
357 # NOTE(xiongkun): why put backend here? because if gloo, we shouldn't set FLAGS_selectedXXX
358 #
359
360 if backend == 'nccl':
361 set_flags({'FLAGS_selected_gpus': env_dict['FLAGS_selected_gpus']})
362 elif backend == 'bkcl':
363 set_flags({'FLAGS_selected_xpus': env_dict['FLAGS_selected_xpus']})
364 else:
365 # NOTE(xiongkun) why not raise Error ?
366 # So far, we added support for CPU parallel, and will be applied when paddle is not
367 # compiled with cuda or xp. just do nothing.
368 pass
369
370 for var_name in env_dict:
371 os.environ[var_name] = env_dict[var_name]
372
373
374def _func_wrapper(func, args, error_queue, return_queue, env_dict, backend):

Callers 1

_func_wrapperFunction · 0.85

Calls 1

set_flagsFunction · 0.90

Tested by

no test coverage detected