RuntimeError: Task <Task pending name='Task-2051' coro=<FutureState.wait() running at /usr/local/lib/python3.8/site-packages/distributed/client.py:482> cb=[WaitIterator._done_callback()]> got Future <Future pending> attached to a different loop

with worker_client()

with worker_client()

def custom_task_function(
    fn: Optional[Callable] = None, **task_kwargs
) -> Union[FunctionTask, Callable[[Callable], FunctionTask]]:
    """
    Wrapped around prefect.task that sets up the Dask client.
    """

    def create_wrapper(fn):
        @wraps(fn)
        def wrapped_fn(*args, **kwargs):
            # do something else in local
            with worker_client():
                return fn(*args, **kwargs)

        return wrapped_fn

    # from @prefect.task
    if fn is None:
        return lambda fn: FunctionTask(
            fn=create_wrapper(fn),
            **task_kwargs,
        )
    return FunctionTask(fn=create_wrapper(fn), **task_kwargs)