role_maker.py 29.1 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
#   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Defination of Role Makers."""
15
import os
16
import time
17
import numpy as np
18
import warnings
19
from multiprocessing import Process, Manager
20

21
import paddle.fluid as fluid
22 23 24 25 26


class Role:
    WORKER = 1
    SERVER = 2
27
    HETER_WORKER = 3
28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171
    ALL = 4


class Gloo(object):
    """
    Gloo is a universal class for barrier and collective communication
    """

    class RENDEZVOUS:
        HDFS = 1
        FILE = 2
        HTTP = 3

    def __init__(self):
        self._worker_comm = None
        self._server_comm = None
        self._nodes_comm = None

        self._comm_world = ["worker", "server", "all"]
        self._err_init = "gloo is not initialized, will not communicator with other nodes"
        self._err_type = "gloo initialized error, please check arguments"
        self._err_world = "argument error, comm_world must in {}".format(
            self._comm_world)

        self._is_initialized = False
        self._init_timeout_seconds = 3600
        self._run_timeout_seconds = 9999999

        self._rendezvous = None
        self._role = None
        self._iface = None

        self._role_id = -1
        self._worker_num = -1
        self._server_num = -1
        self._need_init_all = False

    def init(self,
             rendezvous,
             role,
             role_id,
             worker_num,
             server_num,
             need_init_all=False,
             kwargs=None):

        self._rendezvous = rendezvous
        self._role = role
        self._role_id = role_id
        self._worker_num = worker_num
        self._server_num = server_num
        self._need_init_all = need_init_all
        self._iface = self.__get_default_iface()
        self._prefix = kwargs.get("store.prefix", "")

        if self._rendezvous == Gloo.RENDEZVOUS.HDFS:
            dfs_name = kwargs.get("dfs.name", "")
            dfs_ugi = kwargs.get("dfs.ugi", "")
            dfs_path = kwargs.get("dfs.path", "")

            if not dfs_name or not dfs_ugi or not dfs_path:
                raise ValueError(self._err_type)
            self._init_dfs(dfs_name, dfs_ugi, dfs_path, self._prefix)

        elif self._rendezvous == Gloo.RENDEZVOUS.FILE:
            fs_path = kwargs.get("dfs.path", "")

            if not fs_path:
                raise ValueError(self._err_type)
            self._init_fs(fs_path, self._prefix)

        elif self._rendezvous == Gloo.RENDEZVOUS.HTTP:
            ip = kwargs.get("http.host", "")
            port = kwargs.get("http.port", "")

            if not ip or not port:
                raise ValueError(self._err_type)
            self._init_http(ip, port, self._prefix)

        else:
            raise ValueError(self._err_type)

        self._is_initialized = True

    def _init_fs(self, fs_path, prefix):
        def init(rank, nodes, role):
            gloo = fluid.core.Gloo()
            gloo.set_rank(rank)
            gloo.set_size(nodes)
            gloo.set_prefix(prefix)
            gloo.set_iface(self._iface)
            gloo.set_timeout_seconds(self._init_timeout_seconds,
                                     self._run_timeout_seconds)
            gloo.set_hdfs_store(os.path.join(fs_path, role), "", "")
            gloo.init()
            return gloo

        if self._role == Role.WORKER:
            rank, nodes = self._get_rank_nodes(Role.WORKER)
            gloo = init(rank, nodes, "WORKER")
            self._worker_comm = gloo
        else:
            rank, nodes = self._get_rank_nodes(Role.SERVER)
            gloo = init(rank, nodes, "SERVER")
            self._server_comm = gloo

        if self._need_init_all:
            rank, nodes = self._get_rank_nodes(Role.ALL)
            gloo = init(rank, nodes, "ALL")
            self._nodes_comm = gloo

    def _init_dfs(self, dfs_name, dfs_ugi, dfs_path, prefix):
        def init(rank, nodes, role):
            gloo = fluid.core.Gloo()
            gloo.set_rank(rank)
            gloo.set_size(nodes)
            gloo.set_prefix(prefix)
            gloo.set_iface(self._iface)
            gloo.set_timeout_seconds(self._init_timeout_seconds,
                                     self._run_timeout_seconds)
            gloo.set_hdfs_store(os.path.join(dfs_path, role), dfs_name, dfs_ugi)
            gloo.init()
            return gloo

        if self._role == Role.WORKER:
            rank, nodes = self._get_rank_nodes(Role.WORKER)
            gloo = init(rank, nodes, "WORKER")
            self._worker_comm = gloo
        else:
            rank, nodes = self._get_rank_nodes(Role.SERVER)
            gloo = init(rank, nodes, "SERVER")
            self._server_comm = gloo

        if self._need_init_all:
            rank, nodes = self._get_rank_nodes(Role.ALL)
            gloo = init(rank, nodes, "ALL")
            self._nodes_comm = gloo

    def _init_http(self, ip, port, prefix):
        def __start_kv_server(http_server_d, size_d):
            from paddle.distributed.fleet.utils.http_server import KVServer
            http_server = KVServer(port, size_d)
            http_server.start()
            wait_seconds = 5
172
            while http_server_d.get("running", False):
173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340
                time.sleep(wait_seconds)
            http_server.stop()

        def init_kv_server():
            size_d = {
                "trainer": self._worker_num,
                "pserver": self._server_num,
                "all": self._worker_num + self._server_num
            }

            _http_server_d = {"running": True}
            # child process for http server
            _http_server = Process(
                target=__start_kv_server, args=(_http_server_d, size_d))
            _http_server.daemon = True
            # set running status to True
            # start child process
            _http_server.start()

        def init(rank, nodes, role):
            gloo = fluid.core.Gloo()
            gloo.set_rank(rank)
            gloo.set_size(nodes)
            gloo.set_prefix(prefix)
            gloo.set_iface(self._iface)
            gloo.set_timeout_seconds(self._init_timeout_seconds,
                                     self._run_timeout_seconds)
            gloo.set_http_store(ip, port, role)
            return gloo

        port = int(port)

        if self._role == Role.SERVER and self._role_id == 0:
            init_kv_server()

        if self._role == Role.WORKER:
            rank, nodes = self._get_rank_nodes(Role.WORKER)
            gloo = init(rank, nodes, "WORKER")
            self._worker_comm = gloo
        else:
            rank, nodes = self._get_rank_nodes(Role.SERVER)
            gloo = init(rank, nodes, "SERVER")
            self._server_comm = gloo

        if self._need_init_all:
            rank, nodes = self._get_rank_nodes(Role.ALL)
            gloo = init(rank, nodes, "ALL")
            self._nodes_comm = gloo

    def _get_rank_nodes(self, role):
        nodes = 0
        rank = -1

        if role == Role.WORKER:
            nodes = self._worker_num
            rank = self._role_id
        elif role == Role.SERVER:
            nodes = self._server_num
            rank = self._role_id
        elif role == Role.ALL:
            nodes = self._worker_num + self._server_num

            if self._role == Role.WORKER:
                rank = self._role_id
            else:
                rank = self._worker_num + self._role_id
        else:
            ValueError(self._err_type)

        return rank, nodes

    def __get_default_iface(self):
        """
        get default physical interface
        """
        default1 = self.__get_default_iface_from_gateway()
        default2 = self.__get_default_iface_from_interfaces()
        return default2 if default1 == "lo" else default1

    def __get_default_iface_from_gateway(self):
        """
        get default physical interface
        """
        import netifaces
        gateways = netifaces.gateways()
        if gateways.get(netifaces.AF_INET) != None:
            gateway = gateways[netifaces.AF_INET]
            if len(gateway) > 0 and len(gateway[0]) > 1:
                return gateway[0][1]
        return "lo"

    def __get_default_iface_from_interfaces(self):
        """
        get default physical interface
        """
        import netifaces
        for intf_name in netifaces.interfaces():
            addresses = netifaces.ifaddresses(intf_name)
            if netifaces.AF_INET in addresses:
                ipv4_addresses = addresses[netifaces.AF_INET]
                for ipv4_address in ipv4_addresses:
                    if 'broadcast' in ipv4_address:
                        return intf_name
        return "lo"

    def barrier(self, comm_world):
        """
        dummy barrier, do nothing
        """
        if not self._is_initialized:
            warnings.warn(self._err_init)
            return

        if comm_world not in self._comm_world:
            raise ValueError(self._err_world)

        if comm_world == "worker":
            self._worker_comm.barrier()
        elif comm_world == "server":
            self._server_comm.barrier()
        else:
            self._nodes_comm.barrier()

    def all_reduce(self, input, mode="sum", comm_world="worker"):
        if not self._is_initialized:
            warnings.warn(self._err_init)
            return input

        if comm_world not in self._comm_world:
            raise ValueError(self._err_world)

        input = np.array(input)
        input_shape = input.shape
        input_list = input.reshape(-1).tolist()

        self.barrier(comm_world)

        if comm_world == "worker":
            ans = self._worker_comm.all_reduce(input_list, mode)
        elif comm_world == "server":
            ans = self._server_comm.all_reduce(input_list, mode)
        else:
            ans = self._nodes_comm.all_reduce(input_list, mode)

        output = np.array(ans).reshape(input_shape)
        return output

    def all_gather(self, input, comm_world="worker"):
        """
        dummy all gather, do nothing
        Args:
            obj(any): obj to do all gather
        """
        if not self._is_initialized:
            warnings.warn(self._err_init)
            return input

        if comm_world not in self._comm_world:
            raise ValueError(self._err_world)

        if comm_world == "worker":
            output = self._worker_comm.all_gather(input)
        elif comm_world == "server":
            output = self._server_comm.all_gather(input)
        else:
            output = self._nodes_comm.all_gather(input)

        return output
341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357


class RoleMakerBase(object):
    """
    RoleMakerBase is a base class for assigning a role to current process
    in distributed training.
    A paddle developer can implement RoleMakerBase to design a role maker
    for worker or pserver assignment.
    """

    def __init__(self):
        self._worker_endpoints = []
        self._server_endpoints = []
        self._role_is_generated = False
        self._role = None
        self._current_id = -1

358 359 360 361 362
        # for heter parameter server mode
        self._heter_trainer_endpoints = []
        self._heter_trainer_device = "CPU"
        self._is_heter_parameter_server_mode = False

363
    def _is_worker(self):
364 365 366 367 368
        """
        return is_worker() of current process
        """
        raise NotImplementedError("Please implement this method in child class")

369
    def _is_server(self):
370 371 372 373 374
        """
        return is_server() of current process
        """
        raise NotImplementedError("Please implement this method in child class")

375
    def _is_first_worker(self):
376 377 378 379 380 381 382 383
        """
        Check whether the node is the first instance of worker.
        Returns:
            bool: True if this is the first node of worker,
                  False if not.
        """
        raise NotImplementedError("Please implement this method in child class")

384
    def _worker_num(self):
385 386 387 388 389 390 391 392
        """
        Get current total worker number.

        Returns:
            int: worker number
        """
        raise NotImplementedError("Please implement this method in child class")

393
    def _server_num(self):
394 395 396 397 398 399 400 401
        """
        Get current total server number.

        Returns:
            int: server number
        """
        raise NotImplementedError("Please implement this method in child class")

402
    def _worker_index(self):
403 404 405 406 407 408 409 410
        """
        Get current worker id.

        Returns:
            int: node id
        """
        raise NotImplementedError("Please implement this method in child class")

411
    def _server_index(self):
412 413 414 415 416 417 418 419
        """
        Get current server id.

        Returns:
            int: node id
        """
        raise NotImplementedError("Please implement this method in child class")

420
    def _role_id(self):
421 422 423 424 425 426 427 428
        """
        Get current id.

        Returns:
            int: node id
        """
        raise NotImplementedError("Please implement this method in child class")

429
    def _node_num(self):
430 431 432 433 434 435 436
        """
        Get the training node number
        Returns:
            int: node num
        """
        raise NotImplementedError("Please implement this method in child class")

437
    def _get_trainer_endpoints(self):
438 439 440 441 442
        """
        return trainer endpoints
        """
        return self._worker_endpoints

443
    def _get_pserver_endpoints(self):
444 445 446 447 448 449 450 451 452 453
        """
        return pserver endpoints
        """
        return self._server_endpoints

    def to_string(self):
        return "role: {}, current_id: {}, worker_endpoints: {}, server_endpoints: {}".format(
            self._role, self._current_id, self._worker_endpoints,
            self._server_endpoints)

454 455
    def _all_gather(self, input, comm_world="worker"):
        print("warning: RoleMakerBase does not have all gather worker.")
456 457
        return None

458
    def _all_reduce(self, input, mode="sum", comm_world="worker"):
459 460 461 462 463 464 465 466 467 468 469 470 471 472 473
        """
        Args:
            input(list/numpy.array): array of one dim
            output(list/numpy.array): array of one dim
            mode(str): "sum" or "min" or "max"
        """
        print("warning: RoleMakerBase does not have all reduce worker.")
        return None

    def _barrier(self, comm_world):
        """
        barrier between trainers if current role is TRAINER
        """
        print("warning: RoleMakerBase does not have barrier worker.")

474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496
    def _is_heter_worker(self):
        """
        Return is_heter_worker() of current process
        """
        warnings.warn("RoleMakerBase does not have function: _is_heter_worker.")
        return False

    def _heter_worker_num(self):
        """
        Get current total heter-worker number.

        Returns:
            int: heter_worker number
        """
        warnings.warn(
            "RoleMakerBase does not have function: _heter_worker_num.")
        return 0

    def _get_heter_worker_endpoints(self):
        """
        Returns:
            string: all heter_trainers'endpoints
        """
C
Chengmo 已提交
497
        assert self._heter_trainer_endpoints != [], "Heter Worker Endpoints Not initialized"
498 499 500 501 502 503 504 505 506
        return self._heter_trainer_endpoints

    def _get_heter_worker_endpoint(self):
        """
        Returns:
            int: corresponding heter_trainer's endpoint

        e.g: if we have 4 cpu-trainer(default), 2 gpu-trainer(heter)
             then No.0 and No.2 cpu-trainer will work with No.0 gpu-trainer
C
Chengmo 已提交
507
             and No.1 and No.3 cpu-trainer will work with No.1 gpu-trainer
508
        """
C
Chengmo 已提交
509 510
        assert self._heter_trainer_endpoints != [], "Heter Worker Endpoints Not initialized"
        return self._heter_trainer_endpoints[(self._current_id) %
511 512 513 514 515 516 517 518 519
                                             self._heter_worker_num()]

    def _get_heter_worker_device(self):
        """
        Returns:
            string: heter_trainer's device of current node, e.g: CPU/GPU/XPU
        """
        return self._heter_trainer_device.upper()

520 521

class PaddleCloudRoleMaker(RoleMakerBase):
522
    def __init__(self, is_collective=False, **kwargs):
523 524 525
        super(PaddleCloudRoleMaker, self).__init__()
        self._is_collective = is_collective

526 527 528
        self._non_distributed = False

        self._kwargs = kwargs
529 530 531 532 533
        self._role_is_generated = False

        self._server_endpoints = None
        self._worker_endpoints = None

534
        self._gloo = Gloo()  # gloo instance
535 536

    def _barrier(self, comm_world):
537
        self._gloo.barrier(comm_world)
538

539 540
    def _all_gather(self, input, comm_world="worker"):
        return self._gloo.all_gather(input, comm_world)
541

542 543
    def _all_reduce(self, input, mode="sum", comm_world="worker"):
        return self._gloo.all_reduce(input, mode, comm_world)
544

545
    def _is_worker(self):
546 547 548 549
        """
        whether current process is worker
        """
        if not self._role_is_generated:
550
            self._generate_role()
551 552
        return self._role == Role.WORKER

553
    def _is_server(self):
554 555 556 557
        """
        whether current process is server
        """
        if not self._role_is_generated:
558
            self._generate_role()
559 560
        return self._role == Role.SERVER

561
    def _is_first_worker(self):
562 563 564 565
        """
        whether current process is worker of rank 0
        """
        if not self._role_is_generated:
566
            self._generate_role()
567 568
        return self._role == Role.WORKER and self._current_id == 0

569
    def _worker_index(self):
570 571 572 573
        """
        get index of current worker
        """
        if not self._role_is_generated:
574
            self._generate_role()
575 576
        return self._current_id

577
    def _server_index(self):
578 579 580 581
        """
        get index of current server
        """
        if not self._role_is_generated:
582
            self._generate_role()
583 584
        return self._current_id

585
    def _role_id(self):
586 587 588
        """
        get index of current node
        """
589 590
        if not self._role_is_generated:
            self._generate_role()
591
        return self._current_id
592

593
    def _worker_num(self):
594 595 596 597
        """
        retrun the current number of worker
        """
        if not self._role_is_generated:
598
            self._generate_role()
599 600
        return self._trainers_num

601
    def _server_num(self):
602 603 604 605
        """
        return the current number of server
        """
        if not self._role_is_generated:
606
            self._generate_role()
607 608
        return len(self._get_pserver_endpoints(
        )) if self._get_pserver_endpoints() is not None else 0
609

610
    def _node_num(self):
611 612 613 614
        """
        return the training node number
        """
        if not self._role_is_generated:
615 616
            self._generate_role()
        return self._nodes_num
617

618
    def _get_trainer_endpoints(self):
619 620 621 622
        """
        get endpoint of all trainers
        """
        if not self._role_is_generated:
623
            self._generate_role()
624 625
        return self._worker_endpoints

626
    def _get_pserver_endpoints(self):
627 628 629 630
        """
        get endpoint of all pservers
        """
        if not self._role_is_generated:
631
            self._generate_role()
632 633
        return self._server_endpoints

634 635 636 637 638 639
    def _is_non_distributed(self):
        """
        Return True if indispensable environment for fleetrun is not found
        (use python-run to launch fleet-code directly)
        """
        if not self._role_is_generated:
640
            self._generate_role()
641 642
        return self._non_distributed

643 644 645 646 647
    def _heter_worker_num(self):
        """
        get heter worker nums
        """
        if not self._role_is_generated:
648
            self._generate_role()
649 650 651 652 653 654 655
        return self._heter_trainers_num

    def _is_heter_worker(self):
        """
        whether current process is heter worker
        """
        if not self._role_is_generated:
656
            self._generate_role()
657 658
        return self._role == Role.HETER_WORKER

659 660 661
    def _ps_env(self):
        try:
            # Environment variable PADDLE_PSERVERS_IP_PORT_LIST must be set
662
            # format: string(ip:port,ip:port), eg. 127.0.0.1:6001,127.0.0.1:6002
663
            self._server_endpoints = os.getenv("PADDLE_PSERVERS_IP_PORT_LIST")
664

665 666 667 668 669 670
            if self._server_endpoints is None:
                # back to non_distributed execution.
                self._server_endpoints = ""
                self._trainers_num = 1
                self._role = Role.WORKER
                self._current_id = 0
671
                self._nodes_num = 1
672 673 674 675 676 677
                self._heter_trainers_num = 0
                self._heter_trainer_endpoints = None
                self._non_distributed = True
                return

            self._server_endpoints = self._server_endpoints.split(",")
678 679 680 681 682 683 684

            self._worker_endpoints = os.getenv("PADDLE_TRAINER_ENDPOINTS")
            if self._worker_endpoints:
                self._worker_endpoints = self._worker_endpoints.split(",")
            else:
                self._worker_endpoints = []

685 686 687
            trainers_num = int(os.environ["PADDLE_TRAINERS_NUM"])
            training_role = os.environ["TRAINING_ROLE"]

688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717
            if training_role not in ["TRAINER", "PSERVER", "HETER_TRAINER"]:
                raise ValueError(
                    "TRAINING_ROLE must be PSERVER or TRAINER or HETER_TRAINER, but get {}, please check your environment.".
                    format(training_role))

            # For heter parameter server env setting
            heter_trainer_eplist = os.getenv(
                "PADDLE_HETER_TRAINER_IP_PORT_LIST", None)
            heter_trainer_device = os.getenv("PADDLE_HETER_TRAINER_DEVICE",
                                             None)
            if heter_trainer_eplist and heter_trainer_device:
                try:
                    heter_trainer_eplist = os.environ[
                        "PADDLE_HETER_TRAINER_IP_PORT_LIST"].split(",")
                except:
                    raise ValueError(
                        "Can not Find PADDLE_HETER_TRAINER_IP_PORT_LIST in env or its format doesn't match the requirement: 'IP:PORT,IP:PORT' ."
                    )

                self._is_heter_parameter_server_mode = True
                heter_trainers_num = len(heter_trainer_eplist)
                current_node_device = heter_trainer_device.upper()
                if current_node_device not in ["CPU", "GPU", "XPU"]:
                    raise ValueError(
                        "Heter Trainer doesn't support {} device now, please use CPU / GPU / XPU(KunLun)".
                        format(heter_trainer_device))
                self._heter_trainer_device = current_node_device
            else:
                self._is_heter_parameter_server_mode = False
                heter_trainers_num = 0
718 719 720 721 722 723 724 725 726 727 728 729

            if training_role == "TRAINER":
                role = Role.WORKER
                current_id = int(os.environ["PADDLE_TRAINER_ID"])
                if len(self._worker_endpoints) > 0:
                    self._cur_endpoint = self._worker_endpoints[current_id]
            elif training_role == "PSERVER":
                role = Role.SERVER
                port = os.environ["PADDLE_PORT"]
                ip = os.environ["POD_IP"]
                self._cur_endpoint = ip + ":" + port
                current_id = self._server_endpoints.index(self._cur_endpoint)
730 731 732 733 734 735
            elif training_role == "HETER_TRAINER":
                role = Role.HETER_WORKER
                cur_ip = os.environ["POD_IP"]
                cur_port = os.environ["PADDLE_PORT"]
                curr_endpoint = ":".join([cur_ip, cur_port])
                current_id = heter_trainer_eplist.index(curr_endpoint)
736
            else:
737 738 739
                raise ValueError(
                    "TRAINING_ROLE must be PSERVER or TRAINER or HETER_TRAINER")
        except ValueError as e:
740
            raise ValueError(
741
                "Something wrong with PaddleCloud, please check environment")
742 743 744 745

        self._trainers_num = trainers_num
        self._role = role
        self._current_id = current_id
746
        self._nodes_num = len(
747
            set([x.split(':')[0] for x in self._worker_endpoints]))
748 749
        self._heter_trainers_num = heter_trainers_num
        self._heter_trainer_endpoints = heter_trainer_eplist
750 751 752 753 754

    def _collective_env(self):
        self._current_id = int(os.getenv("PADDLE_TRAINER_ID", "0"))
        self._training_role = os.getenv("PADDLE_TRAINING_ROLE", "TRAINER")
        assert (self._training_role == "TRAINER")
755
        self._role = Role.WORKER
756 757
        self._worker_endpoints = os.getenv("PADDLE_TRAINER_ENDPOINTS")
        self._cur_endpoint = os.getenv("PADDLE_CURRENT_ENDPOINT")
758 759 760 761 762
        if self._worker_endpoints is None:
            # back to non_distributed execution.
            self._worker_endpoints = "127.0.0.1:6170"
            self._cur_endpoint = self._worker_endpoints
            self._non_distributed = True
763 764
        self._worker_endpoints = self._worker_endpoints.split(",")
        self._trainers_num = len(self._worker_endpoints)
765
        self._nodes_num = len(
766
            set([x.split(':')[0] for x in self._worker_endpoints]))
767

768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801
    def _gloo_init(self):
        # PADDLE_WITH_GLOO 1: trainer barrier, 2: all barrier
        use_gloo = int(os.getenv("PADDLE_WITH_GLOO", "0"))
        if use_gloo not in [1, 2]:
            return

        # PADDLE_GLOO_RENDEZVOUS 1: HDFS 2: FILE 3: HTTP
        rendezvous_type = int(os.getenv("PADDLE_GLOO_RENDEZVOUS", "0"))
        prefix = os.getenv("SYS_JOB_ID", "")
        if rendezvous_type not in [
                Gloo.RENDEZVOUS.HDFS, Gloo.RENDEZVOUS.HTTP, Gloo.RENDEZVOUS.FILE
        ]:
            raise ValueError(self._gloo._err_type)

        need_init_all = True if use_gloo == 2 else False

        if rendezvous_type == Gloo.RENDEZVOUS.HDFS:
            dfs_name = os.getenv("PADDLE_GLOO_FS_NAME", "")
            dfs_ugi = os.getenv("PADDLE_GLOO_FS_UGI", "")
            dfs_path = os.getenv("PADDLE_GLOO_FS_PATH", "")
            kwargs = {
                "dfs.name": dfs_name,
                "dfs.ugi": dfs_ugi,
                "dfs.path": dfs_path,
                "store.prefix": prefix,
            }
        elif rendezvous_type == Gloo.RENDEZVOUS.HTTP:
            ip = os.getenv("PADDLE_GLOO_HTTP_HOST", "")
            port = os.getenv("PADDLE_GLOO_HTTP_PORT", "")
            kwargs = {
                "http.host": ip,
                "http.port": port,
                "store.prefix": prefix,
            }
802
        else:
803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820
            dfs_path = os.getenv("PADDLE_GLOO_FS_PATH", "")
            kwargs = {
                "dfs.path": dfs_path,
                "store.prefix": prefix,
            }

        if rendezvous_type == Gloo.RENDEZVOUS.HDFS:
            type = "HDFS"
        elif rendezvous_type == Gloo.RENDEZVOUS.HTTP:
            type = "HTTP"
        else:
            type = "FILE"
        print("Gloo init with {}: need_init_all: {}, args: {}".format(
            type, need_init_all, kwargs))

        self._gloo.init(
            rendezvous=rendezvous_type,
            role=self._role,
821 822 823
            role_id=self._role_id(),
            worker_num=self._worker_num(),
            server_num=self._server_num(),
824 825
            need_init_all=need_init_all,
            kwargs=kwargs)
826

827
    def _generate_role(self):
828 829 830 831 832 833 834 835 836
        """
        generate role for role maker
        """
        if not self._role_is_generated:
            if not self._is_collective:
                self._ps_env()
            else:
                self._collective_env()
            self._role_is_generated = True
837
            self._gloo_init()
838 839 840 841 842 843


class UserDefinedRoleMaker(PaddleCloudRoleMaker):
    def __init__(self, is_collective=False, init_gloo=False, **kwargs):
        super(UserDefinedRoleMaker, self).__init__(
            is_collective=is_collective, init_gloo=init_gloo, **kwargs)
844
        self._init_gloo = init_gloo
845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862

    def _user_defined_ps_env(self):
        self._server_endpoints = self._kwargs.get("server_endpoints")
        self._worker_endpoints = self._kwargs.get("worker_endpoints", [])
        self._trainers_num = self._kwargs.get("worker_num", 0)

        if self._trainers_num == 0:
            assert (len(self._worker_endpoints) > 0)
            self._trainers_num = len(self._worker_endpoints)

        self._role = self._kwargs.get("role")
        self._current_id = self._kwargs.get("current_id")

        if self._role == Role.WORKER and len(
                self._worker_endpoints) > self._current_id:
            self._cur_endpoint = self._worker_endpoints[self._current_id]
        elif self._role == Role.SERVER:
            self._cur_endpoint = self._server_endpoints[self._current_id]
863
        self._nodes_num = len(
864
            set([x.split(':')[0] for x in self._worker_endpoints]))
865 866 867 868 869

    def _user_defined_collective_env(self):
        self._worker_endpoints = self._kwargs.get("worker_endpoints")
        self._current_id = self._kwargs.get("current_id")
        self._trainers_num = len(self._worker_endpoints)
870
        self._training_role = Role.WORKER
871
        self._nodes_num = len(
872
            set([x.split(':')[0] for x in self._worker_endpoints]))
873

874
    def _generate_role(self):
875 876 877 878 879 880 881 882 883
        """
        generate role for role maker
        """
        if not self._role_is_generated:
            if not self._is_collective:
                self._user_defined_ps_env()
            else:
                self._user_defined_collective_env()
            self._role_is_generated = True