Organize outputs directory as nested folders

[zulissimeta]

Summary of Changes

This is a draft PR to organize the RESULTS_DIR as a nested directory based on the DAG of the calling functions. For example, running the EMT example in the docs:

from ase.build import bulk
from quacc.recipes.emt.slabs import bulk_to_slabs_flow

# Define the Atoms object
atoms = bulk("Cu")

# Define the workflow
result = bulk_to_slabs_flow(atoms)

with the new setting QUACC_NESTED_RESULTS_DIR=True will lead to a RESULTS_DIR with layout:

• RESULTS_DIR
  • quacc.recipes.emt.slabs.bulk_to_slabs_flow-d0451dd0-843b-4f45-a6d9-f2de6e79e187
    • quacc.recipes.common.slabs.bulk_to_slabs_subflow-13292e70-7fe9-4dc4-9d01-34088b61cae4
      • quacc.recipes.emt.core.relax_job-b0a60c45-d0e6-4a11-bd42-4f396d3b9b76
      • quacc.recipes.emt.core.static_job-a72fe906-9d19-4bd3-a739-d0433f2eb3fe
      • quacc.recipes.emt.core.relax_job-ab3c836a-66b0-4bc6-a000-1996349d776d
      • quacc.recipes.emt.core.static_job-23df0c5a-5606-454c-8f1a-02861cfb9f2d
      • quacc.recipes.emt.core.relax_job-02717538-3d4b-421d-80ed-ee0f61ffc091
      • quacc.recipes.emt.core.static_job-d7e1d0b5-830b-481a-94c2-2785cbb5b126
      • quacc.recipes.emt.core.relax_job-b6d62a53-0550-445a-b817-5c6c14976bd3
      • quacc.recipes.emt.core.static_job-7302e8cd-5815-4892-a7df-8a3ed5d42a92

This is only a demo to inspire discussion, and is only implemented for prefect right now (though others would be easy to add!).

A couple of potential downsides:

• Passing around wrapped functions increases complexity
• You can only define the task/flow/etc at runtime in the flow, which could lead to some downstream weird logic issues
• Prefect throws some warnings that different functions with the same name are being used (it still works)
• I don't quite understand the logic in change_settings_wrap, so there's probably some overlap / considerations there if it was separately specified for a function (and which should take priority)

[buildbot-princeton]

Can one of the admins verify this patch?

[Andrew-S-Rosen]

Thanks! Let me first address one of your points/questions before proceeding:

I don't quite understand the logic in change_settings_wrap, so there's probably some overlap / considerations there if it was separately specified for a function (and which should take priority)

The function is below.

quacc/src/quacc/settings.py

Lines 589 to 598 in f6bc96c

	original_func = func._original_func if getattr(func, "_changed", False) else func
	@wraps(original_func)
	def wrapper(*args, **kwargs):
	with change_settings(changes):
	return original_func(*args, **kwargs)
	wrapper._changed = True
	wrapper._original_func = original_func
	return wrapper

The important bits are lines 591-594. There, we are basically turning

def add(a, b):
    return a + b

into

from quacc import change_settings

def add(a, b):
    with change_settings(changes):
        return a + b

This is necessary because with change_settings() is modifying the thread-based, in-memory global settings. Since this is a modification to a variable in-memory, it needs to be "shipped" to the remote machine and therefore included within the function definition itself. The other lines in the function definition can be ignored, as they are specific to getting it to work with Parsl.

So, the with change_settings() will create a newly modified function, but where is this called? That is done in the definition of the job function below:

quacc/src/quacc/wflow_tools/decorators.py

Lines 147 to 148 in f6bc96c

	if changes := kwargs.pop("settings_swap", {}):
	return job(change_settings_wrap(_func, changes), **kwargs)

In short, if a user passes @job(settings_swap=changes) where changes: dict[str, Any] (e.g. changes={"RESULTS_DIR": "~/test"}), this will:

1. Remove the keyword argument from the list of keyword arguments to be passed to the workflow engine's decorator (e.g. @task for Prefect).
2. Inject the with change_settings(settings): context manager within the original function definition.
3. Decorate the result with the workflow engine's decorator (e.g. @task for Prefect).

Now we have successfully changed the settings within the @task-decorated function itself.

---

the _changed attribute was a nuance that probably isn't worth thinking much about here tbh. An external contributor wanted to make sure that double-application of change_settings took the outer one as priority.

with change_settings({"RESULTS_DIR": "~/test2"}):
    with change_settings({"RESULTS_DIR": "~/test1"}):
        static_job(atoms)

this came up in practice because they redecorated the same function twice without renaming it:

static_job = redecorate(static_job, job(settings_swap={"RESULTS_DIR": "~/test1"}))
static_job(atoms)

static_job = redecorate(static_job, job(settings_swap={"RESULTS_DIR": "~/test2"}))
static_job(atoms)

this _changed business was to get it to work correctly, with Parsl specifically I believe.

but this is an edge case and a scenario of "the user shouldn't technically do that", so I told them if it caused issues it may be removed. in any case, not much to worry about at the moment.

I'll take care of other workflow engines as needed. if tests pass, then we're fine. I added a lot of tests for the settings handling.

[Andrew-S-Rosen]

@zulissimeta: Thanks for this PR! I will have to try it out, but looking at the code and your writeup, this seems logical to me. I do not have major objections (at this time).

As for this PR, next steps would be:

1. Figure out why the tests are failing and update the code accordingly. We have three tests failing in tests/prefect/test_customizers.py.
2. Address the (minor) comments in my review.
3. Add tests for this newly introduced behavior so we can ensure we don't break it.
4. Extend this to other workflow engines (this is for me to do).

As for your "downsides":

Passing around wrapped functions increases complexity

It is true that wrapping functions increases complexity. That said, we already are wrapping functions once if change_settings_wrap is applied. Additionally, even without the change_settings business, we were already wrapping in a few places depending on the workflow engine (e.g. for Prefect when we have PREFECT_AUTO_SUBMIT as True). So, here we have either single or double wrapping depending on the workflow engine. That should hopefully be okay because we are already doing that in the test suite as-is. You haven't really increased the complexity from a wrapping standpoint in this PR.

You can only define the task/flow/etc at runtime in the flow, which could lead to some downstream weird logic issues

I don't think I understand this point. Could you clarify?

Prefect throws some warnings that different functions with the same name are being used (it still works)

Is this a new warning, or was it always present? Shouldn't the name of the function not change if we use @wraps? I'm curious what happened here to cause that metadata change.

[codecov]

Codecov Report

Attention: Patch coverage is 92.68293% with 3 lines in your changes missing coverage. Please review.

Project coverage is 98.35%. Comparing base (c8e17b5) to head (2215b1c).

Files	Patch %	Lines
src/quacc/wflow_tools/decorators.py	88.88%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2296      +/-   ##
==========================================
- Coverage   98.45%   98.35%   -0.11%     
==========================================
  Files          84       84              
  Lines        3430     3457      +27     
==========================================
+ Hits         3377     3400      +23     
- Misses         53       57       +4     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[zulissimeta]

@zulissimeta: Thanks for this PR! I will have to try it out, but looking at the code and your writeup, this seems logical to me. I do not have major objections (at this time).

As for this PR, next steps would be:

1. Figure out why the tests are failing and update the code accordingly. We have three tests failing in tests/prefect/test_customizers.py.
2. Address the (minor) comments in my review.
3. Add tests for this newly introduced behavior so we can ensure we don't break it.
4. Extend this to other workflow engines (this is for me to do).

Done! (except for 4.)

As for your "downsides":

Passing around wrapped functions increases complexity

It is true that wrapping functions increases complexity. That said, we already are wrapping functions once if change_settings_wrap is applied. Additionally, even without the change_settings business, we were already wrapping in a few places depending on the workflow engine (e.g. for Prefect when we have PREFECT_AUTO_SUBMIT as True). So, here we have either single or double wrapping depending on the workflow engine. That should hopefully be okay because we are already doing that in the test suite as-is. You haven't really increased the complexity from a wrapping standpoint in this PR.

After thinking some more, I made it so the change_settings_wrap keeps a reference to the changes, both so that the nested directory wrapper knows when not to set the new directory (if it was manually specified), and so that it can just be added to the other changes.

You can only define the task/flow/etc at runtime in the flow, which could lead to some downstream weird logic issues

I don't think I understand this point. Could you clarify?

If you don't have this, you can definitely schedule a bunch of future jobs (say a chain of dependent jobs). With this change, the function actually has to be called with each input in the calling flow.

I'm not aware of a concrete problem here, but making it so that the wrapped function 100% has to be called to get it to add the right context manager might lead to problems later. For example, in prefect there's a job.map() that might not explicitly call functions, and instead just schedule them directly on the server.

Prefect throws some warnings that different functions with the same name are being used (it still works)

Is this a new warning, or was it always present? Shouldn't the name of the function not change if we use @wraps? I'm curious what happened here to cause that metadata change.

This is a new warning that will happen for any wrapped function (not just in my PR, but with change_settings too). Wrapped functions are different functions, but have the same name, which prefect is warning about.

[Andrew-S-Rosen]

This seems logical. I did some refactoring to reduce code duplication.

Before this can be merged, the following will need to be addressed:

• It does not currently work when there's no workflow engine being set. We'll need to fix that.
• Tests are needed for the no workflow case.
• We need to test the uncovered lines for the prefect tests. Because NESTED_RESULTS_DIR is True by default, the no-wrapped cases are not being tested anymore.
• Tests are needed for all other workflow engines.

I can work on this soon but might be a bit slower than usual, as I'm getting ready to move across the country this weekend. ✈️ Although, certainly feel free to chip in if you have time and are interested. Either way, hopefully, I can get this merged next week. I quite like the idea. If there is any weirdness we stumble upon, we can set it to False by default for a trial period without much concern.

[Andrew-S-Rosen]

Pinging @honghuikim since you may be interested in this PR based on your prior contribution.

[honghuikim]

@Andrew-S-Rosen Thanks for the tag. I just skimmed this PR, and it looks great! I like this "nested directory" concept, which will be helpful especially for flows.

Hi @zulissimeta, I like the concept of making nested directories, and this quite resembles the purpose of my previous PR (change_settings_wrap). I have one suggestion. Maybe we can try assigning the parent RESULTS_DIR inside customize_funcs? This will be something like,

def customize_funcs(
    names: list[str] | str,
    funcs: list[Callable] | Callable,
    param_defaults: dict[str, dict[str, Any]] | None = None,
    param_swaps: dict[str, dict[str, Any]] | None = None,
    decorators: dict[str, Callable | None] | None = None,
) -> tuple[Callable, ...] | Callable:

    parameters = recursive_dict_merge(param_defaults, param_swaps)
    decorators = decorators or {}
    if QUACC_NESTED_RESULTS_DIR:
        time_now = datetime.now(timezone.utc).strftime("%Y-%m-%d-%H-%M-%S-%f") 
        if prefix is None: 
            prefix = "" 
        flow_dir_dict =  {"all": job(settings_swap={"RESULTS_DIR": Path(f"{prefix}{time_now}-{randint(10000, 99999)}")})}
        decorators = recursive_dict_merge(flow_dir_dict , decorators)
    
    # and the rest of the original code

By using recursive_dict_merge on the flow_dir_dict (i.e. randomly generated directory of flow) and the given decorators argument, we can first assign the parent directory by default when QUACC_NESTED_RESULTS_DIR is True.
If this works, then we can get rid of workflow-specific cases. Nonetheless, this will not work as it is now. I think we need to make sure the flow_dir_dict is kept the same in a flow.

I might have missed some key points to consider, so this might be oversimplifying the task. Please let me know your thoughts on this.

[Andrew-S-Rosen]

Thank you, @honghuikim, for this very interesting suggestion! It would be great if we could have the logic inside something like customize_funcs that way we don't have to deal with the messiness of wrapping functions with the various workflow engines. I do think that is likely going to be a lot better.

One small subtlety to keep track of is in the following proposed set of lines:

flow_dir_dict = {
    "all": job(
        settings_swap={
            "RESULTS_DIR": Path(f"{prefix}{time_now}-{randint(10000, 99999)}")
        }
    )
}
decorators = recursive_dict_merge(flow_dir_dict, decorators)

If the user passes in decorators = {"all": job(executors=["exec1", "exec2"])} to the customize_funcs function for instance, the merged dictionary that results from recursive_dict_merge will still be {"all": job(executors=["exec1", "exec2"])}, without the settings_swaps keyword argument. Of course, that's because it is merging dictionary keys and values, not the keyword arguments of the job call. This doesn't seem too painful to solve. One could, for instance, parse the inspect the function call, copy the arguments, and then merge those if needed. Just wanted to bring that up.

Anyway, in full disclosure, I probably can't tackle this in the very short term. But I would be happy to review such an approach, which I think would further improve this PR.

[honghuikim]

@Andrew-S-Rosen Thanks a lot for the mini-review of it. That indeed is an error to be solved if we use this approach. My short thought is to use an internally defined hidden key (e.g., named _all) that works exclusively with QUACC_NESTED_RESULTS_DIR. We can then merge this with the user-provided inputs as you suggested.

I will also be able to dig into this more later, probably a week later. In the meantime, it would be great if zulissimeta could provide his thoughts on this. Especially regarding potential errors that could arise from using this approach since he has been tackling this.

[zulissimeta]

@Andrew-S-Rosen Thanks a lot for the mini-review of it. That indeed is an error to be solved if we use this approach. My short thought is to use an internally defined hidden key (e.g., named _all) that works exclusively with QUACC_NESTED_RESULTS_DIR. We can then merge this with the user-provided inputs as you suggested.

I will also be able to dig into this more later, probably a week later. In the meantime, it would be great if zulissimeta could provide his thoughts on this. Especially regarding potential errors that could arise from using this approach since he has been tackling this.

I don't have strong feelings on this, and I agree that an approach that is agnostic to the specific workflow would be nice.

However, putting this in customize_funcs means that only flows that explicitly call that will actually update right? I wouldn't be surprised if some flows/recipes don't use that. Something that happens for every job/flow would probably be better.

[Andrew-S-Rosen]

However, putting this in customize_funcs means that only flows that explicitly call that will actually update right? I wouldn't be surprised if some flows/recipes don't use that. Something that happens for every job/flow would probably be better.

That is correct. However, in the quacc codebase itself, there is a strict requirement that all flows call customize_funcs so that the underlying keyword arguments and/or decorators can be modified. Of course, that is a design decision rather than a strict requirements, so someone doing their own thing could very well not do this.

[Andrew-S-Rosen]

Following up the discussion here to state that packing this logic inside customize_funcs is not actually as good of an idea as it may initially seem. While most @flows in quacc call customize_funcs as a convenient way to apply decorator/kwarg updates to individual jobs in a flow, that design decision may not hold for end users. For instance, end users may decide to stitch together their own jobs into a custom flow and never call customize_funcs in their @flow definition. This is actually the expected behavior because customize_funcs is merely for cleanliness of the quacc library. End users will rarely have a need for calling this themselves. They'll just stitch together jobs how they see fit, but they should still be able to benefit from organized output directories. Also, not all flows actually use customize_funcs. @zulissimeta correctly notes that this is the case for quacc.recipes.vasp.core.double_relax_flow.

In short, we want to decouple directory handling from any optional behavior. But we do need it to work for all workflow engines.

[honghuikim]

Following up the discussion here to state that packing this logic inside customize_funcs is not actually as good of an idea as it may initially seem. While most @flows in quacc call customize_funcs as a convenient way to apply decorator/kwarg updates to individual jobs in a flow, that design decision may not hold for end users. For instance, end users may decide to stitch together their own jobs into a custom flow and never call customize_funcs in their @flow definition. This is actually the expected behavior because customize_funcs is merely for cleanliness of the quacc library. End users will rarely have a need for calling this themselves. They'll just stitch together jobs how they see fit, but they should still be able to benefit from organized output directories. Also, not all flows actually use customize_funcs. @zulissimeta correctly notes that this is the case for quacc.recipes.vasp.core.double_relax_flow.

I agree, this is correct. Actually, changing the settings of a flow using settings_swap keyword has the same problem. When we change settings of a flow using the settings_swap keyword, it benefits from customize_funcs to propagate the changes to inner jobs. Which means we cannot change the settings of flows defined by end users (without customize_funcs in it) using the same strategy. Utilizing customize_funcs will cover most pre-defined flows in Quacc, but for more generality, we do need a better approach.