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

prefect/results.py

@@ -38,7 +38,6 @@ from prefect.settings import (
     PREFECT_RESULTS_DEFAULT_SERIALIZER,
     PREFECT_RESULTS_PERSIST_BY_DEFAULT,
     PREFECT_TASK_SCHEDULING_DEFAULT_STORAGE_BLOCK,
-    default_result_storage_block_name,
 )
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import sync_compatible
@@ -62,35 +61,15 @@ logger = get_logger("results")
 P = ParamSpec("P")
 R = TypeVar("R")
 
+_default_storages: Dict[Tuple[str, str], WritableFileSystem] = {}
 
-@sync_compatible
-async def get_default_result_storage() -> ResultStorage:
-    """
-    Generate a default file system for result storage.
-    """
-    try:
-        return await Block.load(PREFECT_DEFAULT_RESULT_STORAGE_BLOCK.value())
-    except ValueError as e:
-        if "Unable to find" not in str(e):
-            raise e
-        elif (
-            PREFECT_DEFAULT_RESULT_STORAGE_BLOCK.value()
-            == default_result_storage_block_name()
-        ):
-            return LocalFileSystem(basepath=PREFECT_LOCAL_STORAGE_PATH.value())
-        else:
-            raise
-
-
-_default_task_scheduling_storages: Dict[Tuple[str, str], WritableFileSystem] = {}
 
-
-async def get_or_create_default_task_scheduling_storage() -> ResultStorage:
+async def _get_or_create_default_storage(block_document_slug: str) -> ResultStorage:
     """
-    Generate a default file system for background task parameter/result storage.
+    Generate a default file system for storage.
     """
     default_storage_name, storage_path = cache_key = (
-        PREFECT_TASK_SCHEDULING_DEFAULT_STORAGE_BLOCK.value(),
+        block_document_slug,
         PREFECT_LOCAL_STORAGE_PATH.value(),
     )
 
@@ -105,8 +84,8 @@ async def get_or_create_default_task_scheduling_storage() -> ResultStorage:
         if block_type_slug == "local-file-system":
             block = LocalFileSystem(basepath=storage_path)
         else:
-            raise Exception(
-                "The default task storage block does not exist, but it is of type "
+            raise ValueError(
+                "The default storage block does not exist, but it is of type "
                 f"'{block_type_slug}' which cannot be created implicitly.  Please create "
                 "the block manually."
             )
@@ -123,13 +102,32 @@ async def get_or_create_default_task_scheduling_storage() -> ResultStorage:
         return block
 
     try:
-        return _default_task_scheduling_storages[cache_key]
+        return _default_storages[cache_key]
     except KeyError:
         storage = await get_storage()
-        _default_task_scheduling_storages[cache_key] = storage
+        _default_storages[cache_key] = storage
         return storage
 
 
+@sync_compatible
+async def get_or_create_default_result_storage() -> ResultStorage:
+    """
+    Generate a default file system for result storage.
+    """
+    return await _get_or_create_default_storage(
+        PREFECT_DEFAULT_RESULT_STORAGE_BLOCK.value()
+    )
+
+
+async def get_or_create_default_task_scheduling_storage() -> ResultStorage:
+    """
+    Generate a default file system for background task parameter/result storage.
+    """
+    return await _get_or_create_default_storage(
+        PREFECT_TASK_SCHEDULING_DEFAULT_STORAGE_BLOCK.value()
+    )
+
+
 def get_default_result_serializer() -> ResultSerializer:
     """
     Generate a default file system for result storage.
@@ -210,7 +208,9 @@ class ResultFactory(BaseModel):
                 kwargs.pop(key)
 
         # Apply defaults
-        kwargs.setdefault("result_storage", await get_default_result_storage())
+        kwargs.setdefault(
+            "result_storage", await get_or_create_default_result_storage()
+        )
         kwargs.setdefault("result_serializer", get_default_result_serializer())
         kwargs.setdefault("persist_result", get_default_persist_setting())
         kwargs.setdefault("cache_result_in_memory", True)
@@ -280,7 +280,9 @@ class ResultFactory(BaseModel):
         """
         Create a new result factory for a task.
         """
-        return await cls._from_task(task, get_default_result_storage, client=client)
+        return await cls._from_task(
+            task, get_or_create_default_result_storage, client=client
+        )
 
     @classmethod
     @inject_client
@@ -429,7 +431,11 @@ class ResultFactory(BaseModel):
 
     @sync_compatible
     async def create_result(
-        self, obj: R, key: Optional[str] = None, expiration: Optional[DateTime] = None
+        self,
+        obj: R,
+        key: Optional[str] = None,
+        expiration: Optional[DateTime] = None,
+        defer_persistence: bool = False,
     ) -> Union[R, "BaseResult[R]"]:
         """
         Create a result type for the given object.
@@ -462,6 +468,7 @@ class ResultFactory(BaseModel):
             serializer=self.serializer,
             cache_object=should_cache_object,
             expiration=expiration,
+            defer_persistence=defer_persistence,
         )
 
     @sync_compatible
@@ -587,6 +594,19 @@ class PersistedResult(BaseResult):
     expiration: Optional[DateTime] = None
 
     _should_cache_object: bool = PrivateAttr(default=True)
+    _persisted: bool = PrivateAttr(default=False)
+    _storage_block: WritableFileSystem = PrivateAttr(default=None)
+    _serializer: Serializer = PrivateAttr(default=None)
+
+    def _cache_object(
+        self,
+        obj: Any,
+        storage_block: WritableFileSystem = None,
+        serializer: Serializer = None,
+    ) -> None:
+        self._cache = obj
+        self._storage_block = storage_block
+        self._serializer = serializer
 
     @sync_compatible
     @inject_client
@@ -599,7 +619,7 @@ class PersistedResult(BaseResult):
             return self._cache
 
         blob = await self._read_blob(client=client)
-        obj = blob.serializer.loads(blob.data)
+        obj = blob.load()
         self.expiration = blob.expiration
 
         if self._should_cache_object:
@@ -630,6 +650,46 @@ class PersistedResult(BaseResult):
         if hasattr(storage_block, "_remote_file_system"):
             return storage_block._remote_file_system._resolve_path(key)
 
+    @sync_compatible
+    @inject_client
+    async def write(self, obj: R = NotSet, client: "PrefectClient" = None) -> None:
+        """
+        Write the result to the storage block.
+        """
+
+        if self._persisted:
+            # don't double write or overwrite
+            return
+
+        # load objects from a cache
+
+        # first the object itself
+        if obj is NotSet and not self.has_cached_object():
+            raise ValueError("Cannot write a result that has no object cached.")
+        obj = obj if obj is not NotSet else self._cache
+
+        # next, the storage block
+        storage_block = self._storage_block
+        if storage_block is None:
+            block_document = await client.read_block_document(self.storage_block_id)
+            storage_block = Block._from_block_document(block_document)
+
+        # finally, the serializer
+        serializer = self._serializer
+        if serializer is None:
+            # this could error if the serializer requires kwargs
+            serializer = Serializer(type=self.serializer_type)
+
+        data = serializer.dumps(obj)
+        blob = PersistedResultBlob(
+            serializer=serializer, data=data, expiration=self.expiration
+        )
+        await storage_block.write_path(self.storage_key, content=blob.to_bytes())
+        self._persisted = True
+
+        if not self._should_cache_object:
+            self._cache = NotSet
+
     @classmethod
     @sync_compatible
     async def create(
@@ -641,6 +701,7 @@ class PersistedResult(BaseResult):
         serializer: Serializer,
         cache_object: bool = True,
         expiration: Optional[DateTime] = None,
+        defer_persistence: bool = False,
     ) -> "PersistedResult[R]":
         """
         Create a new result reference from a user's object.
@@ -650,19 +711,13 @@ class PersistedResult(BaseResult):
         """
         assert (
             storage_block_id is not None
-        ), "Unexpected storage block ID. Was it persisted?"
-        data = serializer.dumps(obj)
-        blob = PersistedResultBlob(
-            serializer=serializer, data=data, expiration=expiration
-        )
+        ), "Unexpected storage block ID. Was it saved?"
 
         key = storage_key_fn()
         if not isinstance(key, str):
             raise TypeError(
                 f"Expected type 'str' for result storage key; got value {key!r}"
             )
-        await storage_block.write_path(key, content=blob.to_bytes())
-
         description = f"Result of type `{type(obj).__name__}`"
         uri = cls._infer_path(storage_block, key)
         if uri:
@@ -682,12 +737,23 @@ class PersistedResult(BaseResult):
             expiration=expiration,
         )
 
-        if cache_object:
+        if cache_object and not defer_persistence:
             # Attach the object to the result so it's available without deserialization
-            result._cache_object(obj)
+            result._cache_object(
+                obj, storage_block=storage_block, serializer=serializer
+            )
 
         object.__setattr__(result, "_should_cache_object", cache_object)
 
+        if not defer_persistence:
+            await result.write(obj=obj)
+        else:
+            # we must cache temporarily to allow for writing later
+            # the cache will be removed on write
+            result._cache_object(
+                obj, storage_block=storage_block, serializer=serializer
+            )
+
         return result
 
 
@@ -703,6 +769,9 @@ class PersistedResultBlob(BaseModel):
     prefect_version: str = Field(default=prefect.__version__)
     expiration: Optional[DateTime] = None
 
+    def load(self) -> Any:
+        return self.serializer.loads(self.data)
+
     def to_bytes(self) -> bytes:
         return self.model_dump_json(serialize_as_any=True).encode()
 

prefect/runner/runner.py

@@ -45,7 +45,7 @@ import threading
 from copy import deepcopy
 from functools import partial
 from pathlib import Path
-from typing import Callable, Dict, Iterable, List, Optional, Set, Union
+from typing import TYPE_CHECKING, Callable, Dict, Iterable, List, Optional, Set, Union
 from uuid import UUID, uuid4
 
 import anyio
@@ -75,7 +75,6 @@ from prefect.deployments.runner import (
     EntrypointType,
     RunnerDeployment,
 )
-from prefect.deployments.schedules import FlexibleScheduleList
 from prefect.events import DeploymentTriggerTypes, TriggerTypes
 from prefect.exceptions import Abort, ObjectNotFound
 from prefect.flows import Flow, load_flow_from_flow_run
@@ -98,6 +97,9 @@ from prefect.utilities.engine import propose_state
 from prefect.utilities.processutils import _register_signal, run_process
 from prefect.utilities.services import critical_service_loop
 
+if TYPE_CHECKING:
+    from prefect.client.types.flexible_schedule_list import FlexibleScheduleList
+
 __all__ = ["Runner"]
 
 
@@ -221,7 +223,7 @@ class Runner:
         cron: Optional[Union[Iterable[str], str]] = None,
         rrule: Optional[Union[Iterable[str], str]] = None,
         paused: Optional[bool] = None,
-        schedules: Optional[FlexibleScheduleList] = None,
+        schedules: Optional["FlexibleScheduleList"] = None,
         schedule: Optional[SCHEDULE_TYPES] = None,
         is_schedule_active: Optional[bool] = None,
         parameters: Optional[dict] = None,

prefect/runner/server.py

@@ -10,7 +10,7 @@ from typing_extensions import Literal
 from prefect._internal.schemas.validators import validate_values_conform_to_schema
 from prefect.client.orchestration import get_client
 from prefect.exceptions import MissingFlowError, ScriptError
-from prefect.flows import Flow, load_flow_from_entrypoint, load_flows_from_script
+from prefect.flows import Flow, load_flow_from_entrypoint
 from prefect.logging import get_logger
 from prefect.runner.utils import (
     inject_schemas_into_openapi,
@@ -24,6 +24,7 @@ from prefect.settings import (
     PREFECT_RUNNER_SERVER_PORT,
 )
 from prefect.utilities.asyncutils import sync_compatible
+from prefect.utilities.importtools import load_script_as_module
 
 if TYPE_CHECKING:
     from prefect.client.schemas.responses import DeploymentResponse
@@ -155,7 +156,10 @@ async def get_subflow_schemas(runner: "Runner") -> Dict[str, Dict]:
                 continue
 
             script = deployment.entrypoint.split(":")[0]
-            subflows = load_flows_from_script(script)
+            module = load_script_as_module(script)
+            subflows = [
+                obj for obj in module.__dict__.values() if isinstance(obj, Flow)
+            ]
             for flow in subflows:
                 schemas[flow.name] = flow.parameters.model_dump()
 

prefect/settings.py

@@ -1541,7 +1541,7 @@ The maximum number of retries to queue for submission.
 
 PREFECT_TASK_SCHEDULING_PENDING_TASK_TIMEOUT = Setting(
     timedelta,
-    default=timedelta(seconds=30),
+    default=timedelta(0),
 )
 """
 How long before a PENDING task are made available to another task worker.  In practice,

prefect/states.py

@@ -209,6 +209,7 @@ async def return_value_to_state(
     result_factory: ResultFactory,
     key: Optional[str] = None,
     expiration: Optional[datetime.datetime] = None,
+    defer_persistence: bool = False,
 ) -> State[R]:
     """
     Given a return value from a user's function, create a `State` the run should
@@ -242,7 +243,10 @@ async def return_value_to_state(
         # to update the data to the correct type
         if not isinstance(state.data, BaseResult):
             state.data = await result_factory.create_result(
-                state.data, key=key, expiration=expiration
+                state.data,
+                key=key,
+                expiration=expiration,
+                defer_persistence=defer_persistence,
             )
 
         return state
@@ -284,7 +288,10 @@ async def return_value_to_state(
             type=new_state_type,
             message=message,
             data=await result_factory.create_result(
-                retval, key=key, expiration=expiration
+                retval,
+                key=key,
+                expiration=expiration,
+                defer_persistence=defer_persistence,
             ),
         )
 
@@ -300,7 +307,10 @@ async def return_value_to_state(
     else:
         return Completed(
             data=await result_factory.create_result(
-                data, key=key, expiration=expiration
+                data,
+                key=key,
+                expiration=expiration,
+                defer_persistence=defer_persistence,
             )
         )
 

prefect/task_engine.py

@@ -310,6 +310,8 @@ class TaskRunEngine(Generic[P, R]):
                 result_factory=result_factory,
                 key=transaction.key,
                 expiration=expiration,
+                # defer persistence to transaction commit
+                defer_persistence=True,
             )
         )
         transaction.stage(
@@ -417,9 +419,7 @@ class TaskRunEngine(Generic[P, R]):
                     log_prints=log_prints,
                     task_run=self.task_run,
                     parameters=self.parameters,
-                    result_factory=run_coro_as_sync(
-                        ResultFactory.from_autonomous_task(self.task)
-                    ),  # type: ignore
+                    result_factory=run_coro_as_sync(ResultFactory.from_task(self.task)),  # type: ignore
                     client=client,
                 )
             )
@@ -467,9 +467,6 @@ class TaskRunEngine(Generic[P, R]):
                                 extra_task_inputs=dependencies,
                             )
                         )
-                        self.logger.info(
-                            f"Created task run {self.task_run.name!r} for task {self.task.name!r}"
-                        )
                     # Emit an event to capture that the task run was in the `PENDING` state.
                     self._last_event = emit_task_run_state_change_event(
                         task_run=self.task_run,
@@ -478,6 +475,10 @@ class TaskRunEngine(Generic[P, R]):
                     )
 
                     with self.setup_run_context():
+                        # setup_run_context might update the task run name, so log creation here
+                        self.logger.info(
+                            f"Created task run {self.task_run.name!r} for task {self.task.name!r}"
+                        )
                         yield self
 
                 except Exception:

prefect/task_runs.py

@@ -92,13 +92,18 @@ class TaskRunWaiter:
             raise RuntimeError("TaskRunWaiter must run on the global loop thread.")
 
         self._loop = loop_thread._loop
-        self._consumer_task = self._loop.create_task(self._consume_events())
+
+        consumer_started = asyncio.Event()
+        self._consumer_task = self._loop.create_task(
+            self._consume_events(consumer_started)
+        )
+        asyncio.run_coroutine_threadsafe(consumer_started.wait(), self._loop)
 
         loop_thread.add_shutdown_call(create_call(self.stop))
         atexit.register(self.stop)
         self._started = True
 
-    async def _consume_events(self):
+    async def _consume_events(self, consumer_started: asyncio.Event):
         async with get_events_subscriber(
             filter=EventFilter(
                 event=EventNameFilter(
@@ -109,6 +114,7 @@ class TaskRunWaiter:
                 )
             )
         ) as subscriber:
+            consumer_started.set()
             async for event in subscriber:
                 try:
                     self.logger.debug(
@@ -119,6 +125,7 @@ class TaskRunWaiter:
                             "prefect.task-run.", ""
                         )
                     )
+
                     with self._observed_completed_task_runs_lock:
                         # Cache the task run ID for a short period of time to avoid
                         # unnecessary waits
@@ -172,14 +179,21 @@ class TaskRunWaiter:
             # when the event is received
             instance._completion_events[task_run_id] = finished_event
 
-        with anyio.move_on_after(delay=timeout):
-            await from_async.wait_for_call_in_loop_thread(
-                create_call(finished_event.wait)
-            )
+        try:
+            # Now check one more time whether the task run arrived before we start to
+            # wait on it, in case it came in while we were setting up the event above.
+            with instance._observed_completed_task_runs_lock:
+                if task_run_id in instance._observed_completed_task_runs:
+                    return
 
-        with instance._completion_events_lock:
-            # Remove the event from the cache after it has been waited on
-            instance._completion_events.pop(task_run_id, None)
+            with anyio.move_on_after(delay=timeout):
+                await from_async.wait_for_call_in_loop_thread(
+                    create_call(finished_event.wait)
+                )
+        finally:
+            with instance._completion_events_lock:
+                # Remove the event from the cache after it has been waited on
+                instance._completion_events.pop(task_run_id, None)
 
     @classmethod
     def instance(cls):