Improve dynamic changing of SETTINGS parameters for decorated functions with active workflow engines

[Andrew-S-Rosen]

What new feature would you like to see?

Recently, we introduced a function change_settings that can be used to temporarily modify the settings of some function calls in a script. See here for details. An example is like that shown below:

from quacc import change_settings

with change_settings({"RESULTS_DIR": "/new/path/to/store/results"}):
    pass  # Your calculation here

This works wonderfully. However, as noted on that same page, if you dispatch the calculation to a remote machine (as is common with the use of a workflow engine), there is no knowledge of changes to in-memory parameters like the global SETTINGS object. We should come up with a smart and user-friendly way to address this.

For instance, let's say you want to change SETTINGS.DEBUG to True for a @job called relax_job. Most likely, it can be solved by:

1. Stripping the @job, @flow, or @subflow decorator (from quacc import strip_decorator) from the imported function. In this case, it'd be a @job.
2. Dynamically redefining the function with an update to SETTINGS.DEBUG inside of the function definition. Important: we must retain the original function's name, signature, type hints, etc. so presumably functools.wraps here.
3. Redecorating it with the original decorator.
4. Returning the redecorated function to the user

The tools to do this are likely already in quacc.wflow_tools.customizers, but the exact logic to get it right in an automated fashion currently escapes me.

For background context, refer to #2145 and my answer therein.

[honghuikim]

I've tried a workflow similar to your suggestion. I wasn't sure how to use quacc.wflow_tools.customizers to change settings for a remote machine, instead, I defined another function for it. In my working environment (latest Quacc and Parsl), it works just as I want, but there might be possible errors.

def change_settings_wf(func: Callable, settings: Dict, decorator: Callable) -> Callable:
    from quacc import change_settings
    def concatenated(*args, **kwargs):
        with change_settings(settings):
            return func.func(*args, **kwargs)
    
    return decorator(concatenated)

from quacc import job
relax_job = change_settings_wf(relax_job, {"DEBUG": True}, job)

I think it would be great if I could directly get the decorator from func so that the change_settings_wf doesn't need to take decorator as a keyword. Let me know your thoughts on this, or if there's a better way.

[Andrew-S-Rosen]

Thanks for the suggestion!

We can replace func.func(*args, **kwargs) with

from quacc import strip_decorator
return strip_decorator(func)(*args, **kwargs)

So that it works for any workflow engine.

If you'd wish to do so, feel free to open a PR and add a test for Parsl! I can add the other tests. Alternatively, if you do not have the time or desire to do so, I can work on implementing this soon.

I think it would be great if I could directly get the decorator from func so that the change_settings_wf doesn't need to take decorator as a keyword

Definitely. This was also a problem in #1865. I tried to solve it recently by appending a .quacc_decorator: Literal["job", "subflow", "flow"] attribute onto the decorated function that could be retrieved later, but some of the workflow engines (Dask) don't support adding new attributes like this. I think I may need to write a custom function that:

1. Detects the workflow engine from SETTINGS
2. Tries to smartly detect which decorator was applied, basically using the reverse of the logic that was within quacc.wflow_tools.customizers.strip_decorator.

Not a huge fan of that approach but perhaps it's the only way. I'll keep thinking on this.

[honghuikim]

No problem!

Sure, I can work on it tomorrow since it's getting late here. But of course you can go ahead if you would like to implement it sooner.

I see, that sounds good. It seems there's no way to know the wrapper from the wrapped function without any modification. I agree that this is the only way.

[Andrew-S-Rosen]

It's no rush from my side. I probably won't have time to seriously devote to this until the middle of next week. Thanks for kicking it off though! Shouldn't be too painful to implement.

[honghuikim]

No problem, just opened a PR. I only added a unit test for Parsl.

FYI, I found that there can be a potential issue in the nested call of change_settings_wf.

static_job1 = change_settings_wf(static_job, {"RESULTS_DIR": "/new/path1"}, job)
static_job2 = change_settings_wf(static_job1, {"RESULTS_DIR": "/new/path2"}, job)
#static_job in static_job2 will have "RESULTS_DIR" of "/new/path1"

This can be problematic when using flow or subflow that calls jobs in it. At the moment, I don't know how to fix it. We might need another approach for change_settings.