diff --git a/emmet-cli/emmet/cli/utils.py b/emmet-cli/emmet/cli/utils.py index 17f73f0732..407397394d 100644 --- a/emmet-cli/emmet/cli/utils.py +++ b/emmet-cli/emmet/cli/utils.py @@ -18,6 +18,7 @@ from mongogrant.client import Client from atomate.vasp.drones import VaspDrone from pymongo.errors import DocumentTooLarge +from dotty_dict import dotty from emmet.core.utils import group_structures from emmet.cli import SETTINGS @@ -389,16 +390,23 @@ def parse_vasp_dirs(vaspdirs, tag, task_ids): try: target.insert_task(task_doc, use_gridfs=True) except DocumentTooLarge: - logger.warning(f"{name} Remove normalmode_eigenvecs and retry ...") - task_doc["calcs_reversed"][0]["output"].pop("normalmode_eigenvecs") - try: - target.insert_task(task_doc, use_gridfs=True) - except DocumentTooLarge: - logger.warning( - f"{name} Also remove force_constants and retry ..." - ) - task_doc["calcs_reversed"][0]["output"].pop("force_constants") - target.insert_task(task_doc, use_gridfs=True) + output = dotty(task_doc["calcs_reversed"][0]["output"]) + pop_keys = ["normalmode_eigenvecs", "force_constants", "outcar.onsite_density_matrices"] + + for k in pop_keys: + if k not in output: + continue + + logger.warning(f"{name} Remove {k} and retry ...") + output.pop(k) + try: + target.insert_task(task_doc, use_gridfs=True) + break + except DocumentTooLarge: + continue + else: + logger.warning(f"{name} failed to reduce document size") + continue if target.collection.count(query): shutil.rmtree(vaspdir) diff --git a/emmet-cli/requirements.txt b/emmet-cli/requirements.txt index b56269462a..4aed08becc 100644 --- a/emmet-cli/requirements.txt +++ b/emmet-cli/requirements.txt @@ -14,3 +14,4 @@ slurmpy==0.0.8 github3.py==1.3.0 hpsspy==0.5.1 multiprocessing-logging==0.3.1 +dotty-dict==1.3.0 diff --git a/emmet-cli/setup.py b/emmet-cli/setup.py index 92bf3d5a1b..5ca1631e91 100644 --- a/emmet-cli/setup.py +++ b/emmet-cli/setup.py @@ -21,6 +21,7 @@ "github3.py", "hpsspy", "multiprocessing-logging", + "dotty-dict", "emmet-core", ], license="modified BSD",