Skip to content

Commit

Permalink
[Fix] PythonJob not need aiida_workgraph on remote computer (#396)
Browse files Browse the repository at this point in the history
In the decorated function of PythonJob, we attached the WorkGraph's `Task` class to the function, this will be pickled together with the function, thus, when unpickling the function on the remote computer, it require `aiida-workgraph`. 

Thsi PR deletes the workgraph related attributes of the decorated func if exist.
  • Loading branch information
superstar54 authored Jan 28, 2025
1 parent 34e32bd commit 9febac0
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion src/aiida_workgraph/engine/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,10 @@ def prepare_for_python_task(task: dict, kwargs: dict, var_kwargs: dict) -> dict:
function_outputs.append(
{"name": output["name"], "identifier": output["identifier"]}
)

# delete workgraph related attributes of the func if exist
for attr in ["task", "tdata", "node"]:
if hasattr(func, attr):
delattr(func, attr)
inputs = prepare_pythonjob_inputs(
function=func,
function_inputs=function_inputs,
Expand Down

0 comments on commit 9febac0

Please sign in to comment.