Use PickledFunction.serialize_function in decorator

aiidateam · Aug 25, 2024 · 6cbf441 · 6cbf441
1 parent 44f4262
commit 6cbf441
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/aiida_workgraph/decorator.py b/aiida_workgraph/decorator.py
@@ -246,7 +246,12 @@ def build_task_from_AiiDA(
             else outputs
         )
         # get the source code of the function
-        tdata["executor"] = PickledFunction(executor).executor
+        tdata["executor"] = {
+            "executor": pickle.dumps(executor),
+            "type": "function",
+            "is_pickle": True,
+        }
+        tdata["executor"].update(PickledFunction.serialize_function(executor))
         # tdata["executor"]["type"] = tdata["task_type"]
     # print("kwargs: ", kwargs)
     # add built-in sockets
@@ -495,9 +500,14 @@ def generate_tdata(
         "properties": properties,
         "inputs": _inputs,
         "outputs": task_outputs,
-        "executor": PickledFunction(func).executor,
         "catalog": catalog,
     }
+    tdata["executor"] = {
+        "executor": pickle.dumps(func),
+        "type": "function",
+        "is_pickle": True,
+    }
+    tdata["executor"].update(PickledFunction.serialize_function(func))
     if additional_data:
         tdata.update(additional_data)
     return tdata