github.com/uber/kraken@v0.1.4/tools/bin/simulation/random_regular_graph.py (about)

     1  # Copyright (c) 2016-2019 Uber Technologies, Inc.
     2  #
     3  # Licensed under the Apache License, Version 2.0 (the "License");
     4  # you may not use this file except in compliance with the License.
     5  # You may obtain a copy of the License at
     6  #
     7  #     http://www.apache.org/licenses/LICENSE-2.0
     8  #
     9  # Unless required by applicable law or agreed to in writing, software
    10  # distributed under the License is distributed on an "AS IS" BASIS,
    11  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  # See the License for the specific language governing permissions and
    13  # limitations under the License.
    14  import json
    15  import random
    16  import sets
    17  
    18  import networkx as nx
    19  
    20  """
    21  Random regular graph with connection limit of 5:
    22   - 5000 peers, 500MB: 17 iterations
    23   - 1000 peers, 10GB: p50 294 iterations, p100 298 iterations (84% ~ 85% speed)
    24  """
    25  PEER_COUNT = 5000
    26  PIECE_COUNT = 125
    27  PIECE_TRANSMIT_LIMIT = 10  # Number of pieces uploaded/downloaded per iteration
    28  DEGREE = 5
    29  
    30  
    31  class Peer(object):
    32      def __init__(self, name, piece_count):
    33          self.name = name
    34          self.neighbors = sets.Set()
    35          self.pieces = [0]*piece_count
    36          self.completed = 0
    37          self.time = 0
    38  
    39          self.uploaded_current_turn = 0
    40          self.downloaded_current_turn = 0
    41  
    42      def connect(self, other):
    43          self.neighbors.add(other)
    44          other.neighbors.add(self)
    45  
    46      def done(self):
    47          return self.completed == len(self.pieces)
    48  
    49      def fetch_step(self, time):
    50          if self.done():
    51              return
    52  
    53          if self.downloaded_current_turn >= PIECE_TRANSMIT_LIMIT:
    54              return
    55  
    56          candidates = []
    57          for n in self.neighbors:
    58              if n.uploaded_current_turn >= PIECE_TRANSMIT_LIMIT:
    59                  continue
    60  
    61              for i in range(0, len(self.pieces)):
    62                  if n.uploaded_current_turn >= PIECE_TRANSMIT_LIMIT:
    63                      continue
    64  
    65                  if n.pieces[i] == 1 and self.pieces[i] == 0:
    66                      candidates.append((n, i))
    67  
    68          if len(candidates) == 0:
    69              return
    70  
    71          c = random.choice(candidates)
    72  
    73          self.pieces[c[1]] = 1
    74          self.completed += 1
    75          self.downloaded_current_turn += 1
    76          c[0].uploaded_current_turn += 1
    77  
    78          print ('Peer %s downloaded one piece from neighbor %s. Total completed: %d.' % (self.name, c[0].name, self.completed))
    79  
    80          if self.completed == len(self.pieces)-1:
    81              self.time = time
    82              print ('Peer %s finished downloading at time %d.' % (self.name, time))
    83  
    84      def fetch_cleanup(self):
    85          self.uploaded_current_turn = 0
    86          self.downloaded_current_turn = 0
    87  
    88  class PeerManager(object):
    89  
    90      def __init__(self):
    91          self.peers = []
    92  
    93          g = nx.random_regular_graph(DEGREE, PEER_COUNT)
    94          for n in g:
    95              peer = Peer(str(n), PIECE_COUNT)
    96              self.peers.append(peer)
    97  
    98          for e in g.edges():
    99              self.peers[e[0]].connect(self.peers[e[1]])
   100  
   101          for peer in self.peers:
   102              neighbors_str = ""
   103              for neighbor in peer.neighbors:
   104                  neighbors_str = neighbors_str + neighbor.name + "; "
   105              print ('Peer %s is connected to peers %s' % (peer.name, neighbors_str))
   106  
   107          # Set peer 0 to be the seeder.
   108          self.peers[0].pieces = [1]*PIECE_COUNT
   109          self.peers[0].completed = len(self.peers[0].pieces)
   110  
   111      def start(self):
   112          time = 0
   113          while True:
   114              print ('current time: %d.' % time)
   115              time += 1
   116  
   117              plan = []
   118              for p in self.peers:
   119                  if not p.done():
   120                      for j in range(0, PIECE_TRANSMIT_LIMIT):
   121                          plan.append(p)
   122              random.shuffle(plan)
   123              for p in plan:
   124                  p.fetch_step(time)
   125  
   126              for p in self.peers:
   127                  p.fetch_cleanup()
   128  
   129              done = True
   130              for p in self.peers:
   131                  if p.completed != len(p.pieces):
   132                      done = False
   133  
   134              if done:
   135                  break
   136  
   137              if time > 1000:
   138                  break
   139  
   140          print ('Done. Total time: %d.' % time)
   141  
   142  
   143  def main():
   144      peer_manager = PeerManager()
   145      peer_manager.start()
   146  
   147  if __name__== "__main__":
   148       main()