chore: import upstream snapshot with attribution

python/paddle/distributed/launch/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2022 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.
+
+__all__ = []

python/paddle/distributed/launch/__main__.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2022 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 .main import launch
+
+launch()

python/paddle/distributed/launch/context/__init__.py

@@ -0,0 +1,107 @@
+# Copyright (c) 2022 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.
+
+import logging
+
+from paddle.distributed.launch import plugins
+
+from .args_envs import env_args_mapping, fetch_envs, parse_args
+from .node import Node
+from .status import Status
+
+
+class Context:
+    def __init__(self, enable_plugin=True):
+        self.args, self.unknown_args = parse_args()
+        self.envs = fetch_envs()
+
+        self.set_env_in_args()
+
+        self.node = Node()
+        self.status = Status()
+
+        self.logger = self.get_logger()
+
+        # design for event queue, later
+        self.events = []
+
+        if enable_plugin:
+            self._enable_plugin()
+        self.max_time_per_task = -1
+        self.run_best = False
+
+    def print(self):
+        self.logger.info("-----------  Configuration  ----------------------")
+        for arg, value in sorted(vars(self.args).items()):
+            self.logger.info(f"{arg}: {value}")
+        self.logger.info("--------------------------------------------------")
+
+    def is_legacy_mode(self):
+        if self.args.legacy:
+            return True
+
+        if self.args.master:
+            return False
+
+        if len(self.unknown_args) > 0:
+            self.logger.warning(
+                f"Compatible mode enable with args {self.unknown_args}"
+            )
+            return True
+
+        return False
+
+    def is_auto_tuner_mode(self):
+        if self.args.auto_tuner_json:
+            return True
+        return False
+
+    def get_envs(self):
+        return self.envs.copy()
+
+    def set_envs(self, env={}):
+        env = {k: v for k, v in env.items() if isinstance(v, str)}
+        self.envs.update(env)
+
+    def _enable_plugin(self):
+        for pl in plugins.enabled_plugins:
+            pl(self)
+
+    def get_logger(self, level=logging.INFO):
+        logger = logging.getLogger("LAUNCH")
+        # forbid the child logger pass on to its parent
+        logger.propagate = False
+        logger.setLevel(self.args.log_level.upper() or level)
+        formatter = logging.Formatter(
+            fmt='%(name)s %(levelname)s %(asctime)s %(message)s'
+        )
+        ch = logging.StreamHandler()
+        ch.setFormatter(formatter)
+        logger.addHandler(ch)
+        return logger
+
+    def continuous_log(self) -> bool:
+        if self.args.log_level.upper() in ['DEBUG', 'ERROR']:
+            return True
+        else:
+            return False
+
+    def set_env_in_args(self):
+        for k, v in env_args_mapping.items():
+            attr, attr_type = v
+            if k in self.envs:
+                print(
+                    f"LAUNCH WARNING args {attr} will be overridden by env: {k} value: {self.envs[k]}"
+                )
+                setattr(self.args, attr, attr_type(self.envs[k]))

python/paddle/distributed/launch/context/args_envs.py

@@ -0,0 +1,246 @@
+# Copyright (c) 2022 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.
+
+import os
+import warnings
+from argparse import REMAINDER, ArgumentParser
+
+from paddle.utils import strtobool
+
+env_args_mapping = {
+    'POD_IP': ('host', str),
+    'PADDLE_MASTER': ('master', str),
+    'PADDLE_DEVICES': ('devices', str),
+    'PADDLE_NNODES': ('nnodes', str),
+    'PADDLE_RUN_MODE': ('run_mode', str),
+    'PADDLE_LOG_LEVEL': ('log_level', str),
+    'PADDLE_LOG_OVERWRITE': ('log_overwrite', strtobool),
+    'PADDLE_SORT_IP': ('sort_ip', strtobool),
+    'PADDLE_NPROC_PER_NODE': ('nproc_per_node', int),
+    'PADDLE_JOB_ID': ('job_id', str),
+    'PADDLE_RANK': ('rank', int),
+    'PADDLE_LOG_DIR': ('log_dir', str),
+    'PADDLE_MAX_RESTART': ('max_restart', int),
+    'PADDLE_ELASTIC_LEVEL': ('elastic_level', int),
+    'PADDLE_ELASTIC_TIMEOUT': ('elastic_timeout', int),
+    'PADDLE_SERVER_NUM': ('server_num', int),
+    'PADDLE_TRAINER_NUM': ('trainer_num', int),
+    'PADDLE_SERVERS_ENDPOINTS': ('servers', str),
+    'PADDLE_TRAINERS_ENDPOINTS': ('trainers', str),
+    'PADDLE_GLOO_PORT': ('gloo_port', int),
+    'PADDLE_WITH_GLOO': ('with_gloo', str),
+    'PADDLE_START_PORT': ('start_port', int),
+    'PADDLE_IPS': ('ips', str),
+    "PADDLE_AUTO_PARALLEL_CONFIG": ('auto_parallel_config', str),
+    'PADDLE_AUTO_CLUSTER': ('auto_cluster_config', strtobool),
+}
+
+
+def fetch_envs():
+    for proxy_key in ("http_proxy", "https_proxy"):
+        if os.environ.get(proxy_key) is not None:
+            os.environ[f"{proxy_key}_original"] = os.environ.pop(proxy_key)
+            warnings.warn(
+                f"Unset '{proxy_key}' to ensure stable NCCL communication in distributed training "
+                f"(backed up as '{proxy_key}_original').",
+                category=UserWarning,
+            )
+
+    return os.environ.copy()
+
+
+def parse_args():
+    parser = ArgumentParser()
+
+    base_group = parser.add_argument_group("Base Parameters")
+
+    base_group.add_argument(
+        "--master",
+        type=str,
+        default=None,
+        help="the master/rendezvous server, ip:port",
+    )
+
+    base_group.add_argument(
+        "--legacy", type=strtobool, default=False, help="use legacy launch"
+    )
+
+    base_group.add_argument(
+        "--rank", type=int, default=-1, help="the node rank"
+    )
+
+    base_group.add_argument(
+        "--log_level", type=str, default="INFO", help="log level. Default INFO"
+    )
+
+    base_group.add_argument(
+        "--log_overwrite",
+        type=strtobool,
+        default=False,
+        help="overwrite exits logfiles. Default False",
+    )
+
+    base_group.add_argument(
+        "--sort_ip",
+        type=strtobool,
+        default=False,
+        help="rank node by ip. Default False",
+    )
+
+    base_group.add_argument(
+        "--enable_gpu_log",
+        type=strtobool,
+        default=True,
+        help="enable capture gpu log while running. Default True",
+    )
+
+    base_group.add_argument(
+        "--nnodes",
+        type=str,
+        default="1",
+        help="the number of nodes, i.e. pod/node number",
+    )
+
+    base_group.add_argument(
+        "--nproc_per_node",
+        type=int,
+        default=None,
+        help="the number of processes in a pod",
+    )
+
+    base_group.add_argument(
+        "--log_dir",
+        type=str,
+        default="log",
+        help="the path for each process's log. Default ./log",
+    )
+    base_group.add_argument(
+        "--run_mode",
+        type=str,
+        default=None,
+        help="run mode of the job, collective/ps/ps-heter",
+    )
+
+    base_group.add_argument(
+        "--job_id",
+        type=str,
+        default="default",
+        help="unique id of the job. Default default",
+    )
+
+    base_group.add_argument(
+        "--devices",
+        "--gpus",
+        "--npus",
+        "--xpus",
+        type=str,
+        default=None,
+        help="accelerate devices. as --gpus,npus,xpus",
+    )
+
+    base_group.add_argument("--host", type=str, default=None, help="host ip")
+
+    base_group.add_argument(
+        "--ips",
+        type=str,
+        default=None,
+        help="nodes ips, e.g. 10.10.1.1,10.10.1.2",
+    )
+
+    base_group.add_argument(
+        "--start_port", type=int, default=6070, help="fix port start with"
+    )
+
+    base_group.add_argument(
+        "--auto_parallel_config",
+        type=str,
+        default=None,
+        help="auto parallel config file absolute path, the file should be json format",
+    )
+
+    base_group.add_argument(
+        "--auto_cluster_config",
+        type=strtobool,
+        default=0,
+        help="auto parallel auto cluster config switch",
+    )
+
+    base_group.add_argument(
+        "training_script",
+        type=str,
+        help="the full path of py script,"
+        "followed by arguments for the "
+        "training script",
+    )
+
+    base_group.add_argument(
+        "--auto_tuner_json",
+        type=str,
+        default=None,
+        help="auto tuner json file path",
+    )
+
+    base_group.add_argument('training_script_args', nargs=REMAINDER)
+
+    ps_group = parser.add_argument_group("Parameter-Server Parameters")
+    # for parameter server
+    ps_group.add_argument(
+        "--servers", type=str, default='', help="servers endpoints full list"
+    )
+    ps_group.add_argument(
+        "--trainers", type=str, default='', help="trainers endpoints full list"
+    )
+
+    ps_group.add_argument(
+        "--trainer_num", type=int, default=None, help="number of trainers"
+    )
+    ps_group.add_argument(
+        "--server_num", type=int, default=None, help="number of servers"
+    )
+    ps_group.add_argument(
+        "--gloo_port", type=int, default=6767, help="gloo http port"
+    )
+    ps_group.add_argument(
+        "--with_gloo", type=str, default="1", help="use gloo or not"
+    )
+
+    # parameter elastic mode
+    elastic_group = parser.add_argument_group("Elastic Parameters")
+    elastic_group.add_argument(
+        "--max_restart",
+        type=int,
+        default=3,
+        help="the times can restart. Default 3",
+    )
+
+    elastic_group.add_argument(
+        "--elastic_level",
+        type=int,
+        default=-1,
+        help="elastic level: -1 disable, 0 failed exit, peers hold, 1 internal restart",
+    )
+
+    elastic_group.add_argument(
+        "--elastic_timeout",
+        type=int,
+        default=30,
+        help="seconds to wait before elastic job begin to train",
+    )
+
+    args = parser.parse_known_args()
+    env_rank = int(os.getenv('PADDLE_TRAINER_ID', -1))
+    if env_rank >= 0:
+        assert hasattr(args[0], "rank")
+        args[0].rank = env_rank
+    return args

python/paddle/distributed/launch/context/device.py

@@ -0,0 +1,181 @@
+# Copyright (c) 2022 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.
+
+import os
+
+# (TODO: GhostScreaming) It will be removed later.
+from paddle.base import core
+from paddle.base.core import get_all_custom_device_type
+from paddle.device import get_available_custom_device
+
+
+class DeviceType:
+    CPU = 'cpu'
+    GPU = 'gpu'
+    XPU = 'xpu'
+    IPU = 'ipu'
+    CUSTOM_DEVICE = 'custom_device'
+
+
+class Device:
+    def __init__(self, dtype=None, memory="", labels=""):
+        self._dtype = dtype
+        self._memory = memory
+        self._labels = labels
+
+    def __str__(self):
+        return ",".join(self._labels)
+
+    @property
+    def dtype(self):
+        return self._dtype
+
+    @property
+    def count(self):
+        return len(self._labels) or 1
+
+    @property
+    def memory(self):
+        return self._memory
+
+    @property
+    def labels(self):
+        return self._labels
+
+    @labels.setter
+    def labels(self, lbs):
+        if isinstance(lbs, str):
+            self._labels = lbs.split(',')
+        elif isinstance(lbs, list):
+            self._labels = lbs
+        else:
+            self._labels = []
+
+    def get_selected_device_key(self):
+        if self._dtype == DeviceType.CPU:
+            return 'FLAGS_selected_cpus'
+        if self._dtype == DeviceType.GPU:
+            return 'FLAGS_selected_gpus'
+        if self._dtype == DeviceType.XPU:
+            return 'FLAGS_selected_xpus'
+        if self._dtype == DeviceType.IPU:
+            return 'FLAGS_selected_ipus'
+        if self._dtype == DeviceType.CUSTOM_DEVICE:
+            custom_device_types = get_all_custom_device_type()
+            device_type = (
+                str(custom_device_types[0]) if custom_device_types else ""
+            )
+            return f'FLAGS_selected_{device_type}s'
+        return 'FLAGS_selected_devices'
+
+    def get_selected_devices(self, devices=''):
+        '''
+        return the device label/id relative to the visible devices
+        '''
+        if not devices:
+            return [str(x) for x in range(0, len(self._labels))]
+        else:
+            devs = [x.strip() for x in devices.split(',')]
+            return [str(self._labels.index(d)) for d in devs]
+
+    def get_custom_device_envs(self):
+        custom_device_types = get_all_custom_device_type()
+        device_type = str(custom_device_types[0]) if custom_device_types else ""
+        return {
+            'PADDLE_DISTRI_BACKEND': 'xccl',
+            'PADDLE_XCCL_BACKEND': device_type,
+        }
+
+    @classmethod
+    def parse_device(self):
+        dev = Device()
+        visible_devices = None
+        custom_device_types = get_all_custom_device_type()
+        if custom_device_types:
+            dev._dtype = DeviceType.CUSTOM_DEVICE
+            device_type = str(custom_device_types[0])
+            visible_devices_str = f'{device_type.upper()}_VISIBLE_DEVICES'
+            if visible_devices_str in os.environ:
+                visible_devices = os.getenv(visible_devices_str)
+        elif 'XPULINK_VISIBLE_DEVICES' in os.environ:
+            dev._dtype = DeviceType.XPU
+            visible_devices = os.getenv("XPULINK_VISIBLE_DEVICES")
+        elif 'XPU_VISIBLE_DEVICES' in os.environ:
+            dev._dtype = DeviceType.XPU
+            visible_devices = os.getenv("XPU_VISIBLE_DEVICES")
+        elif 'CUDA_VISIBLE_DEVICES' in os.environ:
+            if core.is_compiled_with_xpu():
+                dev._dtype = DeviceType.XPU
+            else:
+                dev._dtype = DeviceType.GPU
+            visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
+
+        if visible_devices is not None and visible_devices != 'all':
+            dev._labels = visible_devices.split(',')
+        else:
+            return self.detect_device()
+
+        return dev
+
+    @classmethod
+    def detect_device(self):
+        def get_custom_devices_count(device_type):
+            all_custom_devices = get_available_custom_device()
+            all_custom_devices = [
+                device.split(':')[0] for device in all_custom_devices
+            ]
+            custom_devices_count = all_custom_devices.count(device_type)
+            return custom_devices_count
+
+        dev = Device()
+        num = 0
+        visible_devices = None
+        custom_device_types = get_all_custom_device_type()
+        if custom_device_types:
+            custom_device_type = str(custom_device_types[0])
+            dev._dtype = DeviceType.CUSTOM_DEVICE
+            num = get_custom_devices_count(custom_device_type)
+            visible_devices_str = (
+                f'{custom_device_type.upper()}_VISIBLE_DEVICES'
+            )
+            if visible_devices_str in os.environ:
+                visible_devices = os.getenv(visible_devices_str)
+        elif core.is_compiled_with_cuda():
+            dev._dtype = DeviceType.GPU
+            num = core.get_cuda_device_count()
+            visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
+        elif core.is_compiled_with_xpu():
+            dev._dtype = DeviceType.XPU
+            num = core.get_xpu_device_count()
+            visible_devices = os.getenv("XPU_VISIBLE_DEVICES")
+        elif core.is_compiled_with_ipu():
+            dev._dtype = DeviceType.IPU
+            num = core.get_ipu_device_count()
+            # For IPUs, 'labels' is a list which contains the available numbers of IPU devices.
+            dev._labels = [str(x) for x in range(0, num + 1)]
+            return dev
+
+        if num == 0:
+            dev._dtype = DeviceType.CPU
+        elif visible_devices is None or visible_devices == "all":
+            dev._labels = [str(x) for x in range(0, num)]
+        else:
+            dev._labels = visible_devices.split(',')
+
+        return dev
+
+
+if __name__ == '__main__':
+    d = Device.parse_device()
+    print(d.get_selected_devices())

python/paddle/distributed/launch/context/event.py

@@ -0,0 +1,20 @@
+# Copyright (c) 2022 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.
+
+
+class Event:
+    def __init__(self, kind="status", message="", fatal=False):
+        self.kind = kind
+        self.message = message
+        self.fatal = fatal

python/paddle/distributed/launch/context/node.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2022 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.
+
+import os
+import random
+import socket
+import struct
+from contextlib import closing
+
+from .device import Device
+
+
+class Node:
+    def __init__(self):
+        # self.device = Device.detect_device()
+        self.device = Device.parse_device()
+        self.ip = self.get_host_ip()
+        self.free_ports = []
+        self._allocated_ports = []
+
+        port_range = os.getenv('PORT_RANGE', '35100:64000')
+        port_range = port_range.split(':')
+        self._port_start = int(port_range[0])
+        self._port_end = int(port_range[1])
+        self._port_cur = random.randint(self._port_start, self._port_end)
+
+    def get_host_ip(self):
+        try:
+            self.hostname = socket.gethostname()
+            self.ip = socket.gethostbyname(socket.getfqdn(self.hostname))
+            return self.ip
+        except:
+            return '127.0.0.1'
+
+    def get_free_ports(self, n=1, rank=0):
+        if os.environ.get('FLAGS_FIXED_PORT') is None:
+            free_ports = [self.get_free_port() for i in range(n)]
+            self.free_ports += free_ports
+        else:
+            start_port = int(os.environ.get('FLAGS_FIXED_PORT'))
+            free_ports = list(
+                range(start_port + rank, start_port + rank + n, 1)
+            )
+            self.free_ports += free_ports
+        return free_ports
+
+    def get_ports_occupied(self):
+        return self.free_ports
+
+    def _get_free_port(self, port=0):
+        with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+            s.setsockopt(
+                socket.SOL_SOCKET, socket.SO_LINGER, struct.pack('ii', 1, 0)
+            )
+            try:
+                s.bind(('', port))
+                return s.getsockname()[1]
+            except:
+                return -1
+
+    def _update_port_cur(self):
+        self._port_cur += 1
+        if self._port_cur > self._port_end:
+            self._port_cur = self._port_start
+
+    def get_free_port(self):
+        for _ in range(100):
+            ret = self._get_free_port(self._port_cur)
+            if ret > 0:
+                self._update_port_cur()
+                return ret
+            else:
+                self._update_port_cur()
+
+        return self._port_cur
+
+    @classmethod
+    def is_server_ready(self, ip, port):
+        with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
+            # sock.settimeout(0.01)
+            # sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            if hasattr(socket, 'SO_REUSEPORT'):
+                sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
+            result = sock.connect_ex((ip, int(port)))
+            if result == 0:
+                return True
+            else:
+                return False

python/paddle/distributed/launch/context/resource.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2022 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.
+
+
+class Resource:
+    def __init__(self):
+        self.devices = []

python/paddle/distributed/launch/context/status.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2022 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.
+
+
+class Status:
+    UNINIT = "uninit"
+    READY = "ready"
+    RUNNING = "running"
+    FAILED = "failed"
+    TERMINATING = "terminating"
+    RESTARTING = "restarting"
+    UNKNOWN = "unknown"
+    COMPLETED = "completed"
+    DONE = "done"  # should exit whatever status
+
+    def __init__(self):
+        self._current_status = None
+
+    def current(self):
+        return self._current_status
+
+    def is_running(self):
+        return self._current_status == self.RUNNING
+
+    def is_restarting(self):
+        return self._current_status == self.RESTARTING
+
+    def is_done(self):
+        if self._current_status in [self.DONE, self.COMPLETED, self.FAILED]:
+            return True
+        else:
+            return False
+
+    def run(self):
+        self._current_status = self.RUNNING
+
+    def fail(self):
+        self._current_status = self.FAILED
+
+    def complete(self):
+        self._current_status = self.COMPLETED
+
+    def restart(self):
+        self._current_status = self.RESTARTING
+
+    def done(self):
+        self._current_status = self.DONE

python/paddle/distributed/launch/controllers/__init__.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2022 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.
+
+__all__ = []
+
+from .collective import CollectiveController, CollectiveElasticController
+from .ipu_controller import IPUController
+from .ps import PSController
+from .rpc import RpcController
+
+# the order is extremely important
+_controllers = [
+    IPUController,
+    CollectiveElasticController,
+    PSController,
+    RpcController,
+    CollectiveController,
+]
+
+
+def init(ctx):
+    for c in _controllers:
+        if c.enable(ctx):
+            ctx.print()
+            return c(ctx)

python/paddle/distributed/launch/controllers/collective.py

@@ -0,0 +1,323 @@
+# Copyright (c) 2022 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.
+
+import json
+import os
+
+from ..context.device import DeviceType
+from .controller import Controller, ControllerMode
+
+
+class CollectiveController(Controller):
+    def __init__(self, ctx):
+        self._tuner_run_mode = None  # 'tuner_only', 'run_only', 'tuner_and_run'
+        super().__init__(ctx)
+
+    @classmethod
+    def enable(cls, ctx):
+        # collective is the default mode
+        if ctx:
+            ctx.logger.debug(f"{cls.__name__} enabled")
+            ctx.args.run_mode = ControllerMode.COLLECTIVE
+            return True
+        else:
+            return False
+
+    def build_pod(self):
+        skip_run = self._build_pod_with_tuner()
+        if skip_run:
+            return
+        if (
+            self.ctx.args.master is None
+            and self.ctx.args.start_port
+            and self.ctx.args.ips
+        ):
+            return self._build_pod_with_args()
+        else:
+            if self.ctx.args.auto_parallel_config is None:
+                skip_run = True
+            # only when skip_run is False, should not reset pod
+            return self._build_pod_with_master(skip_run)
+
+    def _build_pod_with_tuner(self):
+        auto_parallel_config = self.ctx.args.auto_parallel_config
+        if auto_parallel_config is not None:
+            if not os.path.exists(auto_parallel_config):
+                self.ctx.logger.warning("auto_parallel_conf not exists!")
+            if not auto_parallel_config.endswith(".json"):
+                self.ctx.logger.warning(
+                    "auto_parallel_config should be a json format file!"
+                )
+
+            with open(auto_parallel_config, 'r') as robj:
+                auto_parallel_data = json.loads(robj.read())
+                self._tuner_run_mode = auto_parallel_data.get(
+                    "tuner_run_mode", 'tuner_and_run'
+                )
+
+            self.ctx.logger.info(f"tuner_run_mode is: {self._tuner_run_mode}")
+            endpoint = f"127.0.0.1:{self.ctx.node.get_free_port()}"
+            pod_replicas = self.pod_replicas()
+            if self._tuner_run_mode in ['tuner_only', 'tuner_and_run']:
+                e = {
+                    "PADDLE_AUTO_PARALLEL_CONFIG": self.ctx.args.auto_parallel_config,
+                    "PADDLE_TRAINERS_NUM": "1",
+                    "PADDLE_TRAINER_ENDPOINTS": endpoint,
+                    "PADDLE_TRAINER_ID": "0",
+                    "PADDLE_CURRENT_ENDPOINT": endpoint,
+                    "FLAGS_selected_gpus": "0",
+                    "PADDLE_AUTO_PARALLEL_STAGE": "tuner",
+                    "PADDLE_GLOBAL_SIZE": f"{pod_replicas * int(self.ctx.args.nnodes)}",
+                    "PADDLE_LOCAL_SIZE": f"{pod_replicas}",
+                }
+                log_file = "tuner.log"
+                self.add_container(envs=e, log_file=log_file, is_init=True)
+
+                if self._tuner_run_mode == 'tuner_only':
+                    return True
+        return False
+
+    def _build_pod_with_args(self):
+        self.pod.replicas = self.pod_replicas()
+
+        start_port = int(self.ctx.args.start_port)
+        ips = self.ctx.args.ips.split(',')
+
+        job_endpoints = [
+            f"{h}:{p + start_port}"
+            for h in ips
+            for p in range(self.pod.replicas)
+        ]
+
+        self.ctx.logger.debug(f"job endpoints: {job_endpoints}")
+
+        self.ctx.logger.warning(
+            f"master is set by args, it will be overwritten by {job_endpoints[0]}."
+        )
+        # this is necessary for tcp store to work when endpoints cannot be passed to sub processes.
+        self.ctx.args.master = job_endpoints[0]
+
+        rank_offset = (
+            ips.index(self.ctx.node.ip) * self.pod.replicas
+            if self.ctx.node.ip in ips
+            else 0
+        )
+
+        self.save_pod_log(job_endpoints)
+
+        selected_dev_key = self.ctx.node.device.get_selected_device_key()
+        selected_dev_list = self.ctx.node.device.get_selected_devices(
+            self.ctx.args.devices
+        )
+
+        for i in range(self.pod.replicas):
+            e = {
+                "PADDLE_MASTER": self.ctx.args.master,
+                "PADDLE_GLOBAL_SIZE": f"{len(job_endpoints)}",
+                "PADDLE_LOCAL_SIZE": f"{self.pod.replicas}",
+                "PADDLE_GLOBAL_RANK": f"{i + rank_offset}",
+                "PADDLE_LOCAL_RANK": f"{i}",
+                "PADDLE_NNODES": f"{len(ips)}",
+                # compatible env
+                "PADDLE_CURRENT_ENDPOINT": job_endpoints[i + rank_offset],
+                "PADDLE_TRAINER_ID": f"{i + rank_offset}",
+                "PADDLE_TRAINERS_NUM": f"{len(job_endpoints)}",
+                "PADDLE_RANK_IN_NODE": str(i),
+                "PADDLE_AUTO_CLUSTER": str(self.ctx.args.auto_cluster_config),
+            }
+            e.update({"PADDLE_TRAINER_ENDPOINTS": ",".join(job_endpoints)})
+
+            if self._tuner_run_mode is not None:
+                e.update(
+                    {
+                        "PADDLE_AUTO_PARALLEL_CONFIG": self.ctx.args.auto_parallel_config,
+                        "PADDLE_AUTO_PARALLEL_STAGE": "run",
+                    }
+                )
+            if len(selected_dev_list) > 0:
+                if self.ctx.node.device.dtype == DeviceType.CUSTOM_DEVICE:
+                    e.update(self.ctx.node.device.get_custom_device_envs())
+                if self.pod.replicas == 1:
+                    e.update({selected_dev_key: ",".join(selected_dev_list)})
+                else:
+                    e.update({selected_dev_key: selected_dev_list[i]})
+            else:
+                e.update({'PADDLE_DISTRI_BACKEND': 'gloo'})
+
+            log_file = f"workerlog.{i}"
+            self.add_container(envs=e, log_file=log_file)
+
+        return True
+
+    def _build_pod_with_master(self, reset_pod=True):
+        self.pod.replicas = self.pod_replicas()
+
+        # rank will be reset when restart
+        self.pod.rank = int(self.ctx.args.rank)
+
+        port = self.ctx.node.get_free_port()
+
+        # compatible
+        endpoints = [
+            f"{self.ctx.node.ip}:{p}"
+            for p in self.ctx.node.get_free_ports(
+                self.pod.replicas, self.pod.rank
+            )
+        ]
+
+        data = json.dumps(
+            {
+                'name': self.pod.name,
+                'rank': self.pod.rank,
+                'replicas': self.pod.replicas,
+                'dtype': self.ctx.node.device.dtype,
+                'candidate': f'{self.ctx.node.ip}:{port}',
+                'endpoints': ",".join(endpoints),
+            }
+        )
+
+        peer_list, rank = self.master.sync_peers(
+            f'/{self.job.id}/info',
+            self.pod.name,
+            data,
+            self.job.replicas,
+            self.pod.rank,
+        )
+        self.pod.rank = rank
+
+        if len(peer_list) < 1:
+            return False
+
+        peer_list = [json.loads(i) for i in peer_list]
+
+        self.ctx.logger.debug(f"sync peers done {peer_list}")
+        self.save_pod_log(peer_list)
+
+        global_size = sum([i['replicas'] for i in peer_list])
+        rank_offset = sum([i['replicas'] for i in peer_list[:rank]])
+        '''
+        The new designed collective need nothing but a master endpoint
+        '''
+        collective_master = peer_list[0]['candidate']
+
+        # get collective master ip
+        collective_master_ip = collective_master.split(':')[0].strip()
+        os.environ["COLLECTIVE_MASTER_IP"] = collective_master_ip
+
+        job_endpoints = [i['endpoints'] for i in peer_list]
+
+        if reset_pod:
+            self.pod.reset()
+        selected_dev_key = self.ctx.node.device.get_selected_device_key()
+        selected_dev_list = self.ctx.node.device.get_selected_devices(
+            self.ctx.args.devices
+        )
+        for i in range(self.pod.replicas):
+            e = {
+                "PADDLE_MASTER": collective_master,
+                "PADDLE_GLOBAL_SIZE": f"{global_size}",
+                "PADDLE_LOCAL_SIZE": f"{self.pod.replicas}",
+                "PADDLE_GLOBAL_RANK": f"{i + rank_offset}",
+                "PADDLE_LOCAL_RANK": f"{i}",
+                "PADDLE_NNODES": f"{self.job.replicas}",
+                # compatible env
+                "PADDLE_CURRENT_ENDPOINT": endpoints[i],
+                "PADDLE_TRAINER_ID": f"{i + rank_offset}",
+                "PADDLE_TRAINERS_NUM": f"{global_size}",
+                "PADDLE_RANK_IN_NODE": str(i),
+                "PADDLE_AUTO_CLUSTER": str(self.ctx.args.auto_cluster_config),
+            }
+            e.update({"PADDLE_TRAINER_ENDPOINTS": ",".join(job_endpoints)})
+
+            if self._tuner_run_mode is not None:
+                e.update(
+                    {
+                        "PADDLE_AUTO_PARALLEL_CONFIG": self.ctx.args.auto_parallel_config,
+                        "PADDLE_AUTO_PARALLEL_STAGE": "run",
+                    }
+                )
+            if len(selected_dev_list) > 0:
+                if self.ctx.node.device.dtype == DeviceType.CUSTOM_DEVICE:
+                    e.update(self.ctx.node.device.get_custom_device_envs())
+                if self.pod.replicas == 1:
+                    e.update({selected_dev_key: ",".join(selected_dev_list)})
+                else:
+                    e.update({selected_dev_key: selected_dev_list[i]})
+            else:
+                e.update({'PADDLE_DISTRI_BACKEND': 'gloo'})
+
+            # log_file = "{}.{}.{}.log".format(self.job.id, self.pod.name, i)
+            log_file = f"workerlog.{i}"
+            self.add_container(envs=e, log_file=log_file)
+
+        return True
+
+
+class CollectiveElasticController(CollectiveController):
+    @classmethod
+    def enable(cls, ctx):
+        if ctx.args.master and ctx.args.master.startswith("etcd://"):
+            ctx.logger.debug(f"{cls.__name__} enabled")
+            ctx.args.run_mode = ControllerMode.COLLECTIVE
+            return True
+        else:
+            return False
+
+    def register(self):
+        if self.job.id == 'default':
+            self.ctx.logger.warning(
+                'Using default job name may cause conflict, add --job_id in args'
+            )
+
+        self.master.register_heartbeat(self.job.id, self.pod.name)
+
+    def run(self):
+        timeout = int(self.ctx.args.elastic_timeout)
+        timeout = timeout if self.job.elastic else timeout * 10
+        self.register()
+
+        while self.pod.restart <= self.ctx.args.max_restart:
+            self.build_job()
+
+            self.ctx.logger.info("Waiting peer ready...")
+
+            ok, replicas = self.master.wait_peer_ready(
+                self.job.replicas_min, self.job.replicas_max, timeout
+            )
+            if ok:
+                self.job.replicas = replicas
+            else:
+                self.ctx.logger.warning(f"peer not ready {self.job}")
+                if self.ctx.is_auto_tuner_mode():
+                    self.ctx.logger.info(
+                        "Failed to start peer, auto tuner exit."
+                    )
+                    import sys
+
+                    sys.exit(-1)
+                break
+
+            self.ctx.logger.debug(f"Run {self.job}")
+
+            if not self.build_pod():
+                continue
+
+            self.master.set_status(self.ctx.status.RUNNING)
+
+            self.deploy_pod()
+
+            if self.watch():
+                break
+
+        self.ctx.logger.debug(f"Job done {self.job}")

python/paddle/distributed/launch/controllers/controller.py

@@ -0,0 +1,341 @@
+# Copyright (c) 2022 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.
+
+import copy
+import os
+import signal
+import sys
+
+from paddle.distributed.launch.job.container import Container
+from paddle.distributed.launch.job.job import Job
+from paddle.distributed.launch.job.pod import Pod
+
+from .master import Master
+from .watcher import Watcher
+
+
+class ControllerMode:
+    COLLECTIVE = "collective"
+    PS = "ps"
+    IPU = "ipu"
+    RPC = "rpc"
+
+
+class ControllerBase:
+    def __init__(self, ctx):
+        signal.signal(signal.SIGTERM, self.signal_handler)
+        signal.signal(signal.SIGABRT, self.signal_handler)
+        signal.signal(signal.SIGINT, self.signal_handler)
+        if ctx.is_auto_tuner_mode():
+            if not ctx.run_best:
+                # set per task timeout
+                signal.signal(signal.SIGALRM, self.not_exit_signal_handler)
+                signal.alarm(ctx.max_time_per_task)
+            else:
+                signal.alarm(0)
+
+        self.ctx = ctx
+        self.master = Master.factory(self.ctx)
+
+        self.watcher = Watcher(self.ctx)
+
+        self.job = Job(
+            nnodes=self.ctx.args.nnodes,
+            mode=self.ctx.args.run_mode,
+            jid=self.ctx.args.job_id,
+        )
+        self.pod = Pod()
+
+        self.ctx.set_envs({"POD_NAME": self.pod.name})
+
+        self.join_server = None
+
+    def deploy_pod(self):
+        assert len(self.pod.containers) + len(self.pod.init_containers) > 0, (
+            "No container in the pod"
+        )
+
+        self.ctx.logger.info(f"Run {self.pod}")
+        if len(self.pod.init_containers) > 0:
+            self.ctx.logger.debug(self.pod.init_containers[0])
+        if len(self.pod.containers) > 0:
+            self.ctx.logger.debug(self.pod.containers[0])
+
+        self.save_pod_env()
+        self.ctx.status.run()
+        self.pod.deploy()
+
+    def run(self):
+        self.build_job()
+        self.build_pod()
+
+        self.deploy_pod()
+
+        self.watch()
+
+    def watch(self) -> bool:
+        '''
+        watch self and peer status, return true to exit
+        '''
+        # TODO(kuizhiqing) unify ctx.status and master status
+
+        self.ctx.logger.info(f"Watching {self.pod}")
+
+        while not self.ctx.status.is_done():
+            status = self.pod.watch(timeout=2)
+
+            # if self.ctx.continuous_log():
+            # default to print log
+            self.pod.logs()
+
+            # completed
+            if status == self.ctx.status.COMPLETED:
+                self.ctx.status.complete()
+
+                self.master.set_status(status)
+
+                while self.pod.logs():
+                    pass
+
+                self.ctx.logger.info(f"Pod {status}")
+                return True
+
+            # self failure
+            elif status == self.ctx.status.FAILED:
+                self.ctx.status.fail()
+
+                self.master.set_status(status)
+                self.master.restart_peer()
+
+                fc = self.pod.failed_container()
+                self.ctx.logger.info(f"Pod {status}")
+                self.ctx.logger.error(f"Container failed !!!\n{fc[0]}")
+                self.ctx.logger.info(
+                    "------------------------- ERROR LOG DETAIL -------------------------"
+                )
+                fc[0].tail()
+
+                if self.ctx.args.elastic_level <= 0:
+                    self.pod.stop(timeout=3)
+                    return True
+                else:
+                    self.pod.stop(timeout=30)
+                    return False
+
+            # peer failure
+            if (
+                self.ctx.status.is_restarting()
+                and self.master.get_status() != self.ctx.status.COMPLETED
+            ):
+                # when peer failure, stop peer
+                if self.ctx.args.elastic_level == -1:
+                    self.pod.stop(timeout=3)
+                    return True
+
+                self.pod.stop(timeout=30)
+                return False
+
+    def stop(self, sigint=None):
+        self.ctx.logger.debug("Controller stop")
+
+        self.watcher.stop()
+
+        self.master.stop()
+        self.pod.stop(timeout=30)
+
+    def finalize(self, exit=True):
+        self.pod.join()
+        self.master.stop()
+
+        self.ctx.logger.info(f"Exit code {self.pod.exit_code}")
+        if exit:
+            sys.exit(self.pod.exit_code)
+
+    def signal_handler(self, sigint, frame):
+        if hasattr(self, 'sigint'):
+            self.ctx.logger.info("Force quit in 10 seconds...")
+            self.pod.stop(timeout=10)
+            sys.exit(sigint)
+
+        self.ctx.logger.info(f"Terminating with signal {sigint}")
+
+        self.sigint = sigint
+        self.ctx.status.done()
+        self.stop(sigint=sigint)
+        self.ctx.logger.info(f"Exit with signal {sigint}")
+        sys.exit(sigint)
+
+    def not_exit_signal_handler(self, sigint, frame):
+        if hasattr(self, 'sigint'):
+            self.ctx.logger.info("Force quit in 10 seconds...")
+            self.pod.stop(timeout=10)
+
+        self.ctx.logger.info(f"Terminating with signal {sigint}")
+
+        self.sigint = sigint
+        self.ctx.status.done()
+        self.stop(sigint=sigint)
+        self.ctx.logger.info(f"Exit with signal {sigint}")
+
+
+class Controller(ControllerBase):
+    '''
+    Controller API for customization
+    '''
+
+    def build_job(self):
+        '''
+        build job fill the job info.
+        '''
+        self.ctx.logger.info(self.job)
+
+    def build_pod(self) -> bool:
+        '''
+        build pod includes creating containers etc.
+
+        Return True if succeed
+        '''
+        raise NotImplementedError
+
+    def _get_entrypoint(self):
+        if self.ctx.args.training_script.endswith('.py'):
+            if os.environ.get("WITH_COVERAGE") == "ON":
+                entrypoint = [
+                    sys.executable,
+                    "-u",
+                    "-m",
+                    "coverage",
+                    "run",
+                    "--branch",
+                    "-p",
+                    self.ctx.args.training_script,
+                ]
+            else:
+                entrypoint = [
+                    sys.executable,
+                    "-u",
+                    self.ctx.args.training_script,
+                ]
+        elif self.ctx.args.training_script.endswith('.pyxes'):
+            entrypoint = [sys.executable, self.ctx.args.training_script]
+        else:
+            entrypoint = [self.ctx.args.training_script]
+
+        entrypoint.extend(self.ctx.args.training_script_args)
+        return entrypoint
+
+    def _get_out_err_file(self, out=None, err=None):
+        if out and self.ctx.args.log_dir != "":
+            out = os.path.join(self.ctx.args.log_dir, out)
+        if err and self.ctx.args.log_dir != "":
+            err = os.path.join(self.ctx.args.log_dir, err)
+        return out, (err or out)
+
+    def new_container(
+        self, entrypoint=None, envs={}, use_ctx_env=True, out=None, err=None
+    ):
+        c = Container(
+            entrypoint=(entrypoint or self._get_entrypoint()),
+            env=(self.ctx.get_envs() if use_ctx_env else {}),
+            overwrite_log=self.ctx.args.log_overwrite,
+        )
+        c.outfile, c.errfile = self._get_out_err_file(out, err)
+        c.update_env(envs)
+        return c
+
+    def add_container(
+        self,
+        container=None,
+        entrypoint=None,
+        envs={},
+        log_file=None,
+        is_init=False,
+    ):
+        if not container:
+            envs = copy.deepcopy(envs)
+            envs['PADDLE_LOG_DIR'] = str(os.path.abspath(self.ctx.args.log_dir))
+            container = self.new_container(
+                entrypoint=entrypoint, envs=envs, out=log_file, err=log_file
+            )
+
+        if is_init:
+            self.pod.add_init_container(container)
+        else:
+            self.pod.add_container(container)
+
+    def pod_replicas(self):
+        '''
+        how many process/container should be run in pod
+        '''
+
+        if self.ctx.args.nproc_per_node:
+            return int(self.ctx.args.nproc_per_node)
+        elif self.ctx.args.devices:
+            return len(self.ctx.args.devices.split(','))
+        else:
+            return self.ctx.node.device.count
+
+    def save_pod_log(self, info):
+        '''
+        save_pod_log append *info* to the log file of pod.name
+        '''
+        if not self.ctx.args.log_dir:
+            return
+
+        f = os.path.join(
+            self.ctx.args.log_dir,
+            f'{self.job.id}.{self.pod.name}.log',
+        )
+        try:
+            os.makedirs(os.path.dirname(f), exist_ok=True)
+            with open(f, 'a+') as fd:
+                if fd.tell() == 0:
+                    fd.write(str(os.environ))
+                    fd.write("\n")
+                fd.write(str(info))
+                fd.write("\n")
+        except Exception as e:
+            self.ctx.logger.error(f"save log failed because {e}")
+
+    def save_pod_env(self):
+        assert len(self.pod.containers) + len(self.pod.init_containers) > 0, (
+            "No container in the pod"
+        )
+
+        if not self.ctx.args.log_dir:
+            return
+
+        for c in self.pod.init_containers:
+            self._save_container_env(c, is_init=True)
+
+        for c in self.pod.containers:
+            self._save_container_env(c)
+
+    def _save_container_env(self, container, is_init=False):
+        f = os.path.join(
+            self.ctx.args.log_dir,
+            (
+                f'envlog.init.{container.rank}'
+                if is_init
+                else f'envlog.{container.rank}'
+            ),
+        )
+        try:
+            os.makedirs(os.path.dirname(f), exist_ok=True)
+            with open(f, container.log_mode) as fd:
+                fd.writelines(
+                    f"{k}={v}\n" for k, v in sorted(container.env.items())
+                )
+        except Exception as e:
+            self.ctx.logger.error(f"save pod env log failed because {e}")

python/paddle/distributed/launch/controllers/ipu_controller.py

@@ -0,0 +1,178 @@
+# Copyright (c) 2022 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.
+
+import argparse
+import os
+import sys
+
+from paddle.distributed.launch.job.container import Container
+
+from .collective import CollectiveController, ControllerMode
+
+
+class IPUController(CollectiveController):
+    @classmethod
+    def enable(cls, ctx):
+        if ctx.args.training_script == "ipu":
+            ctx.logger.debug(f"{cls.__name__} enabled")
+            ctx.args.run_mode = ControllerMode.IPU
+            return True
+        else:
+            return False
+
+    def parse_ipu_args(self, args_list):
+        parser = argparse.ArgumentParser()
+        parser.add_argument(
+            "--hosts", type=str, help="The hosts for IPU distributed training."
+        )
+        parser.add_argument(
+            "--nproc_per_host",
+            type=int,
+            help="The number of processes launched per host.",
+        )
+        parser.add_argument(
+            "--ipus_per_replica",
+            type=int,
+            help="The number of IPUs requested per replica.",
+        )
+        parser.add_argument(
+            "--ipu_partition",
+            type=str,
+            help="The partition name of IPU devices.",
+        )
+        parser.add_argument(
+            "--vipu_server", type=str, help="The ip of the IPU device manager."
+        )
+        parser.add_argument(
+            "training_script",
+            type=str,
+            help="The full path to the IPU distributed training program/script to be launched in parallel. e.g., ``training.py``.",
+        )
+        parser.add_argument('training_script_args', nargs=argparse.REMAINDER)
+        return parser.parse_args(args_list)
+
+    def replace_training_script(self):
+        # IPU distributed computing is based on PopRun which is a wrapper of MPI.
+        self.ctx.args.training_script = "poprun"
+        poprun_args = self.parse_ipu_args(self.ctx.args.training_script_args)
+
+        num_ipus = int(self.ctx.args.devices)
+        # The number of replicas for data parallel
+        assert (num_ipus % poprun_args.ipus_per_replica) == 0, (
+            f"The number of IPUs:{num_ipus} mod the number of IPUs per replica:{poprun_args.ipus_per_replica} must == 0"
+        )
+        num_replicas = num_ipus // poprun_args.ipus_per_replica
+        self.ctx.logger.info(f"The number of total replicas is {num_replicas}.")
+
+        # The number of processes
+        num_nodes = len(poprun_args.hosts.split(','))
+        num_procs = num_nodes * poprun_args.nproc_per_host
+        self.ctx.logger.info(f"The number of total processes is {num_procs}.")
+        assert (num_replicas % num_procs) == 0, (
+            f"The number of replicas:{num_replicas} mod the number of processes:{num_procs} must == 0"
+        )
+
+        # hosts and endpoints
+        hosts = poprun_args.hosts.replace(' ', '').split(',')
+        endpoints = [x + ":8090" for x in hosts]
+
+        # args for poprun
+        poprun_command = []
+
+        poprun_command.append(f'--num-instances={num_procs}')
+        poprun_command.append(f'--num-replicas={num_replicas}')
+        poprun_command.append(
+            f'--ipus-per-replica={poprun_args.ipus_per_replica}'
+        )
+        poprun_command.append('--host={}'.format(','.join(hosts)))
+        poprun_command.append(f'--vipu-partition={poprun_args.ipu_partition}')
+        poprun_command.append(f'--vipu-server-host={poprun_args.vipu_server}')
+
+        poprun_command.extend(
+            [
+                '--update-partition=no',
+                '--vipu-server-timeout=120',
+                '--print-topology=yes',
+                '--numa-aware=yes',
+            ]
+        )
+
+        # global envs
+        global_envs = '--mpi-local-args=\''
+        log_level = os.getenv('POPART_LOG_LEVEL', None)
+        if log_level:
+            global_envs += f'-x POPART_LOG_LEVEL={log_level} '
+        global_envs += (
+            '-x PADDLE_TRAINERS_NUM={} -x PADDLE_TRAINER_ENDPOINTS={}'.format(
+                num_procs, ','.join(endpoints)
+            )
+        )
+        global_envs += '\''
+        poprun_command.append(global_envs)
+
+        # local envs
+        for idx in range(num_procs):
+            cur_endpoint = endpoints[idx // poprun_args.nproc_per_host]
+            rank_in_node = idx % poprun_args.nproc_per_host
+            poprun_command.append(
+                f'--instance-mpi-local-args={idx}:"-x PADDLE_TRAINER_ID={idx} -x PADDLE_CURRENT_ENDPOINT={cur_endpoint} -x PADDLE_RANK_IN_NODE={rank_in_node}"'
+            )
+
+        # executor
+        poprun_command.append(sys.executable)
+
+        # script and script args
+        poprun_command.append(poprun_args.training_script)
+        poprun_command.extend(poprun_args.training_script_args)
+
+        # for debug
+        print("-----------  PopRun Command -----------")
+        print("poprun \\")
+        for i in range(len(poprun_command) - 1):
+            print(f"{poprun_command[i]} \\")
+        print(f"{poprun_command[len(poprun_command) - 1]}")
+        print("---------------------------------------")
+
+        # replace training_script_args
+        self.ctx.args.training_script_args = poprun_command
+
+    def _get_entrypoint(self):
+        entrypoint = [self.ctx.args.training_script]
+        entrypoint.extend(self.ctx.args.training_script_args)
+        entrypoint = [" ".join(entrypoint)]
+        return entrypoint
+
+    def new_container(
+        self, entrypoint=None, envs={}, use_ctx_env=True, out=None, err=None
+    ):
+        c = Container(
+            entrypoint=(entrypoint or self._get_entrypoint()),
+            env=(self.ctx.get_envs() if use_ctx_env else {}),
+        )
+        c.outfile, c.errfile = self._get_out_err_file(out, err)
+        c.update_env(envs)
+        # Need subprocess.Popen(shell=True) for PopRun command
+        c.shell = True
+        return c
+
+    def run(self):
+        # Replace the training script with the PopRun command
+        self.replace_training_script()
+
+        self.build_job()
+        self.build_pod()
+
+        self.deploy_pod()
+
+        self.watch()

python/paddle/distributed/launch/controllers/master.py

@@ -0,0 +1,344 @@
+# Copyright (c) 2022 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.
+
+import copy
+import ipaddress
+import json
+import random
+import sys
+import threading
+import time
+
+from paddle.distributed.launch.utils.kv_client import KVClient
+from paddle.distributed.launch.utils.kv_server import KVServer
+
+ETCD_PROTOCOL = 'etcd://'
+
+
+def _cmp_by_ip(x):
+    x = json.loads(x)
+    ip_x = x.get('candidate', '127.0.0.1:8080').split(':')[0]
+    return int(ipaddress.IPv4Address(ip_x))
+
+
+class Master:
+    '''
+    Master is a distributed store design to exchange info among nodes
+    '''
+
+    MAIN = "main"
+    STANDBY = "standby"
+    PARTICIPANT = "participant"
+
+    def __init__(self, ctx):
+        self.ctx = ctx
+        self.server = None
+        self.initialized = False
+        self.endpoint = None
+
+    def stop(self):
+        raise NotImplementedError
+
+    def set_status(self, status):
+        pass
+
+    def get_status(self):
+        return None
+
+    def restart_peer(self):
+        pass
+
+    def sync_peers(self, prefix, key, value, size, rank=-1) -> (list, int):
+        raise NotImplementedError
+
+    @classmethod
+    def factory(cls, ctx):
+        if ctx.args.master and ctx.args.master.startswith(ETCD_PROTOCOL):
+            return ETCDMaster(ctx)
+        else:
+            return HTTPMaster(ctx)
+
+
+class HTTPMaster(Master):
+    def lazy_init(self):
+        if self.initialized:
+            return
+
+        self.role = Master.PARTICIPANT
+
+        if self.ctx.args.master:
+            self.endpoint = self.ctx.args.master
+            ip, port = self.endpoint.split(':')
+            if ip in ['127.0.0.1', self.ctx.node.ip]:
+                time.sleep(2 * random.random())
+                while not self.ctx.node.is_server_ready(ip, int(port)):
+                    try:
+                        self.server = KVServer(int(port))
+                        self.role = Master.MAIN
+                        break
+                    except Exception as e:
+                        self.ctx.logger.warning(f"start master failed {e}")
+                        time.sleep(0.1)
+                        continue
+        else:
+            port = self.ctx.node.get_free_port()
+            self.endpoint = f"{self.ctx.node.ip}:{port}"
+            self.server = KVServer(port)
+            self.role = Master.MAIN
+
+            print("Copy the following command to other nodes to run.")
+            cmd = [
+                sys.executable.split('/')[-1],
+                "-m",
+                "paddle.distributed.launch",
+            ]
+            cmd.extend(["--master", self.endpoint])
+            cmd.extend(sys.argv[1:])
+            print("-" * 80)
+            print(" ".join(cmd))
+            print("-" * 80)
+
+            if int(self.ctx.args.rank) >= 0:
+                self.ctx.logger.warning(
+                    "--rank set in the command may not compatible in auto mode"
+                )
+
+        if '127.0.0.1' in self.endpoint:
+            self.endpoint = self.endpoint.replace('127.0.0.1', self.ctx.node.ip)
+        self.client = KVClient(self.endpoint)
+
+        self.initialized = True
+
+        self._start_server()
+
+    def _start_server(self):
+        if self.server and not self.server.started:
+            self.server.start()
+            self.ctx.logger.debug(f"KV server start at {self.endpoint}")
+
+    def _stop_server(self):
+        if self.server and not self.server.stopped:
+            self.server.stop()
+            self.ctx.logger.debug("KV server stopped")
+
+    def stop(self):
+        self._stop_server()
+
+    def sync_peers(self, prefix, key, value, size, rank=-1) -> (list, int):
+        if size < 2:
+            return [value], 0
+
+        self.ctx.logger.info("Waiting peer start...")
+
+        self.lazy_init()
+
+        while not self.ctx.status.is_done():
+            if self.client.wait_server_ready(timeout=5):
+                break
+            else:
+                self.ctx.logger.warning("master not ready")
+                time.sleep(0.1)
+
+        # 'aaaaaa' make sure main pod (master server) as rank 0
+        ky = 'aaaaaa' if rank < 0 and self.role == Master.MAIN else key
+        k = f"{prefix}/{ky}/{rank}"
+
+        while not self.ctx.status.is_done():
+            if not self.client.put(k, value):
+                self.ctx.logger.warning("put value failed")
+                time.sleep(0.1)
+                continue
+
+            rjson = self.client.get_prefix(prefix)
+            self.ctx.logger.debug(f"sync peers {rjson}")
+            if rjson and len(rjson) == size:
+                if self.ctx.args.sort_ip:
+                    ret = sorted(rjson.values(), key=_cmp_by_ip)
+                    idx = ret.index(value)
+                    return ret, idx
+                elif rank < 0:
+                    keys = list(rjson.keys())
+                    keys.sort()
+                    ret = [rjson[k] for k in keys]
+                    idx = ret.index(value)
+                    return ret, idx
+                else:
+                    ret = [None] * size
+                    for k, v in rjson.items():
+                        ret[int(k.split('/')[-1])] = v
+                    return ret, rank
+            else:
+                time.sleep(0.5)
+        return [], 0
+
+
+class ETCDMaster(Master):
+    def __init__(self, ctx):
+        super().__init__(ctx)
+
+        if self.ctx.args.master:
+            # etcd://localhost:2379
+            self.endpoint = self.ctx.args.master.removeprefix("etcd://")
+
+        import etcd3
+
+        from ..utils.etcd_client import ETCDClient
+
+        host, port = self.endpoint.split(':')
+        if ctx.is_auto_tuner_mode():
+            self.client = ETCDClient(host=host, port=port)
+        else:
+            self.client = etcd3.client(host=host, port=port)
+
+    def sync_peers(self, prefix, key, value, size, rank=-1) -> (list, int):
+        '''
+        sync_peers gather all value for key under scope prefix
+        result always be sorted either by rank or alphabet of pod.name
+        '''
+
+        if size < 2:
+            return [value], 0
+
+        self.ctx.logger.info("Waiting peer start...")
+
+        path = f"{prefix}/{key}/{rank}"
+
+        self.client.delete_prefix(prefix)
+
+        self.ctx.logger.debug(f"sync path {path} value {value}")
+
+        while not self.ctx.status.is_done():
+            self.client.put(path, value.encode('latin-1'))
+
+            result = list(self.client.get_prefix(prefix))
+            result = copy.deepcopy(result)
+            self.ctx.logger.debug(f"sync peers {result}")
+
+            if len(result) == size:
+                if self.ctx.args.sort_ip:
+                    values = [i[0].decode() for i in result]
+                    ret = sorted(values, key=_cmp_by_ip)
+                    idx = ret.index(value)
+                    return ret, idx
+                elif rank < 0:
+                    keys = [i[1].key.decode() for i in result]
+                    sorted_keys = [i[1].key.decode() for i in result]
+                    sorted_keys.sort()
+                    values = [i[0].decode() for i in result]
+                    ret = [values[keys.index(k)] for k in sorted_keys]
+                    idx = ret.index(value)
+                    return ret, idx
+                else:
+                    ret = [None] * size
+                    for v, k in result:
+                        ii = int(k.key.decode().split('/')[-1])
+                        if ii < 0:
+                            self.ctx.logger.error(f"rank {ii} error in sync")
+                        ret[ii] = v.decode()
+                    return ret, rank
+            else:
+                time.sleep(0.5)
+
+    def register_heartbeat(self, job_id, pod_id, ttl=10):
+        if hasattr(self, 'heartbeat_prefix'):
+            self.ctx.logger.warning("Heartbeat already done")
+            return
+
+        self.job_prefix = f'/paddle/{job_id}'
+        self.heartbeat_prefix = f'{self.job_prefix}/heartbeat'
+        self.client.delete_prefix(self.job_prefix)
+        lease = self.client.lease(ttl)
+
+        # self.client.delete_prefix(self.job_prefix)
+
+        beat_path = f"{self.heartbeat_prefix}/{pod_id}"
+        self.client.put(beat_path, pod_id.encode('latin-1'), lease=lease)
+
+        def _beat_watch(event):
+            self.ctx.status.restart()
+
+        beat_watch = self.client.add_watch_prefix_callback(
+            self.heartbeat_prefix, _beat_watch
+        )
+
+        def _heartbeat():
+            while not self.ctx.status.is_done():
+                try:
+                    lease.refresh()
+                    if pod_id not in self.fetch_peer_alive():
+                        self.client.put(
+                            beat_path, pod_id.encode('latin-1'), lease=lease
+                        )
+                        self.ctx.logger.debug("Heartbeat register again")
+                except Exception as e:
+                    self.ctx.logger.error(f"Heartbeat error {e}")
+                time.sleep(ttl / 2)
+            self.ctx.logger.debug("Heartbeat done")
+            self.client.cancel_watch(beat_watch)
+
+        self.beat_thread = threading.Thread(
+            name='heartbeat', target=_heartbeat, daemon=True
+        )
+        self.beat_thread.start()
+
+    def fetch_peer_alive(self):
+        peer_alive = [
+            i[0].decode() for i in self.client.get_prefix(self.heartbeat_prefix)
+        ]
+        self.ctx.logger.debug(f"peer alive {peer_alive}")
+        return peer_alive
+
+    def wait_peer_ready(self, replicas_min, replicas_max, timeout):
+        timeout = timeout if timeout > 1 else 3
+
+        end = time.time() + timeout
+        np_pre = len(self.fetch_peer_alive())
+        while not self.ctx.status.is_done() and time.time() < end:
+            np = len(self.fetch_peer_alive())
+            if np == replicas_max:
+                # maximum replicas reached, return immediately
+                return (True, replicas_max)
+            elif np != np_pre:
+                # replicas are changing, reset timeout
+                end = time.time() + timeout
+                np_pre = np
+                time.sleep(0.2)
+            else:
+                time.sleep(0.5)
+
+        np = len(self.fetch_peer_alive())
+        if np >= replicas_min and np <= replicas_max:
+            return (True, np)
+        else:
+            return (False, np)
+
+    def restart_peer(self):
+        self.client.delete_prefix(self.heartbeat_prefix)
+
+    def set_status(self, status):
+        assert self.client.put(
+            self.job_prefix,
+            status.encode('latin-1'),
+            lease=self.client.lease(600),
+        ), f"set status failed {status}"
+
+    def get_status(self):
+        value = self.client.get(self.job_prefix)[0]
+        return value.decode() if value is not None else ''
+
+    def stop(self):
+        if hasattr(self, 'beat_thread'):
+            self.ctx.status.done()
+            # daemon thread
+            # self.beat_thread.join()

python/paddle/distributed/launch/controllers/ps.py

@@ -0,0 +1,239 @@
+# Copyright (c) 2022 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.
+
+import json
+import os
+import shutil
+
+from .controller import Controller, ControllerMode
+
+
+class PSController(Controller):
+    @classmethod
+    def enable(cls, ctx):
+        if (
+            ctx.args.run_mode == ControllerMode.PS
+            or ctx.args.server_num
+            or len(ctx.args.servers) > 0
+            or ctx.args.trainer_num
+            or len(ctx.args.trainers) > 0
+        ):
+            ctx.logger.debug(f"{cls.__name__} enabled")
+            ctx.args.run_mode = ControllerMode.PS
+            return True
+        else:
+            return False
+
+    def build_pod(self):
+        if self.ctx.args.servers and self.ctx.args.trainers:
+            self._build_pod_with_args()
+        else:
+            self._build_pod_with_master()
+
+    def _build_pod_with_args(self):
+        if '127.0.0.1' in self.ctx.args.servers:
+            host = '127.0.0.1'
+        else:
+            host = self.ctx.node.ip
+
+        server_endpoints = list(self.ctx.args.servers.split(","))
+        trainer_endpoints = list(self.ctx.args.trainers.split(","))
+        servers = [
+            s for s in self.ctx.args.servers.split(",") if s.startswith(host)
+        ]
+        trainers = [
+            s for s in self.ctx.args.trainers.split(",") if s.startswith(host)
+        ]
+        server_num = len(servers)
+        trainer_num = len(trainers)
+
+        self.pod.replicas = server_num + trainer_num
+
+        self.save_pod_log([server_endpoints, trainer_endpoints])
+
+        import tempfile
+
+        gloo_rendezvous_dir = tempfile.mkdtemp()
+        if os.path.exists(gloo_rendezvous_dir):
+            shutil.rmtree(gloo_rendezvous_dir)
+
+        gloo_port = self.ctx.args.gloo_port
+        gloo_http = "{}:{}".format(server_endpoints[0].split(":")[0], gloo_port)
+
+        _gloo_envs = {
+            "PADDLE_GLOO_RENDEZVOUS": "3",
+            "PADDLE_GLOO_FS_PATH": gloo_rendezvous_dir,
+            "PADDLE_GLOO_HTTP_ENDPOINT": gloo_http,
+            "PADDLE_WITH_GLOO": self.ctx.args.with_gloo,
+        }
+
+        for i in range(server_num):
+            e = {
+                "PADDLE_PSERVERS_IP_PORT_LIST": self.ctx.args.servers,
+                "PADDLE_TRAINER_ENDPOINTS": self.ctx.args.trainers,
+                "PADDLE_PORT": servers[i].split(":")[1],
+                "PADDLE_ROLE": "PSERVER",
+                "TRAINING_ROLE": "PSERVER",
+                "PADDLE_TRAINERS_NUM": f"{len(trainer_endpoints)}",
+                "POD_IP": self.ctx.node.ip,
+            }
+            e.update(_gloo_envs)
+            log_file = f"serverlog.{i}"
+            self.add_container(envs=e, log_file=log_file)
+
+        trainer_rank_offset = 0
+        for s in trainer_endpoints:
+            if s.startswith(host):
+                break
+            else:
+                trainer_rank_offset += 1
+
+        for i in range(trainer_num):
+            e = {
+                "PADDLE_PSERVERS_IP_PORT_LIST": ",".join(server_endpoints),
+                "PADDLE_TRAINER_ENDPOINTS": ",".join(trainer_endpoints),
+                "PADDLE_PORT": trainers[i].split(":")[1],
+                "PADDLE_ROLE": "TRAINER",
+                "TRAINING_ROLE": "TRAINER",
+                "PADDLE_TRAINER_ID": f"{i + trainer_rank_offset}",
+                "PADDLE_TRAINERS_NUM": f"{len(trainer_endpoints)}",
+                "POD_IP": self.ctx.node.ip,
+            }
+            e.update(_gloo_envs)
+            log_file = f"workerlog.{i}"
+            self.add_container(envs=e, log_file=log_file)
+
+    def _build_pod_with_master(self):
+        self.pod.rank = int(self.ctx.args.rank)
+
+        server_num = self.ctx.args.server_num or 1
+        servers = [
+            f"{self.ctx.node.ip}:{p}"
+            for p in self.ctx.node.get_free_ports(server_num)
+        ]
+        trainer_num = self.ctx.args.trainer_num or 1
+        trainers = [
+            f"{self.ctx.node.ip}:{p}"
+            for p in self.ctx.node.get_free_ports(trainer_num)
+        ]
+
+        data = json.dumps(
+            {
+                'name': self.pod.name,
+                'rank': self.pod.rank,
+                'servers': servers,
+                'trainers': trainers,
+                'dtype': self.ctx.node.device.dtype,
+                'gloo_port': self.ctx.node.get_free_port(),
+            }
+        )
+
+        peer_list, rank = self.master.sync_peers(
+            f'/{self.job.id}/info',
+            self.pod.name,
+            data,
+            self.job.replicas,
+            self.pod.rank,
+        )
+
+        self.ctx.logger.debug(f"sync peers done {peer_list}")
+
+        peer_list = [json.loads(i) for i in peer_list]
+
+        self.save_pod_log(peer_list)
+
+        server_endpoints = [j for i in peer_list for j in i['servers']]
+        trainer_endpoints = [j for i in peer_list for j in i['trainers']]
+        # rank_offset = sum([i['replicas'] for i in peer_list[:rank]])
+
+        server_rank_offset = sum([len(i['servers']) for i in peer_list[:rank]])
+        trainer_rank_offset = sum(
+            [len(i['trainers']) for i in peer_list[:rank]]
+        )
+
+        self.pod.rank = rank
+
+        self.pod.replicas = server_num + trainer_num
+
+        import tempfile
+
+        gloo_rendezvous_dir = tempfile.mkdtemp()
+        if os.path.exists(gloo_rendezvous_dir):
+            shutil.rmtree(gloo_rendezvous_dir)
+
+        gloo_port = peer_list[0]['gloo_port']
+        gloo_http = "{}:{}".format(server_endpoints[0].split(":")[0], gloo_port)
+
+        _gloo_envs = {
+            "PADDLE_GLOO_RENDEZVOUS": "3",
+            "PADDLE_GLOO_FS_PATH": gloo_rendezvous_dir,
+            "PADDLE_GLOO_HTTP_ENDPOINT": gloo_http,
+            "PADDLE_WITH_GLOO": self.ctx.args.with_gloo,
+        }
+
+        for i in range(server_num):
+            e = {
+                "PADDLE_NNODES": f"{self.job.replicas}",
+                "PADDLE_PSERVERS_IP_PORT_LIST": ",".join(server_endpoints),
+                "PADDLE_TRAINER_ENDPOINTS": ",".join(trainer_endpoints),
+                "PADDLE_PORT": server_endpoints[i + server_rank_offset].split(
+                    ":"
+                )[1],
+                "PADDLE_ROLE": "PSERVER",
+                "TRAINING_ROLE": "PSERVER",
+                "PADDLE_TRAINERS_NUM": f"{len(trainer_endpoints)}",
+                "POD_IP": self.ctx.node.ip,
+            }
+            e.update(_gloo_envs)
+            log_file = f"serverlog.{i}"
+            self.add_container(envs=e, log_file=log_file)
+
+        for i in range(trainer_num):
+            e = {
+                "PADDLE_NNODES": f"{self.job.replicas}",
+                "PADDLE_PSERVERS_IP_PORT_LIST": ",".join(server_endpoints),
+                "PADDLE_TRAINER_ENDPOINTS": ",".join(trainer_endpoints),
+                "PADDLE_PORT": trainer_endpoints[i + trainer_rank_offset].split(
+                    ":"
+                )[1],
+                "PADDLE_ROLE": "TRAINER",
+                "TRAINING_ROLE": "TRAINER",
+                "PADDLE_TRAINER_ID": f"{i + trainer_rank_offset}",
+                "PADDLE_TRAINERS_NUM": f"{len(trainer_endpoints)}",
+                "POD_IP": self.ctx.node.ip,
+            }
+            e.update(_gloo_envs)
+            log_file = f"workerlog.{i}"
+            self.add_container(envs=e, log_file=log_file)
+        ''' NEW VERSION
+        for i in range(server_num):
+            e = {
+                "PADDLE_PSERVER_ENDPOINTS": ",".join(server_endpoints),
+                "PADDLE_TRAINER_ENDPOINTS": ",".join(trainer_endpoints),
+                "PADDLE_ROLE": "PSERVER",
+                "PADDLE_RANK": "{}".format(i + server_rank_offset),
+            }
+            log_tag = "ps.{}".format(i)
+            self.add_container(envs=e, log_tag=log_tag)
+
+        for i in range(trainer_num):
+            e = {
+                "PADDLE_PSERVER_ENDPOINTS": ",".join(server_endpoints),
+                "PADDLE_TRAINER_ENDPOINTS": ",".join(trainer_endpoints),
+                "PADDLE_ROLE": "TRAINER_CPU",
+                "PADDLE_RANK": "{}".format(i + trainer_rank_offset),
+            }
+            log_tag = "trainer.{}".format(i)
+            self.add_container(envs=e, log_tag=log_tag)
+        '''

python/paddle/distributed/launch/controllers/rpc.py

@@ -0,0 +1,91 @@
+# Copyright (c) 2022 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.
+
+import json
+
+from .controller import Controller, ControllerMode
+
+
+class RpcController(Controller):
+    @classmethod
+    def enable(cls, ctx):
+        if ctx.args.run_mode == ControllerMode.RPC:
+            ctx.logger.debug(f"{cls.__name__} enabled")
+            return True
+        else:
+            return False
+
+    def build_pod(self):
+        assert self.ctx.args.master is not None, (
+            "Master is None, Please set master address!"
+        )
+        self._build_pod_with_master()
+
+    def _build_pod_with_master(self):
+        # nproc_per_node
+        self.pod.replicas = self.pod_replicas()
+
+        # rank will be reset when restart
+        self.pod.rank = int(self.ctx.args.rank)
+
+        port = self.ctx.node.get_free_port()
+
+        # compatible
+        endpoints = [
+            f"{self.ctx.node.ip}:{p}"
+            for p in self.ctx.node.get_free_ports(self.pod.replicas)
+        ]
+
+        data = json.dumps(
+            {
+                "name": self.pod.name,
+                "rank": self.pod.rank,
+                "replicas": self.pod.replicas,
+                "dtype": self.ctx.node.device.dtype,
+                "candidate": f"{self.ctx.node.ip}:{port}",
+                "endpoints": ",".join(endpoints),
+            }
+        )
+        peer_list, rank = self.master.sync_peers(
+            f"/{self.job.id}/info",
+            self.pod.name,
+            data,
+            self.job.replicas,
+            self.pod.rank,
+        )
+        self.pod.rank = rank
+
+        if len(peer_list) < 1:
+            return False
+
+        peer_list = [json.loads(i) for i in peer_list]
+
+        self.ctx.logger.debug(f"sync peers done {peer_list}")
+        self.save_pod_log(peer_list)
+
+        global_size = sum([i["replicas"] for i in peer_list])
+        rank_offset = sum([i["replicas"] for i in peer_list[:rank]])
+
+        rpc_master = peer_list[0]["candidate"]
+        self.pod.reset()
+        for i in range(self.pod.replicas):
+            e = {
+                "PADDLE_MASTER_ENDPOINT": rpc_master,
+                "PADDLE_WORKER_ENDPOINT": endpoints[i],
+                "PADDLE_TRAINER_ID": f"{i + rank_offset}",
+                "PADDLE_TRAINERS_NUM": f"{global_size}",
+            }
+            log_file = f"workerlog.{i + rank_offset}"
+            self.add_container(envs=e, log_file=log_file)
+        return True

python/paddle/distributed/launch/controllers/watcher.py

@@ -0,0 +1,106 @@
+# Copyright (c) 2022 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.
+
+import os
+import time
+from threading import Thread
+
+import paddle
+
+from ..utils.nvsmi import get_gpu_info, get_gpu_process, get_gpu_util
+
+
+class Watcher:
+    def __init__(self, ctx):
+        self.ctx = ctx
+
+        self.interval = 5
+
+        self.gpu_util = []
+
+        if not self.ctx.args.enable_gpu_log:
+            return
+
+        if paddle.is_compiled_with_rocm():
+            return
+
+        # gpu log file
+        self.gpus = self.ctx.args.devices or self.ctx.node.device.labels
+        if len(self.gpus) > 0:
+            fn = os.path.join(
+                self.ctx.args.log_dir, f"{self.ctx.args.job_id}.gpu.log"
+            )
+            os.makedirs(os.path.dirname(fn), exist_ok=True)
+            self.gpu_fd = open(fn, 'w')
+        else:
+            return
+
+        # start
+        self.proc = Thread(target=self.watch)
+        self.proc.daemon = True
+        self.proc.start()
+
+    def watch(self):
+        if not len(self.gpus) > 0:
+            return
+
+        self._print_gpu_info()
+
+        util_key = "index,utilization_gpu,memory_total,memory_used,memory_free,timestamp"
+        self.gpu_fd.write(util_key)
+        self.gpu_fd.write('\n')
+
+        while not self.ctx.status.is_done():
+            self._save_gpu_log(util_key)
+            time.sleep(self.interval)
+
+        if hasattr(self, "gpu_fd"):
+            self.gpu_fd.close()
+
+    def _print_gpu_info(self):
+        try:
+            info_key = "index,uuid,driver_version,name,gpu_serial,display_active,display_mode"
+            self.gpu_fd.write(info_key)
+            self.gpu_fd.write('\n')
+            for line in get_gpu_info(self.gpus):
+                self.gpu_fd.write(line.str(info_key))
+                self.gpu_fd.write('\n')
+            self.gpu_fd.write('\n')
+
+            process_key = "pid,process_name,gpu_uuid,gpu_name,used_memory"
+            self.gpu_fd.write(process_key)
+            self.gpu_fd.write('\n')
+            for line in get_gpu_process(self.gpus):
+                self.gpu_fd.write(line.str(process_key))
+                self.gpu_fd.write('\n')
+            self.gpu_fd.write('\n')
+
+            self.gpu_fd.flush()
+        except Exception as e:
+            self.ctx.logger.warning(f"save gpu info failed: {e!s}")
+
+    def _save_gpu_log(self, util_key):
+        try:
+            for line in get_gpu_util(self.gpus):
+                self.gpu_fd.write(line.str(util_key))
+                self.gpu_fd.write('\n')
+            self.gpu_fd.flush()
+        except Exception as e:
+            self.ctx.logger.warning(f"save gpu log failed: {e!s}")
+
+    def stop(self):
+        if hasattr(self, "proc"):
+            # daemon without join
+            # self.proc.join()
+            pass

python/paddle/distributed/launch/job/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2022 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.

python/paddle/distributed/launch/job/container.py

@@ -0,0 +1,211 @@
+# Copyright (c) 2022 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.
+
+import os
+import sys
+
+from paddle.distributed.launch.utils.process_context import ProcessContext
+
+from .status import Status
+
+
+class Container:
+    '''
+    TODO(kuizhiqing) A container can be run by process/thread or just a callable function
+    '''
+
+    def __init__(self, entrypoint=[], rank=-1, env={}, overwrite_log=False):
+        self._entrypoint = entrypoint
+        self._rank = rank
+        self._out = None
+        self._err = None
+        self._env = env
+        self._proc = None
+
+        self._retry: int = 3
+        self._grace_period = 10
+
+        self._log_handler = None
+        self._shell = False
+
+        self.log_mode = 'w' if overwrite_log else 'a'
+
+    @property
+    def env(self):
+        return self._env
+
+    @property
+    def entrypoint(self):
+        return self._entrypoint
+
+    @entrypoint.setter
+    def entrypoint(self, entry):
+        self._entrypoint = entry
+
+    @property
+    def rank(self):
+        return self._rank
+
+    @rank.setter
+    def rank(self, r):
+        self._rank = r
+
+    @property
+    def outfile(self):
+        return self._out
+
+    @outfile.setter
+    def outfile(self, out):
+        self._out = out
+
+    @property
+    def errfile(self):
+        return self._err
+
+    @errfile.setter
+    def errfile(self, err):
+        self._err = err
+
+    @property
+    def shell(self):
+        return self._shell
+
+    @shell.setter
+    def shell(self, shell):
+        self._shell = shell
+
+    def update_env(self, env={}, **kwargs):
+        env = {k: v for k, v in env.items() if isinstance(v, str)}
+        self._env.update(env)
+
+        kwargs = {k: v for k, v in kwargs.items() if isinstance(v, str)}
+        self._env.update(kwargs)
+
+    def _validate_env(self):
+        for k, v in self._env.items():
+            assert isinstance(k, str) and isinstance(v, str), (
+                f'env {k}:{v} must be str'
+            )
+
+    def _get_fd(self, pth):
+        if not pth:
+            return None
+
+        try:
+            d = os.path.dirname(pth)
+            if not os.path.isdir(d):
+                os.makedirs(d, exist_ok=True)
+            return open(pth, self.log_mode)
+        except:
+            return None
+
+    def start(self):
+        if self._proc and self._proc.alive():
+            return True
+
+        self._validate_env()
+
+        self._stdout = self._get_fd(self._out) or sys.stdout
+        if self._out == self._err:
+            self._stderr = self._stdout
+        elif self._err:
+            self._stderr = self._get_fd(self._err) or sys.stderr
+
+        if self._out and not self._log_handler:
+            self._log_handler = open(self._out)
+            self._log_handler.seek(0, 2)
+            self._log_start_offset = self._log_handler.tell()
+
+        self._proc = ProcessContext(
+            self._entrypoint,
+            env=self._env,
+            out=self._stdout,
+            err=self._stderr,
+            shell=self._shell,
+        )
+
+        self._proc.start()
+
+    def terminate(self, force=False):
+        if self._log_handler:
+            self._log_handler.close()
+            self._log_handler = None
+
+        if self._proc and self._proc.alive():
+            return self._proc.terminate(force)
+
+    def wait(self, timeout=None):
+        try:
+            self._proc.wait(timeout)
+            return True
+        except Exception:
+            return False
+
+    @property
+    def exit_code(self):
+        return self._proc.exit_code() if self._proc else -1
+
+    @property
+    def status(self):
+        if not self._proc:
+            return Status.UNINIT
+        if self._proc.alive():
+            return Status.RUNNING
+        elif self._proc.exit_code() == 0:
+            return Status.COMPLETED
+        else:
+            return Status.FAILED
+
+    def __str__(self):
+        need_print = os.environ.get('FLAGS_print_launcher_env', 'false').lower()
+        if need_print == 'true' or need_print == '1':
+            return f'Container rank {self._rank} status {self.status} cmd {self._entrypoint} code {self.exit_code} log {self.errfile} \nenv {self._env}'
+        return f'Container rank {self._rank} status {self.status} cmd {self._entrypoint} code {self.exit_code} log {self.errfile}'
+
+    def logs(self, fn=None, offset=0, whence=1, limit=1000):
+        if not self._log_handler:
+            return
+
+        if fn is None:
+            fn = sys.stdout
+
+        try:
+            if offset != 0 or whence != 1:
+                if whence == 0 and offset < self._log_start_offset:
+                    offset = self._log_start_offset
+                self._log_handler.seek(offset, whence)
+
+            for _ in range(limit):
+                line = self._log_handler.readline()
+                if not line:
+                    return False
+                fn.write(line)
+            return True
+        except:
+            return
+
+    def tail(self, length=3000):
+        if not self._log_handler:
+            return
+
+        try:
+            self._log_handler.seek(0, 2)
+            ed = self._log_handler.tell()
+        except:
+            pass
+
+        if ed > length:
+            self.logs(offset=ed - length, whence=0)
+        else:
+            self.logs(offset=0, whence=0)

python/paddle/distributed/launch/job/job.py

@@ -0,0 +1,81 @@
+# Copyright (c) 2022 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.
+
+
+class JobMode:
+    COLLECTIVE = 'collective'
+    PS = 'ps'
+    HETER = 'heter'
+
+
+class Job:
+    def __init__(self, jid='default', mode=JobMode.COLLECTIVE, nnodes="1"):
+        self._mode = mode
+        self._id = jid
+
+        self._replicas = 0
+        self._replicas_min = self._replicas
+        self._replicas_max = self._replicas
+        self._elastic = False
+
+        self.set_replicas(str(nnodes))
+
+    def __str__(self):
+        return f"Job: {self.id}, mode {self.mode}, replicas {self._replicas}[{self._replicas_min}:{self._replicas_max}], elastic {self.elastic}"
+
+    @property
+    def mode(self):
+        return self._mode
+
+    @property
+    def id(self):
+        return self._id
+
+    @property
+    def elastic(self):
+        return self._elastic
+
+    @property
+    def replicas(self):
+        return self._replicas
+
+    @property
+    def replicas_min(self):
+        return self._replicas_min
+
+    @property
+    def replicas_max(self):
+        return self._replicas_max
+
+    @replicas.setter
+    def replicas(self, replicas):
+        self._replicas = replicas
+
+    def set_replicas(self, nnodes: str):
+        np = str(nnodes) if nnodes else '1'
+
+        if ':' in np:
+            nps = np.split(':')
+            self._replicas_min, self._replicas_max = int(nps[0]), int(nps[1])
+            self._replicas = self._replicas_max  # default to max
+
+            self._elastic = True
+        else:
+            self._replicas = int(np)
+            self._replicas_min, self._replicas_max = (
+                self._replicas,
+                self._replicas,
+            )
+
+            self._elastic = False

python/paddle/distributed/launch/job/pod.py

@@ -0,0 +1,212 @@
+# Copyright (c) 2022 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 random
+import time
+from typing import TYPE_CHECKING
+
+from .status import Status
+
+if TYPE_CHECKING:
+    from .container import Container
+
+
+class PodSpec:
+    def __init__(self):
+        self._name = ''.join(
+            random.choice('abcdefghijklmnopqrstuvwxyz') for _ in range(6)
+        )
+
+        # by controller
+        self._init_containers: list[Container] = []
+        self._containers: list[Container] = []
+
+        # self.resource: Resource = None
+        # self.status: Status = None
+
+        self._rank = -1
+        self._init_timeout = None
+        self._restart = -1
+        self._replicas = 0  # number of containers
+        self._exit_code = 0
+
+
+class Pod(PodSpec):
+    def __init__(self):
+        super().__init__()
+
+    def __str__(self):
+        return (
+            f"Pod: {self.name}, replicas {self.replicas}, status {self.status}"
+        )
+
+    def failed_container(self):
+        cs = []
+        for c in self._containers:
+            if c.status == Status.FAILED:
+                cs.append(c)
+        return cs
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def replicas(self):
+        return self._replicas
+
+    @replicas.setter
+    def replicas(self, r):
+        self._replicas = max(r, 1)
+
+    @property
+    def rank(self):
+        return self._rank
+
+    @rank.setter
+    def rank(self, r):
+        self._rank = r
+
+    @property
+    def restart(self):
+        return self._restart
+
+    @property
+    def containers(self):
+        return self._containers
+
+    def add_container(self, c):
+        c.rank = len(self._containers)
+        self._containers.append(c)
+
+    @property
+    def init_containers(self):
+        return self._init_containers
+
+    def add_init_container(self, c):
+        c.rank = len(self._init_containers)
+        self._init_containers.append(c)
+
+    @property
+    def exit_code(self):
+        for c in self._containers:
+            if c.exit_code != 0:
+                return c.exit_code
+        return 0
+
+    def deploy(self):
+        # init container should stop before run containers
+        for i in self._init_containers:
+            i.start()
+            i.wait(self._init_timeout)
+        for c in self._containers:
+            c.start()
+
+        self._restart += 1
+
+    def stop(self, sigint=15, timeout=None):
+        for c in self._containers:
+            if isinstance(sigint, int) and timeout is None:
+                c.send_signal(sigint)
+            else:
+                c.terminate()
+
+        if isinstance(timeout, int):
+            if not self.join(timeout):
+                for c in self._containers:
+                    c.terminate(force=True)
+                return False
+            else:
+                return True
+
+    def join(self, timeout=None):
+        for c in self._containers:
+            if not c.wait(timeout):
+                return False
+        return True
+
+    @property
+    def status(self):
+        if self.is_failed():
+            return Status.FAILED
+
+        if self.is_completed():
+            return Status.COMPLETED
+
+        if self.is_running():
+            return Status.RUNNING
+
+        return Status.READY
+
+    def reset(self):
+        self._init_containers = []
+        self._containers = []
+
+    def is_failed(self):
+        for c in self._containers:
+            if c.status == Status.FAILED:
+                return True
+        return False
+
+    def is_completed(self):
+        for c in self._containers:
+            if c.status != Status.COMPLETED:
+                return False
+        return True
+
+    def is_running(self):
+        for c in self._containers:
+            if c.status != Status.RUNNING:
+                return False
+        return True
+
+    def logs(self, idx=None):
+        if idx is None:
+            if len(self._containers) > 0:
+                self._containers[0].logs()
+            if len(self._init_containers) > 0:
+                self._init_containers[0].logs()
+        else:
+            self._containers[idx].logs()
+
+    def tail(self, idx=None):
+        if idx is None:
+            self._containers[0].tail()
+        else:
+            self._containers[idx].tail()
+
+    def watch(
+        self,
+        all_list=[Status.COMPLETED],
+        any_list=[Status.FAILED],
+        interval=1,
+        timeout=-1,
+    ):
+        '''
+        watch return if any container status in any_list
+        or all container status in all_list
+        '''
+        end = time.time() + timeout
+        while timeout < 0 or time.time() < end:
+            for c in self._init_containers + self._containers:
+                if c.status in any_list:
+                    return c.status
+
+            s = [c.status for c in self._init_containers + self._containers]
+            if len(set(s)) == 1 and s[0] in all_list:
+                return s[0]
+
+            time.sleep(interval)

python/paddle/distributed/launch/job/status.py

@@ -0,0 +1,24 @@
+# Copyright (c) 2022 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.
+
+
+class Status:
+    UNINIT = "uninit"
+    READY = "ready"
+    RUNNING = "running"
+    FAILED = "failed"
+    TERMINATING = "terminating"
+    RESTARTING = "restarting"
+    UNKNOWN = "unknown"
+    COMPLETED = "completed"

python/paddle/distributed/launch/plugins/__init__.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2022 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.
+
+import os
+
+__all__ = []
+
+
+# print configuration after args are well filled in controller init
+def log(ctx):
+    ctx.logger.info("-----------  Configuration  ----------------------")
+    for arg, value in sorted(vars(ctx.args).items()):
+        ctx.logger.info(f"{arg}: {value}")
+    ctx.logger.info("--------------------------------------------------")
+
+
+def process_args(ctx):
+    # reset device by args
+    # argdev = ctx.args.gpus or ctx.args.xpus or ctx.args.npus
+    argdev = ctx.args.devices
+    if argdev:
+        for d in argdev.split(','):
+            if d not in ctx.node.device.labels:
+                ctx.logger.error(
+                    f'Device not found {d} from {argdev} for setting {ctx.node.device.labels}'
+                )
+
+    if ctx.args.ips:
+        ips = ctx.args.ips.split(',')
+        if '127.0.0.1' in ips and len(ips) != 1:
+            raise ValueError("127.0.0.1 in ips is not allowed in multi-nodes.")
+
+
+def collective_compatible(ctx):
+    force_use_args = int(os.getenv("PADDLE_LAUNCH_WITH_ARGS", "0"))
+    if 'PADDLE_TRAINER_ENDPOINTS' in ctx.envs:
+        eps = ctx.envs['PADDLE_TRAINER_ENDPOINTS'].split(',')
+        hosts = {h.split(':')[0] for h in eps}
+        if force_use_args:
+            ctx.args.master = None
+        else:
+            ctx.args.master = eps[0] if ':' in eps[0] else f'{eps[0]}:6768'
+        ctx.args.nnodes = len(hosts)
+        ctx.logger.info(f'args reset by env PADDLE_TRAINER_ENDPOINTS\n{eps}')
+
+    if 'DISTRIBUTED_TRAINER_ENDPOINTS' in ctx.envs:
+        eps = ctx.envs['DISTRIBUTED_TRAINER_ENDPOINTS'].split(',')
+        hosts = {h.split(':')[0] for h in eps}
+        if force_use_args:
+            ctx.args.master = None
+        else:
+            ctx.args.master = eps[0]
+        ctx.args.nnodes = len(hosts)
+        ctx.logger.info(
+            f'args reset by env DISTRIBUTED_TRAINER_ENDPOINTS\n{eps}'
+        )
+
+
+def rewrite_host_ip(ctx):
+    if ctx.args.host is not None and "." in ctx.args.host:
+        ctx.logger.warning(f'Host ip reset to {ctx.args.host}')
+        ctx.node.ip = ctx.args.host
+
+
+def test_mode(ctx):
+    if ctx.args.training_script == 'run_check':
+        ctx.logger.info('Paddle Distributed Test begin...')
+        if int(ctx.args.nnodes) < 2:
+            ctx.args.nnodes = 2
+        ctx.args.training_script = f'{os.path.dirname(__file__)}/test.py'
+
+
+enabled_plugins = [
+    test_mode,
+    collective_compatible,
+    rewrite_host_ip,
+    process_args,
+]

python/paddle/distributed/launch/plugins/test.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2022 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.
+
+import numpy as np
+
+import paddle
+from paddle.distributed import fleet
+from paddle.io import DataLoader, Dataset
+from paddle.vision.models import ResNet
+from paddle.vision.models.resnet import BottleneckBlock
+
+base_lr = 0.1
+momentum_rate = 0.9
+l2_decay = 1e-4
+
+epoch = 3
+batch_num = 1
+batch_size = 1
+class_dim = 102
+
+
+# define a random dataset
+class RandomDataset(Dataset):
+    def __init__(self, num_samples):
+        self.num_samples = num_samples
+
+    def __getitem__(self, idx):
+        image = np.random.random([3, 224, 224]).astype('float32')
+        label = np.random.randint(0, class_dim - 1, (1,)).astype('int64')
+        return image, label
+
+    def __len__(self):
+        return self.num_samples
+
+
+def optimizer_setting(parameter_list=None):
+    optimizer = paddle.optimizer.Momentum(
+        learning_rate=base_lr,
+        momentum=momentum_rate,
+        weight_decay=paddle.regularizer.L2Decay(l2_decay),
+        parameters=parameter_list,
+    )
+    return optimizer
+
+
+def train_resnet():
+    fleet.init(is_collective=True)
+
+    resnet = ResNet(BottleneckBlock, 18, num_classes=class_dim)
+    optimizer = optimizer_setting(parameter_list=resnet.parameters())
+    optimizer = fleet.distributed_optimizer(optimizer)
+    resnet = fleet.distributed_model(resnet)
+
+    dataset = RandomDataset(batch_num * batch_size)
+    train_loader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        drop_last=True,
+        num_workers=2,
+    )
+
+    print("Distributed training start...")
+    for eop in range(epoch):
+        resnet.train()
+
+        for batch_id, data in enumerate(train_loader()):
+            img, label = data
+            label.stop_gradient = True
+
+            out = resnet(img)
+            loss = paddle.nn.functional.cross_entropy(input=out, label=label)
+            avg_loss = paddle.mean(x=loss)
+            acc_top1 = paddle.metric.accuracy(input=out, label=label, k=1)
+            acc_top5 = paddle.metric.accuracy(input=out, label=label, k=5)
+
+            avg_loss.backward()
+            optimizer.step()
+            resnet.clear_gradients()
+
+            print(
+                f"[Epoch {eop}, batch {batch_id}] loss: {avg_loss:.5f}, acc1: {acc_top1:.5f}, acc5: {acc_top5:.5f}"
+            )
+
+    print("Distributed training completed")
+
+
+if __name__ == '__main__':
+    import os
+
+    nnodes = os.getenv('PADDLE_NNODES')
+    cn = os.getenv('PADDLE_LOCAL_SIZE')
+    print(f"Prepare distributed training with {nnodes} nodes {cn} cards")
+    train_resnet()

python/paddle/distributed/launch/utils/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2022 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.

python/paddle/distributed/launch/utils/etcd_client.py

@@ -0,0 +1,180 @@
+# Copyright (c) 2023 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.
+
+import logging
+import time
+
+import etcd3
+
+
+class ETCDClient:
+    def __init__(self, host, port, retry_times=20):
+        self.retry_times = retry_times
+        times = 0
+        while times < self.retry_times:
+            try:
+                self.client = etcd3.client(host=host, port=port)
+                break
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Initialize etcd client failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(
+                f"Initialize etcd client failed failed after {self.retry_times} times."
+            )
+
+    def put(self, key, value, lease=None, prev_kv=False):
+        times = 0
+        while times < self.retry_times:
+            try:
+                return self.client.put(key, value, lease, prev_kv)
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Put failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(f"Put failed after {self.retry_times} times.")
+
+    def get(self, key):
+        times = 0
+        while times < self.retry_times:
+            try:
+                return self.client.get(key)
+
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Get {key} failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(
+                f"Get {key} failed after {self.retry_times} times."
+            )
+
+    def delete(self, key, prev_kv=False, return_response=False):
+        times = 0
+        while times < self.retry_times:
+            try:
+                return self.client.delete(key, prev_kv, return_response)
+                break
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Delete {key} failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(
+                f"Delete {key} failed after {self.retry_times} times."
+            )
+
+    def get_prefix(self, key_prefix, sort_order=None, sort_target='key'):
+        times = 0
+        while times < self.retry_times:
+            try:
+                return self.client.get_prefix(key_prefix)
+                break
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Get prefix {key_prefix} failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(
+                f"Get prefix {key_prefix} failed after {self.retry_times} times."
+            )
+
+    def delete_prefix(self, prefix):
+        times = 0
+        while times < self.retry_times:
+            try:
+                return self.client.delete_prefix(prefix)
+                break
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Delete prefix {prefix} failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(
+                f"Delete prefix {prefix} failed after {self.retry_times} times."
+            )
+
+    def lease(self, ttl, lease_id=None):
+        times = 0
+        while times < self.retry_times:
+            try:
+                return self.client.lease(ttl, lease_id)
+                break
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Lease failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(f"Lease failed after {self.retry_times} times.")
+
+    def add_watch_prefix_callback(self, key_prefix, callback, **kwargs):
+        times = 0
+        while times < self.retry_times:
+            try:
+                return self.client.add_watch_prefix_callback(
+                    key_prefix, callback, **kwargs
+                )
+                break
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Add watch prefix callback failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(
+                f"Add watch prefix callback failed after {self.retry_times} times."
+            )
+
+    def cancel_watch(self, watch_id):
+        times = 0
+        while times < self.retry_times:
+            try:
+                return self.client.cancel_watch(watch_id)
+                break
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Cancel watch failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(
+                f"Cancel watch failed after {self.retry_times} times."
+            )

python/paddle/distributed/launch/utils/kv_client.py

@@ -0,0 +1,96 @@
+# Copyright (c) 2022 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.
+
+import time
+
+import httpx
+
+
+class KVClient:
+    def __init__(self, endpoint='localhost:2379'):
+        self.endpoint = (
+            endpoint if endpoint.startswith("http://") else f"http://{endpoint}"
+        )
+
+    def put(self, key, value):
+        key = key if key.startswith('/') else f"/{key}"
+        u = f"{self.endpoint}{key}"
+        try:
+            r = httpx.post(u, data=value, timeout=None, follow_redirects=True)
+            if r.status_code == 200:
+                return True
+            else:
+                return False
+        except:
+            return False
+
+    def get(self, key):
+        key = key if key.startswith('/') else f"/{key}"
+        u = f"{self.endpoint}{key}"
+        try:
+            r = httpx.get(u, timeout=None, follow_redirects=True)
+            if r.status_code == 200:
+                ret = r.json()
+                return ret.get(key, '')
+            else:
+                return "error"
+        except:
+            return ""
+
+    def get_prefix(self, key):
+        key = key if key.startswith('/') else f"/{key}"
+        u = f"{self.endpoint}{key}"
+        try:
+            r = httpx.get(u, timeout=None, follow_redirects=True)
+            if r.status_code == 200:
+                return r.json()
+        except:
+            return ""
+
+    def delete(self, key):
+        key = key if key.startswith('/') else f"/{key}"
+        u = f"{self.endpoint}{key}"
+        try:
+            r = httpx.delete(u, timeout=None, follow_redirects=True)
+            if r.status_code == 200:
+                return True
+            else:
+                return False
+        except:
+            return False
+
+    def wait_server_ready(self, timeout=3):
+        end = time.time() + timeout
+        while time.time() < end:
+            if self.get("/healthy") == "ok":
+                return True
+
+
+if __name__ == '__main__':
+    cli = KVClient("http://localhost:8090")
+    data = {"/workers/1": "rank1", "/workers/2": "rank2"}
+    for k, v in data.items():
+        cli.put(k, v)
+    x = cli.get_prefix("/workers")
+    print(x)
+    for k, v in data.items():
+        assert x[k] == v
+
+    cli.put("key", "value")
+    print(cli.get("key"))
+    assert cli.get("key") == "value"
+    cli.delete("key")
+    print(cli.get("/key"))
+    print(cli.get("/healthy"))
+    assert cli.get("/healthy") == "ok"

python/paddle/distributed/launch/utils/kv_server.py

@@ -0,0 +1,128 @@
+# Copyright (c) 2022 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.
+
+import http.server as SimpleHTTPServer
+import json
+import threading
+from http.server import HTTPServer
+from multiprocessing import Process
+
+from .topology import SingleNodeTopology
+
+
+class KVHandler(SimpleHTTPServer.SimpleHTTPRequestHandler):
+    def do_GET(self):
+        with self.server.kv_lock:
+            ret = {}
+            for k, v in self.server.kv.items():
+                if k.startswith(self.path):
+                    ret[k] = v.decode(encoding="utf-8")
+            if ret:
+                self.output(200, json.dumps(ret).encode("utf-8"))
+            else:
+                self.output(404)
+
+    def do_PUT(self):
+        self.do_POST()
+
+    def do_POST(self):
+        content_length = int(self.headers['Content-Length'] or 0)
+        try:
+            value = self.rfile.read(content_length)
+            with self.server.kv_lock:
+                self.server.kv[self.path] = value
+                self.output(200)
+                return
+        except:
+            self.output(500)
+
+    def do_DELETE(self):
+        with self.server.kv_lock:
+            if self.path in self.server.kv:
+                del self.server.kv[self.path]
+                self.output(200)
+            else:
+                self.output(404)
+
+    def output(self, code, value=''):
+        self.send_response(code)
+        self.send_header("Content-Length", len(value))
+        self.send_header("Content-Type", "application/json; charset=utf8")
+        self.end_headers()
+        if value:
+            self.wfile.write(value)
+
+    def log_message(self, format, *args):
+        return
+
+
+class KVServer(HTTPServer):
+    def __init__(self, port):
+        super().__init__(('', port), KVHandler)
+        self.kv_lock = threading.Lock()
+        self.kv = {'/healthy': b'ok'}
+        self.port = port
+        self.stopped = False
+        self.started = False
+        self.node_topo = None
+
+    def start(self):
+        self.listen_thread = threading.Thread(target=self.serve_forever)
+        self.listen_thread.start()
+        self.started = True
+
+    def stop(self):
+        self.shutdown()
+        self.listen_thread.join()
+        self.server_close()
+        self.stopped = True
+
+    def get_topology(self):
+        if self.node_topo is None:
+            self.node_topo = SingleNodeTopology()
+        self.node_topo.detect()
+        return self.node_topo.json_object
+
+
+class PKVServer:
+    def __init__(self, port):
+        self._server = KVServer(port)
+
+    def start(self):
+        self.proc = Process(target=self._server.start)
+        self.proc.daemon = True
+        self.proc.start()
+
+    def stop(self):
+        self._server.stop()
+        self.proc.join()
+
+    @property
+    def started(self):
+        return self._server.started
+
+    @property
+    def stopped(self):
+        return self._server.stopped
+
+
+if __name__ == '__main__':
+    # kv = PKVServer(8090)
+    kv = KVServer(8090)
+    kv.start()
+    import time
+
+    # print("serve at 8090 for 600 s")
+
+    time.sleep(600)

python/paddle/distributed/launch/utils/nvsmi.py

@@ -0,0 +1,256 @@
+# Copyright (c) 2022 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.
+
+import json
+import os
+import re
+import shutil
+import subprocess
+import time
+
+import paddle
+from paddle.base import core
+
+
+class Info:
+    def __repr__(self):
+        return str(self.__dict__)
+
+    def json(self):
+        return json.dumps(self.__dict__)
+
+    def dict(self):
+        return self.__dict__
+
+    def str(self, keys=None):
+        if keys is None:
+            keys = self.__dict__.keys()
+
+        if isinstance(keys, str):
+            keys = keys.split(',')
+
+        values = [str(self.__dict__.get(k, '')) for k in keys]
+        return ",".join(values)
+
+
+def query_smi(query=None, query_type="gpu", index=None, dtype=None):
+    """
+    query_type: gpu/compute
+    """
+
+    if not has_nvidia_smi():
+        return []
+
+    cmd = ["nvidia-smi", "--format=csv,noheader,nounits"]
+    if isinstance(query, list) and query_type == "gpu":
+        cmd.extend(["--query-gpu={}".format(",".join(query))])
+    elif isinstance(query, list) and query_type.startswith("compute"):
+        cmd.extend(["--query-compute-apps={}".format(",".join(query))])
+    else:
+        return
+
+    if isinstance(index, list) and len(index) > 0:
+        cmd.extend(["--id={}".format(",".join(index))])
+    if not isinstance(dtype, list) or len(dtype) != len(query):
+        dtype = [str] * len(query)
+
+    output = subprocess.check_output(cmd, timeout=3)
+    lines = output.decode("utf-8").split(os.linesep)
+    ret = []
+    for line in lines:
+        if not line:
+            continue
+        info = Info()
+        for k, v, d in zip(query, line.split(", "), dtype):
+            setattr(info, k.replace(".", "_"), d(v))
+        ret.append(info)
+    return ret
+
+
+def query_rocm_smi(query=None, index=None, dtype=None, mem=32150):
+    if not has_rocm_smi():
+        return []
+
+    cmd = ["rocm-smi"]
+
+    if not isinstance(dtype, list) or len(dtype) != len(query):
+        dtype = [str] * len(query)
+
+    output = subprocess.check_output(cmd, timeout=3)
+    lines = output.decode("utf-8").split(os.linesep)
+    ret = []
+    for line in lines:
+        if not line:
+            continue
+        if len(line.split()) != 8 or "DCU" in line.split():
+            continue
+        info = Info()
+        line = line.split()
+        line = [
+            line[0],
+            line[7][: len(line[7]) - 1],
+            mem,
+            mem * float(line[6][: len(line[6]) - 1]) / 100,
+            mem - mem * float(line[6][: len(line[6]) - 1]) / 100,
+            time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
+        ]
+        for k, v, d in zip(query, line, dtype):
+            setattr(info, k.replace(".", "_"), d(v))
+        ret.append(info)
+    return ret
+
+
+def query_npu_smi(query=None, index=None, dtype=None):
+    if not has_npu_smi():
+        return []
+
+    cmd = ["npu-smi", "info"]
+
+    if not isinstance(dtype, list) or len(dtype) != len(query):
+        dtype = [str] * len(query)
+
+    output = subprocess.check_output(cmd, timeout=3)
+    lines = output.decode("utf-8").split(os.linesep)
+    ret = []
+    i = 0
+
+    for line in lines:
+        if not line:
+            continue
+        result = re.split(r',|/|\s+|\|', line)
+        # result = [item for item in result if item]
+        length = len(result)
+        if length not in [18, 19] or "NPU" in result:
+            continue
+        result = [item for item in result if item]
+        info = Info()
+        result = [
+            i,
+            result[2],
+            result[6],
+            float(result[5]),
+            (float(result[6]) - float(result[5])),
+            time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
+        ]
+        i += 1
+        for k, v, d in zip(query, result, dtype):
+            setattr(info, k.replace(".", "_"), d(v))
+        ret.append(info)
+    return ret
+
+
+def query_xpu_smi(query=None, index=None, dtype=None):
+    if (
+        not hasattr(core, "get_xpu_device_count")
+        or core.get_xpu_device_count() == 0
+    ):
+        return []
+    if not isinstance(dtype, list) or len(dtype) != len(query):
+        dtype = [str] * len(query)
+    if not isinstance(index, list) or len(index) == 0:
+        index = list(range(core.get_xpu_device_count()))
+    ret = []
+    for dev_id in index:
+        dev_id = int(dev_id)
+        utilization_xpu = core.get_xpu_device_utilization_rate(dev_id)
+        mem_total = (
+            core.get_xpu_device_total_memory(dev_id) / 1024 / 1024
+        )  # with MB
+        mem_used = (
+            core.get_xpu_device_used_memory(dev_id) / 1024 / 1024
+        )  # with MB
+        result = [
+            dev_id,
+            utilization_xpu,
+            mem_total,
+            mem_used,
+            (mem_total - mem_used),
+            time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
+        ]
+        info = Info()
+        for k, v, d in zip(query, result, dtype):
+            setattr(info, k.replace(".", "_"), d(v))
+        ret.append(info)
+    return ret
+
+
+def get_gpu_info(index=None):
+    q = "index,uuid,driver_version,name,gpu_serial,display_active,display_mode".split(
+        ","
+    )
+    d = [int, str, str, str, str, str, str]
+    index = (
+        index
+        if index is None or isinstance(index, list)
+        else str(index).split(",")
+    )
+
+    return query_smi(q, index=index, dtype=d)
+
+
+def get_gpu_util(index=None):
+    q = "index,utilization.gpu,memory.total,memory.used,memory.free,timestamp".split(
+        ","
+    )
+    d = [int, int, int, int, int, str]
+    index = (
+        index
+        if index is None or isinstance(index, list)
+        else str(index).split(",")
+    )
+    if paddle.device.is_compiled_with_rocm():
+        return query_rocm_smi(q, index=index, dtype=d)
+    elif paddle.device.is_compiled_with_custom_device('npu'):
+        return query_npu_smi(q, index=index, dtype=d)
+    elif paddle.is_compiled_with_xpu():
+        return query_xpu_smi(q, index=index, dtype=d)
+    return query_smi(q, index=index, dtype=d)
+
+
+def get_gpu_process(index=None):
+    q = "pid,process_name,gpu_uuid,gpu_name,used_memory".split(",")
+    d = [int, str, str, str, int]
+    index = (
+        index
+        if index is None or isinstance(index, list)
+        else str(index).split(",")
+    )
+
+    return query_smi(q, index=index, query_type="compute", dtype=d)
+
+
+def has_nvidia_smi():
+    return shutil.which("nvidia-smi")
+
+
+def has_rocm_smi():
+    return shutil.which("rocm-smi")
+
+
+def has_npu_smi():
+    return shutil.which("npu-smi")
+
+
+def has_xpu_smi():
+    return shutil.which("xpu-smi")
+
+
+if __name__ == '__main__':
+    print(get_gpu_info(0))
+    print(get_gpu_util(0))
+    print(get_gpu_process(0))
+
+    u = get_gpu_util()
+    for i in u:
+        print(i.str())

python/paddle/distributed/launch/utils/process_context.py

@@ -0,0 +1,116 @@
+# Copyright (c) 2022 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.
+
+import json
+import os
+import signal
+import subprocess
+import sys
+import time
+
+LIMIT_LEN_ENVS = ["TRAINER_IP_PORT_LIST", "PADDLE_TRAINER_ENDPOINTS"]
+
+
+class ProcessContext:
+    def __init__(
+        self,
+        cmd,
+        env=os.environ,
+        out=sys.stdout,
+        err=sys.stderr,
+        group=True,
+        preexec_fn=None,
+        shell=False,
+    ):
+        self._cmd = cmd
+        self._env = env
+        self._preexec_fn = preexec_fn
+        self._stdout = out
+        self._stderr = err
+        self._group = group if os.name != 'nt' else False
+        self._proc = None
+        self._code = None
+        self._shell = shell
+
+    def _start(self):
+        pre_fn = os.setsid if self._group else None
+        log_dir = self._env["PADDLE_LOG_DIR"]
+        os.makedirs(log_dir, exist_ok=True)
+
+        rank = self._env.get("PADDLE_TRAINER_ID")
+        if rank is not None:
+            rank = int(rank)
+            backup_env_path = str(
+                os.path.join(log_dir, f'backup_env.{rank}.json')
+            )
+            envs = {"PADDLE_BACKUP_ENV_PATH": backup_env_path}
+
+            max_len = int(os.getenv('PADDLE_ENV_LIMIT_LEN', 48000))
+            for k, v in self._env.items():
+                if k not in LIMIT_LEN_ENVS or len(v) < max_len:
+                    envs[k] = v
+
+            with open(backup_env_path, 'w') as f:
+                json.dump(dict(self._env), f, indent=4, sort_keys=True)
+        else:
+            envs = self._env
+
+        self._proc = subprocess.Popen(
+            self._cmd,
+            env=envs,
+            stdout=self._stdout,
+            stderr=self._stderr,
+            preexec_fn=self._preexec_fn or pre_fn,
+            shell=self._shell,
+        )
+
+    def _close_std(self):
+        try:
+            if not self._stdout.isatty():
+                self._stdout.close()
+
+            if not self._stderr.isatty():
+                self._stderr.close()
+        except:
+            pass
+
+    def alive(self):
+        return self._proc and self._proc.poll() is None
+
+    def exit_code(self):
+        return self._proc.poll() if self._proc else None
+
+    def start(self):
+        self._start()
+
+    def terminate(self, force=False, max_retry=3):
+        for i in range(max_retry):
+            if self.alive():
+                if self._group:
+                    os.killpg(os.getpgid(self._proc.pid), signal.SIGTERM)
+                else:
+                    self._proc.terminate()
+                time.sleep(0.2)
+            else:
+                break
+
+        if force and self.alive():
+            self._proc.kill()
+
+        self._close_std()
+
+        return self.alive()
+
+    def wait(self, timeout=None):
+        self._proc.wait(timeout)

python/paddle/distributed/launch/utils/topology.py

@@ -0,0 +1,362 @@
+# Copyright (c) 2023 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.
+
+import json
+import os
+import subprocess
+import warnings
+
+import paddle
+
+
+def call_cmd(cmd, err_msg, default_value):
+    process = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        universal_newlines=True,
+        shell=True,
+    )
+    stdout, stderr = process.communicate()
+    if stderr:
+        warnings.warn(err_msg)
+        stdout = default_value
+
+    return stdout
+
+
+class SingleNodeTopology:
+    def __init__(self):
+        self.pcie_latency = 0.0
+        self.pcie_bandwidth = float('inf')
+        self.nvlink_bandwidth = -1.0
+        self.nb_devices = 8
+
+        self.machine = {}
+        self.devices = []
+        self.links = []
+        self.json_object = None
+
+    def calculate_cpu_flops(self):
+        # Get number sockets
+        cmd = "lscpu | grep 'Socket(s)' | awk '{print $NF}'"
+        err_msg = "Failed to get number of sockets"
+        default_value = 4
+        nb_sockets = call_cmd(cmd, err_msg, default_value)
+
+        # Get number of cores per socket
+        cmd = "lscpu | grep 'Core(s) per socket' | awk '{print $NF}'"
+        err_msg = "Failed to get number of cores per socket"
+        default_value = 20
+        nb_cores_per_socket = call_cmd(cmd, err_msg, default_value)
+
+        # Get clock speed
+        cmd = "lscpu | grep GHz | awk -F '@' '{print $NF}' | awk -F 'G' '{print $1}'"
+        err_msg = "Failed to get cpu clock rate"
+        default_value = 2.4
+        clock_rate = call_cmd(cmd, err_msg, default_value)
+
+        # Get number of FMA units
+        # TODO(changtao02): find a way to detect this value
+        nb_fmas = 2
+
+        # Get SIMD width
+        simd_width_sp = 0
+        simd_width_dp = 0
+
+        cmd = "lscpu | grep sse"
+        err_msg = "Failed to get cpu vector size"
+        default_value = "sse"
+        vector_size = call_cmd(cmd, err_msg, default_value)
+
+        if vector_size:
+            simd_width_sp = 4  # 128 / 32
+            simd_width_dp = 2  # 128 / 64
+
+        cmd = "lscpu | grep avx2"
+        err_msg = "Failed to get cpu vector size"
+        default_value = "avx2"
+        vector_size = call_cmd(cmd, err_msg, default_value)
+
+        if vector_size:
+            simd_width_sp = 8  # 256 / 32
+            simd_width_dp = 4  # 256 / 64
+
+        cmd = "lscpu | grep avx512"
+        err_msg = "Failed to get cpu vector size"
+        default_value = "avx512"
+        vector_size = call_cmd(cmd, err_msg, default_value)
+
+        if vector_size:
+            simd_width_sp = 16  # 512 / 32
+            simd_width_dp = 8  # 512 / 64
+
+        gflops_per_element = (
+            int(nb_sockets)
+            * int(nb_cores_per_socket)
+            * float(clock_rate)
+            * nb_fmas
+        )
+        sp_gflops = gflops_per_element * simd_width_sp
+        dp_gflops = gflops_per_element * simd_width_dp
+
+        self.machine['sp_gflops'] = sp_gflops
+        self.machine['dp_gflops'] = dp_gflops
+
+    def pcie_gen2bandwidth(self, pcie_generation):
+        if pcie_generation == 1:
+            return 0.25
+        elif pcie_generation == 2:
+            return 0.5
+        elif pcie_generation == 3:
+            return 1.0
+        elif pcie_generation == 4:
+            return 2.0
+        elif pcie_generation == 5:
+            return 4.0
+        elif pcie_generation == 6:
+            return 8.0
+
+    def model2gflops(self, model):
+        if "H100" in model and "SXM5" in model:
+            return 60000, 30000
+        elif "H100" in model and "PCIe" in model:
+            return 48000, 24000
+        elif "A100" in model:
+            return 19500, 9700
+        elif "A800" in model:
+            return 19500, 9700
+        elif "V100" in model:
+            return 15700, 7800
+        elif "P100" in model:
+            return 10600, 5300
+
+    def get_link_bandwidth(self, source_id, target_id):
+        # Get link type
+        row_id = 2 + source_id
+        column_id = 2 + target_id
+
+        cmd = (
+            "cat /tmp/matrix.txt | awk 'FNR=="
+            + str(row_id)
+            + " {print $"
+            + str(column_id)
+            + "}'"
+        )
+        err_msg = "Failed to get topo matrix"
+        default_value = "NVL"
+        link_type = call_cmd(cmd, err_msg, default_value)
+
+        link_bandwidth = self.pcie_bandwidth
+
+        if "NV" in link_type:
+            if self.nvlink_bandwidth == -1.0:
+                cmd = "nvidia-smi nvlink -s -i 0 | tail -n 1 | awk '{print $3}'"
+                err_msg = "Failed to get nvlink bandwidth"
+                default_value = "25"
+                self.nvlink_bandwidth = float(
+                    call_cmd(cmd, err_msg, default_value)
+                )
+
+            link_bandwidth = int(link_type[2:]) * self.nvlink_bandwidth
+            link_type = "NVL"
+
+        return link_type, link_bandwidth
+
+    def get_host_info(self):
+        # Get hostname
+        cmd = "hostname -s"
+        err_msg = "Failed to get hostname"
+        default_value = "localhost"
+        hostname = call_cmd(cmd, err_msg, default_value).strip()
+
+        # Get ip address
+        cmd = "hostname -i"
+        err_msg = "Failed to get host ip address"
+        default_value = "127.0.0.1"
+        ip_addr = call_cmd(cmd, err_msg, default_value).strip()
+
+        # Get CPU memory (GB)
+        cmd = "cat /proc/meminfo | grep 'MemAvailable' | awk -F ':' '{print $NF}' | awk '{print $1}'"
+        err_msg = "Failed to get cpu memory"
+        default_value = "41366484"
+        cpu_memory = int(call_cmd(cmd, err_msg, default_value)) // 1e6
+
+        # Get single-point flops and double-point flops (GFLOPs)
+        self.calculate_cpu_flops()
+
+        self.machine['hostname'] = hostname
+        self.machine['addr'] = ip_addr
+        self.machine['memory'] = cpu_memory
+
+    def get_device_info(self):
+        # Get device count
+        cmd = "nvidia-smi -L | wc -l"
+        err_msg = "Failed to get device count"
+        default_value = "8"
+        self.nb_devices = int(call_cmd(cmd, err_msg, default_value))
+
+        local_size = int(os.getenv("PADDLE_LOCAL_SIZE"))
+        if local_size < self.nb_devices:
+            self.nb_devices = local_size
+        # Get PCIe latency and bandwidth (ms, GB/s)
+        for i in range(self.nb_devices):
+            cmd = (
+                "nvidia-smi --id="
+                + str(i)
+                + " --query-gpu=pcie.link.gen.max --format=csv,noheader"
+            )
+            err_msg = "Failed to get max pcie link generation"
+            default_value = "4"
+            pcie_generation = int(call_cmd(cmd, err_msg, default_value))
+
+            cmd = (
+                "nvidia-smi --id="
+                + str(i)
+                + " --query-gpu=pcie.link.width.max --format=csv,noheader"
+            )
+            err_msg = "Failed to get max pcie link width"
+            default_value = "16"
+            pcie_width = int(call_cmd(cmd, err_msg, default_value))
+
+            self.pcie_bandwidth = min(
+                self.pcie_bandwidth,
+                self.pcie_gen2bandwidth(pcie_generation) * pcie_width,
+            )
+
+        dev_global_ids = []
+        dev_local_ids = []
+        dev_types = []
+        dev_models = []
+        dev_memories = []  # GiB
+        dev_sp_gflops = []  # GB/s
+        dev_dp_gflops = []  # GB/s
+
+        # Get device info
+        rank_first = paddle.paddle.distributed.get_rank()
+        for i in range(self.nb_devices):
+            dev_global_ids.append(i + rank_first)
+            dev_local_ids.append(i)
+            dev_types.append("GPU")
+
+            cmd = (
+                "nvidia-smi --id="
+                + str(i)
+                + " --query-gpu=name --format=csv,noheader"
+            )
+            err_msg = "Failed to get device name"
+            default_value = "NVIDIA A100-SXM4-40GB"
+            dev_models.append(call_cmd(cmd, err_msg, default_value).strip())
+
+            cmd = (
+                "nvidia-smi --id="
+                + str(i)
+                + " --query-gpu=memory.free --format=csv,noheader | awk '{print $1}'"
+            )
+            err_msg = "Failed to get device available memory"
+            default_value = "40536"
+            dev_memories.append(
+                int(call_cmd(cmd, err_msg, default_value)) // 1e3
+            )
+
+            sp_gflops, dp_gflops = self.model2gflops(dev_models[i])
+            dev_sp_gflops.append(sp_gflops)
+            dev_dp_gflops.append(dp_gflops)
+
+        for i in range(len(dev_global_ids)):
+            device = {}
+            device['global_id'] = dev_global_ids[i]
+            device['local_id'] = dev_local_ids[i]
+            device['type'] = dev_types[i]
+            device['model'] = dev_models[i]
+            device['memory'] = dev_memories[i]
+            device['sp_gflops'] = dev_sp_gflops[i]
+            device['dp_gflops'] = dev_dp_gflops[i]
+            self.devices.append(device)
+
+        self.machine['latency'] = self.pcie_latency
+        self.machine['bandwidth'] = self.pcie_bandwidth
+        self.machine['device_type'] = dev_types[0]
+        self.machine['device_type_full'] = f"{dev_types[0]}-{dev_models[0]}"
+        self.machine['devices'] = self.devices
+
+    def get_link_info(self):
+        link_source_global_ids = []
+        link_target_global_ids = []
+        link_types = []
+        link_latencies = []  # ms
+        link_bandwidths = []  # GB/s
+
+        cmd = "nvidia-smi topo -m > /tmp/matrix.txt"
+        err_msg = "Failed to get topo matrix"
+        default_value = ""
+        call_cmd(cmd, err_msg, default_value)
+
+        rank_first = paddle.paddle.distributed.get_rank()
+        # Get link info between devices
+        for i in range(self.nb_devices):
+            for j in range(self.nb_devices):
+                if i == j:
+                    link_types.append("X")
+                    link_bandwidths.append(-1.0)
+                else:
+                    link_source_global_ids.append(i + rank_first)
+                    link_target_global_ids.append(j + rank_first)
+                    link_latencies.append(0.0)
+                    if i > j:
+                        index = j * self.nb_devices + i
+                        link_types.append(link_types[index])
+                        link_bandwidths.append(link_bandwidths[index])
+                    elif i < j:
+                        link_type, link_bandwidth = self.get_link_bandwidth(
+                            i, j
+                        )
+                        link_types.append(link_type)
+                        link_bandwidths.append(link_bandwidth)
+
+        for i in reversed(range(self.nb_devices)):
+            link_types.pop(i * self.nb_devices + i)
+            link_bandwidths.pop(i * self.nb_devices + i)
+
+        cmd = "rm /tmp/matrix.txt"
+        err_msg = "Failed to delete matrix.txt"
+        default_value = ""
+        # call_cmd(cmd, err_msg, default_value)
+
+        for i in range(len(link_types)):
+            link = {}
+            link['source_global_id'] = link_source_global_ids[i]
+            link['target_global_id'] = link_target_global_ids[i]
+            link['type'] = link_types[i]
+            link['latency'] = link_latencies[i]
+            link['bandwidth'] = link_bandwidths[i]
+            self.links.append(link)
+
+        self.machine['links'] = self.links
+
+    def detect(self):
+        # Get host info
+        self.get_host_info()
+
+        # Get device info
+        self.get_device_info()
+
+        # Get link info between devices
+        self.get_link_info()
+
+        self.json_object = json.dumps(self.machine, indent=4)
+
+    def dump(self, output_path):
+        with open(output_path, "w") as outfile:
+            json.dump(self.machine, outfile, indent=4)