def clean_ensure_indexes(run, fields, coll):
if run:
created = ensure_indexes(fields, [coll])
if created:
indexes = ", ".join(created[coll.full_name])
logger.info(
f"Created the following index(es) on {coll.full_name}:\n{indexes}"
)
else:
logger.info("All indexes already created.")
else:
fields_list = ", ".join(fields)
logger.info(f"Would create/ensure the following index(es) on "
f"{coll.full_name}:\n{fields_list}")
def parse(task_ids, snl_metas, nproc, store_volumetric_data): # noqa: C901
"""Parse VASP launchers into tasks"""
ctx = click.get_current_context()
if "CLIENT" not in ctx.obj:
raise EmmetCliError("Use --spec to set target DB for tasks!")
run = ctx.parent.parent.params["run"]
nmax = ctx.parent.params["nmax"]
directory = ctx.parent.params["directory"].rstrip(os.sep)
tag = os.path.basename(directory)
target = ctx.obj["CLIENT"]
snl_collection = target.db.snls_user
logger.info(
f"Connected to {target.collection.full_name} with {target.collection.count()} tasks."
)
ensure_indexes(
["task_id", "tags", "dir_name", "retired_task_id"], [target.collection]
)
chunk_size = math.ceil(nmax / nproc)
if nproc > 1 and nmax <= chunk_size:
nproc = 1
logger.warning(
f"nmax = {nmax} but chunk size = {chunk_size} -> sequential parsing."
)
pool = multiprocessing.Pool(processes=nproc)
gen = VaspDirsGenerator()
iterator = iterator_slice(gen, chunk_size) # process in chunks
queue = deque()
count = 0
sep_tid = None
if task_ids:
with open(task_ids, "r") as f:
task_ids = json.load(f)
else:
# reserve list of task_ids to avoid collisions during multiprocessing
# insert empty doc with max ID + 1 into target collection for parallel SLURM jobs
# NOTE use regex first to reduce size of distinct below 16MB
q = {"task_id": {"$regex": r"^mp-\d{7,}$"}}
all_task_ids = [
t["task_id"] for t in target.collection.find(q, {"_id": 0, "task_id": 1})
]
if not all_task_ids:
all_task_ids = target.collection.distinct("task_id")
next_tid = max(int(tid.split("-")[-1]) for tid in all_task_ids) + 1
lst = [f"mp-{next_tid + n}" for n in range(nmax)]
task_ids = chunks(lst, chunk_size)
if run:
sep_tid = f"mp-{next_tid + nmax}"
target.collection.insert({"task_id": sep_tid})
logger.info(f"Inserted separator task with task_id {sep_tid}.")
logger.info(f"Reserved {len(lst)} task ID(s).")
else:
logger.info(f"Would reserve {len(lst)} task ID(s).")
sep_snlid = None
if snl_metas:
with open(snl_metas, "r") as f:
snl_metas = json.load(f)
# reserve list of snl_ids to avoid collisions during multiprocessing
# insert empty doc with max ID + 1 into target collection for parallel SLURM jobs
all_snl_ids = snl_collection.distinct("snl_id")
prefixes = set()
next_snlid = -1
for snlid in all_snl_ids:
prefix, index = snlid.split("-", 1)
index = int(index)
prefixes.add(prefix)
if index > next_snlid:
next_snlid = index
next_snlid += 1
prefix = prefixes.pop() # NOTE use the first prefix found
nsnls = len(snl_metas)
for n, launcher in enumerate(snl_metas):
snl_id = f"{prefix}-{next_snlid + n}"
snl_metas[launcher]["snl_id"] = snl_id
if run:
sep_snlid = f"{prefix}-{next_snlid + nsnls}"
snl_collection.insert({"snl_id": sep_snlid})
logger.info(f"Inserted separator SNL with snl_id {sep_snlid}.")
logger.info(f"Reserved {nsnls} SNL ID(s).")
else:
logger.info(f"Would reserve {nsnls} SNL ID(s).")
while iterator or queue:
try:
args = [next(iterator), tag, task_ids, snl_metas]
queue.append(pool.apply_async(parse_vasp_dirs, args))
except (StopIteration, TypeError):
iterator = None
while queue and (len(queue) >= pool._processes or not iterator):
process = queue.pop()
process.wait(1)
if not process.ready():
queue.append(process)
else:
count += process.get()
pool.close()
if run:
logger.info(
f"Successfully parsed and inserted {count}/{gen.value} tasks in {directory}."
)
if sep_tid:
target.collection.remove({"task_id": sep_tid})
logger.info(f"Removed separator task {sep_tid}.")
if sep_snlid:
snl_collection.remove({"snl_id": sep_snlid})
logger.info(f"Removed separator SNL {sep_snlid}.")
else:
logger.info(f"Would parse and insert {count}/{gen.value} tasks in {directory}.")
return ReturnCodes.SUCCESS if count and gen.value else ReturnCodes.WARNING
def parse(task_ids, nproc):
"""Parse VASP launchers into tasks"""
ctx = click.get_current_context()
if "CLIENT" not in ctx.obj:
raise EmmetCliError(f"Use --spec to set target DB for tasks!")
run = ctx.parent.parent.params["run"]
nmax = ctx.parent.params["nmax"]
directory = ctx.parent.params["directory"].rstrip(os.sep)
tag = os.path.basename(directory)
target = ctx.obj["CLIENT"]
logger.info(
f"Connected to {target.collection.full_name} with {target.collection.count()} tasks."
)
ensure_indexes(["task_id", "tags", "dir_name", "retired_task_id"],
[target.collection])
chunk_size = math.ceil(nmax / nproc)
if nproc > 1 and nmax <= chunk_size:
nproc = 1
logger.warning(
f"nmax = {nmax} but chunk size = {chunk_size} -> sequential parsing."
)
pool = multiprocessing.Pool(processes=nproc)
gen = VaspDirsGenerator()
iterator = iterator_slice(gen, chunk_size) # process in chunks
queue = deque()
count = 0
sep_tid = None
if task_ids:
with open(task_ids, "r") as f:
task_ids = json.load(f)
else:
# reserve list of task_ids to avoid collisions during multiprocessing
# insert empty doc with max ID + 1 into target collection for parallel SLURM jobs
all_task_ids = target.collection.distinct("task_id")
next_tid = max(int(tid.split("-")[-1]) for tid in all_task_ids) + 1
lst = [f"mp-{next_tid + n}" for n in range(nmax)]
if run:
sep_tid = f"mp-{next_tid + nmax}"
target.collection.insert({"task_id": sep_tid})
logger.info(f"Inserted separator task with task_id {sep_tid}.")
task_ids = chunks(lst, chunk_size)
logger.info(f"Reserved {len(lst)} task ID(s).")
while iterator or queue:
try:
args = [next(iterator), tag, task_ids]
queue.append(pool.apply_async(parse_vasp_dirs, args))
except (StopIteration, TypeError):
iterator = None
while queue and (len(queue) >= pool._processes or not iterator):
process = queue.pop()
process.wait(1)
if not process.ready():
queue.append(process)
else:
count += process.get()
pool.close()
if run:
logger.info(
f"Successfully parsed and inserted {count}/{gen.value} tasks in {directory}."
)
if sep_tid:
target.collection.remove({"task_id": sep_tid})
logger.info(f"Removed separator task {sep_tid}.")
else:
logger.info(
f"Would parse and insert {count}/{gen.value} tasks in {directory}."
)
return ReturnCodes.SUCCESS if count and gen.value else ReturnCodes.WARNING