Futures unpickled on Client/Worker are attached to the wrong Client

[crusaderky]

Reopen and rescope of #2336

When a Future is wrapped into a dask collection and then passed as data through client.submit or similar, it is unpickled on the Worker and attached to the first available Client. This does not necessarily match the Client that the user desires to associate it to.

from distributed import Client

client = Client('localhost:8786')
df = pd.DataFrame(data=[[1]], columns=['a'])
df = dd.from_pandas(df, npartitions=1)
df = df.persist()

def remote(df):
    # Crashes with 'Inputs contain futures that were created by another client'
    df.compute() 

future = client.submit(remote, df)

The specific issue is caused by the scheduler being defined by hostname on one side and by IP on another - which has already been addressed by #3729 for Client.get_dataset() specifically.

There is a broader issue however that any given process may legitimately have more than one Client open, or a Client that does not match the default created by dask. For example, I personally use an asynchronous Client running on uvloop on my workers.

Proposed solution:

• In Future.__init__, set Future._client to
  -- the explicit client parameter, if not None,
  -- failing that, the Client from Client.as_current(),
  -- failing that, the default client,
  -- failing that, None.
  Currently:
  -- try using the Client from Client.as_current(),
  -- failing that, Future.client = the latest initialised Client instance, even if it was created with set_as_default=False

• Change Future.client to a property, which returns
  -- Future._client, if not None,
  -- failing that, the Client from Client.as_current(),
  -- failing that, the default client,
  -- failing that, raise ValueError.

• In Future.__getstate__, don't pickle the scheduler address - it causes just too much trouble!

• In Future.__setstate__, implement the same logic as in Future.__init__ (except that you'll never going to have an explicit client parameter).

The above snippet would have to be changed as follows in order to work:

def remote(df):
    distributed.get_client()  # Creates a Client with set_as_default=True)
    df.compute() 

or better, explicitly:

def remote(df):
    with distributed.get_client().as_current():
        df.compute() 

The notable change, which may possibly break stuff for some people, is that manually creating a Future won't hold a reference on cluster-side data anymore when there is no default client, which in turn may cause data to be deallocated. I don't think it should be a big problem as it will only affect people who explicitly created their Client with set_as_default=False, so they probably know what they're doing already.

[jcrist]

I think your proposal above makes sense. Just to clarify:

The notable change, which may possibly break stuff for some people, is that manually creating a Future won't hold a reference on cluster-side data anymore when there is no default client, which in turn may cause data to be deallocated.

This only happens for futures that are created without an explicit client attached (since no global client will be inferred). Things like client.submit, etc... all set the future's client explicitly, so the only operations that deserialize futures into a node without an explicit global client (or create a future by calling Future manually) will have this issue. Correct?

[crusaderky]

correct

[hhuuggoo]

def remote(df):
    with distributed.get_client().as_current():
        df.compute() 

I couldn't get this alternative to work? By the time remote is called, df is already unpickled, and therefore __setstate__ is already called. @crusaderky, how is this working for you?

FWIW - my plan is to use a worker plugin which modifies the worker address - calling resolve_address on get_worker().scheduler.addr. Which doesn't directly address the root of the problem, but should cause the clients to match.

__setstate__ on the Future calls get_client which defaults to resolve_address=True. This is the part which causes my unpickled futures to have a mismatched client from the Worker.

[gimperiale]

@hhuuggoo my proposed change goes beyond the 3 lines you pasted - it's explained in detail in my opening post. Those 3 lines alone won't work for the reason you already figured out.

[mrocklin]

Whoa, name change for @gimperiale . That's going to take some getting used to :)

[crusaderky]

@mrocklin nah I just accidentally commented from my work computer :P

[hhuuggoo]

@hhuuggoo my proposed change goes beyond the 3 lines you pasted - it's explained in detail in my opening post. Those 3 lines alone won't work for the reason you already figured out.

Oh ok apologies. I read that as this is the current work around, but now I understand. Thanks

FWIW on our side we're going to be mitigating this issue by resolving hostnames before passing scheduler addresses anywhere.

[mdwint]

I've been working around this issue in the same way as @hhuuggoo, by resolving the scheduler address before starting workers. This workaround has held up, but it has the obvious downside that workers are unable to reconnect if the scheduler address were to change (after Docker container restarts or deployments).

It would be great if this issue were fixed, and I want to offer my help, or at least get the conversation going again. Following @crusaderky's proposed solution, I created this fork, but I feel like I'm in way over my head already. I'm going to take some time to build a better understanding of the codebase.

@crusaderky Are you still interested in working on this? If anyone else wants to help out, let me know.

[hhuuggoo]

@mdwint we're ended up doing it with worker plugins and it's been pretty good.

worker.scheduler.addr = resolve_address(worker.scheduler.addr)

Feel free to reach out to me directly if more specific information would be useful.

[KaeganCasey]

Hello,
Has there been any other resolutions to this? I am using a managed dask service and trying to run a function in parallel that consumes a dask dataframe and I am getting "ValueError: Inputs contain futures that were created by another client." I don't think I have the ability to use worker plugins since it's a managed service. Is this something that should be theoretically possible?

For background I want to parallelize an operation on every column of a dask dataframe so I created a function that takes a dask dataframe (or series) as an input and then I am looping over every column of the dataframe and submitting it to the client. I can't share the exact code but it would be something like this:

# takes a dask series even though I call it ddf
def myfunc(ddf):
    # do something

for col in ddf.columns:
    client.submit(myfunc, ddf[col])

Any help would be much appreciated!! Thank you in advance!

[crusaderky]

@KaeganCasey passing an actual dask.dataframe through client.submit is not something most people should do.
Can you use https://docs.dask.org/en/stable/generated/dask.dataframe.Series.map_partitions.html#dask.dataframe.Series.map_partitions instead?

[KaeganCasey]

Thank you for your response @crusaderky!! I am still pretty new to dask, but it makes sense to me why that could be problematic. If I use map_partitions() would it look something like this?

# takes a dask series even though I call it ddf
def myfunc(ddf):
    # do something

for col in ddf.columns:
    ddf[col].map_partitions(myfunc)

If so, wouldn't this run sequentially and not in parallel? I think it would have to wait for each function to finish on one column before moving onto the next. Unless there is a way to use map_partitions() to map a function across columns that I am not aware of.

I really appreciate your response to my earlier post! Thank you for any help you can offer!

[crusaderky]

map_partitions will apply myfunc in parallel to each partition. myfunc will receive a pandas dataframe corresponding to the partition.
You need to save the output somewhere.
There's no need to break the columns apart - if you want to increase parallelism, just increase the number of partitions.

def myfunc(df):
    # df is a pandas.DataFrame containing all the columns of ddf, but only some of the rows.
    # Feel free to iterate sequentially on the columns.
    df2 = ...
    return df2

ddf = ddf.map_partitions(myfunc)