Annotate extract_serialize (for Cythonization)

[jakirkham]

This adds some unintrusive type annotations to extract_serialize so that when Cythonizing this function, we can get optimal performance out of it. If run only as pure Python, the performance remains the same.

Pure Python:

In [1]: from distributed.protocol.serialize import extract_serialize

In [2]: data = 1_000_000 * b"abc"
   ...: msg = 11 * [10 * [2 * [5 * [data]]]]

In [3]: %timeit extract_serialize(msg)
789 µs ± 6.68 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Cython:

In [1]: from distributed.protocol.serialize import extract_serialize

In [2]: data = 1_000_000 * b"abc"
   ...: msg = 11 * [10 * [2 * [5 * [data]]]]

In [3]: %timeit extract_serialize(msg)
481 µs ± 5.67 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

[jakirkham]

cc @mrocklin @quasiben

[mrocklin]

As a PR this seems fine to me.

At a larger scale it would obviously be grand if we could avoid having pyx files. My guess is that this will be hard to achieve while still getting optimal performance. I'd be very happy to learn otherwise though.

[jakirkham]

Right. I want to see how far we get with this strategy as...

1. It allows easier collaboration in the near term
2. Assuming this style is agreeable we can merge things in as needed
3. Things still work if someone doesn't have Cython (though they don't get the perf bump then)
4. Installation remains simple for the pure Python case

At least when I looked at the code this generated in C, it was pretty much as good as I would get even if I wrote it in Cython. My guess is pure Python mode in Cython has grown a lot since when Antoine and yourself last tried. So it seems worth exploring again for that reason alone (in addition to the benefits listed above).

If we still find we need .pyx files for a few cases, we can certainly explore that after we have squeezed out as much performance as we can get through this process. My guess is that will only affect a very small amount of code (if any), which will give us a lot of options in terms of how we handle that code (like it could become a small optional dependency for example).

[jrbourbeau]

Let's give this a shot, thanks @jakirkham

[jakirkham]

Thanks all! 😄

[mrocklin]

This approach excites me. I would be very happy if I could continue developing the scheduler in pure Python.