MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / _run_wrapped

Function _run_wrapped

imperative/python/megengine/distributed/launcher.py:20–57  ·  view source on GitHub ↗

r"""Init distributed process group and run wrapped function.

(
    func,
    is_multimachine,
    master_ip,
    port,
    world_size,
    rank,
    dev,
    device_type,
    args,
    kwargs,
    backend,
    queue: mp.Queue,
    machine_ranks: list,
)

Source from the content-addressed store, hash-verified

18
19
20def _run_wrapped(
21 func,
22 is_multimachine,
23 master_ip,
24 port,
25 world_size,
26 rank,
27 dev,
28 device_type,
29 args,
30 kwargs,
31 backend,
32 queue: mp.Queue,
33 machine_ranks: list,
34):
35 r"""Init distributed process group and run wrapped function."""
36 _check_interpreter_status()
37 _check_device_initialized(device_type, dev)
38 init_process_group(
39 master_ip=master_ip,
40 port=port,
41 world_size=world_size,
42 rank=rank,
43 device=dev,
44 backend=backend,
45 device_type=device_type,
46 )
47 # set NCCL_LAUNCH_MODE to avoid deadlock
48 os.environ["NCCL_LAUNCH_MODE"] = "PARALLEL"
49 _set_machine_ranks(machine_ranks)
50 if is_multimachine:
51 group_barrier()
52 ret = func(*args, **kwargs)
53 queue.put((dev, ret))
54 full_sync()
55 if is_multimachine:
56 group_barrier()
57 _exit(0)
58
59
60class launcher:

Callers

nothing calls this directly

Calls 8

init_process_groupFunction · 0.85
_set_machine_ranksFunction · 0.85
group_barrierFunction · 0.85
_exitFunction · 0.85
funcFunction · 0.50
putMethod · 0.45

Tested by

no test coverage detected