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

prefect/_build_info.py

@@ -1,5 +1,5 @@
 # Generated by versioningit
-__version__ = "3.4.0"
-__build_date__ = "2025-05-02 20:01:04.986143+00:00"
-__git_commit__ = "c80e444246c8805f1dfb684267e0d88dbfcc8d38"
+__version__ = "3.4.1.dev1"
+__build_date__ = "2025-05-04 08:07:31.383855+00:00"
+__git_commit__ = "378fc4f6b72aace32343264ee6cea198d949ba6a"
 __dirty__ = False

prefect/events/clients.py

@@ -684,8 +684,8 @@ class PrefectEventSubscriber:
 
     async def __aexit__(
         self,
-        exc_type: Optional[Type[Exception]],
-        exc_val: Optional[Exception],
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
         exc_tb: Optional[TracebackType],
     ) -> None:
         self._websocket = None

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,
@@ -100,6 +95,7 @@ from prefect.events.utilities import emit_event
 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 (
@@ -229,6 +225,7 @@ class Runner:
             raise ValueError("Heartbeat must be 30 seconds or greater.")
         self._heartbeat_task: asyncio.Task[None] | None = None
 
+        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 +498,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 +571,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 +645,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 +956,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)
@@ -1543,43 +1427,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 +1494,18 @@ 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)
 
         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 +1531,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

{prefect_client-3.4.0.dist-info → prefect_client-3.4.1.dev1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prefect-client
-Version: 3.4.0
+Version: 3.4.1.dev1
 Summary: Workflow orchestration and management.
 Project-URL: Changelog, https://github.com/PrefectHQ/prefect/releases
 Project-URL: Documentation, https://docs.prefect.io

{prefect_client-3.4.0.dist-info → prefect_client-3.4.1.dev1.dist-info}/RECORD

@@ -1,7 +1,7 @@
 prefect/.prefectignore,sha256=awSprvKT0vI8a64mEOLrMxhxqcO-b0ERQeYpA2rNKVQ,390
 prefect/__init__.py,sha256=iCdcC5ZmeewikCdnPEP6YBAjPNV5dvfxpYCTpw30Hkw,3685
 prefect/__main__.py,sha256=WFjw3kaYJY6pOTA7WDOgqjsz8zUEUZHCcj3P5wyVa-g,66
-prefect/_build_info.py,sha256=he53eMvLm7NsraWSPtkBQjfJ16fUU8YPCI6caWK9kN0,180
+prefect/_build_info.py,sha256=S2FR6y1xQ3ZKOstDGJZLDW2Sml9MbnhZ9zMjp4hX7qw,185
 prefect/_result_records.py,sha256=S6QmsODkehGVSzbMm6ig022PYbI6gNKz671p_8kBYx4,7789
 prefect/_versioning.py,sha256=YqR5cxXrY4P6LM1Pmhd8iMo7v_G2KJpGNdsf4EvDFQ0,14132
 prefect/_waiters.py,sha256=Ia2ITaXdHzevtyWIgJoOg95lrEXQqNEOquHvw3T33UQ,9026
@@ -148,7 +148,7 @@ prefect/docker/__init__.py,sha256=z6wdc6UFfiBG2jb9Jk64uCWVM04JKVWeVyDWwuuon8M,52
 prefect/docker/docker_image.py,sha256=bR_pEq5-FDxlwTj8CP_7nwZ_MiGK6KxIi8v7DRjy1Kg,3138
 prefect/events/__init__.py,sha256=GtKl2bE--pJduTxelH2xy7SadlLJmmis8WR1EYixhuA,2094
 prefect/events/actions.py,sha256=A7jS8bo4zWGnrt3QfSoQs0uYC1xfKXio3IfU0XtTb5s,9129
-prefect/events/clients.py,sha256=c5ZTt-ZVslvDr6-OrbsWb_7XUocWptGyAuVAVtAzokY,27589
+prefect/events/clients.py,sha256=gp3orepQav99303OC-zK6uz3dpyLlLpQ9ZWJEDol0cs,27597
 prefect/events/filters.py,sha256=2hVfzc3Rdgy0mBHDutWxT__LJY0zpVM8greWX3y6kjM,8233
 prefect/events/related.py,sha256=CTeexYUmmA93V4gsR33GIFmw-SS-X_ouOpRg-oeq-BU,6672
 prefect/events/utilities.py,sha256=ww34bTMENCNwcp6RhhgzG0KgXOvKGe0MKmBdSJ8NpZY,3043
@@ -184,7 +184,8 @@ prefect/logging/highlighters.py,sha256=BCf_LNhFInIfGPqwuu8YVrGa4wVxNc4YXo2pYgftp
 prefect/logging/loggers.py,sha256=rwFJv0i3dhdKr25XX-xUkQy4Vv4dy18bTy366jrC0OQ,12741
 prefect/logging/logging.yml,sha256=tT7gTyC4NmngFSqFkCdHaw7R0GPNPDDsTCGZQByiJAQ,3169
 prefect/runner/__init__.py,sha256=pQBd9wVrUVUDUFJlgiweKSnbahoBZwqnd2O2jkhrULY,158
-prefect/runner/runner.py,sha256=hk-y9Y29rpB7p08u9hlt4V8KmXUGMvkF9DE8uBxuXdI,65118
+prefect/runner/_observers.py,sha256=PpyXQL5bjp86AnDFEzcFPS5ayL6ExqcYgyuBMMQCO9Q,2183
+prefect/runner/runner.py,sha256=DFgZQTkKwmCDMmfA640xY1oTOCURzTOo7HOtwQxRVwA,59443
 prefect/runner/server.py,sha256=YRYFNoYddA9XfiTIYtudxrnD1vCX-PaOLhvyGUOb9AQ,11966
 prefect/runner/storage.py,sha256=n-65YoEf7KNVInnmMPeP5TVFJOa2zOS8w9en9MHi6uo,31328
 prefect/runner/submit.py,sha256=qOEj-NChQ6RYFV35hHEVMTklrNmKwaGs2mR78ku9H0o,9474
@@ -318,7 +319,7 @@ prefect/workers/cloud.py,sha256=dPvG1jDGD5HSH7aM2utwtk6RaJ9qg13XjkA0lAIgQmY,287
 prefect/workers/process.py,sha256=Yi5D0U5AQ51wHT86GdwtImXSefe0gJf3LGq4r4z9zwM,11090
 prefect/workers/server.py,sha256=2pmVeJZiVbEK02SO6BEZaBIvHMsn6G8LzjW8BXyiTtk,1952
 prefect/workers/utilities.py,sha256=VfPfAlGtTuDj0-Kb8WlMgAuOfgXCdrGAnKMapPSBrwc,2483
-prefect_client-3.4.0.dist-info/METADATA,sha256=UbF7RUupdDJ_R4xFe6ybCcDpsKc9d-BwIdHOF4VTVXM,7466
-prefect_client-3.4.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-prefect_client-3.4.0.dist-info/licenses/LICENSE,sha256=MCxsn8osAkzfxKC4CC_dLcUkU8DZLkyihZ8mGs3Ah3Q,11357
-prefect_client-3.4.0.dist-info/RECORD,,
+prefect_client-3.4.1.dev1.dist-info/METADATA,sha256=IFwUWcPPlE7Z-2A_Xn2tZef9WEqSXbfjlRRx0TQnbfI,7471
+prefect_client-3.4.1.dev1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+prefect_client-3.4.1.dev1.dist-info/licenses/LICENSE,sha256=MCxsn8osAkzfxKC4CC_dLcUkU8DZLkyihZ8mGs3Ah3Q,11357
+prefect_client-3.4.1.dev1.dist-info/RECORD,,