role_maker.py 31.8 KB
Newer Older
D
dongdaxiang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13
#   Copyright (c) 2019 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.
X
xujiaqi01 已提交
14
"""Defination of Role Makers."""
D
dongdaxiang 已提交
15

T
tangwei12 已提交
16
from __future__ import print_function
X
xujiaqi01 已提交
17 18 19
import paddle.fluid as fluid
import os
import time
20

T
tangwei12 已提交
21
__all__ = [
22
    'Role', 'RoleMakerBase', 'MPISymetricRoleMaker', 'UserDefinedRoleMaker',
X
xujiaqi01 已提交
23
    'UserDefinedCollectiveRoleMaker', 'PaddleCloudRoleMaker', 'GeneralRoleMaker'
T
tangwei12 已提交
24 25
]

26

T
tangwei12 已提交
27 28
class Role:
    WORKER = 1
29 30
    SERVER = 2

D
dongdaxiang 已提交
31 32

class RoleMakerBase(object):
33 34 35 36 37 38 39
    """
    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.
    """

D
dongdaxiang 已提交
40
    def __init__(self):
T
tangwei12 已提交
41 42
        self._worker_endpoints = []
        self._server_endpoints = []
D
dongdaxiang 已提交
43
        self._role_is_generated = False
T
tangwei12 已提交
44 45
        self._role = None
        self._current_id = -1
D
dongdaxiang 已提交
46

T
tangwei12 已提交
47
    def is_worker(self):
48 49 50
        """
        return is_worker() of current process
        """
D
dongdaxiang 已提交
51 52
        raise NotImplementedError("Please implement this method in child class")

T
tangwei12 已提交
53
    def is_server(self):
54 55 56
        """
        return is_server() of current process
        """
D
dongdaxiang 已提交
57 58
        raise NotImplementedError("Please implement this method in child class")

T
tangwei12 已提交
59
    def is_first_worker(self):
60
        """
T
tangwei12 已提交
61 62 63 64
        Check whether the node is the first instance of worker.
        Returns:
            bool: True if this is the first node of worker,
                  False if not.
65
        """
T
tangwei12 已提交
66
        raise NotImplementedError("Please implement this method in child class")
D
dongdaxiang 已提交
67

68 69 70 71 72 73 74 75 76
    def worker_num(self):
        """
        Get current total worker number.

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

T
tangwei12 已提交
77
    def worker_index(self):
78
        """
T
tangwei12 已提交
79 80 81 82
        Get current worker id.

        Returns:
            int: node id
83
        """
T
tangwei12 已提交
84
        raise NotImplementedError("Please implement this method in child class")
D
dongdaxiang 已提交
85

T
tangwei12 已提交
86
    def server_index(self):
87
        """
T
tangwei12 已提交
88 89 90 91
        Get current server id.

        Returns:
            int: node id
92
        """
T
tangwei12 已提交
93
        raise NotImplementedError("Please implement this method in child class")
D
dongdaxiang 已提交
94

T
tangwei12 已提交
95
    def get_trainer_endpoints(self):
96
        """
T
tangwei12 已提交
97
        return trainer endpoints
98
        """
T
tangwei12 已提交
99 100 101 102 103 104 105
        return self._worker_endpoints

    def get_pserver_endpoints(self):
        """
        return pserver endpoints
        """
        return self._server_endpoints
D
dongdaxiang 已提交
106

T
tangwei12 已提交
107 108 109 110 111
    def to_string(self):
        return "role: {}, current_id: {}, worker_endpoints: {}, server_endpoints: {}".format(
            self._role, self._current_id, self._worker_endpoints,
            self._server_endpoints)

X
xujiaqi01 已提交
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
    def all_gather(self, input):
        """
        all gather between trainers and pservers

        Args:
            input(int|float): input value

        Returns:
            return a list of values
        """
        print("warning: RoleMakerBase does not have all gather.")
        return None

    def all_reduce_worker(self, input, output, mode="sum"):
        """
        all reduce between trainers if current role is TRAINER,
        only support array of one dim.

        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.")

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

    def barrier_all(self):
        """
        barrier between trainers if current role is PSERVER
        """
        print("warning: RoleMakerBase does not have barrier all.")

D
dongdaxiang 已提交
149 150

class MPIRoleMaker(RoleMakerBase):
151 152 153 154 155
    """
    MPIRoleMaker is a MPI-API based role maker which is a counter-part of K8SRoleMaker
    mpi4py will be used if a developer inherits MPIRoleMaker
    """

D
dongdaxiang 已提交
156
    def __init__(self):
X
xujiaqi01 已提交
157
        """Init."""
X
xujiaqi01 已提交
158
        super(MPIRoleMaker, self).__init__()
D
dongdaxiang 已提交
159 160
        from mpi4py import MPI
        self.MPI = MPI
T
tangwei12 已提交
161 162
        self._comm = MPI.COMM_WORLD
        self._node_type_comm = None
D
dongdaxiang 已提交
163
        self._ips = None
T
tangwei12 已提交
164
        self._ip = None
D
dongdaxiang 已提交
165

166
    def _get_rank(self):
X
xujiaqi01 已提交
167
        """Return rank."""
D
dongdaxiang 已提交
168 169
        self._rank = self._comm.Get_rank()
        return self._rank
D
dongdaxiang 已提交
170

171
    def _get_size(self):
X
xujiaqi01 已提交
172
        """Return size."""
D
dongdaxiang 已提交
173 174
        self._size = self._comm.Get_size()
        return self._size
D
dongdaxiang 已提交
175

176
    def _all_gather(self, obj):
177 178 179
        """
        all_gather(obj) will call MPI's allgather function
        """
X
xjqbest 已提交
180
        self._barrier_all()
D
dongdaxiang 已提交
181
        return self._comm.allgather(obj)
D
dongdaxiang 已提交
182

X
xjqbest 已提交
183 184 185 186
    def _worker_gather(self, obj):
        """
        worker_gather(obj) will call MPI's allgather function
        """
T
tangwei12 已提交
187
        if self.is_worker():
D
dongdaxiang 已提交
188 189
            self._node_type_comm.barrier()
            return self._node_type_comm.allgather(obj)
X
xjqbest 已提交
190 191
        return None

192
    def _barrier_all(self):
193 194 195
        """
        barrier_all() will call MPI's barrier_all function
        """
D
dongdaxiang 已提交
196
        self._comm.barrier()
D
dongdaxiang 已提交
197

T
tangwei12 已提交
198 199 200 201
    def _finalize(self):
        """
        finalize the current MPI instance.
        """
202
        self.MPI.Finalize()
T
tangwei12 已提交
203

204
    def _get_ips(self):
205 206 207
        """
        collect current distributed job's ip list
        """
T
tangwei12 已提交
208 209
        if not self._ips:
            self._ips = self._comm.allgather(self.get_local_ip())
D
dongdaxiang 已提交
210
        return self._ips
D
dongdaxiang 已提交
211

T
tangwei12 已提交
212
    def get_local_ip(self):
X
xujiaqi01 已提交
213
        """Return get local ip."""
T
tangwei12 已提交
214 215 216 217 218 219 220 221 222
        import socket
        self._ip = socket.gethostbyname(socket.gethostname())
        return self._ip

    def generate_role(self):
        """
        generate_role() should be called to identify current process's role
        """
        raise NotImplementedError("Please implement this method in child class")
D
dongdaxiang 已提交
223 224 225


class MPISymetricRoleMaker(MPIRoleMaker):
226 227 228 229 230 231
    """
    MPISymetricRoleMaker is designed for worker and server assignment
    under MPI. Typically, a worker and a server node will be appointed
    on each physical node. This role maker can be only used under MPI.
    """

D
dongdaxiang 已提交
232
    def __init__(self):
X
xujiaqi01 已提交
233
        """Init."""
D
dongdaxiang 已提交
234
        super(MPISymetricRoleMaker, self).__init__()
D
dongdaxiang 已提交
235 236
        self._node_type = None
        self._proc_per_node = 2
G
guru4elephant 已提交
237
        self._pserver_rand_port = 0
D
dongdaxiang 已提交
238

239
    def _check_role_generation(self):
X
xujiaqi01 已提交
240
        """Check whether role has been generated."""
D
dongdaxiang 已提交
241
        if not self._role_is_generated:
T
tangwei12 已提交
242
            raise NameError("generate_role() should be called first")
243 244
        return True

X
xujiaqi01 已提交
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
    def all_gather(self, input):
        """
        all gather between trainers and pservers

        Args:
            input(int|float): input value

        Returns:
            return a list of values
        """
        if not self._role_is_generated:
            self.generate_role()
        return self._all_gather(input)

    def all_reduce_worker(self, input, output, mode="sum"):
        """
        all reduce between trainers if current role is TRAINER,
        only support array of one dim.

        Args:
            input(list/numpy.array): array of one dim
            output(list/numpy.array): array of one dim
            mode(str): "sum" or "min" or "max"
        """
        if not self._role_is_generated:
            self.generate_role()
        if not self.is_worker():
            print("warning: current role is not worker in all_reduce_worker")
            return
        self._all_reduce(input, output, mode)

    def barrier_worker(self):
        """
        barrier between trainers if current role is TRAINER
        """
        if not self._role_is_generated:
            self.generate_role()
        if self.is_worker():
            self._node_type_comm.barrier()
        else:
            print("warning: current role is not worker in barrier_worker")

    def barrier_all(self):
        """
        barrier between trainers if current role is PSERVER
        """
        if not self._role_is_generated:
            self.generate_role()
        self._comm.barrier()

T
tangwei12 已提交
295
    def is_first_worker(self):
296 297 298 299
        """
        return whether current process is the first worker assigned by role maker
        """
        if self._check_role_generation():
T
tangwei12 已提交
300
            return self.is_worker() and 0 == self.worker_index()
301
        return False
D
dongdaxiang 已提交
302

G
guru4elephant 已提交
303
    def get_pserver_endpoints(self):
X
xujiaqi01 已提交
304 305 306 307 308 309
        """
        get pserver endpoints
        
        Returns:
            endpoints(list): pserver endpoints
        """
G
guru4elephant 已提交
310 311 312 313 314 315 316 317 318 319 320 321 322
        if self._pserver_rand_port <= 0:
            import random
            random.seed(self._server_num())
            # port will be randomly generated from 60001 to 63999
            # random seed is server num so that all nodes will get
            # the same port
            self._pserver_rand_port = random.randint(60001, 64000)
        endpoints = [
            x + ":" + str(self._pserver_rand_port)
            for x in self._server_endpoints
        ]
        return endpoints

323 324 325
    def worker_num(self):
        return self._worker_num()

T
tangwei12 已提交
326
    def is_worker(self):
327 328 329 330
        """
        return whether current process is worker assigned by role maker
        """
        if self._check_role_generation():
D
dongdaxiang 已提交
331
            return self._node_type == 1
332
        return False
D
dongdaxiang 已提交
333

T
tangwei12 已提交
334
    def is_server(self):
335 336 337 338
        """
        return whether current process is server assigned by role maker
        """
        if self._check_role_generation():
D
dongdaxiang 已提交
339
            return self._node_type == 0
340
        return False
D
dongdaxiang 已提交
341

342
    def _worker_num(self):
343 344 345 346
        """
        return the current number of worker
        """
        if self._check_role_generation():
347
            return self._get_size() / self._proc_per_node
348
        return 0
D
dongdaxiang 已提交
349

350
    def _server_num(self):
351 352 353 354
        """
        return the current number of server
        """
        if self._check_role_generation():
G
guru4elephant 已提交
355 356 357 358
            return self._get_size() / self._proc_per_node
        else:
            self.generate_role()
            return self._get_size() / self._proc_per_node
D
dongdaxiang 已提交
359

T
tangwei12 已提交
360
    def worker_index(self):
361 362 363 364
        """
        return the index of worker
        """
        if self._check_role_generation():
D
dongdaxiang 已提交
365
            return self._rank / self._proc_per_node
G
guru4elephant 已提交
366 367 368
        else:
            self.generate_role()
            return self._get_size() / 2
D
dongdaxiang 已提交
369

T
tangwei12 已提交
370
    def server_index(self):
371 372 373 374
        """
        return the index of server
        """
        if self._check_role_generation():
D
dongdaxiang 已提交
375
            return self._rank / self._proc_per_node
G
guru4elephant 已提交
376 377 378
        else:
            self.generate_role()
            return self._get_size() / self._proc_per_node
D
dongdaxiang 已提交
379

X
xujiaqi01 已提交
380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401
    def _all_reduce(self, input, output, mode="sum"):
        """
        all reduce between trainers if current role is TRAINER,
        only support array of one dim.

        Args:
            input(list/numpy.array): array of one dim
            output(list/numpy.array): array of one dim
            mode(str): "sum" or "min" or "max"
        """
        if not self._role_is_generated:
            self.generate_role()
        if mode == "sum":
            mode = self.MPI.SUM
        elif mode == "max":
            mode = self.MPI.MAX
        elif mode == "min":
            mode = self.MPI.MIN
        else:
            raise ValueError("unknown mode: %s" % mode)
        self._node_type_comm.Allreduce(input, output, op=mode)

402
    def _barrier_worker(self):
403 404 405 406
        """
        barrier all workers in current distributed job
        """
        if self._check_role_generation():
T
tangwei12 已提交
407
            if self.is_worker():
D
dongdaxiang 已提交
408
                self._node_type_comm.barrier()
G
guru4elephant 已提交
409 410
        else:
            raise Exception("You should check role generation first")
D
dongdaxiang 已提交
411

412
    def _barrier_server(self):
413 414 415 416
        """
        barrier all servers in current distributed job
        """
        if self._check_role_generation():
T
tangwei12 已提交
417
            if self.is_server():
D
dongdaxiang 已提交
418
                self._node_type_comm.barrier()
G
guru4elephant 已提交
419 420
        else:
            raise Exception("You should check role generation first")
D
dongdaxiang 已提交
421

T
tangwei12 已提交
422
    def generate_role(self):
423 424 425
        """
        generate currently process's role
        """
D
dongdaxiang 已提交
426
        if not self._role_is_generated:
427
            # TODO(guru4elephant): only allow to be called once
428 429
            self._worker_endpoints = self._get_ips()[1::2]
            self._server_endpoints = self._get_ips()[::2]
430

D
dongdaxiang 已提交
431 432
            if 0 == self._get_rank() % self._proc_per_node % 2:
                self._node_type = 0
433
            else:
D
dongdaxiang 已提交
434 435 436
                self._node_type = 1
            self._node_type_comm = self._comm.Split(self._node_type)
            self._role_is_generated = True
G
guru4elephant 已提交
437 438
        else:
            raise Exception("You should check role generation first")
439 440


441
class PaddleCloudRoleMaker(RoleMakerBase):
X
xujiaqi01 已提交
442 443 444 445 446
    """
    role maker for paddle cloud,
    base class is RoleMakerBase
    """

447
    def __init__(self, is_collective=False):
448
        super(PaddleCloudRoleMaker, self).__init__()
449
        self._role_is_generated = False
450
        self._is_collective = is_collective
451 452

    def generate_role(self):
X
xujiaqi01 已提交
453
        """Generate role."""
454
        if not self._role_is_generated:
455
            if not self._is_collective:
T
tangwei12 已提交
456
                try:
C
Chengmo 已提交
457 458 459 460
                    # Environment variable PADDLE_PSERVERS_IP_PORT_LIST must be set
                    # format: string(ip:port), eg. 127.0.0.1:6001
                    eplist = os.environ["PADDLE_PSERVERS_IP_PORT_LIST"].split(
                        ",")
T
tangwei12 已提交
461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477
                    # note that, we usually assign the same port to different ips
                    # if we run parameter server training in local mode
                    # port should be different in environment variables

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

                    if training_role not in ["TRAINER", "PSERVER"]:
                        raise ValueError(
                            "TRAINING_ROLE must be PSERVER or TRAINER")

                    if training_role == "TRAINER":
                        role = Role.WORKER
                        current_id = int(os.environ["PADDLE_TRAINER_ID"])
                    elif training_role == "PSERVER":
                        role = Role.SERVER
                        cur_ip = os.environ["POD_IP"]
C
Chengmo 已提交
478 479 480
                        curr_port = os.environ["PADDLE_PORT"]
                        curr_endpoint = ":".join([cur_ip, curr_port])
                        current_id = eplist.index(curr_endpoint)
T
tangwei12 已提交
481 482 483 484 485 486 487 488 489 490 491 492
                    else:
                        raise ValueError(
                            "TRAINING_ROLE must be PSERVER or TRAINER")
                except ValueError as ve:
                    raise ValueError(
                        "something wrong with PaddleCloud, please check environment"
                    )

                self._trainers_num = trainers_num
                self._server_endpoints = eplist
                self._role = role
                self._current_id = current_id
493
            else:
494 495 496 497 498 499
                self._current_id = int(os.getenv("PADDLE_TRAINER_ID", "0"))
                self._training_role = os.getenv("PADDLE_TRAINING_ROLE",
                                                "TRAINER")
                assert (self._training_role == "TRAINER")
                self._worker_endpoints = os.getenv("PADDLE_TRAINER_ENDPOINTS")
                self._current_endpoint = os.getenv("PADDLE_CURRENT_ENDPOINT")
500 501 502 503
                assert self._worker_endpoints is not None, "can't find PADDLE_TRAINER_ENDPOINTS"
                self._worker_endpoints = self._worker_endpoints.split(",")
                self._trainers_num = len(self._worker_endpoints)

504 505
            self._role_is_generated = True

506 507 508 509 510
    def get_pserver_endpoints(self):
        if not self._role_is_generated:
            self.generate_role()
        return self._server_endpoints

511 512 513
    def is_worker(self):
        if not self._role_is_generated:
            self.generate_role()
514 515 516
        return self._role == Role.WORKER

    def is_server(self):
517 518
        if not self._role_is_generated:
            self.generate_role()
519 520 521
        return self._role == Role.SERVER

    def is_first_worker(self):
522 523
        if not self._role_is_generated:
            self.generate_role()
524 525 526
        return self._role == Role.WORKER and self._current_id == 0

    def worker_index(self):
527 528
        if not self._role_is_generated:
            self.generate_role()
529 530 531
        return self._current_id

    def server_index(self):
532 533
        if not self._role_is_generated:
            self.generate_role()
534 535 536
        return self._current_id

    def worker_num(self):
537 538
        if not self._role_is_generated:
            self.generate_role()
539
        return self._trainers_num
540 541


X
xujiaqi01 已提交
542 543 544
class GeneralRoleMaker(RoleMakerBase):
    """
    This role maker is for general use, you can set os.environ to customize:
T
tianshuo78520a 已提交
545 546
        PADDLE_PSERVERS_IP_PORT_LIST : all pservers' ip:port, separated by ','
        PADDLE_TRAINER_ENDPOINTS     : all trainers' ip:port, separated by ','
X
xujiaqi01 已提交
547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 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 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 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 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875
        TRAINING_ROLE                : TRAINER or PSERVER
        PADDLE_TRAINER_ID            : current trainer id (only for trainer),
                                       it is index in PADDLE_TRAINER_ENDPOINTS
        PADDLE_PSERVER_ID            : current pserver id (only for pserver)
                                       it is index in PADDLE_PSERVERS_IP_PORT_LIST
    """

    def __init__(self, **kwargs):
        super(RoleMakerBase, self).__init__()
        self._role_is_generated = False
        self._hdfs_name = kwargs.get("hdfs_name", "")
        self._hdfs_ugi = kwargs.get("hdfs_ugi", "")
        self._hdfs_path = kwargs.get("path", "")
        self._iface = self.__get_default_iface()
        # this environment variable can be empty
        self._prefix = os.getenv("SYS_JOB_ID", "")

    def generate_role(self):
        """
        generate role for general role maker
        """
        if not self._role_is_generated:
            eplist = os.environ["PADDLE_PSERVERS_IP_PORT_LIST"].split(",")
            training_role = os.environ["TRAINING_ROLE"]
            worker_endpoints = os.environ["PADDLE_TRAINER_ENDPOINTS"].split(",")
            trainers_num = len(worker_endpoints)
            if training_role not in ["TRAINER", "PSERVER"]:
                raise ValueError("TRAINING_ROLE must be PSERVER or TRAINER")
            if training_role == "TRAINER":
                role = Role.WORKER
                current_id = int(os.environ["PADDLE_TRAINER_ID"])
                self._node_type = 1
                self._cur_endpoint = worker_endpoints[current_id]
                gloo = fluid.core.Gloo()
                gloo.init(current_id,
                          len(worker_endpoints),
                          self._hdfs_path.rstrip("/") + "/trainer",
                          self._hdfs_name, self._hdfs_ugi, self._iface,
                          self._prefix)
                self._node_type_comm = gloo
            elif training_role == "PSERVER":
                role = Role.SERVER
                if os.environ.get("PADDLE_PSERVER_ID") is not None:
                    current_id = int(os.environ["PADDLE_PSERVER_ID"])
                    cur_endpoint = eplist[current_id]
                else:
                    # this is for compatible with paddlecloud
                    cur_ip = os.environ["POD_IP"]
                    cur_port = os.environ["PADDLE_PORT"]
                    cur_endpoint = ":".join([cur_ip, cur_port])
                    current_id = eplist.index(cur_endpoint)
                self._node_type = 0
                self._cur_endpoint = cur_endpoint
                gloo = fluid.core.Gloo()
                gloo.init(current_id,
                          len(eplist),
                          self._hdfs_path.rstrip("/") + "/pserver",
                          self._hdfs_name, self._hdfs_ugi, self._iface,
                          self._prefix)
                self._node_type_comm = gloo

            gloo = fluid.core.Gloo()
            all_list = worker_endpoints + eplist
            gloo.init(
                all_list.index(self._cur_endpoint),
                len(all_list),
                self._hdfs_path.rstrip("/") + "/all", self._hdfs_name,
                self._hdfs_ugi, self._iface, self._prefix)
            self._all_comm = gloo
            self._trainers_num = trainers_num
            self._server_endpoints = eplist
            self._role = role
            self._current_id = current_id
            self._rank = all_list.index(self._cur_endpoint)
            self._size = len(all_list)
            self._worker_endpoints = worker_endpoints
            self._role_is_generated = True

    def all_gather(self, input):
        """
        all gather between trainers and pservers

        Args:
            input(int|float): input value

        Returns:
            return a list of values
        """
        return self._all_gather(input)

    def all_reduce_worker(self, input, output, mode="sum"):
        """
        all reduce between trainers if current role is TRAINER,
        only support array of one dim.

        Args:
            input(list/numpy.array): array of one dim
            output(list/numpy.array): array of one dim
            mode(str): "sum" or "min" or "max"
        """
        if not self.is_worker():
            return
        self._all_reduce(input, output, mode)

    def barrier_worker(self):
        """
        barrier between trainers if current role is TRAINER
        """
        self._barrier_worker()

    def barrier_all(self):
        """
        barrier between trainers if current role is PSERVER
        """
        self._barrier_all()

    def get_local_endpoint(self):
        """
        get local endpoint of current process
        """
        if not self._role_is_generated:
            self.generate_role()
        return self._cur_endpoint

    def get_trainer_endpoints(self):
        """
        get endpoint of all trainers
        """
        if not self._role_is_generated:
            self.generate_role()
        return self._worker_endpoints

    def get_pserver_endpoints(self):
        """
        get endpoint of all pservers
        """
        if not self._role_is_generated:
            self.generate_role()
        return self._server_endpoints

    def is_worker(self):
        """
        whether current process is worker
        """
        if not self._role_is_generated:
            self.generate_role()
        return self._role == Role.WORKER

    def is_server(self):
        """
        whether current process is server
        """
        if not self._role_is_generated:
            self.generate_role()
        return self._role == Role.SERVER

    def is_first_worker(self):
        """
        whether current process is worker of rank 0
        """
        if not self._role_is_generated:
            self.generate_role()
        return self._role == Role.WORKER and self._current_id == 0

    def worker_index(self):
        """
        get index of current worker
        """
        if not self._role_is_generated:
            self.generate_role()
        return self._current_id

    def server_index(self):
        """
        get index of current server
        """
        if not self._role_is_generated:
            self.generate_role()
        return self._current_id

    def worker_num(self):
        """
        retrun the current number of worker
        """
        if not self._role_is_generated:
            self.generate_role()
        return self._worker_num()

    def server_num(self):
        """
        return the current number of server
        """
        if not self._role_is_generated:
            self.generate_role()
        return self._server_num()

    def _barrier_worker(self):
        """
        barrier all workers in current distributed job
        """
        if not self._role_is_generated:
            self.generate_role()
        if self.is_worker():
            self._node_type_comm.barrier()

    def _barrier_all(self):
        """
        barrier all workers and servers in current distributed job
        """
        if not self._role_is_generated:
            self.generate_role()
        self._all_comm.barrier()

    def _barrier_server(self):
        """
        barrier all servers in current distributed job
        """
        if not self._role_is_generated:
            self.generate_role()
        if self.is_server():
            self._node_type_comm.barrier()

    def _worker_num(self):
        """
        return the current number of worker
        """
        if not self._role_is_generated:
            self.generate_role()
        return self._trainers_num

    def _server_num(self):
        """
        return the current number of server
        """
        if not self._role_is_generated:
            self.generate_role()
        return len(self._server_endpoints)

    def _finalize(self):
        """Default do nothing."""
        pass

    def _all_reduce(self, input, output, mode="sum"):
        """
        all reduce between all workers

        Args:
            input(list|numpy.array): array of one dim
            output(list|numpy.array): array of one dim
            mode(str): "sum" or "min" or "max"
        """
        if not self._role_is_generated:
            self.generate_role()
        input_list = [i for i in input]
        ans = self._node_type_comm.all_reduce(input_list, mode)
        for i in range(len(ans)):
            output[i] = ans[i]

    def _all_gather(self, obj):
        """
        gather between all workers and pservers
        """
        if not self._role_is_generated:
            self.generate_role()
        self._barrier_all()
        return self._all_comm.all_gather(obj)

    def _worker_gather(self, obj):
        """
        gather between all workers
        """
        if not self._role_is_generated:
            self.generate_role()
        if not self.is_worker():
            return None
        self._barrier_worker()
        return self._node_type_comm.all_gather(obj)

    def _get_rank(self):
        """
        get current rank in all workers and pservers
        """
        if not self._role_is_generated:
            self.generate_role()
        return self._rank

    def _get_size(self):
        """
        get total num of all workers and pservers
        """
        if not self._role_is_generated:
            self.generate_role()
        return self._size

    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"


876
class UserDefinedRoleMaker(RoleMakerBase):
X
xujiaqi01 已提交
877 878 879 880 881 882
    """
    UserDefinedRoleMaker is designed for worker and server assignment
    under manual. Typically, a worker and a server node will be appointed
    on each physical node, It can be assign by user.
    """

883 884
    def __init__(self,
                 current_id=0,
T
tangwei12 已提交
885 886 887
                 role=Role.WORKER,
                 worker_num=0,
                 server_endpoints=None):
888 889
        super(UserDefinedRoleMaker, self).__init__()

890 891 892 893 894 895 896
        if not isinstance(server_endpoints, list):
            raise TypeError("server_endpoints must be as string list")
        elif len(server_endpoints) <= 0:
            raise ValueError(
                "the length of server_endpoints list must be greater than 0")
        elif len(server_endpoints) != len(set(server_endpoints)):
            raise ValueError("server_endpoints can't have duplicate elements")
897
        else:
898 899 900 901 902 903
            for server_endpoint in server_endpoints:
                if not isinstance(server_endpoint, str):
                    raise TypeError(
                        "every element in server_endpoints list must be as string"
                    )
            self._server_endpoints = server_endpoints
904

T
tangwei12 已提交
905
        if role != Role.WORKER and role != Role.SERVER:
906 907 908 909
            raise TypeError("role must be as Role")
        else:
            self._role = role

910 911 912 913 914 915 916 917 918 919 920 921 922
        if not isinstance(current_id, int):
            raise TypeError("current_id must be as int")
        else:
            if current_id < 0:
                raise ValueError(
                    "current_id must be greater than or equal to 0")
            elif self._role == Role.SERVER and current_id >= len(
                    server_endpoints):
                raise ValueError(
                    "if role is Role.SERVER, current_id must be less than or equal to len(server_endpoints) - 1"
                )
            self._current_id = current_id

923 924 925
        if not isinstance(worker_num, int):
            raise TypeError("worker_num must be as int")
        else:
926 927
            if worker_num <= 0:
                raise ValueError("worker_num must be greater than 0")
928 929
            self._worker_num = worker_num

930 931 932
    def generate_role(self):
        self._role_is_generated = True

T
tangwei12 已提交
933 934 935 936 937
    def is_worker(self):
        return self._role == Role.WORKER

    def is_server(self):
        return self._role == Role.SERVER
938

T
tangwei12 已提交
939 940
    def is_first_worker(self):
        return self._role == Role.WORKER and self._current_id == 0
941

T
tangwei12 已提交
942 943
    def worker_index(self):
        return self._current_id
944

T
tangwei12 已提交
945 946
    def server_index(self):
        return self._current_id
947 948 949

    def worker_num(self):
        return self._worker_num
950 951 952


class UserDefinedCollectiveRoleMaker(RoleMakerBase):
X
xujiaqi01 已提交
953 954 955 956 957
    """
    UserDefinedCollectiveRoleMaker is designed for worker assignment
    under manual for collective mode.
    """

958 959 960
    def __init__(self, current_id=0, worker_endpoints=None):
        super(UserDefinedCollectiveRoleMaker, self).__init__()

961 962 963 964 965 966 967 968 969 970 971 972 973 974 975
        if not isinstance(worker_endpoints, list):
            raise TypeError("worker_endpoints must be as string list")
        elif len(worker_endpoints) <= 0:
            raise ValueError(
                "the length of worker_endpoints list must be greater than 0")
        elif len(worker_endpoints) != len(set(worker_endpoints)):
            raise ValueError("worker_endpoints can't have duplicate elements")
        else:
            for worker_endpoint in worker_endpoints:
                if not isinstance(worker_endpoint, str):
                    raise TypeError(
                        "every element in worker_endpoints list must be as string"
                    )
            self._worker_endpoints = worker_endpoints

976 977 978 979
        if not isinstance(current_id, int):
            raise TypeError("current_id must be as int")
        else:
            if current_id < 0:
980 981 982 983 984 985
                raise ValueError(
                    "current_id must be greater than or equal to 0")
            elif current_id >= len(worker_endpoints):
                raise ValueError(
                    "current_id must be less than or equal to len(worker_endpoints) - 1"
                )
986 987 988 989
            self._current_id = current_id

        self._worker_num = len(self._worker_endpoints)

990 991 992
    def generate_role(self):
        self._role_is_generated = True

993 994 995 996 997 998 999 1000 1001 1002 1003
    def is_worker(self):
        return True

    def is_first_worker(self):
        return self._current_id == 0

    def worker_index(self):
        return self._current_id

    def worker_num(self):
        return self._worker_num