prefect-client 3.2.12__py3-none-any.whl → 3.2.14__py3-none-any.whl

prefect/workers/base.py

@@ -3,31 +3,31 @@ from __future__ import annotations
 import abc
 import asyncio
 import threading
+import warnings
 from contextlib import AsyncExitStack
 from functools import partial
 from typing import (
     TYPE_CHECKING,
     Any,
     Callable,
-    Dict,
     Generic,
-    List,
     Optional,
-    Set,
     Type,
-    Union,
 )
 from uuid import UUID, uuid4
 
 import anyio
 import anyio.abc
 import httpx
-from importlib_metadata import distributions
+from importlib_metadata import (
+    distributions,  # type: ignore[reportUnknownVariableType] incomplete typing
+)
 from pydantic import BaseModel, Field, PrivateAttr, field_validator
 from pydantic.json_schema import GenerateJsonSchema
 from typing_extensions import Literal, Self, TypeVar
 
 import prefect
+from prefect._internal.compatibility.deprecated import PrefectDeprecationWarning
 from prefect._internal.schemas.validators import return_v_or_none
 from prefect.client.base import ServerType
 from prefect.client.orchestration import PrefectClient, get_client
@@ -60,6 +60,7 @@ from prefect.settings import (
     get_current_settings,
 )
 from prefect.states import (
+    Cancelled,
     Crashed,
     Pending,
     exception_to_failed_state,
@@ -94,12 +95,12 @@ class BaseJobConfiguration(BaseModel):
             "will be automatically generated by the worker."
         ),
     )
-    env: Dict[str, Optional[str]] = Field(
+    env: dict[str, Optional[str]] = Field(
         default_factory=dict,
         title="Environment Variables",
         description="Environment variables to set when starting a flow run.",
     )
-    labels: Dict[str, str] = Field(
+    labels: dict[str, str] = Field(
         default_factory=dict,
         description=(
             "Labels applied to infrastructure created by the worker using "
@@ -114,7 +115,7 @@ class BaseJobConfiguration(BaseModel):
         ),
     )
 
-    _related_objects: Dict[str, Any] = PrivateAttr(default_factory=dict)
+    _related_objects: dict[str, Any] = PrivateAttr(default_factory=dict)
 
     @property
     def is_using_a_runner(self) -> bool:
@@ -122,18 +123,18 @@ class BaseJobConfiguration(BaseModel):
 
     @field_validator("command")
     @classmethod
-    def _coerce_command(cls, v):
+    def _coerce_command(cls, v: str | None) -> str | None:
         return return_v_or_none(v)
 
     @field_validator("env", mode="before")
     @classmethod
-    def _coerce_env(cls, v):
+    def _coerce_env(cls, v: dict[str, Any]) -> dict[str, str | None]:
         return {k: str(v) if v is not None else None for k, v in v.items()}
 
     @staticmethod
-    def _get_base_config_defaults(variables: dict) -> dict:
+    def _get_base_config_defaults(variables: dict[str, Any]) -> dict[str, Any]:
         """Get default values from base config for all variables that have them."""
-        defaults = dict()
+        defaults: dict[str, Any] = {}
         for variable_name, attrs in variables.items():
             # We remote `None` values because we don't want to use them in templating.
             # The currently logic depends on keys not existing to populate the correct value
@@ -149,9 +150,9 @@ class BaseJobConfiguration(BaseModel):
     @inject_client
     async def from_template_and_values(
         cls,
-        base_job_template: dict,
-        values: dict,
-        client: Optional["PrefectClient"] = None,
+        base_job_template: dict[str, Any],
+        values: dict[str, Any],
+        client: "PrefectClient | None" = None,
     ):
         """Creates a valid worker configuration object from the provided base
         configuration and overrides.
@@ -159,7 +160,7 @@ class BaseJobConfiguration(BaseModel):
         Important: this method expects that the base_job_template was already
         validated server-side.
         """
-        base_config: Dict[str, Any] = base_job_template["job_configuration"]
+        base_config: dict[str, Any] = base_job_template["job_configuration"]
         variables_schema = base_job_template["variables"]
         variables = cls._get_base_config_defaults(
             variables_schema.get("properties", {})
@@ -213,8 +214,10 @@ class BaseJobConfiguration(BaseModel):
     def prepare_for_flow_run(
         self,
         flow_run: "FlowRun",
-        deployment: Optional["DeploymentResponse"] = None,
-        flow: Optional["Flow"] = None,
+        deployment: "DeploymentResponse | None" = None,
+        flow: "Flow | None" = None,
+        work_pool: "WorkPool | None" = None,
+        worker_name: str | None = None,
     ) -> None:
         """
         Prepare the job configuration for a flow run.
@@ -227,6 +230,8 @@ class BaseJobConfiguration(BaseModel):
             flow_run: The flow run to be executed.
             deployment: The deployment that the flow run is associated with.
             flow: The flow that the flow run is associated with.
+            work_pool: The work pool that the flow run is running in.
+            worker_name: The name of the worker that is submitting the flow run.
         """
 
         self._related_objects = {
@@ -234,26 +239,19 @@ class BaseJobConfiguration(BaseModel):
             "flow": flow,
             "flow-run": flow_run,
         }
-        if deployment is not None:
-            deployment_labels = self._base_deployment_labels(deployment)
-        else:
-            deployment_labels = {}
-
-        if flow is not None:
-            flow_labels = self._base_flow_labels(flow)
-        else:
-            flow_labels = {}
 
         env = {
             **self._base_environment(),
             **self._base_flow_run_environment(flow_run),
-            **(self.env if isinstance(self.env, dict) else {}),
+            **(self.env if isinstance(self.env, dict) else {}),  # pyright: ignore[reportUnnecessaryIsInstance]
         }
         self.env = {key: value for key, value in env.items() if value is not None}
         self.labels = {
             **self._base_flow_run_labels(flow_run),
-            **deployment_labels,
-            **flow_labels,
+            **self._base_work_pool_labels(work_pool),
+            **self._base_worker_name_label(worker_name),
+            **self._base_flow_labels(flow),
+            **self._base_deployment_labels(deployment),
             **self.labels,
         }
         self.name = self.name or flow_run.name
@@ -267,7 +265,7 @@ class BaseJobConfiguration(BaseModel):
         return "prefect flow-run execute"
 
     @staticmethod
-    def _base_flow_run_labels(flow_run: "FlowRun") -> Dict[str, str]:
+    def _base_flow_run_labels(flow_run: "FlowRun") -> dict[str, str]:
         """
         Generate a dictionary of labels for a flow run job.
         """
@@ -278,7 +276,7 @@ class BaseJobConfiguration(BaseModel):
         }
 
     @classmethod
-    def _base_environment(cls) -> Dict[str, str]:
+    def _base_environment(cls) -> dict[str, str]:
         """
         Environment variables that should be passed to all created infrastructure.
 
@@ -287,14 +285,22 @@ class BaseJobConfiguration(BaseModel):
         return get_current_settings().to_environment_variables(exclude_unset=True)
 
     @staticmethod
-    def _base_flow_run_environment(flow_run: "FlowRun") -> Dict[str, str]:
+    def _base_flow_run_environment(flow_run: "FlowRun | None") -> dict[str, str]:
         """
         Generate a dictionary of environment variables for a flow run job.
         """
+        if flow_run is None:
+            return {}
+
         return {"PREFECT__FLOW_RUN_ID": str(flow_run.id)}
 
     @staticmethod
-    def _base_deployment_labels(deployment: "DeploymentResponse") -> Dict[str, str]:
+    def _base_deployment_labels(
+        deployment: "DeploymentResponse | None",
+    ) -> dict[str, str]:
+        if deployment is None:
+            return {}
+
         labels = {
             "prefect.io/deployment-id": str(deployment.id),
             "prefect.io/deployment-name": deployment.name,
@@ -306,15 +312,37 @@ class BaseJobConfiguration(BaseModel):
         return labels
 
     @staticmethod
-    def _base_flow_labels(flow: "Flow") -> Dict[str, str]:
+    def _base_flow_labels(flow: "Flow | None") -> dict[str, str]:
+        if flow is None:
+            return {}
+
         return {
             "prefect.io/flow-id": str(flow.id),
             "prefect.io/flow-name": flow.name,
         }
 
-    def _related_resources(self) -> List[RelatedResource]:
-        tags = set()
-        related = []
+    @staticmethod
+    def _base_work_pool_labels(work_pool: "WorkPool | None") -> dict[str, str]:
+        """Adds the work pool labels to the job manifest."""
+        if work_pool is None:
+            return {}
+
+        return {
+            "prefect.io/work-pool-name": work_pool.name,
+            "prefect.io/work-pool-id": str(work_pool.id),
+        }
+
+    @staticmethod
+    def _base_worker_name_label(worker_name: str | None) -> dict[str, str]:
+        """Adds the worker name label to the job manifest."""
+        if worker_name is None:
+            return {}
+
+        return {"prefect.io/worker-name": worker_name}
+
+    def _related_resources(self) -> list[RelatedResource]:
+        tags: set[str] = set()
+        related: list[RelatedResource] = []
 
         for kind, obj in self._related_objects.items():
             if obj is None:
@@ -331,12 +359,12 @@ class BaseVariables(BaseModel):
         default=None,
         description="Name given to infrastructure created by a worker.",
     )
-    env: Dict[str, Optional[str]] = Field(
+    env: dict[str, Optional[str]] = Field(
         default_factory=dict,
         title="Environment Variables",
         description="Environment variables to set when starting a flow run.",
     )
-    labels: Dict[str, str] = Field(
+    labels: dict[str, str] = Field(
         default_factory=dict,
         description="Labels applied to infrastructure created by a worker.",
     )
@@ -356,7 +384,7 @@ class BaseVariables(BaseModel):
         ref_template: str = "#/definitions/{model}",
         schema_generator: Type[GenerateJsonSchema] = GenerateJsonSchema,
         mode: Literal["validation", "serialization"] = "validation",
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         """TODO: stop overriding this method - use GenerateSchema in ConfigDict instead?"""
         schema = super().model_json_schema(
             by_alias, ref_template, schema_generator, mode
@@ -403,14 +431,14 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
     def __init__(
         self,
         work_pool_name: str,
-        work_queues: Optional[List[str]] = None,
-        name: Optional[str] = None,
-        prefetch_seconds: Optional[float] = None,
+        work_queues: list[str] | None = None,
+        name: str | None = None,
+        prefetch_seconds: float | None = None,
         create_pool_if_not_found: bool = True,
-        limit: Optional[int] = None,
-        heartbeat_interval_seconds: Optional[int] = None,
+        limit: int | None = None,
+        heartbeat_interval_seconds: int | None = None,
         *,
-        base_job_template: Optional[Dict[str, Any]] = None,
+        base_job_template: dict[str, Any] | None = None,
     ):
         """
         Base class for all Prefect workers.
@@ -445,7 +473,7 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
         self._create_pool_if_not_found = create_pool_if_not_found
         self._base_job_template = base_job_template
         self._work_pool_name = work_pool_name
-        self._work_queues: Set[str] = set(work_queues) if work_queues else set()
+        self._work_queues: set[str] = set(work_queues) if work_queues else set()
 
         self._prefetch_seconds: float = (
             prefetch_seconds or PREFECT_WORKER_PREFETCH_SECONDS.value()
@@ -461,11 +489,35 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
         self._last_polled_time: DateTime = DateTime.now("utc")
         self._limit = limit
         self._limiter: Optional[anyio.CapacityLimiter] = None
-        self._submitting_flow_run_ids = set()
-        self._cancelling_flow_run_ids = set()
-        self._scheduled_task_scopes = set()
+        self._submitting_flow_run_ids: set[UUID] = set()
+        self._cancelling_flow_run_ids: set[UUID] = set()
+        self._scheduled_task_scopes: set[anyio.CancelScope] = set()
         self._worker_metadata_sent = False
 
+    @property
+    def client(self) -> PrefectClient:
+        if self._client is None:
+            raise RuntimeError(
+                "Worker has not been correctly initialized. Please use the worker class as an async context manager."
+            )
+        return self._client
+
+    @property
+    def work_pool(self) -> WorkPool:
+        if self._work_pool is None:
+            raise RuntimeError(
+                "Worker has not been correctly initialized. Please use the worker class as an async context manager."
+            )
+        return self._work_pool
+
+    @property
+    def limiter(self) -> anyio.CapacityLimiter:
+        if self._limiter is None:
+            raise RuntimeError(
+                "Worker has not been correctly initialized. Please use the worker class as an async context manager."
+            )
+        return self._limiter
+
     @classmethod
     def get_documentation_url(cls) -> str:
         return cls._documentation_url
@@ -510,7 +562,7 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
             return worker_registry.get(type)
 
     @staticmethod
-    def get_all_available_worker_types() -> List[str]:
+    def get_all_available_worker_types() -> list[str]:
         """
         Returns all worker types available in the local registry.
         """
@@ -790,7 +842,7 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
 
         should_get_worker_id = self._should_get_worker_id()
 
-        params = {
+        params: dict[str, Any] = {
             "work_pool_name": self._work_pool_name,
             "worker_name": self.name,
             "heartbeat_interval_seconds": self.heartbeat_interval_seconds,
@@ -852,7 +904,7 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
 
     async def _get_scheduled_flow_runs(
         self,
-    ) -> List["WorkerFlowRunResponse"]:
+    ) -> list["WorkerFlowRunResponse"]:
         """
         Retrieve scheduled flow runs from the work pool's queues.
         """
@@ -862,7 +914,7 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
         )
         try:
             scheduled_flow_runs = (
-                await self._client.get_scheduled_flow_runs_for_work_pool(
+                await self.client.get_scheduled_flow_runs_for_work_pool(
                     work_pool_name=self._work_pool_name,
                     scheduled_before=scheduled_before,
                     work_queue_names=list(self._work_queues),
@@ -878,8 +930,8 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
             return []
 
     async def _submit_scheduled_flow_runs(
-        self, flow_run_response: List["WorkerFlowRunResponse"]
-    ) -> List["FlowRun"]:
+        self, flow_run_response: list["WorkerFlowRunResponse"]
+    ) -> list["FlowRun"]:
         """
         Takes a list of WorkerFlowRunResponses and submits the referenced flow runs
         for execution by the worker.
@@ -897,7 +949,7 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
                     self._limiter.acquire_on_behalf_of_nowait(flow_run.id)
             except anyio.WouldBlock:
                 self._logger.info(
-                    f"Flow run limit reached; {self._limiter.borrowed_tokens} flow runs"
+                    f"Flow run limit reached; {self.limiter.borrowed_tokens} flow runs"
                     " in progress."
                 )
                 break
@@ -921,6 +973,8 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
                         run_logger.warning(f"Failed to generate worker URL: {ve}")
 
                 self._submitting_flow_run_ids.add(flow_run.id)
+                if TYPE_CHECKING:
+                    assert self._runs_task_group is not None
                 self._runs_task_group.start_soon(
                     self._submit_run,
                     flow_run,
@@ -939,14 +993,9 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
         """
         run_logger = self.get_flow_run_logger(flow_run)
 
-        if flow_run.deployment_id:
-            assert self._client and self._client._started, (
-                "Client must be started to check flow run deployment."
-            )
-
         try:
-            await self._client.read_deployment(flow_run.deployment_id)
-        except ObjectNotFound:
+            await self.client.read_deployment(getattr(flow_run, "deployment_id"))
+        except (ObjectNotFound, AttributeError):
             self._logger.exception(
                 f"Deployment {flow_run.deployment_id} no longer exists. "
                 f"Flow run {flow_run.id} will not be submitted for"
@@ -964,13 +1013,15 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
         ready_to_submit = await self._propose_pending_state(flow_run)
         self._logger.debug(f"Ready to submit {flow_run.id}: {ready_to_submit}")
         if ready_to_submit:
+            if TYPE_CHECKING:
+                assert self._runs_task_group is not None
             readiness_result = await self._runs_task_group.start(
                 self._submit_run_and_capture_errors, flow_run
             )
 
             if readiness_result and not isinstance(readiness_result, Exception):
                 try:
-                    await self._client.update_flow_run(
+                    await self.client.update_flow_run(
                         flow_run_id=flow_run.id,
                         infrastructure_pid=str(readiness_result),
                     )
@@ -991,8 +1042,10 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
         self._submitting_flow_run_ids.remove(flow_run.id)
 
     async def _submit_run_and_capture_errors(
-        self, flow_run: "FlowRun", task_status: Optional[anyio.abc.TaskStatus] = None
-    ) -> Union[BaseWorkerResult, Exception]:
+        self,
+        flow_run: "FlowRun",
+        task_status: anyio.abc.TaskStatus[int | Exception] | None = None,
+    ) -> BaseWorkerResult | Exception:
         run_logger = self.get_flow_run_logger(flow_run)
 
         try:
@@ -1006,7 +1059,7 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
                 configuration=configuration,
             )
         except Exception as exc:
-            if not task_status._future.done():
+            if task_status and not getattr(task_status, "_future").done():
                 # This flow run was being submitted and did not start successfully
                 run_logger.exception(
                     f"Failed to submit flow run '{flow_run.id}' to infrastructure."
@@ -1025,7 +1078,7 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
         finally:
             self._release_limit_slot(flow_run.id)
 
-        if not task_status._future.done():
+        if task_status and not getattr(task_status, "_future").done():
             run_logger.error(
                 f"Infrastructure returned without reporting flow run '{flow_run.id}' "
                 "as started or raising an error. This behavior is not expected and "
@@ -1033,7 +1086,11 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
                 "flow run will not be marked as failed, but an issue may have occurred."
             )
             # Mark the task as started to prevent agent crash
-            task_status.started()
+            task_status.started(
+                RuntimeError(
+                    "Infrastructure returned without reporting flow run as started or raising an error."
+                )
+            )
 
         if result.status_code != 0:
             await self._propose_crashed_state(
@@ -1044,11 +1101,12 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
                 ),
             )
 
-        self._emit_flow_run_executed_event(result, configuration, submitted_event)
+        if submitted_event:
+            self._emit_flow_run_executed_event(result, configuration, submitted_event)
 
         return result
 
-    def _release_limit_slot(self, flow_run_id: str) -> None:
+    def _release_limit_slot(self, flow_run_id: UUID) -> None:
         """
         Frees up a slot taken by the given flow run id.
         """
@@ -1078,14 +1136,12 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
         flow_run: "FlowRun",
         deployment: Optional["DeploymentResponse"] = None,
     ) -> C:
-        deployment = (
-            deployment
-            if deployment
-            else await self._client.read_deployment(flow_run.deployment_id)
-        )
-        flow = await self._client.read_flow(flow_run.flow_id)
+        if not deployment and flow_run.deployment_id:
+            deployment = await self.client.read_deployment(flow_run.deployment_id)
 
-        deployment_vars = deployment.job_variables or {}
+        flow = await self.client.read_flow(flow_run.flow_id)
+
+        deployment_vars = getattr(deployment, "job_variables", {}) or {}
         flow_run_vars = flow_run.job_variables or {}
         job_variables = {**deployment_vars}
 
@@ -1095,22 +1151,37 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
         job_variables.update(flow_run_vars)
 
         configuration = await self.job_configuration.from_template_and_values(
-            base_job_template=self._work_pool.base_job_template,
+            base_job_template=self.work_pool.base_job_template,
             values=job_variables,
-            client=self._client,
-        )
-        configuration.prepare_for_flow_run(
-            flow_run=flow_run, deployment=deployment, flow=flow
+            client=self.client,
         )
+        try:
+            configuration.prepare_for_flow_run(
+                flow_run=flow_run,
+                deployment=deployment,
+                flow=flow,
+                work_pool=self.work_pool,
+                worker_name=self.name,
+            )
+        except TypeError:
+            warnings.warn(
+                "This worker is missing the `work_pool` and `worker_name` arguments "
+                "in its JobConfiguration.prepare_for_flow_run method. Please update "
+                "the worker's JobConfiguration  class to accept these arguments to "
+                "avoid this warning.",
+                category=PrefectDeprecationWarning,
+            )
+            # Handle older subclasses that don't accept work_pool and worker_name
+            configuration.prepare_for_flow_run(
+                flow_run=flow_run, deployment=deployment, flow=flow
+            )
         return configuration
 
     async def _propose_pending_state(self, flow_run: "FlowRun") -> bool:
         run_logger = self.get_flow_run_logger(flow_run)
         state = flow_run.state
         try:
-            state = await propose_state(
-                self._client, Pending(), flow_run_id=flow_run.id
-            )
+            state = await propose_state(self.client, Pending(), flow_run_id=flow_run.id)
         except Abort as exc:
             run_logger.info(
                 (
@@ -1141,7 +1212,7 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
         run_logger = self.get_flow_run_logger(flow_run)
         try:
             await propose_state(
-                self._client,
+                self.client,
                 await exception_to_failed_state(message="Submission failed.", exc=exc),
                 flow_run_id=flow_run.id,
             )
@@ -1159,7 +1230,7 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
         run_logger = self.get_flow_run_logger(flow_run)
         try:
             state = await propose_state(
-                self._client,
+                self.client,
                 Crashed(message=message),
                 flow_run_id=flow_run.id,
             )
@@ -1175,14 +1246,19 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
                 )
 
     async def _mark_flow_run_as_cancelled(
-        self, flow_run: "FlowRun", state_updates: Optional[dict] = None
+        self, flow_run: "FlowRun", state_updates: dict[str, Any] | None = None
     ) -> None:
         state_updates = state_updates or {}
         state_updates.setdefault("name", "Cancelled")
         state_updates.setdefault("type", StateType.CANCELLED)
-        state = flow_run.state.model_copy(update=state_updates)
 
-        await self._client.set_flow_run_state(flow_run.id, state, force=True)
+        if flow_run.state:
+            state = flow_run.state.model_copy(update=state_updates)
+        else:
+            # Unexpectedly when flow run does not have a state, create a new one
+            state = Cancelled(**state_updates)
+
+        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
@@ -1191,16 +1267,21 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
             60 * 10, self._cancelling_flow_run_ids.remove, flow_run.id
         )
 
-    async def _set_work_pool_template(self, work_pool, job_template):
+    async def _set_work_pool_template(
+        self, work_pool: "WorkPool", job_template: dict[str, Any]
+    ):
         """Updates the `base_job_template` for the worker's work pool server side."""
-        await self._client.update_work_pool(
+
+        await self.client.update_work_pool(
             work_pool_name=work_pool.name,
             work_pool=WorkPoolUpdate(
                 base_job_template=job_template,
             ),
         )
 
-    async def _schedule_task(self, __in_seconds: int, fn, *args, **kwargs):
+    async def _schedule_task(
+        self, __in_seconds: int, fn: Callable[..., Any], *args: Any, **kwargs: Any
+    ):
         """
         Schedule a background task to start after some time.
 
@@ -1208,8 +1289,12 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
 
         The function may be async or sync. Async functions will be awaited.
         """
+        if not self._runs_task_group:
+            raise RuntimeError(
+                "Worker has not been correctly initialized. Please use the worker class as an async context manager."
+            )
 
-        async def wrapper(task_status):
+        async def wrapper(task_status: anyio.abc.TaskStatus[Any]):
             # If we are shutting down, do not sleep; otherwise sleep until the scheduled
             # time or shutdown
             if self.is_setup:
@@ -1271,12 +1356,12 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
 
     def _event_related_resources(
         self,
-        configuration: Optional[BaseJobConfiguration] = None,
+        configuration: BaseJobConfiguration | None = None,
         include_self: bool = False,
-    ) -> List[RelatedResource]:
-        related = []
+    ) -> list[RelatedResource]:
+        related: list[RelatedResource] = []
         if configuration:
-            related += configuration._related_resources()
+            related += getattr(configuration, "_related_resources")()
 
         if self._work_pool:
             related.append(
@@ -1294,7 +1379,7 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
 
     def _emit_flow_run_submitted_event(
         self, configuration: BaseJobConfiguration
-    ) -> Event:
+    ) -> Event | None:
         return emit_event(
             event="prefect.worker.submitted-flow-run",
             resource=self._event_resource(),
@@ -1305,7 +1390,7 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
         self,
         result: BaseWorkerResult,
         configuration: BaseJobConfiguration,
-        submitted_event: Event,
+        submitted_event: Event | None = None,
     ):
         related = self._event_related_resources(configuration=configuration)
 
@@ -1321,7 +1406,7 @@ class BaseWorker(abc.ABC, Generic[C, V, R]):
             follows=submitted_event,
         )
 
-    async def _emit_worker_started_event(self) -> Event:
+    async def _emit_worker_started_event(self) -> Event | None:
         return emit_event(
             "prefect.worker.started",
             resource=self._event_resource(),