Efficiently serialize zero strided NumPy arrays by jrbourbeau · Pull Request #3180 · dask/distributed

jrbourbeau · 2019-10-29T19:35:51Z

This PR updates how NumPy arrays with zero-strided axes are serialized in order to avoid transmitting the entire array where possible. Instead dimensions with a zero stride are collapsed, this collapsed array is serialized, and upon deserialization the smaller array is then broadcasted to the appropriate shape.

Taking the following as an example:

In [1]: import numpy as np

In [2]: a = np.arange(3)

In [3]: x = np.broadcast_to(a, (5, 3))

In [4]: x
Out[4]:
array([[0, 1, 2],
       [0, 1, 2],
       [0, 1, 2],
       [0, 1, 2],
       [0, 1, 2]])

Instead of serializing the entire x array (with x.shape = (5,3)), we would serialize the smaller [0, 1, 2] array. Then after deserializing the smaller array, we use np.broadcast_to to reproduce the expected x array with shape (5,3).

TomAugspurger · 2019-10-29T20:17:55Z

That's a neat trick.

Any reason you marked this as a WIP?

No idea what is going on with the CI:

Traceback (most recent call last):

  File "/home/travis/miniconda/envs/test-environment/lib/python3.5/site-packages/_pytest/config/__init__.py", line 451, in _importconftest

    return self._conftestpath2mod[key]

KeyError: PosixPath('/home/travis/build/dask/distributed/conftest.py')

During handling of the above exception, another exception occurred:

Traceback (most recent call last):

  File "/home/travis/miniconda/envs/test-environment/lib/python3.5/site-packages/_pytest/config/__init__.py", line 567, in import_plugin

    __import__(importspec)

  File "/home/travis/build/dask/distributed/distributed/__init__.py", line 3, in <module>

    from .actor import Actor, ActorFuture

  File "/home/travis/build/dask/distributed/distributed/actor.py", line 7, in <module>

    from .client import Future, default_client

  File "/home/travis/build/dask/distributed/distributed/client.py", line 28, in <module>

    from dask.optimization import SubgraphCallable

ImportError: cannot import name 'SubgraphCallable'

During handling of the above exception, another exception occurred:

jrbourbeau · 2019-10-29T21:31:43Z

Any reason you marked this as a WIP?

I wanted to run the full CI with these changes...which is currently experiencing some problems. I've not seen this problem before. Seems to only happen on the 3.5 build, I wonder if it's related to dropping Python 3.5 in dask

mrocklin · 2019-11-04T22:47:35Z

@jrbourbeau should this go in?

mrocklin · 2019-11-04T22:51:11Z

Or rather, in principle this seems fine to me. If you're comfortable merging, then please do.

mrocklin · 2019-11-05T17:38:17Z

Merging this later today if there are no further comments.

jrbourbeau · 2019-11-05T18:14:26Z

Thanks for giving this a look @mrocklin! The only additional change I think we should make here is handling sizeof(x) when x is a zero-strided array. As dask.sizeof.sizeof can be used outside the context of distributed, it probably makes more sense for me to update safe_sizeof here in distributed.

mrocklin · 2019-11-05T18:23:55Z

OK. I'm happy to defer to your judgement here.

mrocklin · 2019-11-06T19:36:55Z

distributed/sizeof.py

    try:
+        # Handle zero-strided NumPy arrays
+        if np is not None and isinstance(obj, np.ndarray) and 0 in obj.strides:
+            from .protocol.numpy import _zero_strided_slice


It seems off to include this special case in this function. Is this a case that we can fix upstream?

I agree this seems a little off. My main motivation for including this special case here instead of upstream in dask.sizeof was because this is tied specifically to how distributed handles zero-strided arrays. In theory, dask.sizeof could be used independently of distributed. Although I'm not sure how common that is (I suspect probably not very common).

We could move this upstream to dask. Does that seem more reasonable than what I have here?

To be honest I'm not sure what you mean by zero-strided numpy array. This is an array that is stored as, say a 2d array, but uses stride tricks to appear to be 3d or 4d or somethiing?

In that case does dask.sizeof report the correct size of the object in memory? I guess I'm not sure how what distributed does here affects things. If we have one of these objects in memory, how much RAM does it take up? I would think that this would be independent of which library is asking for the size.

mrocklin · 2019-11-12T22:24:07Z

Checking in again here @jrbourbeau . I know that you're busy these days so no rush. Should we remove the sizeof change here, merge the rest of this in, and change the sizeof implementation in dask/dask?

jrbourbeau · 2019-11-13T22:13:01Z

Yeah, following our conversation the other day I think that sounds good. Just pushed a commit to revert the sizeof.py changes (although I did leave the added safe_sizeof tests)

mrocklin · 2019-11-14T16:56:29Z

Thanks @jrbourbeau . This is in.

Use code written by jrbourbeau in a private thread. Add tests. Related to dask/distributed#3180

Related to dask/distributed#3180

jrbourbeau added 2 commits October 29, 2019 14:19

Update de/serialize for zero-strided arrays

2dd8846

Adds tests

a796277

Ensure serialized has same writeable flag

31e96a1

jrbourbeau changed the title ~~[WIP] Efficiently serialize zero strided NumPy arrays~~ Efficiently serialize zero strided NumPy arrays Oct 31, 2019

jrbourbeau added 2 commits November 5, 2019 15:10

Update safe_sizeof for zero strided arrays

e3c9a24

Adds basic tests for sizeof_safe

c423c31

mrocklin reviewed Nov 6, 2019

View reviewed changes

jrbourbeau added 2 commits November 13, 2019 16:05

Consolidate tests

d231ce1

Revert sizeof modifications

8239c94

mrocklin merged commit fd98e30 into dask:master Nov 14, 2019

jrbourbeau deleted the serialize-zero-strided branch November 14, 2019 17:07

dcherian mentioned this pull request Nov 20, 2019

cannot set WRITEABLE flag to True of this array #3252

Closed

jrbourbeau mentioned this pull request Nov 21, 2019

Fix NumPy writeable serialization bug #3253

Merged

Carreau added a commit to Carreau/dask that referenced this pull request Jun 23, 2020

Register sizeof for numpy zero-strided arrays

b6e22ef

Use code written by jrbourbeau in a private thread. Add tests. Related to dask/distributed#3180

Carreau mentioned this pull request Jun 23, 2020

Register sizeof for numpy zero-strided arrays dask/dask#6343

Merged

2 tasks

Carreau added a commit to Carreau/dask that referenced this pull request Jun 23, 2020

Register sizeof for numpy zero-strided arrays

d432260

Use code written by jrbourbeau in a private thread. Add tests. Related to dask/distributed#3180

TomAugspurger pushed a commit to dask/dask that referenced this pull request Jun 24, 2020

Register sizeof for numpy zero-strided arrays (#6343)

326a12a

Related to dask/distributed#3180

kumarprabhu1988 pushed a commit to kumarprabhu1988/dask that referenced this pull request Oct 29, 2020

Register sizeof for numpy zero-strided arrays (dask#6343)

2df4613

Related to dask/distributed#3180

Uh oh!

Conversation

jrbourbeau commented Oct 29, 2019

Uh oh!

TomAugspurger commented Oct 29, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jrbourbeau commented Oct 29, 2019

Uh oh!

mrocklin commented Nov 4, 2019

Uh oh!

mrocklin commented Nov 4, 2019

Uh oh!

mrocklin commented Nov 5, 2019

Uh oh!

jrbourbeau commented Nov 5, 2019

Uh oh!

mrocklin commented Nov 5, 2019

Uh oh!

mrocklin Nov 6, 2019

Choose a reason for hiding this comment

Uh oh!

jrbourbeau Nov 6, 2019

Choose a reason for hiding this comment

Uh oh!

mrocklin Nov 6, 2019

Choose a reason for hiding this comment

Uh oh!

mrocklin commented Nov 12, 2019

Uh oh!

jrbourbeau commented Nov 13, 2019

Uh oh!

mrocklin commented Nov 14, 2019

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

TomAugspurger commented Oct 29, 2019 •

edited

Loading