github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/scripts/meta_transfer_as_rename2workload.py (about) 1 #!/usr/bin/env python3 2 # -*- coding: utf-8 -*- 3 import argparse 4 import functools 5 import json 6 import os 7 8 import etcd3 9 import sys 10 11 check_conflict = False 12 dry_run = False 13 override = False 14 15 16 def remove_prefix(s, prefix): 17 return s[len(prefix):].lstrip('/') if s.startswith(prefix) else s 18 19 20 def range_prefix(meta, obj_prefix, fn): 21 etcd = meta.etcd 22 orig_prefix = os.path.join(meta.orig_root_prefix, obj_prefix) 23 range_start = orig_prefix 24 range_end = etcd3.utils.increment_last_byte( 25 etcd3.utils.to_bytes(range_start) 26 ) 27 28 while True: 29 range_request = etcd3.etcdrpc.RangeRequest() 30 range_request.key = etcd3.utils.to_bytes(range_start) 31 range_request.keys_only = False 32 range_request.range_end = etcd3.utils.to_bytes(range_end) 33 range_request.sort_order = etcd3.etcdrpc.RangeRequest.ASCEND 34 range_request.sort_target = etcd3.etcdrpc.RangeRequest.KEY 35 range_request.serializable = True 36 range_request.limit = 1000 37 38 range_response = etcd.kvstub.Range( 39 range_request, 40 etcd.timeout, 41 credentials=etcd.call_credentials, 42 metadata=etcd.metadata, 43 ) 44 45 for kv in range_response.kvs: 46 orig_key = kv.key.decode('utf-8') 47 objname = remove_prefix(orig_key, orig_prefix) 48 new_key = fn(objname, kv.value.decode('utf-8')) 49 if not new_key: 50 continue 51 52 if check_conflict: 53 if etcd.get(new_key)[0]: 54 print('conflict: new key exists! original: %s, new: %s' % (orig_key, new_key)) 55 else: 56 print('convert %s to %s' % (orig_key, new_key)) 57 58 if not range_response.more: 59 break 60 61 range_start = etcd3.utils.increment_last_byte(kv.key) 62 63 64 def exists(client, key): 65 if override: 66 return False 67 return client.get(key)[0] is not None 68 69 70 class Pod(object): 71 72 def __init__(self, meta, podname=None): 73 """Initializes a pod transfer.""" 74 self.meta = meta 75 self.pod_prefix = 'pod/info' 76 self.range_prefix = functools.partial(range_prefix, self.meta) 77 self.podname = podname 78 self.exists = functools.partial(exists, self.meta.etcd) 79 80 def trans(self): 81 self.range_prefix(self.pod_prefix, self._trans) 82 83 def _trans(self, podname, orig_value): 84 # Only trans the specific pod if any. 85 if self.podname and self.podname != podname: 86 return 87 88 new_key = os.path.join(self.meta.new_root_prefix, self.pod_prefix, podname) 89 if not dry_run and not self.exists(new_key): 90 self.meta.etcd.put(new_key, orig_value) 91 return new_key 92 93 94 class Node(object): 95 96 def __init__(self, meta, podname): 97 """Initializes a node transfer.""" 98 self.meta = meta 99 self.info_prefix = 'node' 100 self.range_prefix = functools.partial(range_prefix, self.meta, self.info_prefix) 101 self.pod_nodes = {} 102 self.podname = podname 103 self.exists = functools.partial(exists, self.meta.etcd) 104 105 def trans(self): 106 self.range_prefix(self._trans_pod) 107 self.range_prefix(self._trans_info) 108 self.range_prefix(self._trans_cert) 109 self.range_prefix(self._trans_workload) 110 return self.pod_nodes.keys() 111 112 def _trans_info(self, nodename, orig_value): 113 # skipping extra info. 114 if ':' in nodename: 115 return 116 117 if not self._belong_pod(nodename): 118 return 119 120 self.pod_nodes[nodename] = json.loads(orig_value) 121 122 new_key = os.path.join(self.meta.new_root_prefix, self.info_prefix, nodename) 123 if not dry_run and not self.exists(new_key): 124 self.meta.etcd.put(new_key, orig_value) 125 return new_key 126 127 def _trans_pod(self, node_pod_pair, orig_value): 128 # parsering node-pod-pair itself only 129 if ':pod/' not in node_pod_pair: 130 return 131 132 podname, _, nodename = node_pod_pair.partition(':pod/') 133 if not (podname and nodename): 134 raise ValueError('invalid podname or nodename for %s' % node_pod_pair) 135 136 # Only trans the specific pod if any. 137 if self.podname and self.podname != podname: 138 return 139 140 self.pod_nodes[nodename] = {} 141 142 new_key = os.path.join(self.meta.new_root_prefix, self.info_prefix, '%s:pod' % podname, nodename) 143 if not dry_run and not self.exists(new_key): 144 self.meta.etcd.put(new_key, orig_value) 145 return new_key 146 147 def _trans_cert(self, cert_key, orig_value): 148 nodename, _, cert_type = cert_key.partition(':') 149 if not self._belong_pod(nodename): 150 return 151 152 # parsering orig_key which ends with :ca, :cert, :key only 153 if cert_type not in ('ca', 'cert', 'key'): 154 return 155 156 new_key = os.path.join(self.meta.new_root_prefix, self.info_prefix, '%s:%s' % (nodename, cert_type)) 157 if not dry_run and not self.exists(new_key): 158 self.meta.etcd.put(new_key, orig_value) 159 return new_key 160 161 def _trans_workload(self, node_wrk_pair, orig_value): 162 nodename, _, wrk_id = node_wrk_pair.partition(':containers/') 163 # parsering orig_key which belongs node-workload pair only. 164 if not (nodename and wrk_id): 165 return 166 167 if not self._belong_pod(nodename): 168 return 169 170 new_key = os.path.join(self.meta.new_root_prefix, self.info_prefix, '%s:workloads' % nodename, wrk_id) 171 wrk = Workload.conv(orig_value, self) 172 if not dry_run and not self.exists(new_key): 173 self.meta.etcd.put(new_key, json.dumps(wrk)) 174 return new_key 175 176 def _belong_pod(self, nodename): 177 if not self.podname: 178 return True 179 return nodename in self.pod_nodes.keys() 180 181 def get_numa_node(self, cpumap, nodename): 182 """Ref core types/core.go GetNUMANode func.""" 183 numa_node_id = "" 184 185 node = self.pod_nodes.get(nodename) 186 if not node: 187 raise ValueError('invalid nodename %s' % nodename) 188 189 numa = node.get('numa') 190 if not numa: 191 return numa_node_id 192 193 for cpu_id in cpumap: 194 mem_node = numa.get(cpu_id) 195 if not mem_node: 196 continue 197 198 if numa_node_id == '': 199 numa_node_id = mem_node 200 elif numa_node_id != mem_node: 201 numa_node_id = '' 202 203 return numa_node_id 204 205 206 class Workload(object): 207 208 def __init__(self, meta, node_transfer, podname=None, pod_nodes=None): 209 """Initializes a workload transfer.""" 210 self.meta = meta 211 self.container_prefix = 'containers' 212 self.wrk_prefix = 'workloads' 213 self.deploy_prefix = 'deploy' 214 self.range_prefix = functools.partial(range_prefix, self.meta) 215 self.node_transfer = node_transfer 216 self.podname = podname 217 self.pod_nodes = pod_nodes 218 self.workloads = set() 219 self.exists = functools.partial(exists, self.meta.etcd) 220 221 def trans(self): 222 self.range_prefix(self.deploy_prefix, self._trans_deploy) 223 self.range_prefix(self.container_prefix, self._trans_container) 224 225 def _trans_container(self, wrk_id, orig_value): 226 # skipping if the workload doesn't belong the specific pod. 227 if not self.podname or wrk_id not in self.workloads: 228 return 229 230 new_key = os.path.join(self.meta.new_root_prefix, self.wrk_prefix, wrk_id) 231 wrk = self.conv(orig_value, self.node_transfer) 232 if not dry_run and not self.exists(new_key): 233 self.meta.etcd.put(new_key, json.dumps(wrk)) 234 return new_key 235 236 def _trans_deploy(self, deploy_key, orig_value): 237 parts = deploy_key.split('/') 238 if len(parts) != 4: 239 print('invalid deploy key: %s' % deploy_key) 240 return 241 242 appname, entrypoint, nodename, wrk_id = parts 243 if not self._belong_pod(nodename): 244 return 245 246 self.workloads.update({wrk_id}) 247 248 new_key = os.path.join(self.meta.new_root_prefix, self.deploy_prefix, appname, entrypoint, nodename, wrk_id) 249 wrk = self.conv(orig_value, self.node_transfer) 250 if not dry_run and not self.exists(new_key): 251 self.meta.etcd.put(new_key, json.dumps(wrk)) 252 253 return new_key 254 255 def _belong_pod(self, nodename): 256 if not self.podname: 257 return True 258 return nodename in self.pod_nodes 259 260 @classmethod 261 def conv(cls, orig_value, node_transfer): 262 263 def delete(*keys): 264 for k in keys: 265 try: 266 del dic[k] 267 except KeyError: 268 pass 269 270 del_keys = set() 271 272 def get(new_field, orig_field, transit_field, default=None): 273 value = None 274 275 # don't use dic.get(new_field, dic[orig_field]), 276 # due to there isn't orig_field but has new_field. 277 if new_field in dic: 278 value = dic[new_field] 279 elif transit_field in dic: 280 value = dic[transit_field] 281 else: 282 if default is None: 283 value = dic[orig_field] 284 else: 285 value = dic.get(orig_field, default) 286 287 del_keys.update({orig_field, transit_field}) 288 289 return value 290 291 dic = json.loads(orig_value) 292 dic['cpu'] = get('CPU', 'cpu', 'CPU') 293 294 dic.update(dict( 295 create_time=1553990400, 296 cpu_quota_request=get('cpu_quota_request', 'quota', 'CPUQuotaRequest'), 297 cpu_quota_limit=get('cpu_quota_limit', 'quota', 'CPUQuotaLimit'), 298 memory_request=get('memory_request', 'memory', 'MemoryRequest'), 299 memory_limit=get('memory_limit', 'memory', 'MemoryLimit'), 300 volume_request=get('volume_request', 'volumes', 'VolumeRequest', default=[]), 301 volume_limit=get('volume_limit', 'volumes', 'VolumeLimit', default=[]), 302 volume_plan_request=get('volume_plan_request', 'volume_plan', 'VolumePlanRequest', default={}), 303 volume_plan_limit=get('volume_plan_limit', 'volume_plan', 'VolumePlanLimit', default={}), 304 volume_changed=dic.get('volume_changed', False), 305 storage_request=get('storage_request', 'storage', 'StorageRequest'), 306 storage_limit=get('storage_limit', 'storage', 'StorageLimit'), 307 )) 308 309 dic['numa_node'] = '' 310 if dic['cpu'] and node_transfer: 311 numa_node = node_transfer.get_numa_node(dic['cpu'], dic['nodename']) 312 dic['numa_node'] = dic.get('NUMANode', numa_node) 313 314 # don't removing *cpu* from the original dict. 315 try: 316 del_keys.remove('cpu') 317 except KeyError: 318 pass 319 320 del_keys.update({'softlimit', 'VolumeChanged', 'NUMANode'}) 321 delete(*list(del_keys)) 322 323 return dic 324 325 326 class Transfer(object): 327 328 def __init__(self, etcd, orig_root_prefix, new_root_prefix): 329 """Initializes a transfer which includes common utilities.""" 330 self.etcd = etcd 331 self.orig_root_prefix = orig_root_prefix 332 self.new_root_prefix = new_root_prefix 333 334 def trans(self, podname=None): 335 Pod(self, podname).trans() 336 337 node_transfer = Node(self, podname) 338 nodes = node_transfer.trans() 339 340 Workload(self, node_transfer, podname, nodes).trans() 341 342 343 def getargs(): 344 ap = argparse.ArgumentParser() 345 ap.add_argument('-o', '--orig', dest='orig_root_prefix', help='original prefix', default='/eru') 346 ap.add_argument('-n', '--new', dest='new_root_prefix', help='new prefix', default='/eru2') 347 ap.add_argument('-p', '--pod', dest='podname') 348 ap.add_argument('--etcd-host', default='127.0.0.1') 349 ap.add_argument('--etcd-port', type=int, default=2379) 350 ap.add_argument('--dry-run', dest='dry_run', action='store_true', help='dry run, will not actually migrate') 351 ap.add_argument('--check-conflict', dest='check_conflict', action='store_true', 352 help='check conflict, checks if destination already has the key') 353 ap.add_argument('--override', dest='override', action='store_true', 354 help='if set, the value will be migrated anyway no matter if it already exists') 355 return ap.parse_args() 356 357 358 def connect_etcd(host, port): 359 return etcd3.client(host=host, port=port) 360 361 362 def main(): 363 args = getargs() 364 365 global dry_run, check_conflict, override 366 dry_run = args.dry_run 367 check_conflict = args.check_conflict 368 if check_conflict: 369 dry_run = True 370 override = args.override 371 372 etcd = connect_etcd(args.etcd_host, args.etcd_port) 373 trans = Transfer(etcd, args.orig_root_prefix, args.new_root_prefix) 374 trans.trans(args.podname) 375 return 0 376 377 378 if __name__ == '__main__': 379 sys.exit(main())