def _do_alloc_ids(self, count, id_alloc_state):
"""Helper to allocate unique ids using the key/value store"""
batch = Config().batch
if not id_alloc_state:
id_alloc_state = []
else:
id_alloc_state = yaml.safe_load(id_alloc_state)
jsonschema.validate(id_alloc_state,
yaml.safe_load(id_allocation_schema))
num_ids_preclean = len(id_alloc_state)
# Cleanup completed jobs
try:
joblist = batch.list_all_jobs()
id_alloc_state = [
pk for pk in id_alloc_state if int(pk['batchid']) in joblist
]
except PcoccError:
pass
num_ids = len(id_alloc_state)
stray_ids = num_ids_preclean - num_ids
if stray_ids > 0:
logging.warning('Found %s leftover Ids, will try to cleanup',
stray_ids)
if num_ids + count > self.num_ids:
raise PcoccError('Not enough free ids in %s' % self.key_path)
id_indexes = []
i = 0
for allocated_id in sorted(id_alloc_state,
key=lambda x: x['pkey_index']):
while i < allocated_id['pkey_index'] and count > 0:
id_indexes.append(i)
i += 1
count -= 1
if count == 0:
break
i += 1
else:
id_indexes += [i for i in xrange(i, i + count)]
for i in id_indexes:
id_alloc_state.append({'pkey_index': i, 'batchid': batch.batchid})
return yaml.dump(id_alloc_state), id_indexes
def coll_alloc(self, count, master, key):
# Master allocates the ids and broadcasts to the others
coll_path = os.path.join('coll_alloc', self.key_path, key)
if Config().batch.node_rank == master:
ids = Config().batch.atom_update_key('global', self.key_path,
self._do_alloc_ids, count)
Config().batch.write_key('cluster', coll_path, yaml.dump(ids))
else:
ids = Config().batch.read_key('cluster',
coll_path,
blocking=True,
timeout=30)
ids = yaml.safe_load(ids)
return ids
def _do_free_ids(self, id_indexes, id_alloc_state):
"""Helper to free unique ids using the key/value store"""
id_alloc_state = yaml.safe_load(id_alloc_state)
jsonschema.validate(id_alloc_state,
yaml.safe_load(id_allocation_schema))
batchid = Config().batch.batchid
id_alloc_state[:] = [
allocated_id for allocated_id in id_alloc_state
if allocated_id['batchid'] != batchid
or allocated_id['pkey_index'] not in id_indexes
]
return yaml.dump(id_alloc_state), None
def free(self, ids):
return Config().batch.atom_update_key('global', self.key_path,
self._do_free_ids, ids)
def alloc(self, count):
return Config().batch.atom_update_key('global', self.key_path,
self._do_alloc_ids, count)