materialsproject / jobflow

jobflow is a library for writing computational workflows.
https://materialsproject.github.io/jobflow
Other
90 stars 25 forks source link

Dynamic update for `update_kwargs` and `update_maker_kwargs` #588

Open gpetretto opened 2 months ago

gpetretto commented 2 months ago

I have realized that the functions update_kwargs and update_maker_kwargs do not apply to Jobs that are generated dynamically, unless the Maker that generates them is already instantiated when the Flow is created.

Here is a minimal example:

from dataclasses import dataclass
from jobflow import job, Flow, Maker, Response
from jobflow.core.job import Job
from jobflow.managers.local import run_locally
from dataclasses import dataclass, field

@dataclass
class TestMaker(Maker):
    name: str = "Test Maker"
    prop: str = "X"

    @job
    def make(self):
        return self.prop

@job()
def replace():
    return Response(replace=TestMaker().make())

@dataclass
class ReplaceMaker(Maker):
    test_maker: TestMaker = field(default_factory=TestMaker)
    name: str = "Replace maker"

    @job
    def make(self):
        return Response(replace=self.test_maker.make())

def run_job(job):
    flow = Flow([job])
    flow.update_maker_kwargs({"prop": "Y"}, name_filter="Test Maker")
    responses = run_locally(flow, log=False)
    print(responses[job.uuid][max(responses[job.uuid].keys())].output)

run_job(TestMaker().make())
run_job(replace())
run_job(ReplaceMaker().make())

whose output is

Y X Y

I am not sure if this is the intended behaviour, but I would have somehow expected this to work. Maybe a dynamic argument could be added to the two functions, similar to the one implemented for update_metadata and update_config (https://github.com/materialsproject/jobflow/pull/198)? I did not check if the procedure could be completely equivalent or adaptations will be needed.

utf commented 2 months ago

I'm not sure this behaviour is unexpected, since the Maker does not exist at the time update_maker_kwargs is applied.

I'd support adding a dynamic argument as you suggested. Although I feel like it should be set to False by default since for some reason this feels more unexpected to me than the config/metadata case.