From 34e1a36d0125216f27bfb513eb7479b758453ba9 Mon Sep 17 00:00:00 2001 From: "Andrew S. Rosen" Date: Sat, 4 Jan 2025 20:05:35 -0500 Subject: [PATCH 1/2] Fix Prefect --- src/quacc/wflow_tools/prefect_utils.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/src/quacc/wflow_tools/prefect_utils.py b/src/quacc/wflow_tools/prefect_utils.py index 5ec1923700..67066a88eb 100644 --- a/src/quacc/wflow_tools/prefect_utils.py +++ b/src/quacc/wflow_tools/prefect_utils.py @@ -35,7 +35,7 @@ def resolve_futures_to_results(expr: PrefectFuture | Any) -> State | Any: def _collect_futures(futures, expr, context): # Expressions inside quotes should not be traversed if isinstance(context.get("annotation"), quote): - raise StopVisiting + raise StopVisiting() if isinstance(expr, PrefectFuture): futures.add(expr) @@ -56,20 +56,18 @@ def _collect_futures(futures, expr, context): future.wait() if future.state.is_completed(): result = future.state.result() - if isinstance(result, BaseResult): - result = result.get() - elif isinstance(result, ResultRecord): + if isinstance(result, ResultRecord): result = result.result results.append(result) else: raise BaseException("At least one result did not complete successfully") - states_by_future = dict(zip(futures, results, strict=False)) + states_by_future = dict(zip(futures, results)) def replace_futures_with_states(expr, context): # Expressions inside quotes should not be modified if isinstance(context.get("annotation"), quote): - raise StopVisiting + raise StopVisiting() if isinstance(expr, PrefectFuture): return states_by_future[expr] @@ -121,20 +119,18 @@ def _collect_futures(futures, expr, context): future.wait() if future.state.is_completed(): result = future.state.result() - if isinstance(result, BaseResult): - result = result.get() - elif isinstance(result, ResultRecord): + if isinstance(result, ResultRecord): result = result.result results.append(result) else: raise BaseException("At least one result did not complete successfully") - states_by_future = dict(zip(futures, results, strict=False)) + states_by_future = dict(zip(futures, results)) def replace_futures_with_states(expr, context): # Expressions inside quotes should not be modified if isinstance(context.get("annotation"), quote): - raise StopVisiting + raise StopVisiting() if isinstance(expr, PrefectFuture): return states_by_future[expr] From ec506e5401c3f1ab188bd8f7cde9778ae9a262d7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 5 Jan 2025 01:05:45 +0000 Subject: [PATCH 2/2] pre-commit auto-fixes --- src/quacc/wflow_tools/prefect_utils.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/quacc/wflow_tools/prefect_utils.py b/src/quacc/wflow_tools/prefect_utils.py index 67066a88eb..155827fdd2 100644 --- a/src/quacc/wflow_tools/prefect_utils.py +++ b/src/quacc/wflow_tools/prefect_utils.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Any from prefect.futures import PrefectFuture -from prefect.results import BaseResult, ResultRecord +from prefect.results import ResultRecord from prefect.utilities.annotations import quote from prefect.utilities.collections import StopVisiting, visit_collection from typing_extensions import TypeVar @@ -35,7 +35,7 @@ def resolve_futures_to_results(expr: PrefectFuture | Any) -> State | Any: def _collect_futures(futures, expr, context): # Expressions inside quotes should not be traversed if isinstance(context.get("annotation"), quote): - raise StopVisiting() + raise StopVisiting if isinstance(expr, PrefectFuture): futures.add(expr) @@ -62,12 +62,12 @@ def _collect_futures(futures, expr, context): else: raise BaseException("At least one result did not complete successfully") - states_by_future = dict(zip(futures, results)) + states_by_future = dict(zip(futures, results, strict=False)) def replace_futures_with_states(expr, context): # Expressions inside quotes should not be modified if isinstance(context.get("annotation"), quote): - raise StopVisiting() + raise StopVisiting if isinstance(expr, PrefectFuture): return states_by_future[expr] @@ -125,12 +125,12 @@ def _collect_futures(futures, expr, context): else: raise BaseException("At least one result did not complete successfully") - states_by_future = dict(zip(futures, results)) + states_by_future = dict(zip(futures, results, strict=False)) def replace_futures_with_states(expr, context): # Expressions inside quotes should not be modified if isinstance(context.get("annotation"), quote): - raise StopVisiting() + raise StopVisiting if isinstance(expr, PrefectFuture): return states_by_future[expr]