prefect-client 3.2.7__py3-none-any.whl → 3.2.8__py3-none-any.whl

prefect/runner/runner.py

@@ -39,6 +39,7 @@ import datetime
 import inspect
 import logging
 import os
+import shlex
 import shutil
 import signal
 import subprocess
@@ -227,6 +228,7 @@ class Runner:
         self._scheduled_task_scopes: set[anyio.abc.CancelScope] = set()
         self._deployment_ids: set[UUID] = set()
         self._flow_run_process_map: dict[UUID, ProcessMapEntry] = dict()
+        self.__flow_run_process_map_lock: asyncio.Lock | None = None
         self._flow_run_bundle_map: dict[UUID, SerializedBundle] = dict()
 
         self._tmp_dir: Path = (
@@ -243,6 +245,12 @@ class Runner:
         )
         self._flow_cache: LRUCache[UUID, "APIFlow"] = LRUCache(maxsize=100)
 
+    @property
+    def _flow_run_process_map_lock(self) -> asyncio.Lock:
+        if self.__flow_run_process_map_lock is None:
+            self.__flow_run_process_map_lock = asyncio.Lock()
+        return self.__flow_run_process_map_lock
+
     @sync_compatible
     async def add_deployment(
         self,
@@ -550,13 +558,23 @@ class Runner:
             )
 
     async def execute_flow_run(
-        self, flow_run_id: UUID, entrypoint: Optional[str] = None
-    ) -> None:
+        self,
+        flow_run_id: UUID,
+        entrypoint: str | None = None,
+        command: str | None = None,
+        cwd: Path | None = None,
+        env: dict[str, str | None] | None = None,
+        task_status: anyio.abc.TaskStatus[int] | None = None,
+        stream_output: bool = True,
+    ) -> anyio.abc.Process | None:
         """
         Executes a single flow run with the given ID.
 
         Execution will wait to monitor for cancellation requests. Exits once
         the flow run process has exited.
+
+        Returns:
+            The flow run process.
         """
         self.pause_on_shutdown = False
         context = self if not self.started else asyncnullcontext()
@@ -570,17 +588,26 @@ class Runner:
                     self._submitting_flow_run_ids.add(flow_run_id)
                     flow_run = await self._client.read_flow_run(flow_run_id)
 
-                    pid = await self._runs_task_group.start(
+                    process: anyio.abc.Process = 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,
                         ),
                     )
-
-                    self._flow_run_process_map[flow_run.id] = ProcessMapEntry(
-                        pid=pid, flow_run=flow_run
-                    )
+                    if task_status:
+                        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
+                            )
 
                     # 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
@@ -609,6 +636,8 @@ class Runner:
                             )
                         )
 
+                    return process
+
     async def execute_bundle(self, bundle: SerializedBundle) -> None:
         """
         Executes a bundle in a subprocess.
@@ -696,7 +725,7 @@ class Runner:
                     )
                 elif (
                     sys.platform == "win32"
-                    and process.returncode == STATUS_CONTROL_C_EXIT
+                    and process.exitcode == STATUS_CONTROL_C_EXIT
                 ):
                     level = logging.INFO
                     help_message = (
@@ -733,9 +762,13 @@ class Runner:
     async def _run_process(
         self,
         flow_run: "FlowRun",
-        task_status: Optional[anyio.abc.TaskStatus[int]] = None,
-        entrypoint: Optional[str] = None,
-    ) -> int:
+        task_status: anyio.abc.TaskStatus[anyio.abc.Process] | None = None,
+        entrypoint: str | None = None,
+        command: str | None = None,
+        cwd: Path | None = None,
+        env: dict[str, str | None] | None = None,
+        stream_output: bool = True,
+    ) -> anyio.abc.Process:
         """
         Runs the given flow run in a subprocess.
 
@@ -747,7 +780,10 @@ class Runner:
             task_status: anyio task status used to send a message to the caller
                 than the flow run process has started.
         """
-        command = [get_sys_executable(), "-m", "prefect.engine"]
+        if command is None:
+            runner_command = [get_sys_executable(), "-m", "prefect.engine"]
+        else:
+            runner_command = shlex.split(command)
 
         flow_run_logger = self._get_flow_run_logger(flow_run)
 
@@ -760,7 +796,9 @@ class Runner:
 
         flow_run_logger.info("Opening process...")
 
-        env = get_current_settings().to_environment_variables(exclude_unset=True)
+        if env is None:
+            env = {}
+        env.update(get_current_settings().to_environment_variables(exclude_unset=True))
         env.update(
             {
                 **{
@@ -798,12 +836,12 @@ class Runner:
                 setattr(storage, "last_adhoc_pull", datetime.datetime.now())
 
         process = await run_process(
-            command=command,
-            stream_output=True,
+            command=runner_command,
+            stream_output=stream_output,
             task_status=task_status,
-            task_status_handler=None,
+            task_status_handler=lambda process: process,
             env=env,
-            cwd=storage.destination if storage else None,
+            cwd=storage.destination if storage else cwd,
             **kwargs,
         )
 
@@ -852,7 +890,7 @@ class Runner:
                 f"Process for flow run {flow_run.name!r} exited cleanly."
             )
 
-        return process.returncode
+        return process
 
     async def _kill_process(
         self,
@@ -1316,9 +1354,10 @@ class Runner:
             )
 
             if readiness_result and not isinstance(readiness_result, Exception):
-                self._flow_run_process_map[flow_run.id] = ProcessMapEntry(
-                    pid=readiness_result, flow_run=flow_run
-                )
+                async with self._flow_run_process_map_lock:
+                    self._flow_run_process_map[flow_run.id] = ProcessMapEntry(
+                        pid=readiness_result, flow_run=flow_run
+                    )
             # Heartbeats are opt-in and only emitted if a heartbeat frequency is set
             if self.heartbeat_seconds is not None:
                 await self._emit_flow_run_heartbeat(flow_run)
@@ -1333,17 +1372,26 @@ class Runner:
     async def _submit_run_and_capture_errors(
         self,
         flow_run: "FlowRun",
-        task_status: anyio.abc.TaskStatus[int | Exception],
-        entrypoint: Optional[str] = None,
+        task_status: anyio.abc.TaskStatus[anyio.abc.Process | Exception],
+        entrypoint: str | None = None,
+        command: str | None = None,
+        cwd: Path | None = None,
+        env: dict[str, str | None] | None = None,
+        stream_output: bool = True,
     ) -> Union[Optional[int], Exception]:
         run_logger = self._get_flow_run_logger(flow_run)
 
         try:
-            status_code = await self._run_process(
+            process = await self._run_process(
                 flow_run=flow_run,
                 task_status=task_status,
                 entrypoint=entrypoint,
+                command=command,
+                cwd=cwd,
+                env=env,
+                stream_output=stream_output,
             )
+            status_code = process.returncode
         except Exception as exc:
             if task_status:
                 # This flow run was being submitted and did not start successfully
@@ -1363,7 +1411,9 @@ class Runner:
             return exc
         finally:
             self._release_limit_slot(flow_run.id)
-            self._flow_run_process_map.pop(flow_run.id, None)
+
+            async with self._flow_run_process_map_lock:
+                self._flow_run_process_map.pop(flow_run.id, None)
 
         if status_code != 0:
             await self._propose_crashed_state(
@@ -1513,6 +1563,7 @@ class Runner:
         """
         Run the hooks for a flow.
         """
+        run_logger = self._get_flow_run_logger(flow_run)
         if state.is_cancelling():
             try:
                 if flow_run.id in self._flow_run_bundle_map:
@@ -1520,6 +1571,7 @@ class Runner:
                         self._flow_run_bundle_map[flow_run.id]
                     )
                 else:
+                    run_logger.info("Loading flow to check for on_cancellation hooks")
                     flow = await load_flow_from_flow_run(
                         flow_run, storage_base_path=str(self._tmp_dir)
                     )
@@ -1529,7 +1581,7 @@ class Runner:
             except ObjectNotFound:
                 run_logger = self._get_flow_run_logger(flow_run)
                 run_logger.warning(
-                    f"Runner cannot retrieve flow to execute cancellation hooks for flow run {flow_run.id!r}."
+                    f"Runner failed to retrieve flow to execute on_cancellation hooks for flow run {flow_run.id!r}."
                 )
 
     async def _run_on_crashed_hooks(
@@ -1540,16 +1592,25 @@ class Runner:
         """
         Run the hooks for a flow.
         """
+        run_logger = self._get_flow_run_logger(flow_run)
         if state.is_crashed():
-            if flow_run.id in self._flow_run_bundle_map:
-                flow = extract_flow_from_bundle(self._flow_run_bundle_map[flow_run.id])
-            else:
-                flow = await load_flow_from_flow_run(
-                    flow_run, storage_base_path=str(self._tmp_dir)
-                )
-            hooks = flow.on_crashed_hooks or []
+            try:
+                if flow_run.id in self._flow_run_bundle_map:
+                    flow = extract_flow_from_bundle(
+                        self._flow_run_bundle_map[flow_run.id]
+                    )
+                else:
+                    run_logger.info("Loading flow to check for on_crashed hooks")
+                    flow = await load_flow_from_flow_run(
+                        flow_run, storage_base_path=str(self._tmp_dir)
+                    )
+                hooks = flow.on_crashed_hooks or []
 
-            await _run_hooks(hooks, flow_run, flow, state)
+                await _run_hooks(hooks, flow_run, flow, state)
+            except ObjectNotFound:
+                run_logger.warning(
+                    f"Runner failed to retrieve flow to execute on_crashed hooks for flow run {flow_run.id!r}."
+                )
 
     async def __aenter__(self) -> Self:
         self._logger.debug("Starting runner...")

prefect/server/api/artifacts.py

@@ -26,6 +26,11 @@ async def create_artifact(
     response: Response,
     db: PrefectDBInterface = Depends(provide_database_interface),
 ) -> core.Artifact:
+    """
+    Create an artifact.
+
+    For more information, see https://docs.prefect.io/v3/develop/artifacts.
+    """
     artifact = core.Artifact(**artifact.model_dump())
 
     right_now = now("UTC")

prefect/server/api/automations.py

@@ -39,6 +39,11 @@ async def create_automation(
     automation: AutomationCreate,
     db: PrefectDBInterface = Depends(provide_database_interface),
 ) -> Automation:
+    """
+    Create an automation.
+
+    For more information, see https://docs.prefect.io/v3/automate.
+    """
     # reset any client-provided IDs on the provided triggers
     automation.trigger.reset_ids()
 

prefect/server/api/block_capabilities.py

@@ -19,6 +19,11 @@ router: PrefectRouter = PrefectRouter(
 async def read_available_block_capabilities(
     db: PrefectDBInterface = Depends(provide_database_interface),
 ) -> List[str]:
+    """
+    Get available block capabilities.
+
+    For more information, see https://docs.prefect.io/v3/develop/blocks.
+    """
     async with db.session_context() as session:
         return await models.block_schemas.read_available_block_capabilities(
             session=session

prefect/server/api/block_documents.py

@@ -24,6 +24,8 @@ async def create_block_document(
 ) -> schemas.core.BlockDocument:
     """
     Create a new block document.
+
+    For more information, see https://docs.prefect.io/v3/develop/blocks.
     """
     async with db.session_context(begin_transaction=True) as session:
         if block_document.name is not None:

prefect/server/api/block_schemas.py

@@ -30,6 +30,11 @@ async def create_block_schema(
     response: Response,
     db: PrefectDBInterface = Depends(provide_database_interface),
 ) -> schemas.core.BlockSchema:
+    """
+    Create a block schema.
+
+    For more information, see https://docs.prefect.io/v3/develop/blocks.
+    """
     from prefect.blocks.core import Block
 
     async with db.session_context(begin_transaction=True) as session:

prefect/server/api/block_types.py

@@ -19,7 +19,9 @@ async def create_block_type(
     db: PrefectDBInterface = Depends(provide_database_interface),
 ) -> schemas.core.BlockType:
     """
-    Create a new block type
+    Create a new block type.
+
+    For more information, see https://docs.prefect.io/v3/develop/blocks.
     """
     # API-created blocks cannot start with the word "Prefect"
     # as it is reserved for system use

prefect/server/api/concurrency_limits.py

@@ -28,6 +28,11 @@ async def create_concurrency_limit(
     response: Response,
     db: PrefectDBInterface = Depends(provide_database_interface),
 ) -> schemas.core.ConcurrencyLimit:
+    """
+    Create a task run concurrency limit.
+
+    For more information, see https://docs.prefect.io/v3/develop/task-run-limits.
+    """
     # hydrate the input model into a full model
     concurrency_limit_model = schemas.core.ConcurrencyLimit(
         **concurrency_limit.model_dump()

prefect/server/api/concurrency_limits_v2.py

@@ -21,6 +21,11 @@ async def create_concurrency_limit_v2(
     concurrency_limit: actions.ConcurrencyLimitV2Create,
     db: PrefectDBInterface = Depends(provide_database_interface),
 ) -> schemas.core.ConcurrencyLimitV2:
+    """
+    Create a task run concurrency limit.
+
+    For more information, see https://docs.prefect.io/v3/develop/global-concurrency-limits.
+    """
     async with db.session_context(begin_transaction=True) as session:
         model = await models.concurrency_limits_v2.create_concurrency_limit(
             session=session, concurrency_limit=concurrency_limit

prefect/server/api/deployments.py

@@ -68,6 +68,8 @@ async def create_deployment(
 
     If the deployment has an active schedule, flow runs will be scheduled.
     When upserting, any scheduled runs from the existing deployment will be deleted.
+
+    For more information, see https://docs.prefect.io/v3/deploy.
     """
 
     data = deployment.model_dump(exclude_unset=True)

prefect/server/api/events.py

@@ -48,7 +48,11 @@ async def create_events(
     events: List[Event],
     ephemeral_request: bool = Depends(is_ephemeral_request),
 ) -> None:
-    """Record a batch of Events"""
+    """
+    Record a batch of Events.
+
+    For more information, see https://docs.prefect.io/v3/automate/events/events.
+    """
     if ephemeral_request:
         await EventsPipeline().process_events(events)
     else:

prefect/server/api/flow_run_notification_policies.py

@@ -25,6 +25,8 @@ async def create_flow_run_notification_policy(
 ) -> schemas.core.FlowRunNotificationPolicy:
     """
     Creates a new flow run notification policy.
+
+    For more information, see https://docs.prefect.io/v3/automate/events/automations-triggers#sending-notifications-with-automations.
     """
     async with db.session_context(begin_transaction=True) as session:
         return await models.flow_run_notification_policies.create_flow_run_notification_policy(

prefect/server/api/flow_run_states.py

@@ -26,6 +26,8 @@ async def read_flow_run_state(
 ) -> schemas.states.State:
     """
     Get a flow run state by id.
+
+    For more information, see https://docs.prefect.io/v3/develop/write-flows#final-state-determination.
     """
     async with db.session_context() as session:
         flow_run_state = await models.flow_run_states.read_flow_run_state(

prefect/server/api/flow_runs.py

@@ -74,6 +74,8 @@ async def create_flow_run(
     idempotency key already exists, the existing flow run will be returned.
 
     If no state is provided, the flow run will be created in a PENDING state.
+
+    For more information, see https://docs.prefect.io/v3/develop/write-flows.
     """
     # hydrate the input model into a full flow run / state model
     flow_run_object = schemas.core.FlowRun(

prefect/server/api/flows.py

@@ -27,6 +27,8 @@ async def create_flow(
 ) -> schemas.core.Flow:
     """Gracefully creates a new flow from the provided schema. If a flow with the
     same name already exists, the existing flow is returned.
+
+    For more information, see https://docs.prefect.io/v3/develop/write-flows.
     """
     # hydrate the input model into a full flow model
     flow = schemas.core.Flow(**flow.model_dump())

prefect/server/api/logs.py

@@ -20,7 +20,11 @@ async def create_logs(
     logs: List[schemas.actions.LogCreate],
     db: PrefectDBInterface = Depends(provide_database_interface),
 ) -> None:
-    """Create new logs from the provided schema."""
+    """
+    Create new logs from the provided schema.
+
+    For more information, see https://docs.prefect.io/v3/develop/logging.
+    """
     for batch in models.logs.split_logs_into_batches(logs):
         async with db.session_context(begin_transaction=True) as session:
             await models.logs.create_logs(session=session, logs=batch)

prefect/server/api/task_run_states.py

@@ -26,6 +26,8 @@ async def read_task_run_state(
 ) -> schemas.states.State:
     """
     Get a task run state by id.
+
+    For more information, see https://docs.prefect.io/v3/develop/write-tasks.
     """
     async with db.session_context() as session:
         task_run_state = await models.task_run_states.read_task_run_state(

prefect/server/api/task_runs.py

@@ -57,6 +57,8 @@ async def create_task_run(
     run will be returned.
 
     If no state is provided, the task run will be created in a PENDING state.
+
+    For more information, see https://docs.prefect.io/v3/develop/write-tasks.
     """
     # hydrate the input model into a full task run / state model
     task_run_dict = task_run.model_dump()

prefect/server/api/task_workers.py

@@ -20,7 +20,11 @@ async def read_task_workers(
         default=None, description="The task worker filter", embed=True
     ),
 ) -> List[TaskWorkerResponse]:
-    """Read active task workers. Optionally filter by task keys."""
+    """
+    Read active task workers. Optionally filter by task keys.
+
+    For more information, see https://docs.prefect.io/v3/develop/deferred-tasks.
+    """
 
     if task_worker_filter and task_worker_filter.task_keys:
         return await models.task_workers.get_workers_for_task_keys(

prefect/server/api/variables.py

@@ -57,6 +57,11 @@ async def create_variable(
     variable: actions.VariableCreate,
     db: PrefectDBInterface = Depends(provide_database_interface),
 ) -> core.Variable:
+    """
+    Create a variable.
+
+    For more information, see https://docs.prefect.io/v3/develop/variables.
+    """
     async with db.session_context(begin_transaction=True) as session:
         try:
             model = await models.variables.create_variable(

prefect/server/api/work_queues.py

@@ -45,6 +45,8 @@ async def create_work_queue(
 
     If a work queue with the same name already exists, an error
     will be raised.
+
+    For more information, see https://docs.prefect.io/v3/deploy/infrastructure-concepts/work-pools#work-queues.
     """
 
     try:

prefect/server/api/workers.py

@@ -161,6 +161,8 @@ async def create_work_pool(
     """
     Creates a new work pool. If a work pool with the same
     name already exists, an error will be raised.
+
+    For more information, see https://docs.prefect.io/v3/deploy/infrastructure-concepts/work-pools.
     """
     if work_pool.name.lower().startswith("prefect"):
         raise HTTPException(
@@ -408,6 +410,8 @@ async def create_work_queue(
     """
     Creates a new work pool queue. If a work pool queue with the same
     name already exists, an error will be raised.
+
+    For more information, see https://docs.prefect.io/v3/deploy/infrastructure-concepts/work-pools#work-queues.
     """
 
     try:

prefect/settings/profiles.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import inspect
 import warnings
 from pathlib import Path
@@ -9,7 +11,6 @@ from typing import (
     Iterable,
     Iterator,
     Optional,
-    Set,
     Union,
 )
 
@@ -32,8 +33,8 @@ from prefect.settings.models.root import Settings
 
 
 def _cast_settings(
-    settings: Union[Dict[Union[str, Setting], Any], Any],
-) -> Dict[Setting, Any]:
+    settings: dict[str | Setting, Any] | Any,
+) -> dict[Setting, Any]:
     """For backwards compatibility, allow either Settings objects as keys or string references to settings."""
     if not isinstance(settings, dict):
         raise ValueError("Settings must be a dictionary.")
@@ -63,7 +64,7 @@ class Profile(BaseModel):
     )
 
     name: str
-    settings: Annotated[Dict[Setting, Any], BeforeValidator(_cast_settings)] = Field(
+    settings: Annotated[dict[Setting, Any], BeforeValidator(_cast_settings)] = Field(
         default_factory=dict
     )
     source: Optional[Path] = None
@@ -114,7 +115,7 @@ class ProfilesCollection:
         self.active_name = active
 
     @property
-    def names(self) -> Set[str]:
+    def names(self) -> set[str]:
         """
         Return a set of profile names in this collection.
         """

prefect/task_worker.py

@@ -18,7 +18,7 @@ import uvicorn
 from exceptiongroup import BaseExceptionGroup  # novermin
 from fastapi import FastAPI
 from typing_extensions import ParamSpec, Self, TypeVar
-from websockets.exceptions import InvalidStatusCode
+from websockets.exceptions import InvalidStatus
 
 from prefect import Task
 from prefect._internal.concurrency.api import create_call, from_sync
@@ -181,8 +181,8 @@ class TaskWorker:
             logger.info("Starting task worker...")
             try:
                 await self._subscribe_to_task_scheduling()
-            except InvalidStatusCode as exc:
-                if exc.status_code == 403:
+            except InvalidStatus as exc:
+                if exc.response.status_code == 403:
                     logger.error(
                         "403: Could not establish a connection to the `/task_runs/subscriptions/scheduled`"
                         f" endpoint found at:\n\n {PREFECT_API_URL.value()}"

prefect/telemetry/instrumentation.py

@@ -95,7 +95,7 @@ def _setup_meter_provider(
     resource: Resource, headers: dict[str, str], telemetry_url: str
 ) -> MeterProvider:
     metric_reader = PeriodicExportingMetricReader(
-        OTLPMetricExporter(
+        OTLPMetricExporter(  # pyright: ignore[reportArgumentType] `preferred_temporality` and `preferred_aggregation` default to `None`, but otel's typing doesn't include it
             endpoint=_url_join(telemetry_url, "v1/metrics"),
             headers=headers,
         )
@@ -109,7 +109,7 @@ def _setup_meter_provider(
 def _setup_logger_provider(
     resource: Resource, headers: dict[str, str], telemetry_url: str
 ) -> LoggerProvider:
-    logger_provider = LoggerProvider(resource=resource)
+    logger_provider = LoggerProvider(resource=resource)  # pyright: ignore[reportArgumentType] `multi_log_record_processor` defaults to `None` but otel's typing doesn't include it
     queueing_log_exporter = QueueingLogExporter.instance(
         _url_join(telemetry_url, "v1/logs"), tuple(headers.items())
     )