diff --git a/src/prefect/main.py b/src/prefect/main.py index afbba1334ce2..72f1490980b6 100644 --- a/src/prefect/main.py +++ b/src/prefect/main.py @@ -25,10 +25,14 @@ # Perform any forward-ref updates needed for Pydantic models import prefect.client.schemas -prefect.context.FlowRunContext.model_rebuild() -prefect.context.TaskRunContext.model_rebuild() -prefect.client.schemas.State.model_rebuild() -prefect.client.schemas.StateCreate.model_rebuild() +prefect.context.FlowRunContext.model_rebuild( + _types_namespace={"Flow": Flow, "BaseResult": BaseResult} +) +prefect.context.TaskRunContext.model_rebuild(_types_namespace={"Task": Task}) +prefect.client.schemas.State.model_rebuild(_types_namespace={"BaseResult": BaseResult}) +prefect.client.schemas.StateCreate.model_rebuild( + _types_namespace={"BaseResult": BaseResult} +) Transaction.model_rebuild() # Configure logging diff --git a/src/prefect/utilities/dispatch.py b/src/prefect/utilities/dispatch.py index 999f80319d1a..599a9afafda4 100644 --- a/src/prefect/utilities/dispatch.py +++ b/src/prefect/utilities/dispatch.py @@ -162,17 +162,22 @@ def register_type(cls: T) -> T: key = get_dispatch_key(cls) existing_value = registry.get(key) if existing_value is not None and id(existing_value) != id(cls): - # Get line numbers for debugging - file = inspect.getsourcefile(cls) - line_number = inspect.getsourcelines(cls)[1] - existing_file = inspect.getsourcefile(existing_value) - existing_line_number = inspect.getsourcelines(existing_value)[1] - warnings.warn( - f"Type {cls.__name__!r} at {file}:{line_number} has key {key!r} that " - f"matches existing registered type {existing_value.__name__!r} from " - f"{existing_file}:{existing_line_number}. The existing type will be " - "overridden." - ) + try: + # Get line numbers for debugging + file = inspect.getsourcefile(cls) + line_number = inspect.getsourcelines(cls)[1] + existing_file = inspect.getsourcefile(existing_value) + existing_line_number = inspect.getsourcelines(existing_value)[1] + warnings.warn( + f"Type {cls.__name__!r} at {file}:{line_number} has key {key!r} that " + f"matches existing registered type {existing_value.__name__!r} from " + f"{existing_file}:{existing_line_number}. The existing type will be " + "overridden." + ) + except OSError: + # If we can't get the source, another actor is loading this class via eval + # and we shouldn't update the registry + return cls # Add to the registry registry[key] = cls