Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

The buffer of embedded numpy variables is deep-copied in client->scheduler comms #8608

Open
crusaderky opened this issue Apr 3, 2024 · 2 comments
Labels
p3 Affects a small number of users or is largely cosmetic

Comments

@crusaderky
Copy link
Collaborator

import distributed
import numpy as np

if __name__ == "__main__":
    with distributed.Client(n_workers=2, processes=False, protocol="tcp") as client:
        a, b = client.has_what()
        x = client.submit(np.random.random, 1024, key="x", workers=[a])
        y = client.submit(lambda x: x, x, key="y", workers=[b])
        y.result()

When the buffer reaches distributed.protocol.serialize.pickle_loads, buffers[0] is a bytes object.
This causes pickle_loads to deep-copy the buffer in order to honour the writeable flag of the original.

To verify, add at the top of pickle_loads:

print(header["writeable"], [ensure_memoryview(b).readonly for b in buffers])

What's causing me a migraine is:

  • if you replace
        x = client.submit(np.random.random, 1024, key="x", workers=[a])

with

        x = client.submit(lambda: np.random.random(1024), key="x", workers=[a])

then the numpy object is no longer deserialized by distributed.protocol.serialize.pickle_loads, but it's instead processed by distributed.protocol.numpy.deserialize_numpy_array, which receives a writeable buffer

  • if you replace the submit API with dask.array:
x = da.random.random(1024).persist(workers=[a])
y = x.map_blocks(lambda x: x).persist(workers=[b])
y.compute()

then we are using again distributed.protocol.serialize.pickle_loads, which receives a read-only buffer but this time the writeable flag is False so no deep copy happens.

@crusaderky
Copy link
Collaborator Author

Ok, found the difference. the deep-copy is NOT tripped on the transfer of the task output from a to b; it's the random seed that's sent from the client to scheduler. 🤦

@crusaderky crusaderky changed the title Buffer is deep-copied in weird edge case random seed is deep-copied in edge case Apr 3, 2024
@crusaderky crusaderky added p3 Affects a small number of users or is largely cosmetic and removed needs triage labels Apr 3, 2024
@crusaderky crusaderky changed the title random seed is deep-copied in edge case np.random.RandomState buffer is deep-copied in edge case Apr 3, 2024
@crusaderky crusaderky changed the title np.random.RandomState buffer is deep-copied in edge case The buffer of embedded numpy variables is deep-copied from client to scheduler Apr 3, 2024
@crusaderky
Copy link
Collaborator Author

crusaderky commented Apr 3, 2024

This issue applies to all embedded variables that are sent from client to scheduler, e.g.

x = client.submit(lambda x: x + 1, np.random.random(1024), key="x", workers=[a])

@crusaderky crusaderky changed the title The buffer of embedded numpy variables is deep-copied from client to scheduler The buffer of embedded numpy variables is deep-copied in client->scheduler comms Apr 4, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
p3 Affects a small number of users or is largely cosmetic
Projects
None yet
Development

No branches or pull requests

1 participant