fix loading pythonjob task:

1) serialize and deserialize the properties when saving and loading the task
2) update the outputs of the task when loading the task.

aiida_workgraph/engine/workgraph.py

@@ -537,12 +537,23 @@ def get_task(self, name: str):
  task = PythonJob.from_dict(self.ctx._tasks[name])
  else:
  task = Task.from_dict(self.ctx._tasks[name])
+ # update task results
+ for output in task.outputs:
+ output.value = get_nested_dict(
+ self.ctx._tasks[name]["results"],
+ output.name,
+ default=output.value,
+ )
  return task
 
  def update_task(self, task: Task):
  """Update task in the context.
  This is used in error handlers to update the task parameters."""
- self.ctx._tasks[task.name]["properties"] = task.properties_to_dict()
+ from aiida_workgraph.utils import serialize_pythonjob_properties
+
+ tdata = task.to_dict()
+ serialize_pythonjob_properties(tdata)
+ self.ctx._tasks[task.name]["properties"] = tdata["properties"]
  self.reset_task(task.name)
 
  def get_task_state_info(self, name: str, key: str) -> str:

aiida_workgraph/task.py

@@ -112,11 +112,15 @@ def from_dict(cls, data: Dict[str, Any], task_pool: Optional[Any] = None) -> "Ta
  Returns:
  Node: An instance of Node initialized with the provided data."""
  from aiida_workgraph.tasks import task_pool
+ from aiida.orm.utils.serialize import deserialize_unsafe
 
  task = super().from_dict(data, node_pool=task_pool)
  task.context_mapping = data.get("context_mapping", {})
  task.waiting_on.add(data.get("wait", []))
- task.process = data.get("process", None)
+ process = data.get("process", None)
+ if process and isinstance(process, str):
+ process = deserialize_unsafe(process)
+ task.process = process
 
  return task
 

aiida_workgraph/utils/__init__.py

@@ -244,9 +244,12 @@ def get_pythonjob_data(tdata: Dict[str, Any]) -> Dict[str, Any]:
  if name == "computer":
  break
  if name in tdata["properties"]:
- tdata["properties"][name]["value"] = tdata["properties"][name][
- "value"
- ].value
+ value = tdata["properties"][name]["value"]
+ if isinstance(value, orm.Data):
+ value = value.value
+ elif value != {}:
+ raise ValueError(f"There something wrong with the input {name}")
+ tdata["properties"][name]["value"] = value
  return tdata
 
 
@@ -382,6 +385,27 @@ def get_or_create_code(
  return code
 
 
+def serialize_pythonjob_properties(task):
+ """Serialize the properties for PythonJob."""
+
+ from aiida_workgraph.orm.serializer import general_serializer
+
+ input_kwargs = []
+ for input in task["inputs"]:
+ if input["name"] == "_wait":
+ break
+ input_kwargs.append(input["name"])
+ for name in input_kwargs:
+ prop = task["properties"][name]
+ # if value is not None, not {}
+ if not (
+ prop["value"] is None
+ or isinstance(prop["value"], dict)
+ and prop["value"] == {}
+ ):
+ prop["value"] = general_serializer(prop["value"])
+
+
 def serialize_properties(wgdata):
  """Serialize the properties.
  Because we use yaml (aiida's serialize) to serialize the data and
@@ -390,27 +414,13 @@ def serialize_properties(wgdata):
  So, if a function is used as input, we needt to serialize the function.
 
  For PythonJob, serialize the function inputs."""
- from aiida_workgraph.orm.serializer import general_serializer
  from aiida_workgraph.orm.function_data import PickledLocalFunction
  import inspect
 
  for _, task in wgdata["tasks"].items():
  if task["metadata"]["node_type"].upper() == "PYTHONJOB":
  # get the names kwargs for the PythonJob, which are the inputs before _wait
- input_kwargs = []
- for input in task["inputs"]:
- if input["name"] == "_wait":
- break
- input_kwargs.append(input["name"])
- for name in input_kwargs:
- prop = task["properties"][name]
- # if value is not None, not {}
- if not (
- prop["value"] is None
- or isinstance(prop["value"], dict)
- and prop["value"] == {}
- ):
- prop["value"] = general_serializer(prop["value"])
+ serialize_pythonjob_properties(task)
  else:
  for _, prop in task["properties"].items():
  if inspect.isfunction(prop["value"]):