prefect-client 3.0.0rc13__py3-none-any.whl → 3.0.0rc15__py3-none-any.whl

prefect/events/clients.py

@@ -15,10 +15,10 @@ from typing import (
 )
 from uuid import UUID
 
-import httpx
 import orjson
 import pendulum
 from cachetools import TTLCache
+from prometheus_client import Counter
 from typing_extensions import Self
 from websockets import Subprotocol
 from websockets.client import WebSocketClientProtocol, connect
@@ -28,20 +28,48 @@ from websockets.exceptions import (
     ConnectionClosedOK,
 )
 
-from prefect.client.base import PrefectHttpxAsyncClient
 from prefect.events import Event
 from prefect.logging import get_logger
-from prefect.settings import PREFECT_API_KEY, PREFECT_API_URL, PREFECT_CLOUD_API_URL
+from prefect.settings import (
+    PREFECT_API_KEY,
+    PREFECT_API_URL,
+    PREFECT_CLOUD_API_URL,
+    PREFECT_SERVER_ALLOW_EPHEMERAL_MODE,
+)
 
 if TYPE_CHECKING:
     from prefect.events.filters import EventFilter
 
+EVENTS_EMITTED = Counter(
+    "prefect_events_emitted",
+    "The number of events emitted by Prefect event clients",
+    labelnames=["client"],
+)
+EVENTS_OBSERVED = Counter(
+    "prefect_events_observed",
+    "The number of events observed by Prefect event subscribers",
+    labelnames=["client"],
+)
+EVENT_WEBSOCKET_CONNECTIONS = Counter(
+    "prefect_event_websocket_connections",
+    (
+        "The number of times Prefect event clients have connected to an event stream, "
+        "broken down by direction (in/out) and connection (initial/reconnect)"
+    ),
+    labelnames=["client", "direction", "connection"],
+)
+EVENT_WEBSOCKET_CHECKPOINTS = Counter(
+    "prefect_event_websocket_checkpoints",
+    "The number of checkpoints performed by Prefect event clients",
+    labelnames=["client"],
+)
+
 logger = get_logger(__name__)
 
 
 def get_events_client(
     reconnection_attempts: int = 10,
-    checkpoint_every: int = 20,
+    checkpoint_every: int = 700,
 ) -> "EventsClient":
     api_url = PREFECT_API_URL.value()
     if isinstance(api_url, str) and api_url.startswith(PREFECT_CLOUD_API_URL.value()):
@@ -49,13 +77,25 @@ def get_events_client(
             reconnection_attempts=reconnection_attempts,
             checkpoint_every=checkpoint_every,
         )
-    elif PREFECT_API_URL:
+    elif api_url:
+        return PrefectEventsClient(
+            reconnection_attempts=reconnection_attempts,
+            checkpoint_every=checkpoint_every,
+        )
+    elif PREFECT_SERVER_ALLOW_EPHEMERAL_MODE:
+        from prefect.server.api.server import SubprocessASGIServer
+
+        server = SubprocessASGIServer()
+        server.start()
         return PrefectEventsClient(
+            api_url=server.api_url,
             reconnection_attempts=reconnection_attempts,
             checkpoint_every=checkpoint_every,
         )
     else:
-        return PrefectEphemeralEventsClient()
+        raise ValueError(
+            "No Prefect API URL provided. Please set PREFECT_API_URL to the address of a running Prefect server."
+        )
 
 
 def get_events_subscriber(
@@ -63,25 +103,38 @@ def get_events_subscriber(
     reconnection_attempts: int = 10,
 ) -> "PrefectEventSubscriber":
     api_url = PREFECT_API_URL.value()
-    if not api_url:
-        raise ValueError(
-            "A Prefect server or Prefect Cloud is required to start an event "
-            "subscriber.  Please check the PREFECT_API_URL setting in your profile."
-        )
 
     if isinstance(api_url, str) and api_url.startswith(PREFECT_CLOUD_API_URL.value()):
         return PrefectCloudEventSubscriber(
             filter=filter, reconnection_attempts=reconnection_attempts
         )
-    else:
+    elif api_url:
         return PrefectEventSubscriber(
             filter=filter, reconnection_attempts=reconnection_attempts
         )
+    elif PREFECT_SERVER_ALLOW_EPHEMERAL_MODE:
+        from prefect.server.api.server import SubprocessASGIServer
+
+        server = SubprocessASGIServer()
+        server.start()
+        return PrefectEventSubscriber(
+            api_url=server.api_url,
+            filter=filter,
+            reconnection_attempts=reconnection_attempts,
+        )
+    else:
+        raise ValueError(
+            "No Prefect API URL provided. Please set PREFECT_API_URL to the address of a running Prefect server."
+        )
 
 
 class EventsClient(abc.ABC):
     """The abstract interface for all Prefect Events clients"""
 
+    @property
+    def client_name(self) -> str:
+        return self.__class__.__name__
+
     async def emit(self, event: Event) -> None:
         """Emit a single event"""
         if not hasattr(self, "_in_context"):
@@ -89,7 +142,11 @@ class EventsClient(abc.ABC):
                 "Events may only be emitted while this client is being used as a "
                 "context manager"
             )
-        return await self._emit(event)
+
+        try:
+            return await self._emit(event)
+        finally:
+            EVENTS_EMITTED.labels(self.client_name).inc()
 
     @abc.abstractmethod
     async def _emit(self, event: Event) -> None:  # pragma: no cover
@@ -168,47 +225,6 @@ def _get_api_url_and_key(
     return api_url, api_key
 
 
-class PrefectEphemeralEventsClient(EventsClient):
-    """A Prefect Events client that sends events to an ephemeral Prefect server"""
-
-    def __init__(self):
-        if PREFECT_API_KEY.value():
-            raise ValueError(
-                "PrefectEphemeralEventsClient cannot be used when PREFECT_API_KEY is set."
-                " Please use PrefectEventsClient or PrefectCloudEventsClient instead."
-            )
-        from prefect.server.api.server import create_app
-
-        app = create_app(ephemeral=True)
-
-        self._http_client = PrefectHttpxAsyncClient(
-            transport=httpx.ASGITransport(app=app, raise_app_exceptions=False),
-            base_url="http://ephemeral-prefect/api",
-            enable_csrf_support=False,
-        )
-
-    async def __aenter__(self) -> Self:
-        await super().__aenter__()
-        await self._http_client.__aenter__()
-        return self
-
-    async def __aexit__(
-        self,
-        exc_type: Optional[Type[Exception]],
-        exc_val: Optional[Exception],
-        exc_tb: Optional[TracebackType],
-    ) -> None:
-        self._websocket = None
-        await self._http_client.__aexit__(exc_type, exc_val, exc_tb)
-        return await super().__aexit__(exc_type, exc_val, exc_tb)
-
-    async def _emit(self, event: Event) -> None:
-        await self._http_client.post(
-            "/events",
-            json=[event.model_dump(mode="json")],
-        )
-
-
 class PrefectEventsClient(EventsClient):
     """A Prefect Events client that streams events to a Prefect server"""
 
@@ -219,7 +235,7 @@ class PrefectEventsClient(EventsClient):
         self,
         api_url: Optional[str] = None,
         reconnection_attempts: int = 10,
-        checkpoint_every: int = 20,
+        checkpoint_every: int = 700,
     ):
         """
         Args:
@@ -299,6 +315,8 @@ class PrefectEventsClient(EventsClient):
         # don't clear the list, just the ones that we are sure of.
         self._unconfirmed_events = self._unconfirmed_events[unconfirmed_count:]
 
+        EVENT_WEBSOCKET_CHECKPOINTS.labels(self.client_name).inc()
+
     async def _emit(self, event: Event) -> None:
         for i in range(self._reconnection_attempts + 1):
             try:
@@ -336,7 +354,7 @@ class PrefectCloudEventsClient(PrefectEventsClient):
         api_url: Optional[str] = None,
         api_key: Optional[str] = None,
         reconnection_attempts: int = 10,
-        checkpoint_every: int = 20,
+        checkpoint_every: int = 700,
     ):
         """
         Args:
@@ -400,9 +418,9 @@ class PrefectEventSubscriber:
             reconnection_attempts: When the client is disconnected, how many times
                 the client should attempt to reconnect
         """
+        self._api_key = None
         if not api_url:
             api_url = cast(str, PREFECT_API_URL.value())
-            self._api_key = None
 
         from prefect.events.filters import EventFilter
 
@@ -426,10 +444,17 @@ class PrefectEventSubscriber:
         if self._reconnection_attempts < 0:
             raise ValueError("reconnection_attempts must be a non-negative integer")
 
+    @property
+    def client_name(self) -> str:
+        return self.__class__.__name__
+
     async def __aenter__(self) -> Self:
         # Don't handle any errors in the initial connection, because these are most
         # likely a permission or configuration issue that should propagate
-        await self._reconnect()
+        try:
+            await self._reconnect()
+        finally:
+            EVENT_WEBSOCKET_CONNECTIONS.labels(self.client_name, "out", "initial")
         return self
 
     async def _reconnect(self) -> None:
@@ -503,7 +528,12 @@ class PrefectEventSubscriber:
                 # Otherwise, after the first time through this loop, we're recovering
                 # from a ConnectionClosed, so reconnect now.
                 if not self._websocket or i > 0:
-                    await self._reconnect()
+                    try:
+                        await self._reconnect()
+                    finally:
+                        EVENT_WEBSOCKET_CONNECTIONS.labels(
+                            self.client_name, "out", "reconnect"
+                        )
                     assert self._websocket
 
                 while True:
@@ -514,7 +544,10 @@ class PrefectEventSubscriber:
                         continue
                     self._seen_events[event.id] = True
 
-                    return event
+                    try:
+                        return event
+                    finally:
+                        EVENTS_OBSERVED.labels(self.client_name).inc()
             except ConnectionClosedOK:
                 logger.debug('Connection closed with "OK" status')
                 raise StopAsyncIteration

prefect/events/utilities.py

@@ -8,7 +8,6 @@ from pydantic_extra_types.pendulum_dt import DateTime
 from .clients import (
     AssertingEventsClient,
     PrefectCloudEventsClient,
-    PrefectEphemeralEventsClient,
     PrefectEventsClient,
 )
 from .schemas.events import Event, RelatedResource
@@ -53,7 +52,6 @@ def emit_event(
         AssertingEventsClient,
         PrefectCloudEventsClient,
         PrefectEventsClient,
-        PrefectEphemeralEventsClient,
     ]
     worker_instance = EventsWorker.instance()
 

prefect/events/worker.py

@@ -17,7 +17,6 @@ from .clients import (
     EventsClient,
     NullEventsClient,
     PrefectCloudEventsClient,
-    PrefectEphemeralEventsClient,
     PrefectEventsClient,
 )
 from .related import related_resources_from_run_context
@@ -97,7 +96,15 @@ class EventsWorker(QueueService[Event]):
             elif should_emit_events_to_running_server():
                 client_type = PrefectEventsClient
             elif should_emit_events_to_ephemeral_server():
-                client_type = PrefectEphemeralEventsClient
+                # create an ephemeral API if none was provided
+                from prefect.server.api.server import SubprocessASGIServer
+
+                server = SubprocessASGIServer()
+                server.start()
+                assert server.server_process is not None, "Server process did not start"
+
+                client_kwargs = {"api_url": server.api_url}
+                client_type = PrefectEventsClient
             else:
                 client_type = NullEventsClient
 

prefect/flow_engine.py

@@ -29,7 +29,8 @@ from prefect.client.orchestration import SyncPrefectClient, get_client
 from prefect.client.schemas import FlowRun, TaskRun
 from prefect.client.schemas.filters import FlowRunFilter
 from prefect.client.schemas.sorting import FlowRunSort
-from prefect.context import ClientContext, FlowRunContext, TagsContext
+from prefect.concurrency.context import ConcurrencyContext
+from prefect.context import FlowRunContext, SyncClientContext, TagsContext
 from prefect.exceptions import (
     Abort,
     Pause,
@@ -505,6 +506,8 @@ class FlowRunEngine(Generic[P, R]):
                     task_runner=task_runner,
                 )
             )
+            stack.enter_context(ConcurrencyContext())
+
             # set the logger to the flow run logger
             self.logger = flow_run_logger(flow_run=self.flow_run, flow=self.flow)
 
@@ -529,8 +532,8 @@ class FlowRunEngine(Generic[P, R]):
         """
         Enters a client context and creates a flow run if needed.
         """
-        with ClientContext.get_or_create() as client_ctx:
-            self._client = client_ctx.sync_client
+        with SyncClientContext.get_or_create() as client_ctx:
+            self._client = client_ctx.client
             self._is_started = True
 
             if not self.flow_run:

prefect/flows.py

@@ -66,6 +66,7 @@ from prefect.exceptions import (
     ObjectNotFound,
     ParameterTypeError,
     ScriptError,
+    TerminationSignal,
     UnspecifiedFlowError,
 )
 from prefect.filesystems import LocalFileSystem, ReadableDeploymentStorage
@@ -95,7 +96,7 @@ from prefect.utilities.callables import (
     parameters_to_args_kwargs,
     raise_for_reserved_arguments,
 )
-from prefect.utilities.collections import listrepr
+from prefect.utilities.collections import listrepr, visit_collection
 from prefect.utilities.filesystem import relative_path_to_current_platform
 from prefect.utilities.hashing import file_hash
 from prefect.utilities.importtools import import_object, safe_load_namespace
@@ -535,6 +536,21 @@ class Flow(Generic[P, R]):
         Raises:
             ParameterTypeError: if the provided parameters are not valid
         """
+
+        def resolve_block_reference(data: Any) -> Any:
+            if isinstance(data, dict) and "$ref" in data:
+                return Block.load_from_ref(data["$ref"])
+            return data
+
+        try:
+            parameters = visit_collection(
+                parameters, resolve_block_reference, return_data=True
+            )
+        except (ValueError, RuntimeError) as exc:
+            raise ParameterTypeError(
+                "Failed to resolve block references in parameters."
+            ) from exc
+
         args, kwargs = parameters_to_args_kwargs(self.fn, parameters)
 
         with warnings.catch_warnings():
@@ -931,10 +947,15 @@ class Flow(Generic[P, R]):
             else:
                 raise
 
-        if loop is not None:
-            loop.run_until_complete(runner.start(webserver=webserver))
-        else:
-            asyncio.run(runner.start(webserver=webserver))
+        try:
+            if loop is not None:
+                loop.run_until_complete(runner.start(webserver=webserver))
+            else:
+                asyncio.run(runner.start(webserver=webserver))
+        except (KeyboardInterrupt, TerminationSignal) as exc:
+            logger.info(f"Received {type(exc).__name__}, shutting down...")
+            if loop is not None:
+                loop.stop()
 
     @classmethod
     @sync_compatible
@@ -1734,14 +1755,13 @@ def load_flow_from_entrypoint(
         raise MissingFlowError(
             f"Flow function with name {func_name!r} not found in {path!r}. "
         ) from exc
-    except ScriptError as exc:
+    except ScriptError:
         # If the flow has dependencies that are not installed in the current
-        # environment, fallback to loading the flow via AST parsing. The
-        # drawback of this approach is that we're unable to actually load the
-        # function, so we create a placeholder flow that will re-raise this
-        # exception when called.
+        # environment, fallback to loading the flow via AST parsing.
         if use_placeholder_flow:
-            flow = load_placeholder_flow(entrypoint=entrypoint, raises=exc)
+            flow = safe_load_flow_from_entrypoint(entrypoint)
+            if flow is None:
+                raise
         else:
             raise
 
@@ -1976,6 +1996,147 @@ def load_placeholder_flow(entrypoint: str, raises: Exception):
     return Flow(**arguments)
 
 
+def safe_load_flow_from_entrypoint(entrypoint: str) -> Optional[Flow]:
+    """
+    Load a flow from an entrypoint and return None if an exception is raised.
+
+    Args:
+        entrypoint: a string in the format `<path_to_script>:<flow_func_name>`
+          or a module path to a flow function
+    """
+    func_def, source_code = _entrypoint_definition_and_source(entrypoint)
+    path = None
+    if ":" in entrypoint:
+        path = entrypoint.rsplit(":")[0]
+    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)
+
+
+def _sanitize_and_load_flow(
+    func_def: Union[ast.FunctionDef, ast.AsyncFunctionDef], namespace: Dict[str, Any]
+) -> Optional[Flow]:
+    """
+    Attempt to load a flow from the function definition after sanitizing the annotations
+    and defaults that can't be compiled.
+
+    Args:
+        func_def: the function definition
+        namespace: the namespace to load the function into
+
+    Returns:
+        The loaded function or None if the function can't be loaded
+        after sanitizing the annotations and defaults.
+    """
+    args = func_def.args.posonlyargs + func_def.args.args + func_def.args.kwonlyargs
+    if func_def.args.vararg:
+        args.append(func_def.args.vararg)
+    if func_def.args.kwarg:
+        args.append(func_def.args.kwarg)
+    # Remove annotations that can't be compiled
+    for arg in args:
+        if arg.annotation is not None:
+            try:
+                code = compile(
+                    ast.Expression(arg.annotation),
+                    filename="<ast>",
+                    mode="eval",
+                )
+                exec(code, namespace)
+            except Exception as e:
+                logger.debug(
+                    "Failed to evaluate annotation for argument %s due to the following error. Ignoring annotation.",
+                    arg.arg,
+                    exc_info=e,
+                )
+                arg.annotation = None
+
+    # Remove defaults that can't be compiled
+    new_defaults = []
+    for default in func_def.args.defaults:
+        try:
+            code = compile(ast.Expression(default), "<ast>", "eval")
+            exec(code, namespace)
+            new_defaults.append(default)
+        except Exception as e:
+            logger.debug(
+                "Failed to evaluate default value %s due to the following error. Ignoring default.",
+                default,
+                exc_info=e,
+            )
+            new_defaults.append(
+                ast.Constant(
+                    value=None, lineno=default.lineno, col_offset=default.col_offset
+                )
+            )
+    func_def.args.defaults = new_defaults
+
+    # Remove kw_defaults that can't be compiled
+    new_kw_defaults = []
+    for default in func_def.args.kw_defaults:
+        if default is not None:
+            try:
+                code = compile(ast.Expression(default), "<ast>", "eval")
+                exec(code, namespace)
+                new_kw_defaults.append(default)
+            except Exception as e:
+                logger.debug(
+                    "Failed to evaluate default value %s due to the following error. Ignoring default.",
+                    default,
+                    exc_info=e,
+                )
+                new_kw_defaults.append(
+                    ast.Constant(
+                        value=None,
+                        lineno=default.lineno,
+                        col_offset=default.col_offset,
+                    )
+                )
+        else:
+            new_kw_defaults.append(
+                ast.Constant(
+                    value=None,
+                    lineno=func_def.lineno,
+                    col_offset=func_def.col_offset,
+                )
+            )
+    func_def.args.kw_defaults = new_kw_defaults
+
+    if func_def.returns is not None:
+        try:
+            code = compile(
+                ast.Expression(func_def.returns), filename="<ast>", mode="eval"
+            )
+            exec(code, namespace)
+        except Exception as e:
+            logger.debug(
+                "Failed to evaluate return annotation due to the following error. Ignoring annotation.",
+                exc_info=e,
+            )
+            func_def.returns = None
+
+    # Attempt to compile the function without annotations and defaults that
+    # can't be compiled
+    try:
+        code = compile(
+            ast.Module(body=[func_def], type_ignores=[]),
+            filename="<ast>",
+            mode="exec",
+        )
+        exec(code, namespace)
+    except Exception as e:
+        logger.debug("Failed to compile: %s", e)
+    else:
+        return namespace.get(func_def.name)
+
+
 def load_flow_arguments_from_entrypoint(
     entrypoint: str, arguments: Optional[Union[List[str], Set[str]]] = None
 ) -> dict[str, Any]:
@@ -1991,6 +2152,9 @@ def load_flow_arguments_from_entrypoint(
     """
 
     func_def, source_code = _entrypoint_definition_and_source(entrypoint)
+    path = None
+    if ":" in entrypoint:
+        path = entrypoint.rsplit(":")[0]
 
     if arguments is None:
         # If no arguments are provided default to known arguments that are of
@@ -2026,7 +2190,7 @@ def load_flow_arguments_from_entrypoint(
 
                 # if the arg value is not a raw str (i.e. a variable or expression),
                 # then attempt to evaluate it
-                namespace = safe_load_namespace(source_code)
+                namespace = safe_load_namespace(source_code, filepath=path)
                 literal_arg_value = ast.get_source_segment(source_code, keyword.value)
                 cleaned_value = (
                     literal_arg_value.replace("\n", "") if literal_arg_value else ""