From 8afb69e033f9f5ea1d7e6fb6f58ccac6767cefeb Mon Sep 17 00:00:00 2001 From: Zack Ulissi Date: Fri, 21 Jun 2024 15:49:07 +0000 Subject: [PATCH 01/17] first commit --- src/quacc/settings.py | 51 ++++++++++++++++++++++++++++ src/quacc/wflow_tools/decorators.py | 52 ++++++++++++++++++++++++++--- 2 files changed, 98 insertions(+), 5 deletions(-) diff --git a/src/quacc/settings.py b/src/quacc/settings.py index db4b778006..cf3a1cecd2 100644 --- a/src/quacc/settings.py +++ b/src/quacc/settings.py @@ -8,6 +8,8 @@ from pathlib import Path from shutil import which from typing import TYPE_CHECKING, Literal, Optional, Union +import inspect +import uuid import psutil from maggma.core import Store @@ -97,6 +99,14 @@ class QuaccSettings(BaseSettings): """ ), ) + NESTED_RESULTS_DIR: bool = Field( + True, + description=( + """ + Whether to nest the results dir by the calling flow/subflow/etc + """ + ), + ) GZIP_FILES: bool = Field( True, description="Whether generated files should be gzip'd." ) @@ -596,3 +606,44 @@ def wrapper(*args, **kwargs): wrapper._changed = True wrapper._original_func = original_func return wrapper + + +def nest_results_dir_wrap(func: Callable) -> Callable: + """ + Wraps a function with the change_settings context manager if not already wrapped. + + Parameters + ---------- + func + The function to wrap. + changes + The settings to apply within the context manager. + + Returns + ------- + Callable + The wrapped function. + """ + original_func = func._original_func if getattr(func, "_changed", False) else func + + from quacc import get_settings + + # Get the settings from the calling function's context + results_parent_dir = get_settings().RESULTS_DIR + + @wraps(original_func) + def wrapper(*args, **kwargs): + + # Set the settings within the new function's context to be a subdirectory + # of the parent's folder + with change_settings( + { + "RESULTS_DIR": results_parent_dir + / f"{inspect.getmodule(func).__name__}.{func.__name__}-{uuid.uuid4()}" + } + ): + return original_func(*args, **kwargs) + + wrapper._changed = True + wrapper._original_func = original_func + return wrapper diff --git a/src/quacc/wflow_tools/decorators.py b/src/quacc/wflow_tools/decorators.py index 5370ad11b1..d844968486 100644 --- a/src/quacc/wflow_tools/decorators.py +++ b/src/quacc/wflow_tools/decorators.py @@ -4,8 +4,10 @@ from functools import partial, wraps from typing import TYPE_CHECKING, TypeVar +import inspect +import uuid -from quacc.settings import change_settings_wrap +from quacc.settings import change_settings_wrap, nest_results_dir_wrap Job = TypeVar("Job") Flow = TypeVar("Flow") @@ -147,6 +149,9 @@ def add(a, b): if changes := kwargs.pop("settings_swap", {}): return job(change_settings_wrap(_func, changes), **kwargs) + # if settings.NESTED_RESULTS_DIR: + # _func = nest_results_dir_wrap(_func) + if settings.WORKFLOW_ENGINE == "covalent": import covalent as ct @@ -182,12 +187,19 @@ def wrapper(*f_args, **f_kwargs): @wraps(_func) def wrapper(*f_args, **f_kwargs): - decorated = task(_func, **kwargs) + adjusted_results_func = nest_results_dir_wrap(_func) + decorated = task(adjusted_results_func, **kwargs) return decorated.submit(*f_args, **f_kwargs) return wrapper else: - return task(_func, **kwargs) + + @wraps(_func) + def wrapper(*f_args, **f_kwargs): + adjusted_results_func = nest_results_dir_wrap(_func) + return task(adjusted_results_func, **kwargs)(*f_args, **f_kwargs) + + return wrapper else: return _func @@ -338,6 +350,9 @@ def workflow(a, b, c): settings = get_settings() + if settings.NESTED_RESULTS_DIR: + _func = nest_results_dir_wrap(_func) + if _func is None: return partial(flow, **kwargs) @@ -352,7 +367,20 @@ def workflow(a, b, c): elif settings.WORKFLOW_ENGINE == "prefect": from prefect import flow as prefect_flow - return prefect_flow(_func, validate_parameters=False, **kwargs) + if settings.NESTED_RESULTS_DIR: + + @wraps(_func) + def wrapper(*f_args, **f_kwargs): + adjusted_results_func = nest_results_dir_wrap(_func) + return prefect_flow( + adjusted_results_func, validate_parameters=False, **kwargs + )(*f_args, **f_kwargs) + + return wrapper + else: + return prefect_flow(_func, validate_parameters=False, **kwargs) + + # return prefect_flow(_func, validate_parameters=False, **kwargs) else: return _func @@ -554,6 +582,9 @@ def workflow(a, b, c): settings = get_settings() + # if settings.NESTED_RESULTS_DIR: + # _func = nest_results_dir_wrap(_func) + if _func is None: return partial(subflow, **kwargs) @@ -583,7 +614,18 @@ def wrapper(*f_args, **f_kwargs): elif settings.WORKFLOW_ENGINE == "prefect": from prefect import flow as prefect_flow - return prefect_flow(_func, validate_parameters=False, **kwargs) + if settings.NESTED_RESULTS_DIR: + + @wraps(_func) + def wrapper(*f_args, **f_kwargs): + adjusted_results_func = nest_results_dir_wrap(_func) + return prefect_flow( + adjusted_results_func, validate_parameters=False, **kwargs + )(*f_args, **f_kwargs) + + return wrapper + else: + return prefect_flow(_func, validate_parameters=False, **kwargs) elif settings.WORKFLOW_ENGINE == "redun": from redun import task From 93d02a70ce1fd63c668596d5e0991aa02ba793ec Mon Sep 17 00:00:00 2001 From: Zack Ulissi Date: Fri, 21 Jun 2024 15:49:31 +0000 Subject: [PATCH 02/17] remove commented lines --- src/quacc/wflow_tools/decorators.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/quacc/wflow_tools/decorators.py b/src/quacc/wflow_tools/decorators.py index d844968486..9ed72d28e1 100644 --- a/src/quacc/wflow_tools/decorators.py +++ b/src/quacc/wflow_tools/decorators.py @@ -149,9 +149,6 @@ def add(a, b): if changes := kwargs.pop("settings_swap", {}): return job(change_settings_wrap(_func, changes), **kwargs) - # if settings.NESTED_RESULTS_DIR: - # _func = nest_results_dir_wrap(_func) - if settings.WORKFLOW_ENGINE == "covalent": import covalent as ct From 58666be42483646a04d7f78f803698b236f98e77 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 21 Jun 2024 16:06:45 +0000 Subject: [PATCH 03/17] pre-commit auto-fixes --- src/quacc/settings.py | 5 ++--- src/quacc/wflow_tools/decorators.py | 2 -- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/src/quacc/settings.py b/src/quacc/settings.py index cf3a1cecd2..676c88ddcf 100644 --- a/src/quacc/settings.py +++ b/src/quacc/settings.py @@ -2,14 +2,14 @@ from __future__ import annotations +import inspect import os +import uuid from contextlib import contextmanager from functools import wraps from pathlib import Path from shutil import which from typing import TYPE_CHECKING, Literal, Optional, Union -import inspect -import uuid import psutil from maggma.core import Store @@ -633,7 +633,6 @@ def nest_results_dir_wrap(func: Callable) -> Callable: @wraps(original_func) def wrapper(*args, **kwargs): - # Set the settings within the new function's context to be a subdirectory # of the parent's folder with change_settings( diff --git a/src/quacc/wflow_tools/decorators.py b/src/quacc/wflow_tools/decorators.py index 9ed72d28e1..bf51625eaf 100644 --- a/src/quacc/wflow_tools/decorators.py +++ b/src/quacc/wflow_tools/decorators.py @@ -4,8 +4,6 @@ from functools import partial, wraps from typing import TYPE_CHECKING, TypeVar -import inspect -import uuid from quacc.settings import change_settings_wrap, nest_results_dir_wrap From 1bed53ea3197fc8f2e01b4dee2deeb6210955359 Mon Sep 17 00:00:00 2001 From: Zack Ulissi Date: Fri, 21 Jun 2024 16:06:55 +0000 Subject: [PATCH 04/17] clean up nested dir function --- src/quacc/settings.py | 30 +++++++++--------------------- 1 file changed, 9 insertions(+), 21 deletions(-) diff --git a/src/quacc/settings.py b/src/quacc/settings.py index cf3a1cecd2..8d8871c4f8 100644 --- a/src/quacc/settings.py +++ b/src/quacc/settings.py @@ -610,40 +610,28 @@ def wrapper(*args, **kwargs): def nest_results_dir_wrap(func: Callable) -> Callable: """ - Wraps a function with the change_settings context manager if not already wrapped. + Wraps a function with the change_settings context manager using a nested RESULTS_DIR + Parameters ---------- func The function to wrap. - changes - The settings to apply within the context manager. Returns ------- Callable The wrapped function. """ - original_func = func._original_func if getattr(func, "_changed", False) else func - from quacc import get_settings # Get the settings from the calling function's context results_parent_dir = get_settings().RESULTS_DIR - @wraps(original_func) - def wrapper(*args, **kwargs): - - # Set the settings within the new function's context to be a subdirectory - # of the parent's folder - with change_settings( - { - "RESULTS_DIR": results_parent_dir - / f"{inspect.getmodule(func).__name__}.{func.__name__}-{uuid.uuid4()}" - } - ): - return original_func(*args, **kwargs) - - wrapper._changed = True - wrapper._original_func = original_func - return wrapper + return change_settings_wrap( + func, + { + "RESULTS_DIR": results_parent_dir + / f"{inspect.getmodule(func).__name__}.{func.__name__}-{uuid.uuid4()}" + }, + ) From ae8c71a9e7c18171390003e9ba79f79b9bf5e540 Mon Sep 17 00:00:00 2001 From: Zack Ulissi Date: Fri, 21 Jun 2024 16:22:47 +0000 Subject: [PATCH 05/17] clean up decorators --- src/quacc/wflow_tools/decorators.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/src/quacc/wflow_tools/decorators.py b/src/quacc/wflow_tools/decorators.py index bf51625eaf..22461b9dbc 100644 --- a/src/quacc/wflow_tools/decorators.py +++ b/src/quacc/wflow_tools/decorators.py @@ -178,7 +178,7 @@ def wrapper(*f_args, **f_kwargs): elif settings.WORKFLOW_ENGINE == "prefect": from prefect import task - if settings.PREFECT_AUTO_SUBMIT: + if settings.PREFECT_AUTO_SUBMIT and settings.NESTED_RESULTS_DIR: @wraps(_func) def wrapper(*f_args, **f_kwargs): @@ -187,7 +187,7 @@ def wrapper(*f_args, **f_kwargs): return decorated.submit(*f_args, **f_kwargs) return wrapper - else: + elif (not settings.PREFECT_AUTO_SUBMIT) and settings.NESTED_RESULTS_DIR: @wraps(_func) def wrapper(*f_args, **f_kwargs): @@ -195,6 +195,17 @@ def wrapper(*f_args, **f_kwargs): return task(adjusted_results_func, **kwargs)(*f_args, **f_kwargs) return wrapper + elif settings.PREFECT_AUTO_SUBMIT and (not settings.NESTED_RESULTS_DIR): + + @wraps(_func) + def wrapper(*f_args, **f_kwargs): + decorated = task(_func, **kwargs) + return decorated.submit(*f_args, **f_kwargs) + + return wrapper + else: + return task(_func, **kwargs) + else: return _func @@ -345,9 +356,6 @@ def workflow(a, b, c): settings = get_settings() - if settings.NESTED_RESULTS_DIR: - _func = nest_results_dir_wrap(_func) - if _func is None: return partial(flow, **kwargs) @@ -374,8 +382,6 @@ def wrapper(*f_args, **f_kwargs): return wrapper else: return prefect_flow(_func, validate_parameters=False, **kwargs) - - # return prefect_flow(_func, validate_parameters=False, **kwargs) else: return _func @@ -577,9 +583,6 @@ def workflow(a, b, c): settings = get_settings() - # if settings.NESTED_RESULTS_DIR: - # _func = nest_results_dir_wrap(_func) - if _func is None: return partial(subflow, **kwargs) From cbb07dbc195f28ea361e5ecb134a0cae22030f9a Mon Sep 17 00:00:00 2001 From: Zack Ulissi Date: Fri, 21 Jun 2024 23:30:06 +0000 Subject: [PATCH 06/17] use datetime and random number instead of uuid --- src/quacc/settings.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/quacc/settings.py b/src/quacc/settings.py index 1e319382a5..03160f8100 100644 --- a/src/quacc/settings.py +++ b/src/quacc/settings.py @@ -4,7 +4,8 @@ import inspect import os -import uuid +from datetime import datetime, timezone +from random import randint from contextlib import contextmanager from functools import wraps from pathlib import Path @@ -628,10 +629,9 @@ def nest_results_dir_wrap(func: Callable) -> Callable: # Get the settings from the calling function's context results_parent_dir = get_settings().RESULTS_DIR - return change_settings_wrap( - func, - { - "RESULTS_DIR": results_parent_dir - / f"{inspect.getmodule(func).__name__}.{func.__name__}-{uuid.uuid4()}" - }, + time_now = datetime.now(timezone.utc).strftime("%Y-%m-%d-%H-%M-%S-%f") + + nested_results_dir = results_parent_dir / Path( + f"{inspect.getmodule(func).__name__}.{func.__name__}-{time_now}-{randint(10000, 99999)}" ) + return change_settings_wrap(func, {"RESULTS_DIR": nested_results_dir}) From 01ce6b50eb545c6ede63690072c8842d749de9c5 Mon Sep 17 00:00:00 2001 From: Zack Ulissi Date: Sat, 22 Jun 2024 00:13:50 +0000 Subject: [PATCH 07/17] datetime, dont overwrite changes, and add a test --- src/quacc/settings.py | 17 ++++++++++++++--- tests/prefect/test_customizers.py | 27 +++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 3 deletions(-) diff --git a/src/quacc/settings.py b/src/quacc/settings.py index 03160f8100..ffb916c7ad 100644 --- a/src/quacc/settings.py +++ b/src/quacc/settings.py @@ -605,6 +605,7 @@ def wrapper(*args, **kwargs): return original_func(*args, **kwargs) wrapper._changed = True + wrapper._changes = changes wrapper._original_func = original_func return wrapper @@ -626,12 +627,22 @@ def nest_results_dir_wrap(func: Callable) -> Callable: """ from quacc import get_settings + if getattr(func, "_changed", False): + changes = func._changes + else: + changes = [] + # Get the settings from the calling function's context results_parent_dir = get_settings().RESULTS_DIR - time_now = datetime.now(timezone.utc).strftime("%Y-%m-%d-%H-%M-%S-%f") - nested_results_dir = results_parent_dir / Path( f"{inspect.getmodule(func).__name__}.{func.__name__}-{time_now}-{randint(10000, 99999)}" ) - return change_settings_wrap(func, {"RESULTS_DIR": nested_results_dir}) + + # If someone explcitly set RESULTS_DIR in change_settings already + # they probably know what they're doing and really want a specific + # folder! + if "RESULTS_DIR" not in changes: + changes["RESULTS_DIR"] = nested_results_dir + + return change_settings_wrap(func, changes) diff --git a/tests/prefect/test_customizers.py b/tests/prefect/test_customizers.py index b80d34728f..217371a16a 100644 --- a/tests/prefect/test_customizers.py +++ b/tests/prefect/test_customizers.py @@ -98,3 +98,30 @@ def my_flow(): my_flow().result() assert not Path(tmp_dir1 / "job.txt").exists() assert Path(tmp_dir2 / "job.txt").exists() + + +def test_nested_output_directory(tmp_path_factory): + + @job + def write_file_job(name="job.txt"): + results_file_path = Path(get_settings().RESULTS_DIR, name) + with open(results_file_path, "w") as f: + f.write("test job file") + + return results_file_path + + @flow + def write_file_flow(name="flow.txt", job_decorators=None): + job_results_file_path = write_file_job() + + flow_results_file_path = Path(get_settings().RESULTS_DIR, name) + with open(flow_results_file_path, "w") as f: + f.write("test flow file") + + return job_results_file_path, flow_results_file_path + + # Test with redecorating a job in a flow + job_results_file_path, flow_results_file_path = write_file_flow().result() + assert Path(job_results_file_path).exists() + assert Path(flow_results_file_path).exists() + assert Path(job_results_file_path).parent == Path(flow_results_file_path) From c6aaa69ac1312cd1eadf82b40ea7c4a167d0fd7f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 22 Jun 2024 00:14:04 +0000 Subject: [PATCH 08/17] pre-commit auto-fixes --- src/quacc/settings.py | 9 +++------ tests/prefect/test_customizers.py | 1 - 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/src/quacc/settings.py b/src/quacc/settings.py index ffb916c7ad..7a586e15ce 100644 --- a/src/quacc/settings.py +++ b/src/quacc/settings.py @@ -4,11 +4,11 @@ import inspect import os -from datetime import datetime, timezone -from random import randint from contextlib import contextmanager +from datetime import datetime, timezone from functools import wraps from pathlib import Path +from random import randint from shutil import which from typing import TYPE_CHECKING, Literal, Optional, Union @@ -627,10 +627,7 @@ def nest_results_dir_wrap(func: Callable) -> Callable: """ from quacc import get_settings - if getattr(func, "_changed", False): - changes = func._changes - else: - changes = [] + changes = func._changes if getattr(func, "_changed", False) else [] # Get the settings from the calling function's context results_parent_dir = get_settings().RESULTS_DIR diff --git a/tests/prefect/test_customizers.py b/tests/prefect/test_customizers.py index 217371a16a..75a2d7a356 100644 --- a/tests/prefect/test_customizers.py +++ b/tests/prefect/test_customizers.py @@ -101,7 +101,6 @@ def my_flow(): def test_nested_output_directory(tmp_path_factory): - @job def write_file_job(name="job.txt"): results_file_path = Path(get_settings().RESULTS_DIR, name) From ea726557060bc84c209a5517e8669744a5a6f1a8 Mon Sep 17 00:00:00 2001 From: Zack Ulissi Date: Sat, 22 Jun 2024 00:24:54 +0000 Subject: [PATCH 09/17] typo in code and settings --- src/quacc/settings.py | 2 +- tests/prefect/test_customizers.py | 9 +++++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/quacc/settings.py b/src/quacc/settings.py index ffb916c7ad..4396eec3ee 100644 --- a/src/quacc/settings.py +++ b/src/quacc/settings.py @@ -630,7 +630,7 @@ def nest_results_dir_wrap(func: Callable) -> Callable: if getattr(func, "_changed", False): changes = func._changes else: - changes = [] + changes = {} # Get the settings from the calling function's context results_parent_dir = get_settings().RESULTS_DIR diff --git a/tests/prefect/test_customizers.py b/tests/prefect/test_customizers.py index 217371a16a..589821e3aa 100644 --- a/tests/prefect/test_customizers.py +++ b/tests/prefect/test_customizers.py @@ -121,7 +121,12 @@ def write_file_flow(name="flow.txt", job_decorators=None): return job_results_file_path, flow_results_file_path # Test with redecorating a job in a flow - job_results_file_path, flow_results_file_path = write_file_flow().result() + job_results_file_path, flow_results_file_path = write_file_flow() + job_results_file_path = job_results_file_path.result() assert Path(job_results_file_path).exists() assert Path(flow_results_file_path).exists() - assert Path(job_results_file_path).parent == Path(flow_results_file_path) + + # Check the job output is in a subfolder of the flow folder + assert ( + Path(job_results_file_path).parent.parent == Path(flow_results_file_path).parent + ) From 58ac8994f70f22700f6c97ea5b25c7e2d8cc68b4 Mon Sep 17 00:00:00 2001 From: Zack Ulissi Date: Sat, 22 Jun 2024 00:35:10 +0000 Subject: [PATCH 10/17] refactor prefect decoration code --- src/quacc/wflow_tools/decorators.py | 106 ++++++++++++++-------------- 1 file changed, 52 insertions(+), 54 deletions(-) diff --git a/src/quacc/wflow_tools/decorators.py b/src/quacc/wflow_tools/decorators.py index 22461b9dbc..0571ffa283 100644 --- a/src/quacc/wflow_tools/decorators.py +++ b/src/quacc/wflow_tools/decorators.py @@ -176,36 +176,7 @@ def wrapper(*f_args, **f_kwargs): return task(_func, namespace=_func.__module__, **kwargs) elif settings.WORKFLOW_ENGINE == "prefect": - from prefect import task - - if settings.PREFECT_AUTO_SUBMIT and settings.NESTED_RESULTS_DIR: - - @wraps(_func) - def wrapper(*f_args, **f_kwargs): - adjusted_results_func = nest_results_dir_wrap(_func) - decorated = task(adjusted_results_func, **kwargs) - return decorated.submit(*f_args, **f_kwargs) - - return wrapper - elif (not settings.PREFECT_AUTO_SUBMIT) and settings.NESTED_RESULTS_DIR: - - @wraps(_func) - def wrapper(*f_args, **f_kwargs): - adjusted_results_func = nest_results_dir_wrap(_func) - return task(adjusted_results_func, **kwargs)(*f_args, **f_kwargs) - - return wrapper - elif settings.PREFECT_AUTO_SUBMIT and (not settings.NESTED_RESULTS_DIR): - - @wraps(_func) - def wrapper(*f_args, **f_kwargs): - decorated = task(_func, **kwargs) - return decorated.submit(*f_args, **f_kwargs) - - return wrapper - else: - return task(_func, **kwargs) - + return _decorate_prefect_job(_func, kwargs, settings) else: return _func @@ -368,20 +339,9 @@ def workflow(a, b, c): return task(_func, namespace=_func.__module__, **kwargs) elif settings.WORKFLOW_ENGINE == "prefect": - from prefect import flow as prefect_flow - if settings.NESTED_RESULTS_DIR: + return _decorate_prefect_flow_subflow(_func, kwargs, settings) - @wraps(_func) - def wrapper(*f_args, **f_kwargs): - adjusted_results_func = nest_results_dir_wrap(_func) - return prefect_flow( - adjusted_results_func, validate_parameters=False, **kwargs - )(*f_args, **f_kwargs) - - return wrapper - else: - return prefect_flow(_func, validate_parameters=False, **kwargs) else: return _func @@ -610,20 +570,9 @@ def wrapper(*f_args, **f_kwargs): return join_app(wrapped_fn, **kwargs) elif settings.WORKFLOW_ENGINE == "prefect": - from prefect import flow as prefect_flow - - if settings.NESTED_RESULTS_DIR: - @wraps(_func) - def wrapper(*f_args, **f_kwargs): - adjusted_results_func = nest_results_dir_wrap(_func) - return prefect_flow( - adjusted_results_func, validate_parameters=False, **kwargs - )(*f_args, **f_kwargs) + return _decorate_prefect_flow_subflow(_func, kwargs, settings) - return wrapper - else: - return prefect_flow(_func, validate_parameters=False, **kwargs) elif settings.WORKFLOW_ENGINE == "redun": from redun import task @@ -632,6 +581,55 @@ def wrapper(*f_args, **f_kwargs): return _func +def _decorate_prefect_job(_func, kwargs, settings): + from prefect import task + + if settings.PREFECT_AUTO_SUBMIT and settings.NESTED_RESULTS_DIR: + + @wraps(_func) + def wrapper(*f_args, **f_kwargs): + adjusted_results_func = nest_results_dir_wrap(_func) + decorated = task(adjusted_results_func, **kwargs) + return decorated.submit(*f_args, **f_kwargs) + + return wrapper + elif (not settings.PREFECT_AUTO_SUBMIT) and settings.NESTED_RESULTS_DIR: + + @wraps(_func) + def wrapper(*f_args, **f_kwargs): + adjusted_results_func = nest_results_dir_wrap(_func) + return task(adjusted_results_func, **kwargs)(*f_args, **f_kwargs) + + return wrapper + elif settings.PREFECT_AUTO_SUBMIT and (not settings.NESTED_RESULTS_DIR): + + @wraps(_func) + def wrapper(*f_args, **f_kwargs): + decorated = task(_func, **kwargs) + return decorated.submit(*f_args, **f_kwargs) + + return wrapper + else: + return task(_func, **kwargs) + + +def _decorate_prefect_flow_subflow(_func, kwargs, settings): + from prefect import flow as prefect_flow + + if settings.NESTED_RESULTS_DIR: + + @wraps(_func) + def wrapper(*f_args, **f_kwargs): + adjusted_results_func = nest_results_dir_wrap(_func) + return prefect_flow( + adjusted_results_func, validate_parameters=False, **kwargs + )(*f_args, **f_kwargs) + + return wrapper + else: + return prefect_flow(_func, validate_parameters=False, **kwargs) + + def _get_parsl_wrapped_func( func: Callable, decorator_kwargs: dict[str, Any] ) -> Callable: From 63e0224c3f3cddd883b058aece4b55fe1f0b3985 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 22 Jun 2024 00:35:23 +0000 Subject: [PATCH 11/17] pre-commit auto-fixes --- src/quacc/wflow_tools/decorators.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/quacc/wflow_tools/decorators.py b/src/quacc/wflow_tools/decorators.py index 0571ffa283..9ce40e2e27 100644 --- a/src/quacc/wflow_tools/decorators.py +++ b/src/quacc/wflow_tools/decorators.py @@ -339,7 +339,6 @@ def workflow(a, b, c): return task(_func, namespace=_func.__module__, **kwargs) elif settings.WORKFLOW_ENGINE == "prefect": - return _decorate_prefect_flow_subflow(_func, kwargs, settings) else: @@ -570,7 +569,6 @@ def wrapper(*f_args, **f_kwargs): return join_app(wrapped_fn, **kwargs) elif settings.WORKFLOW_ENGINE == "prefect": - return _decorate_prefect_flow_subflow(_func, kwargs, settings) elif settings.WORKFLOW_ENGINE == "redun": From 28224da8774943b646c35bf1f863358df43e9f87 Mon Sep 17 00:00:00 2001 From: Andrew Rosen Date: Wed, 26 Jun 2024 23:19:54 -0700 Subject: [PATCH 12/17] Minor refactoring --- src/quacc/settings.py | 4 +-- src/quacc/wflow_tools/decorators.py | 44 ++++++++++++++--------------- tests/prefect/test_customizers.py | 6 ++-- 3 files changed, 28 insertions(+), 26 deletions(-) diff --git a/src/quacc/settings.py b/src/quacc/settings.py index c397a30029..6623eb1f06 100644 --- a/src/quacc/settings.py +++ b/src/quacc/settings.py @@ -104,7 +104,7 @@ class QuaccSettings(BaseSettings): True, description=( """ - Whether to nest the results dir by the calling flow/subflow/etc + Whether to automatically nest the results directories by the calling flow, subflow, etc. """ ), ) @@ -627,7 +627,7 @@ def nest_results_dir_wrap(func: Callable) -> Callable: """ from quacc import get_settings - changes = func._changes if getattr(func, "_changed", False) else {} + changes = getattr(func, "_changes", {}) # Get the settings from the calling function's context results_parent_dir = get_settings().RESULTS_DIR diff --git a/src/quacc/wflow_tools/decorators.py b/src/quacc/wflow_tools/decorators.py index 9ce40e2e27..3c439ec5c6 100644 --- a/src/quacc/wflow_tools/decorators.py +++ b/src/quacc/wflow_tools/decorators.py @@ -3,6 +3,7 @@ from __future__ import annotations from functools import partial, wraps +from importlib.util import find_spec from typing import TYPE_CHECKING, TypeVar from quacc.settings import change_settings_wrap, nest_results_dir_wrap @@ -14,6 +15,12 @@ if TYPE_CHECKING: from typing import Any, Callable + from quacc.settings import QuaccSettings + + if find_spec("prefect"): + from prefect import Flow as PrefectFlow + from prefect import Task + def job(_func: Callable | None = None, **kwargs) -> Job: """ @@ -579,39 +586,32 @@ def wrapper(*f_args, **f_kwargs): return _func -def _decorate_prefect_job(_func, kwargs, settings): +def _decorate_prefect_job( + _func: Callable, kwargs: dict[str, Any], settings: QuaccSettings +) -> Task: from prefect import task - if settings.PREFECT_AUTO_SUBMIT and settings.NESTED_RESULTS_DIR: - - @wraps(_func) - def wrapper(*f_args, **f_kwargs): - adjusted_results_func = nest_results_dir_wrap(_func) - decorated = task(adjusted_results_func, **kwargs) - return decorated.submit(*f_args, **f_kwargs) - - return wrapper - elif (not settings.PREFECT_AUTO_SUBMIT) and settings.NESTED_RESULTS_DIR: - - @wraps(_func) - def wrapper(*f_args, **f_kwargs): - adjusted_results_func = nest_results_dir_wrap(_func) - return task(adjusted_results_func, **kwargs)(*f_args, **f_kwargs) - - return wrapper - elif settings.PREFECT_AUTO_SUBMIT and (not settings.NESTED_RESULTS_DIR): + if settings.PREFECT_AUTO_SUBMIT or settings.NESTED_RESULTS_DIR: @wraps(_func) def wrapper(*f_args, **f_kwargs): - decorated = task(_func, **kwargs) - return decorated.submit(*f_args, **f_kwargs) + if settings.NESTED_RESULTS_DIR: + decorated = task(nest_results_dir_wrap(_func),**kwargs) + else: + decorated = task(_func, **kwargs) + if settings.PREFECT_AUTO_SUBMIT: + return decorated.submit(*f_args, **f_kwargs) + else: + return decorated(*f_args, **f_kwargs) return wrapper else: return task(_func, **kwargs) -def _decorate_prefect_flow_subflow(_func, kwargs, settings): +def _decorate_prefect_flow_subflow( + _func: Callable, kwargs: dict[str, Any], settings: QuaccSettings +) -> PrefectFlow: from prefect import flow as prefect_flow if settings.NESTED_RESULTS_DIR: diff --git a/tests/prefect/test_customizers.py b/tests/prefect/test_customizers.py index 5680169b6c..338c69c5a0 100644 --- a/tests/prefect/test_customizers.py +++ b/tests/prefect/test_customizers.py @@ -100,7 +100,9 @@ def my_flow(): assert Path(tmp_dir2 / "job.txt").exists() -def test_nested_output_directory(tmp_path_factory): +def test_nested_output_directory(tmp_path,monkeypatch): + monkeypatch.chdir(tmp_path) + @job def write_file_job(name="job.txt"): results_file_path = Path(get_settings().RESULTS_DIR, name) @@ -110,7 +112,7 @@ def write_file_job(name="job.txt"): return results_file_path @flow - def write_file_flow(name="flow.txt", job_decorators=None): + def write_file_flow(name="flow.txt"): job_results_file_path = write_file_job() flow_results_file_path = Path(get_settings().RESULTS_DIR, name) From 094ce69d6d61b95ee4b8fdd6b92830b2fd2fc5c1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 27 Jun 2024 06:20:09 +0000 Subject: [PATCH 13/17] pre-commit auto-fixes --- src/quacc/wflow_tools/decorators.py | 4 ++-- tests/prefect/test_customizers.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/quacc/wflow_tools/decorators.py b/src/quacc/wflow_tools/decorators.py index 3c439ec5c6..7e888d75f9 100644 --- a/src/quacc/wflow_tools/decorators.py +++ b/src/quacc/wflow_tools/decorators.py @@ -3,7 +3,6 @@ from __future__ import annotations from functools import partial, wraps -from importlib.util import find_spec from typing import TYPE_CHECKING, TypeVar from quacc.settings import change_settings_wrap, nest_results_dir_wrap @@ -13,6 +12,7 @@ Subflow = TypeVar("Subflow") if TYPE_CHECKING: + from importlib.util import find_spec from typing import Any, Callable from quacc.settings import QuaccSettings @@ -596,7 +596,7 @@ def _decorate_prefect_job( @wraps(_func) def wrapper(*f_args, **f_kwargs): if settings.NESTED_RESULTS_DIR: - decorated = task(nest_results_dir_wrap(_func),**kwargs) + decorated = task(nest_results_dir_wrap(_func), **kwargs) else: decorated = task(_func, **kwargs) if settings.PREFECT_AUTO_SUBMIT: diff --git a/tests/prefect/test_customizers.py b/tests/prefect/test_customizers.py index 338c69c5a0..b932c95035 100644 --- a/tests/prefect/test_customizers.py +++ b/tests/prefect/test_customizers.py @@ -100,7 +100,7 @@ def my_flow(): assert Path(tmp_dir2 / "job.txt").exists() -def test_nested_output_directory(tmp_path,monkeypatch): +def test_nested_output_directory(tmp_path, monkeypatch): monkeypatch.chdir(tmp_path) @job From 4bfbc747e7884adbf096f1c77973ff72b2b628b1 Mon Sep 17 00:00:00 2001 From: Andrew Rosen Date: Wed, 26 Jun 2024 23:20:42 -0700 Subject: [PATCH 14/17] Fix type hitn --- src/quacc/wflow_tools/decorators.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/quacc/wflow_tools/decorators.py b/src/quacc/wflow_tools/decorators.py index 3c439ec5c6..57aeb833b0 100644 --- a/src/quacc/wflow_tools/decorators.py +++ b/src/quacc/wflow_tools/decorators.py @@ -17,7 +17,7 @@ from quacc.settings import QuaccSettings - if find_spec("prefect"): + if bool(find_spec("prefect")): from prefect import Flow as PrefectFlow from prefect import Task @@ -596,7 +596,7 @@ def _decorate_prefect_job( @wraps(_func) def wrapper(*f_args, **f_kwargs): if settings.NESTED_RESULTS_DIR: - decorated = task(nest_results_dir_wrap(_func),**kwargs) + decorated = task(nest_results_dir_wrap(_func), **kwargs) else: decorated = task(_func, **kwargs) if settings.PREFECT_AUTO_SUBMIT: @@ -618,10 +618,10 @@ def _decorate_prefect_flow_subflow( @wraps(_func) def wrapper(*f_args, **f_kwargs): - adjusted_results_func = nest_results_dir_wrap(_func) - return prefect_flow( - adjusted_results_func, validate_parameters=False, **kwargs - )(*f_args, **f_kwargs) + decorated = prefect_flow( + nest_results_dir_wrap(_func), validate_parameters=False, **kwargs + ) + return decorated(*f_args, **f_kwargs) return wrapper else: From 7f7a690373c38a2df4095426a808f040c6975d69 Mon Sep 17 00:00:00 2001 From: Andrew Rosen Date: Wed, 26 Jun 2024 23:21:06 -0700 Subject: [PATCH 15/17] Fix imports --- src/quacc/wflow_tools/decorators.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/quacc/wflow_tools/decorators.py b/src/quacc/wflow_tools/decorators.py index 57aeb833b0..e64a69964e 100644 --- a/src/quacc/wflow_tools/decorators.py +++ b/src/quacc/wflow_tools/decorators.py @@ -3,7 +3,6 @@ from __future__ import annotations from functools import partial, wraps -from importlib.util import find_spec from typing import TYPE_CHECKING, TypeVar from quacc.settings import change_settings_wrap, nest_results_dir_wrap @@ -13,6 +12,7 @@ Subflow = TypeVar("Subflow") if TYPE_CHECKING: + from importlib.util import find_spec from typing import Any, Callable from quacc.settings import QuaccSettings From 75d249c0b05c21d37e31856ffd61818e1b91f440 Mon Sep 17 00:00:00 2001 From: Andrew Rosen Date: Wed, 26 Jun 2024 23:41:20 -0700 Subject: [PATCH 16/17] Fix spacing --- src/quacc/settings.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/quacc/settings.py b/src/quacc/settings.py index 6623eb1f06..83f766ff3b 100644 --- a/src/quacc/settings.py +++ b/src/quacc/settings.py @@ -614,7 +614,6 @@ def nest_results_dir_wrap(func: Callable) -> Callable: """ Wraps a function with the change_settings context manager using a nested RESULTS_DIR - Parameters ---------- func From b67ddb8015ad4a3d40f9bd529db0a1b25fa20ac4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 27 Jul 2024 15:33:10 +0000 Subject: [PATCH 17/17] pre-commit auto-fixes --- src/quacc/wflow_tools/decorators.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/quacc/wflow_tools/decorators.py b/src/quacc/wflow_tools/decorators.py index 2d3f2c1dc3..0da0a24ccc 100644 --- a/src/quacc/wflow_tools/decorators.py +++ b/src/quacc/wflow_tools/decorators.py @@ -3,22 +3,23 @@ from __future__ import annotations from functools import partial, wraps -from typing import Any, Callable +from typing import TYPE_CHECKING, Any, Callable from quacc.settings import change_settings_wrap, nest_results_dir_wrap +if TYPE_CHECKING: + from prefect import Flow as PrefectFlow + Job = Callable[..., Any] Flow = Callable[..., Any] Subflow = Callable[..., Any] if TYPE_CHECKING: from importlib.util import find_spec - from typing import Any, Callable from quacc.settings import QuaccSettings if bool(find_spec("prefect")): - from prefect import Flow as PrefectFlow from prefect import Task