prefect-client 3.4.0__py3-none-any.whl → 3.4.1__py3-none-any.whl

prefect/events/schemas/automations.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import abc
 import textwrap
 from datetime import timedelta
@@ -103,7 +105,7 @@ class ResourceTrigger(Trigger, abc.ABC):
         default_factory=lambda: ResourceSpecification.model_validate({}),
         description="Labels for resources which this trigger will match.",
     )
-    match_related: ResourceSpecification = Field(
+    match_related: Union[ResourceSpecification, list[ResourceSpecification]] = Field(
         default_factory=lambda: ResourceSpecification.model_validate({}),
         description="Labels for related resources which this trigger will match.",
     )

prefect/events/schemas/events.py

@@ -13,7 +13,7 @@ from typing import (
     Tuple,
     Union,
 )
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from pydantic import (
     AfterValidator,
@@ -26,6 +26,7 @@ from typing_extensions import Annotated, Self
 
 import prefect.types._datetime
 from prefect._internal.schemas.bases import PrefectBaseModel
+from prefect._internal.uuid7 import uuid7
 from prefect.logging import get_logger
 from prefect.settings import (
     PREFECT_EVENTS_MAXIMUM_LABELS_PER_RESOURCE,
@@ -135,7 +136,7 @@ class Event(PrefectBaseModel):
         description="An open-ended set of data describing what happened",
     )
     id: UUID = Field(
-        default_factory=uuid4,
+        default_factory=uuid7,
         description="The client-provided identifier of this event",
     )
     follows: Optional[UUID] = Field(

prefect/flows.py

@@ -27,6 +27,7 @@ from typing import (
     Coroutine,
     Generic,
     Iterable,
+    List,
     NoReturn,
     Optional,
     Protocol,
@@ -2309,8 +2310,9 @@ def load_flow_from_entrypoint(
     Extract a flow object from a script at an entrypoint by running all of the code in the file.
 
     Args:
-        entrypoint: a string in the format `<path_to_script>:<flow_func_name>` or a module path
-            to a flow function
+        entrypoint: a string in the format `<path_to_script>:<flow_func_name>`
+            or a string in the format `<path_to_script>:<class_name>.<flow_method_name>`
+            or a module path to a flow function
         use_placeholder_flow: if True, use a placeholder Flow object if the actual flow object
             cannot be loaded from the entrypoint (e.g. dependencies are missing)
 
@@ -2700,26 +2702,55 @@ def load_placeholder_flow(entrypoint: str, raises: Exception) -> Flow[P, Any]:
 
 def safe_load_flow_from_entrypoint(entrypoint: str) -> Optional[Flow[P, Any]]:
     """
-    Load a flow from an entrypoint and return None if an exception is raised.
+    Safely load a Prefect flow from an entrypoint string. Returns None if loading fails.
 
     Args:
-        entrypoint: a string in the format `<path_to_script>:<flow_func_name>`
-          or a module path to a flow function
+        entrypoint (str): A string identifying the flow to load. Can be in one of the following formats:
+            - `<path_to_script>:<flow_func_name>`
+            - `<path_to_script>:<class_name>.<flow_method_name>`
+            - `<module_path>.<flow_func_name>`
+
+    Returns:
+        Optional[Flow]: The loaded Prefect flow object, or None if loading fails due to errors
+        (e.g. unresolved dependencies, syntax errors, or missing objects).
     """
-    func_def, source_code = _entrypoint_definition_and_source(entrypoint)
-    path = None
-    if ":" in entrypoint:
-        path = entrypoint.rsplit(":")[0]
+    func_or_cls_def, source_code, parts = _entrypoint_definition_and_source(entrypoint)
+
+    path = entrypoint.rsplit(":", maxsplit=1)[0] if ":" in entrypoint else None
     namespace = safe_load_namespace(source_code, filepath=path)
-    if func_def.name in namespace:
-        return namespace[func_def.name]
-    else:
-        # If the function is not in the namespace, if may be due to missing dependencies
-        # for the function. We will attempt to compile each annotation and default value
-        # and remove them from the function definition to see if the function can be
-        # compiled without them.
 
-        return _sanitize_and_load_flow(func_def, namespace)
+    if parts[0] not in namespace:
+        # If the object is not in the namespace, it may be due to missing dependencies
+        # in annotations or default values. We will attempt to sanitize them by removing
+        # anything that cannot be compiled, and then recompile the function or class.
+        if isinstance(func_or_cls_def, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            return _sanitize_and_load_flow(func_or_cls_def, namespace)
+        elif (
+            isinstance(func_or_cls_def, ast.ClassDef)
+            and len(parts) >= 2
+            and func_or_cls_def.name == parts[0]
+        ):
+            method_name = parts[1]
+            method_def = next(
+                (
+                    stmt
+                    for stmt in func_or_cls_def.body
+                    if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef))
+                    and stmt.name == method_name
+                ),
+                None,
+            )
+            if method_def is not None:
+                return _sanitize_and_load_flow(method_def, namespace)
+        else:
+            return None
+
+    obj = namespace.get(parts[0])
+    for part in parts[1:]:
+        obj = getattr(obj, part, None)
+        if obj is None:
+            return None
+    return obj
 
 
 def _sanitize_and_load_flow(
@@ -2853,7 +2884,7 @@ def load_flow_arguments_from_entrypoint(
           or a module path to a flow function
     """
 
-    func_def, source_code = _entrypoint_definition_and_source(entrypoint)
+    func_def, source_code, _ = _entrypoint_definition_and_source(entrypoint)
     path = None
     if ":" in entrypoint:
         path = entrypoint.rsplit(":")[0]
@@ -2930,26 +2961,45 @@ def is_entrypoint_async(entrypoint: str) -> bool:
     Returns:
         True if the function is asynchronous, False otherwise.
     """
-    func_def, _ = _entrypoint_definition_and_source(entrypoint)
+    func_def, _, _ = _entrypoint_definition_and_source(entrypoint)
     return isinstance(func_def, ast.AsyncFunctionDef)
 
 
 def _entrypoint_definition_and_source(
     entrypoint: str,
-) -> Tuple[Union[ast.FunctionDef, ast.AsyncFunctionDef], str]:
+) -> Tuple[Union[ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef], str, List[str]]:
+    """
+    Resolves and parses the source definition of a given entrypoint.
+
+    The entrypoint can be provided in one of the following formats:
+        - '<path_to_script>:<flow_func_name>'
+        - '<path_to_script>:<class_name>.<flow_method_name>'
+        - '<module_path.to.flow_function>'
+
+    Returns:
+        A tuple containing:
+        - The AST node (FunctionDef, AsyncFunctionDef, or ClassDef) of the base object.
+        - The full source code of the file or module as a string.
+        - A list of attribute access parts from the object path (e.g., ['MyFlowClass', 'run']).
+
+    Raises:
+        ValueError: If the module or target object cannot be found.
+    """
     if ":" in entrypoint:
-        # Split by the last colon once to handle Windows paths with drive letters i.e C:\path\to\file.py:do_stuff
-        path, func_name = entrypoint.rsplit(":", maxsplit=1)
+        path, object_path = entrypoint.rsplit(":", maxsplit=1)
         source_code = Path(path).read_text()
     else:
-        path, func_name = entrypoint.rsplit(".", maxsplit=1)
+        path, object_path = entrypoint.rsplit(".", maxsplit=1)
         spec = importlib.util.find_spec(path)
         if not spec or not spec.origin:
             raise ValueError(f"Could not find module {path!r}")
         source_code = Path(spec.origin).read_text()
 
     parsed_code = ast.parse(source_code)
-    func_def = next(
+    parts = object_path.split(".")
+    base_name = parts[0]
+
+    base_def = next(
         (
             node
             for node in ast.walk(parsed_code)
@@ -2958,14 +3008,15 @@ def _entrypoint_definition_and_source(
                 (
                     ast.FunctionDef,
                     ast.AsyncFunctionDef,
+                    ast.ClassDef,  # flow can be staticmethod/classmethod
                 ),
             )
-            and node.name == func_name
+            and node.name == base_name
         ),
         None,
     )
 
-    if not func_def:
-        raise ValueError(f"Could not find flow {func_name!r} in {path!r}")
+    if not base_def:
+        raise ValueError(f"Could not find object {base_name!r} in {path!r}")
 
-    return func_def, source_code
+    return base_def, source_code, parts

prefect/runner/_observers.py

@@ -0,0 +1,60 @@
+import asyncio
+import uuid
+from contextlib import AsyncExitStack
+from typing import Any, Protocol
+
+from prefect.events.clients import PrefectEventSubscriber, get_events_subscriber
+from prefect.events.filters import EventFilter, EventNameFilter
+from prefect.logging.loggers import get_logger
+
+
+class OnCancellingCallback(Protocol):
+    def __call__(self, flow_run_id: uuid.UUID) -> None: ...
+
+
+class FlowRunCancellingObserver:
+    def __init__(self, on_cancelling: OnCancellingCallback):
+        self.logger = get_logger("FlowRunCancellingObserver")
+        self.on_cancelling = on_cancelling
+        self._events_subscriber: PrefectEventSubscriber | None
+        self._exit_stack = AsyncExitStack()
+
+    async def _consume_events(self):
+        if self._events_subscriber is None:
+            raise RuntimeError(
+                "Events subscriber not initialized. Please use `async with` to initialize the observer."
+            )
+        async for event in self._events_subscriber:
+            try:
+                flow_run_id = uuid.UUID(
+                    event.resource["prefect.resource.id"].replace(
+                        "prefect.flow-run.", ""
+                    )
+                )
+                self.on_cancelling(flow_run_id)
+            except ValueError:
+                self.logger.debug(
+                    "Received event with invalid flow run ID: %s",
+                    event.resource["prefect.resource.id"],
+                )
+
+    async def __aenter__(self):
+        self._events_subscriber = await self._exit_stack.enter_async_context(
+            get_events_subscriber(
+                filter=EventFilter(
+                    event=EventNameFilter(name=["prefect.flow-run.Cancelling"])
+                )
+            )
+        )
+        self._consumer_task = asyncio.create_task(self._consume_events())
+        return self
+
+    async def __aexit__(self, *exc_info: Any):
+        await self._exit_stack.__aexit__(*exc_info)
+        self._consumer_task.cancel()
+        try:
+            await self._consumer_task
+        except asyncio.CancelledError:
+            pass
+        except Exception:
+            self.logger.exception("Error consuming events")

prefect/runner/runner.py

@@ -46,6 +46,8 @@ import subprocess
 import sys
 import tempfile
 import threading
+import uuid
+from contextlib import AsyncExitStack
 from copy import deepcopy
 from functools import partial
 from pathlib import Path
@@ -80,13 +82,6 @@ from prefect._internal.concurrency.api import (
     from_sync,
 )
 from prefect.client.orchestration import PrefectClient, get_client
-from prefect.client.schemas.filters import (
-    FlowRunFilter,
-    FlowRunFilterId,
-    FlowRunFilterState,
-    FlowRunFilterStateName,
-    FlowRunFilterStateType,
-)
 from prefect.client.schemas.objects import (
     ConcurrencyLimitConfig,
     State,
@@ -94,12 +89,13 @@ from prefect.client.schemas.objects import (
 )
 from prefect.client.schemas.objects import Flow as APIFlow
 from prefect.events import DeploymentTriggerTypes, TriggerTypes
+from prefect.events.clients import EventsClient, get_events_client
 from prefect.events.related import tags_as_related_resources
-from prefect.events.schemas.events import RelatedResource
-from prefect.events.utilities import emit_event
+from prefect.events.schemas.events import Event, RelatedResource, Resource
 from prefect.exceptions import Abort, ObjectNotFound
 from prefect.flows import Flow, FlowStateHook, load_flow_from_flow_run
 from prefect.logging.loggers import PrefectLogAdapter, flow_run_logger, get_logger
+from prefect.runner._observers import FlowRunCancellingObserver
 from prefect.runner.storage import RunnerStorage
 from prefect.schedules import Schedule
 from prefect.settings import (
@@ -228,7 +224,9 @@ class Runner:
         if self.heartbeat_seconds is not None and self.heartbeat_seconds < 30:
             raise ValueError("Heartbeat must be 30 seconds or greater.")
         self._heartbeat_task: asyncio.Task[None] | None = None
+        self._events_client: EventsClient = get_events_client(checkpoint_every=1)
 
+        self._exit_stack = AsyncExitStack()
         self._limiter: anyio.CapacityLimiter | None = None
         self._client: PrefectClient = get_client()
         self._submitting_flow_run_ids: set[UUID] = set()
@@ -501,15 +499,6 @@ class Runner:
                         jitter_range=0.3,
                     )
                 )
-                loops_task_group.start_soon(
-                    partial(
-                        critical_service_loop,
-                        workload=runner._check_for_cancelled_flow_runs,
-                        interval=self.query_seconds * 2,
-                        run_once=run_once,
-                        jitter_range=0.3,
-                    )
-                )
 
     def execute_in_background(
         self, func: Callable[..., Any], *args: Any, **kwargs: Any
@@ -583,58 +572,42 @@ class Runner:
             if not self._acquire_limit_slot(flow_run_id):
                 return
 
-            async with anyio.create_task_group() as tg:
-                with anyio.CancelScope():
-                    self._submitting_flow_run_ids.add(flow_run_id)
-                    flow_run = await self._client.read_flow_run(flow_run_id)
-
-                    process: (
-                        anyio.abc.Process | Exception
-                    ) = await self._runs_task_group.start(
-                        partial(
-                            self._submit_run_and_capture_errors,
-                            flow_run=flow_run,
-                            entrypoint=entrypoint,
-                            command=command,
-                            cwd=cwd,
-                            env=env,
-                            stream_output=stream_output,
-                        ),
-                    )
-                    if isinstance(process, Exception):
-                        return
+            self._submitting_flow_run_ids.add(flow_run_id)
+            flow_run = await self._client.read_flow_run(flow_run_id)
 
-                    task_status.started(process.pid)
+            process: anyio.abc.Process | Exception = await self._runs_task_group.start(
+                partial(
+                    self._submit_run_and_capture_errors,
+                    flow_run=flow_run,
+                    entrypoint=entrypoint,
+                    command=command,
+                    cwd=cwd,
+                    env=env,
+                    stream_output=stream_output,
+                ),
+            )
+            if isinstance(process, Exception):
+                return
 
-                    if self.heartbeat_seconds is not None:
-                        await self._emit_flow_run_heartbeat(flow_run)
+            task_status.started(process.pid)
 
-                    async with self._flow_run_process_map_lock:
-                        # Only add the process to the map if it is still running
-                        if process.returncode is None:
-                            self._flow_run_process_map[flow_run.id] = ProcessMapEntry(
-                                pid=process.pid, flow_run=flow_run
-                            )
+            if self.heartbeat_seconds is not None:
+                await self._emit_flow_run_heartbeat(flow_run)
 
-                    # We want this loop to stop when the flow run process exits
-                    # so we'll check if the flow run process is still alive on
-                    # each iteration and cancel the task group if it is not.
-                    workload = partial(
-                        self._check_for_cancelled_flow_runs,
-                        should_stop=lambda: not self._flow_run_process_map,
-                        on_stop=tg.cancel_scope.cancel,
+            async with self._flow_run_process_map_lock:
+                # Only add the process to the map if it is still running
+                if process.returncode is None:
+                    self._flow_run_process_map[flow_run.id] = ProcessMapEntry(
+                        pid=process.pid, flow_run=flow_run
                     )
 
-                    tg.start_soon(
-                        partial(
-                            critical_service_loop,
-                            workload=workload,
-                            interval=self.query_seconds,
-                            jitter_range=0.3,
-                        )
-                    )
+            while True:
+                # Wait until flow run execution is complete and the process has been removed from the map
+                await anyio.sleep(0.1)
+                if self._flow_run_process_map.get(flow_run.id) is None:
+                    break
 
-                    return process
+            return process
 
     async def execute_bundle(
         self,
@@ -673,24 +646,8 @@ class Runner:
             )
             self._flow_run_bundle_map[flow_run.id] = bundle
 
-            tasks: list[asyncio.Task[None]] = []
-            tasks.append(
-                asyncio.create_task(
-                    critical_service_loop(
-                        workload=self._check_for_cancelled_flow_runs,
-                        interval=self.query_seconds,
-                        jitter_range=0.3,
-                    )
-                )
-            )
-
             await anyio.to_thread.run_sync(process.join)
 
-            for task in tasks:
-                task.cancel()
-
-            await asyncio.gather(*tasks, return_exceptions=True)
-
             self._flow_run_process_map.pop(flow_run.id)
 
             flow_run_logger = self._get_flow_run_logger(flow_run)
@@ -1000,83 +957,11 @@ class Runner:
         self.last_polled: datetime.datetime = now("UTC")
         return await self._submit_scheduled_flow_runs(flow_run_response=runs_response)
 
-    async def _check_for_cancelled_flow_runs(
-        self,
-        should_stop: Callable[[], bool] = lambda: False,
-        on_stop: Callable[[], None] = lambda: None,
+    async def _cancel_run(
+        self, flow_run: "FlowRun | uuid.UUID", state_msg: Optional[str] = None
     ):
-        """
-        Checks for flow runs with CANCELLING a cancelling state and attempts to
-        cancel them.
-
-        Args:
-            should_stop: A callable that returns a boolean indicating whether or not
-                the runner should stop checking for cancelled flow runs.
-            on_stop: A callable that is called when the runner should stop checking
-                for cancelled flow runs.
-        """
-        if self.stopping:
-            return
-        if not self.started:
-            raise RuntimeError(
-                "Runner is not set up. Please make sure you are running this runner "
-                "as an async context manager."
-            )
-
-        if should_stop():
-            self._logger.debug(
-                "Runner has no active flow runs or deployments. Sending message to loop"
-                " service that no further cancellation checks are needed."
-            )
-            on_stop()
-
-        self._logger.debug("Checking for cancelled flow runs...")
-
-        named_cancelling_flow_runs = await self._client.read_flow_runs(
-            flow_run_filter=FlowRunFilter(
-                state=FlowRunFilterState(
-                    type=FlowRunFilterStateType(any_=[StateType.CANCELLED]),
-                    name=FlowRunFilterStateName(any_=["Cancelling"]),
-                ),
-                # Avoid duplicate cancellation calls
-                id=FlowRunFilterId(
-                    any_=list(
-                        self._flow_run_process_map.keys()
-                        - self._cancelling_flow_run_ids
-                    )
-                ),
-            ),
-        )
-
-        typed_cancelling_flow_runs = await self._client.read_flow_runs(
-            flow_run_filter=FlowRunFilter(
-                state=FlowRunFilterState(
-                    type=FlowRunFilterStateType(any_=[StateType.CANCELLING]),
-                ),
-                # Avoid duplicate cancellation calls
-                id=FlowRunFilterId(
-                    any_=list(
-                        self._flow_run_process_map.keys()
-                        - self._cancelling_flow_run_ids
-                    )
-                ),
-            ),
-        )
-
-        cancelling_flow_runs = named_cancelling_flow_runs + typed_cancelling_flow_runs
-
-        if cancelling_flow_runs:
-            self._logger.info(
-                f"Found {len(cancelling_flow_runs)} flow runs awaiting cancellation."
-            )
-
-        for flow_run in cancelling_flow_runs:
-            self._cancelling_flow_run_ids.add(flow_run.id)
-            self._runs_task_group.start_soon(self._cancel_run, flow_run)
-
-        return cancelling_flow_runs
-
-    async def _cancel_run(self, flow_run: "FlowRun", state_msg: Optional[str] = None):
+        if isinstance(flow_run, uuid.UUID):
+            flow_run = await self._client.read_flow_run(flow_run)
         run_logger = self._get_flow_run_logger(flow_run)
 
         process_map_entry = self._flow_run_process_map.get(flow_run.id)
@@ -1121,7 +1006,7 @@ class Runner:
             )
 
             flow, deployment = await self._get_flow_and_deployment(flow_run)
-            self._emit_flow_run_cancelled_event(
+            await self._emit_flow_run_cancelled_event(
                 flow_run=flow_run, flow=flow, deployment=deployment
             )
             run_logger.info(f"Cancelled flow run '{flow_run.name}'!")
@@ -1180,14 +1065,18 @@ class Runner:
         related = [RelatedResource.model_validate(r) for r in related]
         related += tags_as_related_resources(set(tags))
 
-        emit_event(
-            event="prefect.flow-run.heartbeat",
-            resource={
-                "prefect.resource.id": f"prefect.flow-run.{flow_run.id}",
-                "prefect.resource.name": flow_run.name,
-                "prefect.version": __version__,
-            },
-            related=related,
+        await self._events_client.emit(
+            Event(
+                event="prefect.flow-run.heartbeat",
+                resource=Resource(
+                    {
+                        "prefect.resource.id": f"prefect.flow-run.{flow_run.id}",
+                        "prefect.resource.name": flow_run.name,
+                        "prefect.version": __version__,
+                    }
+                ),
+                related=related,
+            )
         )
 
     def _event_resource(self):
@@ -1199,7 +1088,7 @@ class Runner:
             "prefect.version": __version__,
         }
 
-    def _emit_flow_run_cancelled_event(
+    async def _emit_flow_run_cancelled_event(
         self,
         flow_run: "FlowRun",
         flow: "Optional[APIFlow]",
@@ -1234,10 +1123,12 @@ class Runner:
         related = [RelatedResource.model_validate(r) for r in related]
         related += tags_as_related_resources(set(tags))
 
-        emit_event(
-            event="prefect.runner.cancelled-flow-run",
-            resource=self._event_resource(),
-            related=related,
+        await self._events_client.emit(
+            Event(
+                event="prefect.runner.cancelled-flow-run",
+                resource=Resource(self._event_resource()),
+                related=related,
+            )
         )
         self._logger.debug(f"Emitted flow run heartbeat event for {flow_run.id}")
 
@@ -1543,43 +1434,6 @@ class Runner:
 
         await self._client.set_flow_run_state(flow_run.id, state, force=True)
 
-        # Do not remove the flow run from the cancelling set immediately because
-        # the API caches responses for the `read_flow_runs` and we do not want to
-        # duplicate cancellations.
-        await self._schedule_task(
-            60 * 10, self._cancelling_flow_run_ids.remove, flow_run.id
-        )
-
-    async def _schedule_task(
-        self, __in_seconds: int, fn: Callable[..., Any], *args: Any, **kwargs: Any
-    ) -> None:
-        """
-        Schedule a background task to start after some time.
-
-        These tasks will be run immediately when the runner exits instead of waiting.
-
-        The function may be async or sync. Async functions will be awaited.
-        """
-
-        async def wrapper(task_status: anyio.abc.TaskStatus[None]) -> None:
-            # If we are shutting down, do not sleep; otherwise sleep until the scheduled
-            # time or shutdown
-            if self.started:
-                with anyio.CancelScope() as scope:
-                    self._scheduled_task_scopes.add(scope)
-                    task_status.started()
-                    await anyio.sleep(__in_seconds)
-
-                self._scheduled_task_scopes.remove(scope)
-            else:
-                task_status.started()
-
-            result = fn(*args, **kwargs)
-            if asyncio.iscoroutine(result):
-                await result
-
-        await self._runs_task_group.start(wrapper)
-
     async def _run_on_cancellation_hooks(
         self,
         flow_run: "FlowRun",
@@ -1647,11 +1501,19 @@ class Runner:
         if not hasattr(self, "_loop") or not self._loop:
             self._loop = asyncio.get_event_loop()
 
-        await self._client.__aenter__()
+        await self._exit_stack.enter_async_context(
+            FlowRunCancellingObserver(
+                on_cancelling=lambda flow_run_id: self._runs_task_group.start_soon(
+                    self._cancel_run, flow_run_id
+                )
+            )
+        )
+        await self._exit_stack.enter_async_context(self._client)
+        await self._exit_stack.enter_async_context(self._events_client)
 
         if not hasattr(self, "_runs_task_group") or not self._runs_task_group:
             self._runs_task_group: anyio.abc.TaskGroup = anyio.create_task_group()
-        await self._runs_task_group.__aenter__()
+        await self._exit_stack.enter_async_context(self._runs_task_group)
 
         if not hasattr(self, "_loops_task_group") or not self._loops_task_group:
             self._loops_task_group: anyio.abc.TaskGroup = anyio.create_task_group()
@@ -1677,11 +1539,7 @@ class Runner:
         for scope in self._scheduled_task_scopes:
             scope.cancel()
 
-        if self._runs_task_group:
-            await self._runs_task_group.__aexit__(*exc_info)
-
-        if self._client:
-            await self._client.__aexit__(*exc_info)
+        await self._exit_stack.__aexit__(*exc_info)
 
         shutil.rmtree(str(self._tmp_dir))
         del self._runs_task_group, self._loops_task_group