Dynamically modifying the `.electron_object.executor` from an imported electron does not always work as expected

[Andrew-S-Rosen]

Dynamically modifying the `.electron_object.executor` from an imported electron does not always work as expected

[santoshkumarradha]

Thanks for this @Andrew-S-Rosen ideally it should be so that all kwargs of electrons are properties as well that can be set post creation. I will update the issue as a bug/feature-enhancement as well.

Cc: @kessler-frost / @jimmylism

[santoshkumarradha]

CC: @cjao , any quick idea on why this would be the case for packages?

[wjcunningham7]

@Andrew-S-Rosen I think this may be due to how the objects are imported: https://stackoverflow.com/a/3536638

[cjao]

@Andrew-S-Rosen This is a fun question and boils down to how cloudpickle (Covalent's serialization mechanism) works. The key difference between your examples

(1) non-working

import covalent as ct
from mypackage import increment, subflow

@ct.lattice
def workflow(x, op):
    return subflow(x, op)

increment.electron_object.executor = "local"

ct.dispatch(workflow)(1, increment)

(2) working

import covalent as ct
from mypackage import increment

@ct.lattice
def workflow(x, op):
    return op(x)

increment.electron_object.executor = "local"

ct.dispatch(workflow)(1, increment)

is that (1) pickles the wrapped increment function -- as a parameter of the sub_workflow electron -- whereas in (2) only the underlying function is serialized, not the electron_object attribute. Because increment lives in an external module instead of __main__, Cloudpickle will serialize it by reference, not by value. Thus when the increment electron is unpickled in the executor that builds out sub_workflow, its attributes -- such as electron_object -- are whatever they are defined to be in mypackage.py.

Example:

mypackage.py

class Electron:
    def __init__(self):
        self.executor = "dask"

def my_func():
    pass

my_func.electron_object = Electron()

main.py

import cloudpickle
import mypackage

mypackage.my_func.electron_object.executor = "local"

with open("my_func.pkl", "wb") as f:
    cloudpickle.dump(mypackage.my_func, f)

exec.py

import pickle

with open("my_func.pkl", "rb") as f:
    my_func = pickle.load(f)
    print(my_func.electron_object.executor)

casey@AgnostiqHQ tests$ python main.py  && python exec.py
dask

But now if we instruct cloudpickle to serialize the mypackage module by value:

main.py

import cloudpickle
import mypackage

mypackage.my_func.electron_object.executor = "local"

cloudpickle.register_pickle_by_value(mypackage)

with open("my_func.pkl", "wb") as f:
    cloudpickle.dump(mypackage.my_func, f)

we get

casey@AgnostiqHQ tests$ python main.py  && python exec.py
local

[Andrew-S-Rosen]

@wjcunningham7, @cjao: Thank you both for the very insightful replies!

Because increment lives in an external module instead of main, Cloudpickle will serialize it by reference, not by value.

That's the main takeaway for sure. I didn't realize that --- subtle! In any case, this isn't a major issue, but I did want to flag it as a potential "gotcha." With this response, I feel pretty comfortable knowing what's going on now. Feel free to close if you wish!