blob: 303427ac9536f72fbf250c95c65b68b7d54b0cfd [file] [log] [blame]
Jon Hall9ebd1bd2016-04-19 01:37:17 -07001#!/usr/bin/env python
2"""
3 Generate the partitions json file from the $OC* environment variables
4
5 Usage: onos-gen-partitions output_file [num_nodes] [-e]
6 If output file is not provided, the json is written to stdout.
7"""
8
9from os import environ
10from collections import deque, OrderedDict
11import re
12import json
13import sys
14import hashlib
15
16convert = lambda text: int(text) if text.isdigit() else text.lower()
17alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key)]
18
19def get_OC_vars():
20 vars = []
21 for var in environ:
22 if re.match(r"OC[0-9]+", var):
23 vars.append(var)
24 return sorted(vars, key=alphanum_key)
25
26def get_nodes(vars, port=9876):
27 node = lambda k: { 'id': k, 'ip': k, 'port': port }
28 return [ node(environ[v]) for v in vars ]
29
Jon Hall9ebd1bd2016-04-19 01:37:17 -070030def generate_extended_partitions_scaling(nodes, k, partitions=3, equal=False):
31 l = deque(nodes)
32 perms = []
33 for i in range(1, partitions + 1):
34 if equal:
35 members = list(l)
36 else:
37 members = list(l)[:k]
38
39 part = {
40 'id': i,
41 'members': members
42 }
43 perms.append(part)
44 l.rotate(-2)
45 return perms
46
47if __name__ == '__main__':
48 vars = get_OC_vars()
49 # NOTE: likely prone to errors
50 nodes = get_nodes(vars)
51 num = None
52 if len(sys.argv) >= 3:
53 num = int(sys.argv[2])
54 try:
55 equal = "-e" in sys.argv[3]
56 except:
57 equal = False
58 if num:
59 nodes = nodes[:num]
60
Jon Hall9ebd1bd2016-04-19 01:37:17 -070061 extended_partitions = generate_extended_partitions_scaling([v.get('id') for v in nodes],
62 3, equal=equal)
63 partitions = []
Jon Hall9ebd1bd2016-04-19 01:37:17 -070064 partitions.extend(extended_partitions)
65 name = hash("HAScaling")
66 data = {
67 'name': name,
68 'nodes': nodes,
69 'partitions': partitions
70 }
71 output = json.dumps(data, indent=4)
72
73 if len(sys.argv) >= 2:
74 filename = sys.argv[1]
75 with open(filename, 'w') as f:
76 f.write(output)
77 else:
78 print output