Question about mixing classes and lattices

[Andrew-S-Rosen]

Consider the following example:

from __future__ import annotations
from dataclasses import dataclass
import covalent as ct

@dataclass
class MyClass:
    executor: str | None = None
    
    @ct.electron(executor=executor)
    def add(self, val1: int, val2: int) -> int:
        return val1 + val2

@ct.lattice
def job(val1: int, val2: int) -> int:
    return MyClass(executor="cow").add(val1, val2)

dispatch_id = ct.dispatch(job)(1, 2)
result = ct.get_result(dispatch_id)
print(result)

My intuition would be that passing executor="cow" would throw an error because it's not one of the supported executors. However, the code runs successfully and uses the default Dask executor. Does anyone have a suggestion as to why this is? If I instead set executor: str | None = "cow" it fails, and if I instead do @ct.electron(executor="cow") it also fails, both as expected.

You may also be asking why I'm making not just modifying the electron's executor by specifying the executor in the lattice decorator. The reason for this is that, in my actual test case, I will have many electrons in the class and want to be able to programmatically control where each one will execute without hard-coding it in the electron decorator (but rather via a series of class arguments).

---

A more representative use case might look like the following, with the same issue as above in that here, the "local" executor in the second electron does not get activated.

from __future__ import annotations
from dataclasses import dataclass
import covalent as ct

@dataclass
class MyClass:
    executor1: str | None = None
    executor2: str | None = None
    
    @ct.electron(executor=executor1)
    def add(self, val1: int, val2: int) -> int:
        return val1 + val2
    
    @ct.electron(executor=executor2)
    def mult(self, val1: int, val2: int) -> int:
        return val1 * val2

@ct.lattice
def job(val1: int, val2: int) -> int:
    mc = MyClass(executor2="local")
    out1 = mc.add(val1, val2)
    out2 = mc.mult(out1, val2)
    return out2

dispatch_id = ct.dispatch(job)(1, 2)
result = ct.get_result(dispatch_id)
print(result)

[santoshkumarradha]

Hi @arosen93 !

Great question! I'll address your concerns in two parts: first, I'll explain why your current implementation isn't working and how to fix it; second, I'll suggest a better pattern for your use case.

Why isnt this working and how to solve it ?

The issue arises because the default Dask executor is used for both add and mult methods, as they are already defined with their respective executors when you define the class functions. When you change the executor2 value during class instantiation, the executor2 value is updated, but the parameter for the electron's executor isn't replaced in-place. The following code demonstrates this:

from covalent._shared_files.context_managers import active_lattice_manager
mc = MyClass(executor2="local")
print(mc.executor2)
with active_lattice_manager.claim(job):
    a=mc.mult(1,2)
    print(a.metadata['executor'])
    a=mc.add(1,2)
    print(a.metadata['executor'])

Note that here active_lattice_manager makes the "electron" objects behave as if they are electrons (if you call electrons outside this context, they just act as function). This will print

local
dask
dask

Which says that the executor's electrons is the default one.

To get clarity, lets slightly modify and add 2 new class methods as shown bellow and dispatch the workflow

from __future__ import annotations
from dataclasses import dataclass
import covalent as ct

@dataclass
class MyClass:
    executor1: str | None = None
    executor2: str | None = None
    
    @ct.electron(executor=executor1)
    def add(self, val1: int, val2: int) -> int:
        return val1 + val2
    
    @ct.electron(executor=executor2)
    def mult(self, val1: int, val2: int) -> int:
        return val1 * val2
    
    def mult_electron(self, val1: int, val2: int) -> int:
        @ct.electron(executor=self.executor2)
        def mult_electron(val1: int, val2: int) -> int:
            return val1 * val2
        return mult_electron(val1, val2)
    
    def mult_electron_default(self, val1: int, val2: int,executor=executor2) -> int:
        @ct.electron(executor=executor)
        def mult_electron_default(val1: int, val2: int) -> int:
            return val1 * val2
        return mult_electron_default(val1, val2)

@ct.lattice
def job(val1: int, val2: int) -> int:
    mc = MyClass(executor2="local")
    out1 = mc.add(val1, val2)
    out2 = mc.mult(out1, val2)
    out3=mc.mult_electron(out1, val2)
    out4=mc.mult_electron_default(out1, val2)
    return out2

dispatch_id = ct.dispatch(job)(1, 2)
result = ct.get_result(dispatch_id)
print(result)

This is what we see in the graph !

First, here we see that mult_electron_default which returns an electron of that is using the executor2 uses the default executor (as its only be changed by reference when defining the class and the electron function is already created) But the mult_electron, which uses the self uses the defined executor. Hope this answers your question !

Better (and recommended) pattern

Covalent primarily focuses on functional programming patterns rather than object-oriented ones. The primary reason is that the lattice in Covalent can be thought of as a pseudo-compiler, which determines the calling order of Electron objects by calculating the graph displayed in the UI. As a result, although manipulating class objects within a lattice is allowed (such as changing an attribute with class_object.some_attribute = 10 or invoking a class method with class_object.some_classmethod()), it's highly recommended to only call Electron objects (functions decorated with electrons) within a lattice.

Based on your example, there could be two potential reasons why you're exploring object-oriented patterns rather than functional ones:

1. You might be interested in "dynamic" executor selection.
2. You want to store your code or parameters within a class object.

Dynamic executors

For dynamic executors, we recommend using the pattern where a normal function takes in the executor and calls an electron internally to return an electron object. This is whats shown in above class example, for example you want to have a pattern like

def return_some_electron_function(executor,**kwargs):
    @ct.electron(executor=executor)
    def actual_function(**kwargs):
        pass
    return actual_function(**kwargs) # this will return an Electron object

Using class methods

If you have a legacy code that has high compute functions as methods inside a class object (For instance ase objects usually do), then the ideal pattern is something like this (essentially passing the instance as input/output of a function

#Legacy code you usually dont have access to (like ase or pymatgen object)
class A:
    def __init__(self,a) -> None:
        self.a=a
    def add(self,b):
        self.a=self.a+b
        
    def subtract(self,b):
        self.a=self.a-b
        
    def __repr__(self): #Lets add this just to see the object better in UI
        return f"A({self.a})"
#--------------------------------------------

@ct.electron
def add(A,b):
    A.add(b)
    return A

@ct.electron
def subtract(A,b):
    A.subtract(b)
    return A

@ct.lattice
def job(A,b,c):
    a_out=add(A,b)
    a_out=subtract(a_out,c)
    return a_out

a=A(10)
runid=ct.dispatch(job)(a,2,3)
result=ct.get_result(runid,wait=True).result
print(result.a) # 9

This should print 9=10+2-3. Note that we are passing class objects instead of defining electrons inside a class and also storing the data class variable throughout electrons which is automatically cached at each stage.

---

I hope this information is helpful to you! Please don't hesitate to ask more questions or seek clarification on any aspect. We're always here to support you and explore new use cases together. Let's keep the conversation going!

[Andrew-S-Rosen]

Thank you so much for providing what was probably the clearest and most thorough answer to a coding question I've ever had. I greatly appreciate it. The example you provided with the four different functions in the class was extremely illustrative!

For context, I was mostly asking because I was hoping to be able to set the executor on-the-fly, which I believe is what you're calling a "dynamic executor." I describe my motivation more in #1587. I thought using a class might be a way to achieve this, but I see that it's quite messy to do so in a way that actually works. The alternate suggestion you gave is definitely more concise. My suggestion in #1587 of having something like Lattice.set_electron_executors(["local", "slurm"]) to update electron 1 and 2 to "local" and "slurm," respectively, would be pretty sweet. I don't know enough about the codebase yet to make a PR, but I'll be keeping it in mind.

Also, super thankful for your example of dealing with objects in other codes. While that wasn't the specific use case I had in mind here, you're right that I am very often thinking about things like ASE Atoms objects and Pymatgen Structure objects!