chore: import upstream snapshot with attribution

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