prefect-client 3.0.0rc1__py3-none-any.whl → 3.0.0rc3__py3-none-any.whl

prefect/tasks.py

@@ -22,6 +22,7 @@ from typing import (
     Optional,
     Set,
     Tuple,
+    Type,
     TypeVar,
     Union,
     cast,
@@ -43,6 +44,7 @@ from prefect.context import (
 )
 from prefect.futures import PrefectDistributedFuture, PrefectFuture
 from prefect.logging.loggers import get_logger
+from prefect.records.cache_policies import DEFAULT, CachePolicy
 from prefect.results import ResultFactory, ResultSerializer, ResultStorage
 from prefect.settings import (
     PREFECT_TASK_DEFAULT_RETRIES,
@@ -62,7 +64,6 @@ from prefect.utilities.importtools import to_qualified_name
 if TYPE_CHECKING:
     from prefect.client.orchestration import PrefectClient
     from prefect.context import TaskRunContext
-    from prefect.task_runners import BaseTaskRunner
     from prefect.transactions import Transaction
 
 T = TypeVar("T")  # Generic type var for capturing the inner return type of async funcs
@@ -122,6 +123,57 @@ def exponential_backoff(backoff_factor: float) -> Callable[[int], List[float]]:
     return retry_backoff_callable
 
 
+def _infer_parent_task_runs(
+    flow_run_context: Optional[FlowRunContext],
+    task_run_context: Optional[TaskRunContext],
+    parameters: Dict[str, Any],
+):
+    """
+    Attempt to infer the parent task runs for this task run based on the
+    provided flow run and task run contexts, as well as any parameters. It is
+    assumed that the task run is running within those contexts.
+    If any parameter comes from a running task run, that task run is considered
+    a parent. This is expected to happen when task inputs are yielded from
+    generator tasks.
+    """
+    parents = []
+
+    # check if this task has a parent task run based on running in another
+    # task run's existing context. A task run is only considered a parent if
+    # it is in the same flow run (because otherwise presumably the child is
+    # in a subflow, so the subflow serves as the parent) or if there is no
+    # flow run
+    if task_run_context:
+        # there is no flow run
+        if not flow_run_context:
+            parents.append(TaskRunResult(id=task_run_context.task_run.id))
+        # there is a flow run and the task run is in the same flow run
+        elif flow_run_context and task_run_context.task_run.flow_run_id == getattr(
+            flow_run_context.flow_run, "id", None
+        ):
+            parents.append(TaskRunResult(id=task_run_context.task_run.id))
+
+    # parent dependency tracking: for every provided parameter value, try to
+    # load the corresponding task run state. If the task run state is still
+    # running, we consider it a parent task run. Note this is only done if
+    # there is an active flow run context because dependencies are only
+    # tracked within the same flow run.
+    if flow_run_context:
+        for v in parameters.values():
+            if isinstance(v, State):
+                upstream_state = v
+            elif isinstance(v, PrefectFuture):
+                upstream_state = v.state
+            else:
+                upstream_state = flow_run_context.task_run_results.get(id(v))
+            if upstream_state and upstream_state.is_running():
+                parents.append(
+                    TaskRunResult(id=upstream_state.state_details.task_run_id)
+                )
+
+    return parents
+
+
 @PrefectObjectRegistry.register_instances
 class Task(Generic[P, R]):
     """
@@ -145,6 +197,7 @@ class Task(Generic[P, R]):
             tags are combined with any tags defined by a `prefect.tags` context at
             task runtime.
         version: An optional string specifying the version of this task definition
+        cache_policy: A cache policy that determines the level of caching for this task
         cache_key_fn: An optional callable that, given the task run context and call
             parameters, generates a string key; if the key matches a previous completed
             state, that state result will be restored instead of running the task again.
@@ -204,6 +257,7 @@ class Task(Generic[P, R]):
         description: Optional[str] = None,
         tags: Optional[Iterable[str]] = None,
         version: Optional[str] = None,
+        cache_policy: Optional[CachePolicy] = NotSet,
         cache_key_fn: Optional[
             Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
         ] = None,
@@ -266,7 +320,18 @@ class Task(Generic[P, R]):
         self.description = description or inspect.getdoc(fn)
         update_wrapper(self, fn)
         self.fn = fn
-        self.isasync = inspect.iscoroutinefunction(self.fn)
+
+        # the task is considered async if its function is async or an async
+        # generator
+        self.isasync = inspect.iscoroutinefunction(
+            self.fn
+        ) or inspect.isasyncgenfunction(self.fn)
+
+        # the task is considered a generator if its function is a generator or
+        # an async generator
+        self.isgenerator = inspect.isgeneratorfunction(
+            self.fn
+        ) or inspect.isasyncgenfunction(self.fn)
 
         if not name:
             if not hasattr(self.fn, "__name__"):
@@ -303,10 +368,23 @@ class Task(Generic[P, R]):
 
             self.task_key = f"{self.fn.__qualname__}-{task_origin_hash}"
 
+        # TODO: warn of precedence of cache policies and cache key fn if both provided?
+        if cache_key_fn:
+            cache_policy = CachePolicy.from_cache_key_fn(cache_key_fn)
+
+        # TODO: manage expiration and cache refresh
         self.cache_key_fn = cache_key_fn
         self.cache_expiration = cache_expiration
         self.refresh_cache = refresh_cache
 
+        if cache_policy is NotSet and result_storage_key is None:
+            self.cache_policy = DEFAULT
+        elif result_storage_key:
+            # TODO: handle this situation with double storage
+            self.cache_policy = None
+        else:
+            self.cache_policy = cache_policy
+
         # TaskRunPolicy settings
         # TODO: We can instantiate a `TaskRunPolicy` and add Pydantic bound checks to
         #       validate that the user passes positive numbers here
@@ -352,33 +430,57 @@ class Task(Generic[P, R]):
         self.retry_condition_fn = retry_condition_fn
         self.viz_return_value = viz_return_value
 
+    @property
+    def ismethod(self) -> bool:
+        return hasattr(self.fn, "__prefect_self__")
+
+    def __get__(self, instance, owner):
+        """
+        Implement the descriptor protocol so that the task can be used as an instance method.
+        When an instance method is loaded, this method is called with the "self" instance as
+        an argument. We return a copy of the task with that instance bound to the task's function.
+        """
+
+        # if no instance is provided, it's being accessed on the class
+        if instance is None:
+            return self
+
+        # if the task is being accessed on an instance, bind the instance to the __prefect_self__ attribute
+        # of the task's function. This will allow it to be automatically added to the task's parameters
+        else:
+            bound_task = copy(self)
+            bound_task.fn.__prefect_self__ = instance
+            return bound_task
+
     def with_options(
         self,
         *,
-        name: str = None,
-        description: str = None,
-        tags: Iterable[str] = None,
-        cache_key_fn: Callable[
-            ["TaskRunContext", Dict[str, Any]], Optional[str]
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        tags: Optional[Iterable[str]] = None,
+        cache_policy: Union[CachePolicy, Type[NotSet]] = NotSet,
+        cache_key_fn: Optional[
+            Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
         ] = None,
         task_run_name: Optional[Union[Callable[[], str], str]] = None,
-        cache_expiration: datetime.timedelta = None,
-        retries: Optional[int] = NotSet,
+        cache_expiration: Optional[datetime.timedelta] = None,
+        retries: Union[int, Type[NotSet]] = NotSet,
         retry_delay_seconds: Union[
             float,
             int,
             List[float],
             Callable[[int], List[float]],
+            Type[NotSet],
         ] = NotSet,
-        retry_jitter_factor: Optional[float] = NotSet,
-        persist_result: Optional[bool] = NotSet,
-        result_storage: Optional[ResultStorage] = NotSet,
-        result_serializer: Optional[ResultSerializer] = NotSet,
-        result_storage_key: Optional[str] = NotSet,
+        retry_jitter_factor: Union[float, Type[NotSet]] = NotSet,
+        persist_result: Union[bool, Type[NotSet]] = NotSet,
+        result_storage: Union[ResultStorage, Type[NotSet]] = NotSet,
+        result_serializer: Union[ResultSerializer, Type[NotSet]] = NotSet,
+        result_storage_key: Union[str, Type[NotSet]] = NotSet,
         cache_result_in_memory: Optional[bool] = None,
-        timeout_seconds: Union[int, float] = None,
-        log_prints: Optional[bool] = NotSet,
-        refresh_cache: Optional[bool] = NotSet,
+        timeout_seconds: Union[int, float, None] = None,
+        log_prints: Union[bool, Type[NotSet]] = NotSet,
+        refresh_cache: Union[bool, Type[NotSet]] = NotSet,
         on_completion: Optional[
             List[Callable[["Task", TaskRun, State], Union[Awaitable[None], None]]]
         ] = None,
@@ -469,6 +571,9 @@ class Task(Generic[P, R]):
             name=name or self.name,
             description=description or self.description,
             tags=tags or copy(self.tags),
+            cache_policy=cache_policy
+            if cache_policy is not NotSet
+            else self.cache_policy,
             cache_key_fn=cache_key_fn or self.cache_key_fn,
             cache_expiration=cache_expiration or self.cache_expiration,
             task_run_name=task_run_name,
@@ -569,7 +674,7 @@ class Task(Generic[P, R]):
         async with client:
             if not flow_run_context:
                 dynamic_key = f"{self.task_key}-{str(uuid4().hex)}"
-                task_run_name = f"{self.name}-{dynamic_key[:NUM_CHARS_DYNAMIC_KEY]}"
+                task_run_name = self.name
             else:
                 dynamic_key = _dynamic_key_for_task_run(
                     context=flow_run_context, task=self
@@ -582,7 +687,7 @@ class Task(Generic[P, R]):
             else:
                 state = Pending()
 
-            # store parameters for background tasks so that task servers
+            # store parameters for background tasks so that task worker
             # can retrieve them at runtime
             if deferred and (parameters or wait_for):
                 parameters_id = uuid4()
@@ -605,27 +710,15 @@ class Task(Generic[P, R]):
                 k: collect_task_run_inputs_sync(v) for k, v in parameters.items()
             }
 
-            # check if this task has a parent task run based on running in another
-            # task run's existing context. A task run is only considered a parent if
-            # it is in the same flow run (because otherwise presumably the child is
-            # in a subflow, so the subflow serves as the parent) or if there is no
-            # flow run
-            if parent_task_run_context:
-                # there is no flow run
-                if not flow_run_context:
-                    task_inputs["__parents__"] = [
-                        TaskRunResult(id=parent_task_run_context.task_run.id)
-                    ]
-                # there is a flow run and the task run is in the same flow run
-                elif (
-                    flow_run_context
-                    and parent_task_run_context.task_run.flow_run_id
-                    == getattr(flow_run_context.flow_run, "id", None)
-                ):
-                    task_inputs["__parents__"] = [
-                        TaskRunResult(id=parent_task_run_context.task_run.id)
-                    ]
+            # collect all parent dependencies
+            if task_parents := _infer_parent_task_runs(
+                flow_run_context=flow_run_context,
+                task_run_context=parent_task_run_context,
+                parameters=parameters,
+            ):
+                task_inputs["__parents__"] = task_parents
 
+            # check wait for dependencies
             if wait_for:
                 task_inputs["wait_for"] = collect_task_run_inputs_sync(wait_for)
 
@@ -755,8 +848,6 @@ class Task(Generic[P, R]):
         """
         Submit a run of the task to the engine.
 
-        If writing an async task, this call must be awaited.
-
         Will create a new task run in the backing API and submit the task to the flow's
         task runner. This call only blocks execution while the task is being submitted,
         once it is submitted, the flow function will continue executing.
@@ -849,7 +940,11 @@ class Task(Generic[P, R]):
         flow_run_context = FlowRunContext.get()
 
         if not flow_run_context:
-            raise ValueError("Task.submit() must be called within a flow")
+            raise RuntimeError(
+                "Unable to determine task runner to use for submission. If you are"
+                " submitting a task outside of a flow, please use `.delay`"
+                " to submit the task run for deferred execution."
+            )
 
         task_viz_tracker = get_task_viz_tracker()
         if task_viz_tracker:
@@ -897,6 +992,7 @@ class Task(Generic[P, R]):
         *args: Any,
         return_state: bool = False,
         wait_for: Optional[Iterable[PrefectFuture]] = None,
+        deferred: bool = False,
         **kwargs: Any,
     ):
         """
@@ -1010,6 +1106,7 @@ class Task(Generic[P, R]):
             [[11, 21], [12, 22], [13, 23]]
         """
 
+        from prefect.task_runners import TaskRunner
         from prefect.utilities.visualization import (
             VisualizationUnsupportedError,
             get_task_viz_tracker,
@@ -1026,22 +1123,22 @@ class Task(Generic[P, R]):
                 "`task.map()` is not currently supported by `flow.visualize()`"
             )
 
-        if not flow_run_context:
-            # TODO: Should we split out background task mapping into a separate method
-            # like we do for the `submit`/`apply_async` split?
+        if deferred:
             parameters_list = expand_mapping_parameters(self.fn, parameters)
-            # TODO: Make this non-blocking once we can return a list of futures
-            # instead of a list of task runs
-            return [
-                run_coro_as_sync(self.create_run(parameters=parameters, deferred=True))
+            futures = [
+                self.apply_async(kwargs=parameters, wait_for=wait_for)
                 for parameters in parameters_list
             ]
-
-        from prefect.task_runners import TaskRunner
-
-        task_runner = flow_run_context.task_runner
-        assert isinstance(task_runner, TaskRunner)
-        futures = task_runner.map(self, parameters, wait_for)
+        elif task_runner := getattr(flow_run_context, "task_runner", None):
+            assert isinstance(task_runner, TaskRunner)
+            futures = task_runner.map(self, parameters, wait_for)
+        else:
+            raise RuntimeError(
+                "Unable to determine task runner to use for mapped task runs. If"
+                " you are mapping a task outside of a flow, please provide"
+                " `deferred=True` to submit the mapped task runs for deferred"
+                " execution."
+            )
         if return_state:
             states = []
             for future in futures:
@@ -1059,7 +1156,7 @@ class Task(Generic[P, R]):
         dependencies: Optional[Dict[str, Set[TaskRunInput]]] = None,
     ) -> PrefectDistributedFuture:
         """
-        Create a pending task run for a task server to execute.
+        Create a pending task run for a task worker to execute.
 
         Args:
             args: Arguments to run the task with
@@ -1181,7 +1278,7 @@ class Task(Generic[P, R]):
         """
         return self.apply_async(args=args, kwargs=kwargs)
 
-    def serve(self, task_runner: Optional["BaseTaskRunner"] = None) -> "Task":
+    def serve(self) -> "Task":
         """Serve the task using the provided task runner. This method is used to
         establish a websocket connection with the Prefect server and listen for
         submitted task runs to execute.
@@ -1198,9 +1295,9 @@ class Task(Generic[P, R]):
 
             >>> my_task.serve()
         """
-        from prefect.task_server import serve
+        from prefect.task_worker import serve
 
-        serve(self, task_runner=task_runner)
+        serve(self)
 
 
 @overload
@@ -1211,12 +1308,15 @@ def task(__fn: Callable[P, R]) -> Task[P, R]:
 @overload
 def task(
     *,
-    name: str = None,
-    description: str = None,
-    tags: Iterable[str] = None,
-    version: str = None,
-    cache_key_fn: Callable[["TaskRunContext", Dict[str, Any]], Optional[str]] = None,
-    cache_expiration: datetime.timedelta = None,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    tags: Optional[Iterable[str]] = None,
+    version: Optional[str] = None,
+    cache_policy: CachePolicy = NotSet,
+    cache_key_fn: Optional[
+        Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
+    ] = None,
+    cache_expiration: Optional[datetime.timedelta] = None,
     task_run_name: Optional[Union[Callable[[], str], str]] = None,
     retries: int = 0,
     retry_delay_seconds: Union[
@@ -1231,7 +1331,7 @@ def task(
     result_storage_key: Optional[str] = None,
     result_serializer: Optional[ResultSerializer] = None,
     cache_result_in_memory: bool = True,
-    timeout_seconds: Union[int, float] = None,
+    timeout_seconds: Union[int, float, None] = None,
     log_prints: Optional[bool] = None,
     refresh_cache: Optional[bool] = None,
     on_completion: Optional[List[Callable[["Task", TaskRun, State], None]]] = None,
@@ -1245,19 +1345,17 @@ def task(
 def task(
     __fn=None,
     *,
-    name: str = None,
-    description: str = None,
-    tags: Iterable[str] = None,
-    version: str = None,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    tags: Optional[Iterable[str]] = None,
+    version: Optional[str] = None,
+    cache_policy: Union[CachePolicy, Type[NotSet]] = NotSet,
     cache_key_fn: Callable[["TaskRunContext", Dict[str, Any]], Optional[str]] = None,
-    cache_expiration: datetime.timedelta = None,
+    cache_expiration: Optional[datetime.timedelta] = None,
     task_run_name: Optional[Union[Callable[[], str], str]] = None,
-    retries: int = None,
+    retries: Optional[int] = None,
     retry_delay_seconds: Union[
-        float,
-        int,
-        List[float],
-        Callable[[int], List[float]],
+        float, int, List[float], Callable[[int], List[float]], None
     ] = None,
     retry_jitter_factor: Optional[float] = None,
     persist_result: Optional[bool] = None,
@@ -1265,7 +1363,7 @@ def task(
     result_storage_key: Optional[str] = None,
     result_serializer: Optional[ResultSerializer] = None,
     cache_result_in_memory: bool = True,
-    timeout_seconds: Union[int, float] = None,
+    timeout_seconds: Union[int, float, None] = None,
     log_prints: Optional[bool] = None,
     refresh_cache: Optional[bool] = None,
     on_completion: Optional[List[Callable[["Task", TaskRun, State], None]]] = None,
@@ -1383,6 +1481,9 @@ def task(
     """
 
     if __fn:
+        if isinstance(__fn, (classmethod, staticmethod)):
+            method_decorator = type(__fn).__name__
+            raise TypeError(f"@{method_decorator} should be applied on top of @task")
         return cast(
             Task[P, R],
             Task(
@@ -1391,6 +1492,7 @@ def task(
                 description=description,
                 tags=tags,
                 version=version,
+                cache_policy=cache_policy,
                 cache_key_fn=cache_key_fn,
                 cache_expiration=cache_expiration,
                 task_run_name=task_run_name,
@@ -1420,6 +1522,7 @@ def task(
                 description=description,
                 tags=tags,
                 version=version,
+                cache_policy=cache_policy,
                 cache_key_fn=cache_key_fn,
                 cache_expiration=cache_expiration,
                 task_run_name=task_run_name,

prefect/transactions.py

@@ -7,17 +7,19 @@ from typing import (
     List,
     Optional,
     Type,
-    TypeVar,
 )
 
 from pydantic import Field
+from typing_extensions import Self
 
-from prefect.context import ContextModel
+from prefect.context import ContextModel, FlowRunContext, TaskRunContext
 from prefect.records import RecordStore
+from prefect.records.result_store import ResultFactoryStore
+from prefect.results import BaseResult, ResultFactory, get_default_result_storage
+from prefect.settings import PREFECT_DEFAULT_RESULT_STORAGE_BLOCK
+from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.collections import AutoEnum
 
-T = TypeVar("T")
-
 
 class IsolationLevel(AutoEnum):
     READ_COMMITTED = AutoEnum.auto()
@@ -52,8 +54,9 @@ class Transaction(ContextModel):
     on_rollback_hooks: List[Callable[["Transaction"], None]] = Field(
         default_factory=list
     )
+    overwrite: bool = False
     _staged_value: Any = None
-    __var__ = ContextVar("transaction")
+    __var__: ContextVar = ContextVar("transaction")
 
     def is_committed(self) -> bool:
         return self.state == TransactionState.COMMITTED
@@ -91,7 +94,8 @@ class Transaction(ContextModel):
         self._token = self.__var__.set(self)
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, *exc_info):
+        exc_type, exc_val, _ = exc_info
         if not self._token:
             raise RuntimeError(
                 "Asymmetric use of context. Context exit called without an enter."
@@ -122,11 +126,19 @@ class Transaction(ContextModel):
     def begin(self):
         # currently we only support READ_COMMITTED isolation
         # i.e., no locking behavior
-        if self.store and self.store.exists(key=self.key):
+        if (
+            not self.overwrite
+            and self.store
+            and self.key
+            and self.store.exists(key=self.key)
+        ):
             self.state = TransactionState.COMMITTED
 
-    def read(self) -> dict:
-        return self.store.read(key=self.key)
+    def read(self) -> BaseResult:
+        if self.store and self.key:
+            return self.store.read(key=self.key)
+        else:
+            return {}  # TODO: Determine what this should be
 
     def reset(self) -> None:
         parent = self.get_parent()
@@ -135,8 +147,9 @@ class Transaction(ContextModel):
             # parent takes responsibility
             parent.add_child(self)
 
-        self.__var__.reset(self._token)
-        self._token = None
+        if self._token:
+            self.__var__.reset(self._token)
+            self._token = None
 
         # do this below reset so that get_transaction() returns the relevant txn
         if parent and self.state == TransactionState.ROLLED_BACK:
@@ -164,7 +177,7 @@ class Transaction(ContextModel):
             for hook in self.on_commit_hooks:
                 hook(self)
 
-            if self.store:
+            if self.store and self.key:
                 self.store.write(key=self.key, value=self._staged_value)
             self.state = TransactionState.COMMITTED
             return True
@@ -173,11 +186,17 @@ class Transaction(ContextModel):
             return False
 
     def stage(
-        self, value: dict, on_rollback_hooks: list, on_commit_hooks: list
+        self,
+        value: BaseResult,
+        on_rollback_hooks: Optional[List] = None,
+        on_commit_hooks: Optional[List] = None,
     ) -> None:
         """
         Stage a value to be committed later.
         """
+        on_commit_hooks = on_commit_hooks or []
+        on_rollback_hooks = on_rollback_hooks or []
+
         if self.state != TransactionState.COMMITTED:
             self._staged_value = value
             self.on_rollback_hooks += on_rollback_hooks
@@ -202,11 +221,11 @@ class Transaction(ContextModel):
             return False
 
     @classmethod
-    def get_active(cls: Type[T]) -> Optional[T]:
+    def get_active(cls: Type[Self]) -> Optional[Self]:
         return cls.__var__.get(None)
 
 
-def get_transaction() -> Transaction:
+def get_transaction() -> Optional[Transaction]:
     return Transaction.get_active()
 
 
@@ -215,6 +234,63 @@ def transaction(
     key: Optional[str] = None,
     store: Optional[RecordStore] = None,
     commit_mode: CommitMode = CommitMode.LAZY,
+    overwrite: bool = False,
 ) -> Generator[Transaction, None, None]:
-    with Transaction(key=key, store=store, commit_mode=commit_mode) as txn:
+    """
+    A context manager for opening and managing a transaction.
+
+    Args:
+        - key: An identifier to use for the transaction
+        - store: The store to use for persisting the transaction result. If not provided,
+            a default store will be used based on the current run context.
+        - commit_mode: The commit mode controlling when the transaction and
+            child transactions are committed
+        - overwrite: Whether to overwrite an existing transaction record in the store
+
+    Yields:
+        - Transaction: An object representing the transaction state
+    """
+    # if there is no key, we won't persist a record
+    if key and not store:
+        flow_run_context = FlowRunContext.get()
+        task_run_context = TaskRunContext.get()
+        existing_factory = getattr(task_run_context, "result_factory", None) or getattr(
+            flow_run_context, "result_factory", None
+        )
+
+        if existing_factory and existing_factory.storage_block_id:
+            new_factory = existing_factory.model_copy(
+                update={
+                    "persist_result": True,
+                }
+            )
+        else:
+            default_storage = get_default_result_storage(_sync=True)
+            if not default_storage._block_document_id:
+                default_name = PREFECT_DEFAULT_RESULT_STORAGE_BLOCK.value().split("/")[
+                    -1
+                ]
+                default_storage.save(default_name, overwrite=True, _sync=True)
+            if existing_factory:
+                new_factory = existing_factory.model_copy(
+                    update={
+                        "persist_result": True,
+                        "storage_block": default_storage,
+                        "storage_block_id": default_storage._block_document_id,
+                    }
+                )
+            else:
+                new_factory = run_coro_as_sync(
+                    ResultFactory.default_factory(
+                        persist_result=True,
+                        result_storage=default_storage,
+                    )
+                )
+        store = ResultFactoryStore(
+            result_factory=new_factory,
+        )
+
+    with Transaction(
+        key=key, store=store, commit_mode=commit_mode, overwrite=overwrite
+    ) as txn:
         yield txn

prefect/types/__init__.py

@@ -20,7 +20,7 @@ timezone_set = available_timezones()
 NonNegativeInteger = Annotated[int, Field(ge=0)]
 PositiveInteger = Annotated[int, Field(gt=0)]
 NonNegativeFloat = Annotated[float, Field(ge=0.0)]
-TimeZone = Annotated[str, Field(default="UTC", pattern="|".join(timezone_set))]
+TimeZone = Annotated[str, Field(default="UTC", pattern="|".join(sorted(timezone_set)))]
 
 
 BANNED_CHARACTERS = ["/", "%", "&", ">", "<"]

prefect/utilities/asyncutils.py

@@ -314,7 +314,7 @@ def sync_compatible(async_fn: T, force_sync: bool = False) -> T:
     """
 
     @wraps(async_fn)
-    def coroutine_wrapper(*args, _sync: bool = None, **kwargs):
+    def coroutine_wrapper(*args, _sync: Optional[bool] = None, **kwargs):
         from prefect.context import MissingContextError, get_run_context
         from prefect.settings import (
             PREFECT_EXPERIMENTAL_DISABLE_SYNC_COMPAT,
@@ -376,8 +376,8 @@ def sync_compatible(async_fn: T, force_sync: bool = False) -> T:
 
 
 @asynccontextmanager
-async def asyncnullcontext():
-    yield
+async def asyncnullcontext(value=None):
+    yield value
 
 
 def sync(__async_fn: Callable[P, Awaitable[T]], *args: P.args, **kwargs: P.kwargs) -> T: