How to use the tasks collection in Mongodb to build the materials collection with MaterialsBuilder

[LuckRookie]

First, I optimized a structure using VASP. Then I converted it into a TaskDoc using atomate2’s VaspDrone. Finally, I used model_dump to convert it into a dict and inserted it into a MongoDB database. The following is my insertion code

from pathlib import Path
import uuid
from maggma.stores import MongoStore
from maggma.stores import GridFSStore
from emmet.core.vasp.calculation import StoreTrajectoryOption
from atomate2.vasp.drones import VaspDrone

calypso_store = MongoStore(
    database="emmet",
    collection_name="tasks",
    host="127.0.0.1",
    port=27017,
    username="user",
    password="password",
    key="task_id", 
)
calypso_store.connect()


vasp_dir = Path("E:\\database\\xxx") 
set_task = {
"store_trajectory":StoreTrajectoryOption.FULL,
"parse_dos":False,
"parse_bandstructure":False,
}
drone = VaspDrone(**set_task)

doc = drone.assimilate(vasp_dir)

if not getattr(doc, "task_id", None):
    doc = doc.model_copy(update={"task_id": f"{uuid.uuid4()}"})

payload = doc.model_dump()
payload["state"] = payload["state"]["value"]
calypso_store.update([payload])

This gave me a task store containing the TaskDoc. I then connected to this TaskDoc with MaterialsBuilder and ran run. However, MaterialsBuilder could not find any structures when executing the get_items function. I eventually discovered that MaterialsBuilder’s query uses temp_query["state"] = "successful", but the 'state' values in the task store I inserted are dictionaries rather than strings. Therefore, MaterialsBuilder cannot query them and cannot build the materials store.

I considered manually flattening the state value after model_dump, turning state into a string instead of a dictionary. However, I realized that if I have to manually make changes all the way through to the final SummaryBuilder, there would be an overwhelming number of manual adjustments required.

Therefore, I would like to ask: is my method of using model_dump incorrect? The data in the task store cannot interface with the subsequent pipeline.

Finally, I started building the materials store from the task store as suggested by the MaterialsProject

[tsmathis]

Hi @LuckRookie, I'm a little confused as to why your state field is a dictionary rather than just a string. TaskState is just a simple enum that should be serialized into the corresponding string prior to inserting into your mongo collection

emmet/emmet-core/emmet/core/vasp/task_valid.py

Lines 19 to 26 in 8a4fc37

	class TaskState(ValueEnum):
	"""
	VASP Calculation State
	"""
	SUCCESS = "successful"
	FAILED = "failed"
	ERROR = "error"

your snippet is ostensibly doing that here, right:
payload["state"] = payload["state"]["value"]

I am going to assume all of your workflows were successful since you are inserting them into mongo. If that's the case, can you just update your tasks collection and set all the state fields to the string "successful"?

[tsmathis]

Although you may need to be doing payload["state"] = payload["state"].value instead (attr access rather than dict-like access on the enum)

[LuckRookie]

Thank you for your reply. I'm very sorry that I didn't express it well. payload["state"] = payload["state"].value is the method I'm going to explore next. It was a mistake for me to put it in when pasting the code. The problem with my original text is that it runs without the code payload["state"] = payload["state"].value.

But after I used model_dump for the TaskDoc class, the 'state' field was dumped as a dictionary instead of a string. The dictionary type of 'state' caused the query of MaterialsBuilder to become invalid. This is my current problem.

The solution I can do at present, of course, is to serialize 'state', just like the code payload["state"] = payload["state"].value. But I think the huge system made by the Materials Project is sensitive to the modification of this kind of pinion. If I modify the 'state' field in this way, do I need to string all the enumeration types? This is a very serious problem that will affect the subsequent operation of the Builder, which is why I came to seek help.

My question is very simple: Is this modification after using model_dump reasonable? If it is reasonable, it indicates that the operation of the entire Builder must not only require the modification of 'state', but also other fields that need to be modified to meet the operation requirements of various builders. Where should I look for this part of information? If it is unreasonable, then what is the correct way to insert TaskDoc into MongoDB and ensure that the subsequent Builder does not report errors? I'm here to seek a suitable way to handle it.

Finally, thank you for your reply.

[tsmathis]

Yes it is reasonable to modify the dictionaries that result from model_dumping a pydantic model. There are actually quite a few data types that pydantic doesn't know how to serialize, and simply skips (this is an irrelevant detail for your case, but this is because those data types don't define a __get_pydantic_core_schema__ classmethod). Enums are a good example, there are also quite a few pymatgen classes as well that won't serialize correctly when calling model_dump.

Emmet has a utility function that can help with this though:

emmet/emmet-core/emmet/core/utils.py

Line 352 in 8a4fc37

def jsanitize(obj, strict=False, allow_bson=False):

Rather than mess with manually serializing every field, just do this:
payload = jsanitize(doc.model_dump())

[LuckRookie]

Wow, thank you for providing me with such a great idea. I wrote a code myself to serialize model_dump. However, before writing, I had doubts. Doesn't the official emmet require model_dump serialization? So far, it seems that I haven't found their code. Thank you very much for your help

[tsmathis]

if you look at some of the builder code you'll see jsanitize sprinkled everywhere to handle these types of things:

emmet/emmet-builders/emmet/builders/vasp/materials.py

Line 280 in 8a4fc37

return jsanitize([mat.model_dump() for mat in materials], allow_bson=True)

[LuckRookie]

It seems that there is still a lot for me to learn

[tsmathis]

(de)serialization can be tricky on its own. Adding in all the different data types in the different repositories in the Materials Project ecosystem doesn't make it easier.

Happy to help where we can with emmet related stuff