# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

from __future__ import annotations

import multiprocessing

import os

import signal

import sys

import warnings

from typing import TYPE_CHECKING, Any, Literal, TypedDict

# deprecated module import

# (TODO: GhostScreaming) It will be removed later.

from paddle.base import core

from paddle.device import get_device

from paddle.distributed.cloud_utils import (

_get_trainers_num,

get_cluster_and_pod,

)

from paddle.distributed.fleet.cloud_utils import use_paddlecloud

from paddle.distributed.fleet.launch import get_cluster_from_args

from paddle.distributed.fleet.launch_utils import (

DeviceMode,

block_windows_and_macos,

check_backend,

)

from paddle.distributed.utils.launch_utils import (

_prepare_trainer_env,

_print_arguments,

get_host_name_ip,

)

from paddle.framework import set_flags

if TYPE_CHECKING:

from collections.abc import Callable, Iterable

from typing_extensions import NotRequired, Unpack

class _SpawnOptions(TypedDict):

start_method: NotRequired[Literal['spawn', 'fork', 'forkserver']]

gpus: NotRequired[str | None]

xpus: NotRequired[str | None]

ips: NotRequired[str]

__all__ = []

class ParallelEnvArgs:

def __init__(self):

# Paddle cluster nodes ips, such as 192.168.0.16,192.168.0.17..

self.cluster_node_ips = None

# The current node ip.

self.node_ip = None

# whether to use paddlecloud platform to run your multi-process job.

# If false, no need to set this argument.

self.use_paddlecloud = None

# The trainer's started port on a single node

self.started_port = None

# Print the config or not

self.print_config = True

# It's for gpu training and the training process will run

# on the selected_devices, each process is bound to a single GPU.

# And if it's not set, this module will use all the gpu cards

# for training.

self.selected_devices = None

def _options_valid_check(options):

# `print_config` keeped as a debug options, not show to users

supported_options = [

'start_method',

'ips',

'gpus',

'xpus',

'print_config',

'backend',

]

deprecated_options = [

'selected_devices',

'started_port',

'cluster_node_ips',

'node_ip',

'use_paddlecloud',

]

for key in options:

if key not in supported_options:

if key in deprecated_options:

warnings.warn(

f"The config option ({key}) of `paddle.distributed.spawn` is deprecated. "

"Please use the latest config options stated in the `spawn` API documentation.",

DeprecationWarning,

)

else:

raise ValueError(

f"The config option ({key}) of `paddle.distributed.spawn` is not supported."

)

def _get_default_nprocs():

device = get_device()

if device in core.get_available_custom_device():

return core.get_custom_device_count(device.split(":")[0])

elif 'gpu' in device:

return core.get_cuda_device_count()

elif 'xpu' in device:

return core.get_xpu_device_count()

elif 'cpu' in device:

return multiprocessing.cpu_count()

else:

raise RuntimeError(

f"`paddle.distributed.spawn` does not support parallel training on device `{device}` now."

)

def _get_default_backend():

device = get_device()

if device in core.get_available_custom_device():

return 'xccl'

elif 'gpu' in device:

return 'nccl'

elif 'xpu' in device:

return 'bkcl'

elif 'cpu' in device:

return 'gloo'

else:

raise RuntimeError(

f"`paddle.distributed.spawn` does not support parallel training on device `{device}` now."

)

def _get_node_ip(ips):

node_ip = None

node_ips = [x.strip() for x in ips.split(',')]

if len(node_ips) == 1:

node_ip = node_ips[0]

else:

_, node_ip = get_host_name_ip()

return node_ip

def _get_subprocess_env_list(nprocs, options):

# NOTE (xiongkun03) Why put backend deduction here ?

# Because _get_subprocess_env_list is used by many testcases.

# So for compatibility, we put backend deduction here

# logic for handle backend option

if 'backend' not in options or options['backend'] == 'auto':

options['backend'] = _get_default_backend()

check_backend(options['backend'])

block_windows_and_macos(options['backend'])

# construct processes env list

processes_env_list = []

# get args from kwargs

args = ParallelEnvArgs()

# deal with `ips`

args.cluster_node_ips = options.get('ips', None)

if args.cluster_node_ips is None:

args.cluster_node_ips = options.get('cluster_node_ips', None)

if args.cluster_node_ips is None:

args.cluster_node_ips = "127.0.0.1"

# deal with `gpus` or `xpus`

# set default selected devices(gpus or xpus)

# e.g. if the nprocs is 4, the selected gpus is "0,1,2,3"

# NOTE(chenweihang): [ why not use FLAGS_selected_gpus or FLAGS_selected_xpus directly? ]

# because the FLAGS_selected_gpus or FLAGS_selected_xpus may be used in other place,

# if we set FLAGS_selected_gpus or FLAGS_selected_xpus to be `0,1,2,3`, it may cause error

# when using `ParallelEnv`

# NOTE(chenweihang): use absolute gpu or xpu card id

if options['backend'] == 'nccl':

args.selected_devices = options.get('gpus', None)

if args.selected_devices is None:

args.selected_devices = options.get('selected_devices', None)

env_devices = os.getenv("CUDA_VISIBLE_DEVICES", None)

if env_devices is None or env_devices == "":

env_devices_list = [

str(x) for x in range(core.get_cuda_device_count())

]

else:

env_devices_list = env_devices.split(',')

if args.selected_devices is None:

if len(env_devices_list) < nprocs:

raise RuntimeError(

f"the number of visible devices({len(env_devices_list)}) is less than the number "

f"of spawn processes({nprocs}), please ensure that the correct "

"`nprocs` argument is passed or the environment variable "

"`CUDA_VISIBLE_DEVICES` is correctly configured."

)

args.selected_devices = ",".join(

[str(env_devices_list[x]) for x in range(0, nprocs)]

)

else:

selected_device_list = args.selected_devices.split(',')

if len(selected_device_list) != nprocs:

raise ValueError(

f"The number of selected devices({len(selected_device_list)}) is not equal to "

f"the number of spawn processes({nprocs}), please ensure that the "

"correct `nprocs` and `gpus` arguments are passed."

)

for card_id in selected_device_list:

if card_id not in env_devices_list:

raise ValueError(

"The selected gpu card {} cannot found in "

"CUDA_VISIBLE_DEVICES ({}).".format(

card_id, ",".join(env_devices_list)

)

)

elif options['backend'] == 'bkcl':

args.selected_devices = options.get('xpus', None)

if args.selected_devices is None:

args.selected_devices = options.get('selected_devices', None)

env_devices = os.getenv("XPU_VISIBLE_DEVICES", None)

if env_devices is None or env_devices == "":

env_devices_list = [

str(x) for x in range(core.get_xpu_device_count())

]

else:

env_devices_list = env_devices.split(',')

if args.selected_devices is None:

if len(env_devices_list) < nprocs:

raise RuntimeError(

f"the number of visible devices({len(env_devices_list)}) is less than the number "

f"of spawn processes({nprocs}), please ensure that the correct "

"`nprocs` argument is passed or the environment variable "

"`XPU_VISIBLE_DEVICES` is correctly configured."

)

args.selected_devices = ",".join(

[str(env_devices_list[x]) for x in range(0, nprocs)]

)

else:

selected_device_list = args.selected_devices.split(',')

if len(selected_device_list) != nprocs:

raise ValueError(

f"The number of selected devices({len(selected_device_list)}) is not equal to "

f"the number of spawn processes({nprocs}), please ensure that the "

"correct `nprocs` and `xpus` arguments are passed."

)

for card_id in selected_device_list:

if card_id not in env_devices_list:

raise ValueError(

"The selected xpu card {} cannot found in "

"XPU_VISIBLE_DEVICES ({}).".format(

card_id, ",".join(env_devices_list)

)

)

elif options['backend'] == 'gloo':

# TODO check gpu / xpu flag must not exist

warnings.warn(

"Your model will be trained under CPUONLY mode by using GLOO,"

"because CPUPlace is specified manually or your installed PaddlePaddle only support CPU Device."

)

args.paddle_cpuonly = True

args.selected_devices = None

args.ips = args.cluster_node_ips

assert options.get('use_paddlecloud', None) is None, (

"CPUONLY spawn doesn't support use paddle cloud"

)

assert len(args.cluster_node_ips.split(',')) <= 1, (

"CPUONLY spawn only support single trainer, that is len(ips)=1, but got %s."

)

assert _get_trainers_num() == 1, (

"CPUONLY spawn doesn't support multi-trainer"

)

elif options['backend'] == 'xccl':

args.selected_devices = None

custom_device_name = core.get_all_custom_device_type()[0]

env_devices = os.getenv(f"FLAGS_selected_{custom_device_name}s", None)

if env_devices is None or env_devices == "":

env_devices_list = [

str(x)

for x in range(core.get_custom_device_count(custom_device_name))

]

else:

env_devices_list = env_devices.split(',')

if len(env_devices_list) < nprocs:

raise RuntimeError(

f"the number of visible devices({len(env_devices_list)}) is less than the number "

f"of spawn processes({nprocs}), please ensure that the correct "

"`nprocs` argument is passed or the environment variable "

f"`FLAGS_selected_{custom_device_name}s` is correctly configured."

)

args.selected_devices = ",".join(

[str(env_devices_list[x]) for x in range(0, nprocs)]

)

# set other inner args

args.node_ip = options.get('node_ip', None)

if args.node_ip is None:

args.node_ip = _get_node_ip(args.cluster_node_ips)

args.started_port = options.get('started_port', None)

args.use_paddlecloud = options.get('use_paddlecloud', None)

if args.use_paddlecloud is None:

args.use_paddlecloud = use_paddlecloud()

# get cluster and pod config

if options['backend'] == 'gloo':

devices_per_proc = list(range(0, nprocs))

cluster, pod = get_cluster_from_args(

args, DeviceMode.CPU, devices_per_proc

)

else:

cluster, pod = get_cluster_and_pod(args)

# prepare subprocess env list

for trainer in pod.trainers:

processes_env_list.append(

_prepare_trainer_env(cluster, trainer, options['backend'])

)

# [Debug] print config

args.print_config = options.get('print_config', False)

if args.print_config:

_print_arguments(args)

return processes_env_list

def _remove_risky_env():

# remove useless env vars

# no copy, each process will hold env vars itself

os.environ.pop("http_proxy", None)

os.environ.pop("https_proxy", None)

def _set_trainer_env(env_dict, backend):

# NOTE(chenweihang): [ Why need set FLAGS_selected_gpus or FLAGS_selected_xpus here? ]

# When the child process starts, it will inherit the configuration of the

# main process and set the FLAGS once, but the environment variable has

# not been set at this time, which leads to the FLAGS_selected_gpus or FLAGS_selected_xpus

# is keep same with mainprocess(usually empty), so manually update the flags here

# NOTE(xiongkun): why put backend here? because if gloo, we shouldn't set FLAGS_selectedXXX

#

if backend == 'nccl':

set_flags({'FLAGS_selected_gpus': env_dict['FLAGS_selected_gpus']})

elif backend == 'bkcl':

set_flags({'FLAGS_selected_xpus': env_dict['FLAGS_selected_xpus']})

else:

# NOTE(xiongkun) why not raise Error ?

# So far, we added support for CPU parallel, and will be applied when paddle is not

# compiled with cuda or xp. just do nothing.

pass

for var_name in env_dict:

os.environ[var_name] = env_dict[var_name]

def _func_wrapper(func, args, error_queue, return_queue, env_dict, backend):

try:

# config subprocess environment variables

_remove_risky_env()

_set_trainer_env(env_dict, backend)

# execute function

result = func(*args)

# record function return value

return_queue.put(result)

except KeyboardInterrupt:

pass

except Exception:

import traceback

error_queue.put(traceback.format_exc())

sys.exit(1)

class MultiprocessContext:

def __init__(self, processes, error_queues, return_queues):

self.error_queues = error_queues

# NOTE(chenweihang): The `spawn` method is mainly used

# to wrap the outermost execution function of the program for

# parallel execution. Generally, the return value is not concerned,

# but if the user needs to obtain the return value, users can get

# the return result of each process from context.return_queues

self.return_queues = return_queues

self.processes = processes

self.sentinels = {

process.sentinel: index for index, process in enumerate(processes)

}

def join(self, timeout=None):

if len(self.sentinels) == 0:

return True

ready = multiprocessing.connection.wait(

self.sentinels.keys(), timeout=timeout

)

error_index = None

for sentinel in ready:

index = self.sentinels.pop(sentinel)

process = self.processes[index]

process.join()

if process.exitcode != 0:

error_index = index

break

if error_index is None:

return len(self.sentinels) == 0

for process in self.processes:

if process.is_alive():

process.terminate()

process.join()

self._throw_exception(error_index)

def _throw_exception(self, error_index):

if self.error_queues[error_index].empty():

exitcode = self.processes[error_index].exitcode

if exitcode < 0:

name = signal.Signals(-exitcode).name

raise Exception(

f"Process {error_index} terminated with signal {name}."

)

else:

raise Exception(

f"Process {error_index} terminated with exit code {exitcode}."

)

original_trace = self.error_queues[error_index].get()

msg = (

"\n\n----------------------------------------------\n"

f"Process {error_index} terminated with the following error:\n"

"----------------------------------------------\n\n"

)

msg += original_trace

raise Exception(msg)

def spawn(

func: Callable[..., None],

args: Iterable[Any] = (),

nprocs: int = -1,

join: bool = True,

daemon: bool = False,

**options: Unpack[_SpawnOptions],

) -> MultiprocessContext:

"""

Start multiple processes with ``spawn`` method for parallel training.

.. note::

``spawn`` now only supports GPU or XPU collective mode. The collective mode

of GPU and XPU cannot be started at the same time, so the option `gpus` and

`xpus` cannot be configured at the same time.

Args:

func (function): The target function is called by spawned process.

This function need to be able to pickled, so it must be defined

at the top level of a module.

args (list|tuple, optional): Arguments passed to ``func``.

nprocs (int, optional): Number of processed to start. Default: -1.

when nprocs is -1, the available device will be obtained from

the environment variable when the model is executed: If use GPU,

the currently available device ID is obtained from the environment

variable CUDA_VISIBLE_DEVICES; If use XPU, the currently available

device ID is obtained from the environment variable XPU_VISIBLE_DEVICES.

join (bool, optional): Perform a blocking join on all spawned processes.

Default: True.

daemon (bool, optional): The spawned processes' daemon flag. Default: False.

**options(dict, optional): Other initial parallel execution environment

configuration options. The following options are currently supported:

(1) start_method (string): the way to start a process.

The start method can be ``spawn`` , ``fork`` , ``forkserver`` .

Because the CUDA runtime does not support the ``fork`` start method,

when use CUDA in subprocesses, we should start process by ``spawn``

or ``forkserver`` method. Default: "spawn" ;

(2) gpus (string): The training process will run on the

selected gpus, such as "0,1,2,3". Default: None;

(3) xpus (string): The training process will run on the

selected xpus, such as "0,1,2,3". Default: None;

(5) ips (string): Paddle cluster nodes ips, such as

"192.168.0.16,192.168.0.17". Default: "127.0.0.1" .

Returns:

``MultiprocessContext`` object, it hold the spawned processes.

Examples:

.. code-block:: python

>>> # doctest: +REQUIRES(env:DISTRIBUTED)

>>> import paddle

>>> import paddle.nn as nn

>>> import paddle.optimizer as opt

>>> import paddle.distributed as dist

>>> class LinearNet(nn.Layer):

... def __init__(self):

... super().__init__()

... self._linear1 = nn.Linear(10, 10)

... self._linear2 = nn.Linear(10, 1)

... def forward(self, x):

... return self._linear2(self._linear1(x))

>>> def train(print_result=False):

... # 1. initialize parallel environment

... group = dist.init_parallel_env()

... process_group = group.process_group if group else None

... # 2. create data parallel layer & optimizer

... layer = LinearNet()

... dp_layer = paddle.DataParallel(layer, group = process_group) # type: ignore[arg-type]

... loss_fn = nn.MSELoss()

... adam = opt.Adam(

... learning_rate=0.001, parameters=dp_layer.parameters())

... # 3. run layer

... inputs = paddle.randn([10, 10], 'float32')

... outputs = dp_layer(inputs)

... labels = paddle.randn([10, 1], 'float32')

... loss = loss_fn(outputs, labels)

... if print_result is True:

... print("loss:", loss.numpy())

... loss.backward()

... adam.step()

... adam.clear_grad()

>>> # Usage 1: only pass function.

>>> # If your training method no need any argument, and

>>> # use all visible devices for parallel training.

>>> if __name__ == '__main__':

... dist.spawn(train)

>>> # Usage 2: pass function and arguments.

>>> # If your training method need some arguments, and

>>> # use all visible devices for parallel training.

>>> if __name__ == '__main__':

... dist.spawn(train, args=(True,))

>>> # Usage 3: pass function, arguments and nprocs.

>>> # If your training method need some arguments, and

>>> # only use part of visible devices for parallel training.

>>> # If your machine hold 8 cards {0,1,2,3,4,5,6,7},

>>> # this case will use cards {0,1}; If you set

>>> # CUDA_VISIBLE_DEVICES=4,5,6,7, this case will use

>>> # cards {4,5}

>>> if __name__ == '__main__':

... dist.spawn(train, args=(True,), nprocs=2)

>>> # Usage 4: pass function, arguments, nprocs and gpus.

>>> # If your training method need some arguments, and

>>> # only use part of visible devices for parallel training,

>>> # but you can't set your machine's environment variable

>>> # CUDA_VISIBLE_DEVICES, such as it is None or all cards

>>> # {0,1,2,3,4,5,6,7}, you can pass `gpus` to

>>> # select the GPU cards you want to use. For example,

>>> # this case will use cards {4,5} if your machine hold 8 cards.

>>> if __name__ == '__main__':

... dist.spawn(train, args=(True,), nprocs=2, gpus='4,5')

"""

# Give an error hint when the users enter a configuration option

# that does not exist

_options_valid_check(options)

# get default nprocs

if nprocs == -1:

nprocs = _get_default_nprocs()

# NOTE(chenweihang): [ why need get cluster info before run? ]

# when using `paddle.distributed.spawn` start parallel training,

# we should get cluster info before starting subprocess, and pass

# correct info to each subprocess

procs_env_list = _get_subprocess_env_list(nprocs, options)

# start processes

# NOTE(chenweihang): [ why default start method is spawn? ]

# The CUDA runtime does not support the fork start method,

# either the spawn or forkserver start method are required

# to use CUDA in subprocesses.

start_method = options.get('start_method', None)

if start_method is None:

start_method = 'spawn'

mp = multiprocessing.get_context(start_method)

error_queues = []

return_queues = []

processes = []

for i in range(nprocs):

error_queue = mp.SimpleQueue()

return_queue = mp.SimpleQueue()

process = mp.Process(

target=_func_wrapper,

args=(

func,

args,

error_queue,

return_queue,

procs_env_list[i],

options['backend'],

),

)

process.daemon = daemon

process.start()

error_queues.append(error_queue)

return_queues.append(return_queue)

processes.append(process)

context = MultiprocessContext(processes, error_queues, return_queues)

if not join:

return context

# loop until all process end

while not context.join():

pass

# finally return context

return context