cli: improve doc size reduction @utf

materialsproject · Nov 5, 2020 · e65daae · e65daae
1 parent 3d069e0
commit e65daae
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 10 deletions.
diff --git 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
@@ -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
@@ -21,6 +21,7 @@
         "github3.py",
         "hpsspy",
         "multiprocessing-logging",
+        "dotty-dict",
         "emmet-core",
     ],
     license="modified BSD",