Pydantic serialization error when using DaskTaskRunner and prefect >3.1.10

[thomashenn]

Bug summary

When using the latest version of Prefect (3.1.12) and prefect-dask (0.3.2), running the flow fails with this error:

'MockValSer' object cannot be converted to 'SchemaSerializer'

However:

• running the same code without DaskTaskRunner works without issues.
• running the same code with DaskTaskRunner and an older version of prefect (3.1.10) (or 3.1.4 as suggested here by @rama-reddy) works without issues.

Example code:

from pathlib import Path

from prefect import flow, task
from prefect_dask.task_runners import DaskTaskRunner

@flow(task_runner=DaskTaskRunner(address="tcp://dask-scheduler-service:8786"), log_prints=True)
def demo_dask_flow():
    message = demo_dask_task.submit("Hello!")
    print(f"message: {message.result()}")
    
@task
def demo_dask_task(message: str):
    return message

if __name__ == "__main__":
    demo_dask_flow.from_source(
        source=str(Path(__file__).parent),
        entrypoint="demo_dask.py:demo_dask_flow",
    ).deploy(
        name="demo-dask-flow",
        work_pool_name="demo-work-pool",
    )

Ran like this:

python3.12 demo_dask.py
prefect deployment run 'demo-dask-flow/demo-dask-flow'

will result in this error:

Encountered exception during execution: TypeError("'MockValSer' object cannot be converted to 'SchemaSerializer'")
Traceback (most recent call last):
  File "/usr/local/lib/python3.12/dist-packages/prefect/flow_engine.py", line 737, in run_context
    yield self
  File "/usr/local/lib/python3.12/dist-packages/prefect/flow_engine.py", line 1336, in run_flow_sync
    engine.call_flow_fn()
  File "/usr/local/lib/python3.12/dist-packages/prefect/flow_engine.py", line 757, in call_flow_fn
    result = call_with_parameters(self.flow.fn, self.parameters)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/prefect/utilities/callables.py", line 208, in call_with_parameters
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/ml_worker/demo_dask.py", line 8, in demo_dask_flow
    message = demo_dask_task.submit("Hello!")
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/prefect/tasks.py", line 1201, in submit
    future = task_runner.submit(self, parameters, wait_for)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/prefect_dask/task_runners.py", line 353, in submit
    future = self._client.submit(
             ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/prefect_dask/client.py", line 35, in submit
    run_task_kwargs["context"] = serialize_context()
                                 ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/prefect/context.py", line 61, in serialize_context
    "flow_run_context": flow_run_context.serialize() if flow_run_context else {},
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/prefect/context.py", line 376, in serialize
    return self.model_dump(
           ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/pydantic/main.py", line 426, in model_dump
    return self.__pydantic_serializer__.to_python(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: 'MockValSer' object cannot be converted to 'SchemaSerializer'

Thank you

Version info

Version:             3.1.12
API version:         0.8.4
Python version:      3.12.3
Git commit:          e299e5a7
Built:               Thu, Jan 9, 2025 10:09 AM
OS/Arch:             linux/x86_64
Profile:             ephemeral
Server type:         server
Pydantic version:    2.10.5
Integrations:
  prefect-dask:      0.3.2

Additional context

I believe it is related to this issue #16422 but not limited to launching tasks from other tasks.