prefect-client 3.4.5.dev4__py3-none-any.whl → 3.4.6__py3-none-any.whl

prefect/settings/profiles.py

@@ -1,17 +1,14 @@
 from __future__ import annotations
 
-import inspect
 import warnings
 from pathlib import Path
 from typing import (
     Annotated,
     Any,
     ClassVar,
-    Dict,
     Iterable,
     Iterator,
     Optional,
-    Union,
 )
 
 import toml
@@ -20,16 +17,15 @@ from pydantic import (
     BeforeValidator,
     ConfigDict,
     Field,
-    TypeAdapter,
     ValidationError,
 )
-from pydantic_settings import BaseSettings
 
 from prefect.exceptions import ProfileSettingsValidationError
 from prefect.settings.constants import DEFAULT_PROFILES_PATH
 from prefect.settings.context import get_current_settings
 from prefect.settings.legacy import Setting, _get_settings_fields
 from prefect.settings.models.root import Settings
+from prefect.utilities.collections import set_in_dict
 
 
 def _cast_settings(
@@ -69,7 +65,7 @@ class Profile(BaseModel):
     )
     source: Optional[Path] = None
 
-    def to_environment_variables(self) -> Dict[str, str]:
+    def to_environment_variables(self) -> dict[str, str]:
         """Convert the profile settings to a dictionary of environment variables."""
         return {
             setting.name: str(value)
@@ -78,23 +74,40 @@ class Profile(BaseModel):
         }
 
     def validate_settings(self) -> None:
-        errors: list[tuple[Setting, ValidationError]] = []
+        """
+        Validate all settings in this profile by creating a partial Settings object
+        with the nested structure properly constructed using accessor paths.
+        """
+        if not self.settings:
+            return
+
+        nested_settings: dict[str, Any] = {}
+
         for setting, value in self.settings.items():
-            try:
-                model_fields = Settings.model_fields
-                annotation = None
-                for section in setting.accessor.split("."):
-                    annotation = model_fields[section].annotation
-                    if inspect.isclass(annotation) and issubclass(
-                        annotation, BaseSettings
-                    ):
-                        model_fields = annotation.model_fields
-
-                TypeAdapter(annotation).validate_python(value)
-            except ValidationError as e:
-                errors.append((setting, e))
-        if errors:
-            raise ProfileSettingsValidationError(errors)
+            set_in_dict(nested_settings, setting.accessor, value)
+
+        try:
+            Settings.model_validate(nested_settings)
+        except ValidationError as e:
+            errors: list[tuple[Setting, ValidationError]] = []
+
+            for error in e.errors():
+                error_path = ".".join(str(loc) for loc in error["loc"])
+
+                for setting in self.settings.keys():
+                    if setting.accessor == error_path:
+                        errors.append(
+                            (
+                                setting,
+                                ValidationError.from_exception_data(
+                                    "ValidationError", [error]
+                                ),
+                            )
+                        )
+                        break
+
+            if errors:
+                raise ProfileSettingsValidationError(errors)
 
 
 class ProfilesCollection:
@@ -106,9 +119,7 @@ class ProfilesCollection:
     The collection may store the name of the active profile.
     """
 
-    def __init__(
-        self, profiles: Iterable[Profile], active: Optional[str] = None
-    ) -> None:
+    def __init__(self, profiles: Iterable[Profile], active: str | None = None) -> None:
         self.profiles_by_name: dict[str, Profile] = {
             profile.name: profile for profile in profiles
         }
@@ -122,7 +133,7 @@ class ProfilesCollection:
         return set(self.profiles_by_name.keys())
 
     @property
-    def active_profile(self) -> Optional[Profile]:
+    def active_profile(self) -> Profile | None:
         """
         Retrieve the active profile in this collection.
         """
@@ -130,7 +141,7 @@ class ProfilesCollection:
             return None
         return self[self.active_name]
 
-    def set_active(self, name: Optional[str], check: bool = True) -> None:
+    def set_active(self, name: str | None, check: bool = True) -> None:
         """
         Set the active profile name in the collection.
 
@@ -145,7 +156,7 @@ class ProfilesCollection:
         self,
         name: str,
         settings: dict[Setting, Any],
-        source: Optional[Path] = None,
+        source: Path | None = None,
     ) -> Profile:
         """
         Add a profile to the collection or update the existing on if the name is already
@@ -201,7 +212,7 @@ class ProfilesCollection:
         """
         self.profiles_by_name.pop(name)
 
-    def without_profile_source(self, path: Optional[Path]) -> "ProfilesCollection":
+    def without_profile_source(self, path: Path | None) -> "ProfilesCollection":
         """
         Remove profiles that were loaded from a given path.
 
@@ -367,7 +378,7 @@ def load_profile(name: str) -> Profile:
 
 
 def update_current_profile(
-    settings: Dict[Union[str, Setting], Any],
+    settings: dict[str | Setting, Any],
 ) -> Profile:
     """
     Update the persisted data for the profile currently in-use.

prefect/task_engine.py

@@ -43,6 +43,7 @@ from prefect.concurrency.v1.asyncio import concurrency as aconcurrency
 from prefect.concurrency.v1.context import ConcurrencyContext as ConcurrencyContextV1
 from prefect.concurrency.v1.sync import concurrency
 from prefect.context import (
+    AssetContext,
     AsyncClientContext,
     FlowRunContext,
     SyncClientContext,
@@ -314,10 +315,13 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             raise RuntimeError("Engine has not started.")
         return self._client
 
-    def can_retry(self, exc: Exception) -> bool:
+    def can_retry(self, exc_or_state: Exception | State[R]) -> bool:
         retry_condition: Optional[
-            Callable[["Task[P, Coroutine[Any, Any, R]]", TaskRun, State], bool]
+            Callable[["Task[P, Coroutine[Any, Any, R]]", TaskRun, State[R]], bool]
         ] = self.task.retry_condition_fn
+
+        failure_type = "exception" if isinstance(exc_or_state, Exception) else "state"
+
         if not self.task_run:
             raise ValueError("Task run is not set")
         try:
@@ -326,8 +330,8 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 f" {self.task.name!r}"
             )
             state = Failed(
-                data=exc,
-                message=f"Task run encountered unexpected exception: {repr(exc)}",
+                data=exc_or_state,
+                message=f"Task run encountered unexpected {failure_type}: {repr(exc_or_state)}",
             )
             if asyncio.iscoroutinefunction(retry_condition):
                 should_retry = run_coro_as_sync(
@@ -449,7 +453,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             else:
                 result = state.data
 
-            link_state_to_result(state, result)
+            link_state_to_result(new_state, result)
 
         # emit a state change event
         self._last_event = emit_task_run_state_change_event(
@@ -476,7 +480,15 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             # otherwise, return the exception
             return self._raised
 
-    def handle_success(self, result: R, transaction: Transaction) -> R:
+    def handle_success(
+        self, result: R, transaction: Transaction
+    ) -> Union[ResultRecord[R], None, Coroutine[Any, Any, R], R]:
+        # Handle the case where the task explicitly returns a failed state, in
+        # which case we should retry the task if it has retries left.
+        if isinstance(result, State) and result.is_failed():
+            if self.handle_retry(result):
+                return None
+
         if self.task.cache_expiration is not None:
             expiration = prefect.types._datetime.now("UTC") + self.task.cache_expiration
         else:
@@ -508,16 +520,16 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self._return_value = result
 
         self._telemetry.end_span_on_success()
-        return result
 
-    def handle_retry(self, exc: Exception) -> bool:
+    def handle_retry(self, exc_or_state: Exception | State[R]) -> bool:
         """Handle any task run retries.
 
         - If the task has retries left, and the retry condition is met, set the task to retrying and return True.
         - If the task has a retry delay, place in AwaitingRetry state with a delayed scheduled time.
         - If the task has no retries left, or the retry condition is not met, return False.
         """
-        if self.retries < self.task.retries and self.can_retry(exc):
+        failure_type = "exception" if isinstance(exc_or_state, Exception) else "state"
+        if self.retries < self.task.retries and self.can_retry(exc_or_state):
             if self.task.retry_delay_seconds:
                 delay = (
                     self.task.retry_delay_seconds[
@@ -535,8 +547,9 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 new_state = Retrying()
 
             self.logger.info(
-                "Task run failed with exception: %r - Retry %s/%s will start %s",
-                exc,
+                "Task run failed with %s: %r - Retry %s/%s will start %s",
+                failure_type,
+                exc_or_state,
                 self.retries + 1,
                 self.task.retries,
                 str(delay) + " second(s) from now" if delay else "immediately",
@@ -552,7 +565,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 else "No retries configured for this task."
             )
             self.logger.error(
-                f"Task run failed with exception: {exc!r} - {retry_message_suffix}",
+                f"Task run failed with {failure_type}: {exc_or_state!r} - {retry_message_suffix}",
                 exc_info=True,
             )
             return False
@@ -625,6 +638,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 persist_result = settings.tasks.default_persist_result
             else:
                 persist_result = should_persist_result()
+
             stack.enter_context(
                 TaskRunContext(
                     task=self.task,
@@ -647,6 +661,24 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
             yield
 
+    @contextmanager
+    def asset_context(self):
+        parent_asset_ctx = AssetContext.get()
+
+        if parent_asset_ctx and parent_asset_ctx.copy_to_child_ctx:
+            asset_ctx = parent_asset_ctx.model_copy()
+            asset_ctx.copy_to_child_ctx = False
+        else:
+            asset_ctx = AssetContext.from_task_and_inputs(
+                self.task, self.task_run.id, self.task_run.task_inputs
+            )
+
+        with asset_ctx as ctx:
+            try:
+                yield
+            finally:
+                ctx.emit_events(self.state)
+
     @contextmanager
     def initialize_run(
         self,
@@ -830,7 +862,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
     def call_task_fn(
         self, transaction: Transaction
-    ) -> Union[R, Coroutine[Any, Any, R]]:
+    ) -> Union[ResultRecord[Any], None, Coroutine[Any, Any, R], R]:
         """
         Convenience method to call the task function. Returns a coroutine if the
         task is async.
@@ -855,10 +887,13 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             raise RuntimeError("Engine has not started.")
         return self._client
 
-    async def can_retry(self, exc: Exception) -> bool:
+    async def can_retry(self, exc_or_state: Exception | State[R]) -> bool:
         retry_condition: Optional[
-            Callable[["Task[P, Coroutine[Any, Any, R]]", TaskRun, State], bool]
+            Callable[["Task[P, Coroutine[Any, Any, R]]", TaskRun, State[R]], bool]
         ] = self.task.retry_condition_fn
+
+        failure_type = "exception" if isinstance(exc_or_state, Exception) else "state"
+
         if not self.task_run:
             raise ValueError("Task run is not set")
         try:
@@ -867,8 +902,8 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 f" {self.task.name!r}"
             )
             state = Failed(
-                data=exc,
-                message=f"Task run encountered unexpected exception: {repr(exc)}",
+                data=exc_or_state,
+                message=f"Task run encountered unexpected {failure_type}: {repr(exc_or_state)}",
             )
             if asyncio.iscoroutinefunction(retry_condition):
                 should_retry = await retry_condition(self.task, self.task_run, state)
@@ -1031,7 +1066,13 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             # otherwise, return the exception
             return self._raised
 
-    async def handle_success(self, result: R, transaction: AsyncTransaction) -> R:
+    async def handle_success(
+        self, result: R, transaction: AsyncTransaction
+    ) -> Union[ResultRecord[R], None, Coroutine[Any, Any, R], R]:
+        if isinstance(result, State) and result.is_failed():
+            if await self.handle_retry(result):
+                return None
+
         if self.task.cache_expiration is not None:
             expiration = prefect.types._datetime.now("UTC") + self.task.cache_expiration
         else:
@@ -1059,19 +1100,20 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self.record_terminal_state_timing(terminal_state)
         await self.set_state(terminal_state)
         self._return_value = result
-
         self._telemetry.end_span_on_success()
 
         return result
 
-    async def handle_retry(self, exc: Exception) -> bool:
+    async def handle_retry(self, exc_or_state: Exception | State[R]) -> bool:
         """Handle any task run retries.
 
         - If the task has retries left, and the retry condition is met, set the task to retrying and return True.
         - If the task has a retry delay, place in AwaitingRetry state with a delayed scheduled time.
         - If the task has no retries left, or the retry condition is not met, return False.
         """
-        if self.retries < self.task.retries and await self.can_retry(exc):
+        failure_type = "exception" if isinstance(exc_or_state, Exception) else "state"
+
+        if self.retries < self.task.retries and await self.can_retry(exc_or_state):
             if self.task.retry_delay_seconds:
                 delay = (
                     self.task.retry_delay_seconds[
@@ -1089,8 +1131,9 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 new_state = Retrying()
 
             self.logger.info(
-                "Task run failed with exception: %r - Retry %s/%s will start %s",
-                exc,
+                "Task run failed with %s: %r - Retry %s/%s will start %s",
+                failure_type,
+                exc_or_state,
                 self.retries + 1,
                 self.task.retries,
                 str(delay) + " second(s) from now" if delay else "immediately",
@@ -1106,7 +1149,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 else "No retries configured for this task."
             )
             self.logger.error(
-                f"Task run failed with exception: {exc!r} - {retry_message_suffix}",
+                f"Task run failed with {failure_type}: {exc_or_state!r} - {retry_message_suffix}",
                 exc_info=True,
             )
             return False
@@ -1180,6 +1223,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 persist_result = settings.tasks.default_persist_result
             else:
                 persist_result = should_persist_result()
+
             stack.enter_context(
                 TaskRunContext(
                     task=self.task,
@@ -1201,6 +1245,24 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
             yield
 
+    @asynccontextmanager
+    async def asset_context(self):
+        parent_asset_ctx = AssetContext.get()
+
+        if parent_asset_ctx and parent_asset_ctx.copy_to_child_ctx:
+            asset_ctx = parent_asset_ctx.model_copy()
+            asset_ctx.copy_to_child_ctx = False
+        else:
+            asset_ctx = AssetContext.from_task_and_inputs(
+                self.task, self.task_run.id, self.task_run.task_inputs
+            )
+
+        with asset_ctx as ctx:
+            try:
+                yield
+            finally:
+                ctx.emit_events(self.state)
+
     @asynccontextmanager
     async def initialize_run(
         self,
@@ -1382,7 +1444,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
     async def call_task_fn(
         self, transaction: AsyncTransaction
-    ) -> Union[R, Coroutine[Any, Any, R]]:
+    ) -> Union[ResultRecord[Any], None, Coroutine[Any, Any, R], R]:
         """
         Convenience method to call the task function. Returns a coroutine if the
         task is async.
@@ -1417,7 +1479,11 @@ def run_task_sync(
     with engine.start(task_run_id=task_run_id, dependencies=dependencies):
         while engine.is_running():
             run_coro_as_sync(engine.wait_until_ready())
-            with engine.run_context(), engine.transaction_context() as txn:
+            with (
+                engine.asset_context(),
+                engine.run_context(),
+                engine.transaction_context() as txn,
+            ):
                 engine.call_task_fn(txn)
 
     return engine.state if return_type == "state" else engine.result()
@@ -1444,7 +1510,11 @@ async def run_task_async(
     async with engine.start(task_run_id=task_run_id, dependencies=dependencies):
         while engine.is_running():
             await engine.wait_until_ready()
-            async with engine.run_context(), engine.transaction_context() as txn:
+            async with (
+                engine.asset_context(),
+                engine.run_context(),
+                engine.transaction_context() as txn,
+            ):
                 await engine.call_task_fn(txn)
 
     return engine.state if return_type == "state" else await engine.result()
@@ -1474,7 +1544,11 @@ def run_generator_task_sync(
     with engine.start(task_run_id=task_run_id, dependencies=dependencies):
         while engine.is_running():
             run_coro_as_sync(engine.wait_until_ready())
-            with engine.run_context(), engine.transaction_context() as txn:
+            with (
+                engine.asset_context(),
+                engine.run_context(),
+                engine.transaction_context() as txn,
+            ):
                 # TODO: generators should default to commit_mode=OFF
                 # because they are dynamic by definition
                 # for now we just prevent this branch explicitly
@@ -1528,7 +1602,11 @@ async def run_generator_task_async(
     async with engine.start(task_run_id=task_run_id, dependencies=dependencies):
         while engine.is_running():
             await engine.wait_until_ready()
-            async with engine.run_context(), engine.transaction_context() as txn:
+            async with (
+                engine.asset_context(),
+                engine.run_context(),
+                engine.transaction_context() as txn,
+            ):
                 # TODO: generators should default to commit_mode=OFF
                 # because they are dynamic by definition
                 # for now we just prevent this branch explicitly