diff --git a/continuous_integration/environment-3.10.yaml b/continuous_integration/environment-3.10.yaml index 7279c6b84d8..d92f3fb902c 100644 --- a/continuous_integration/environment-3.10.yaml +++ b/continuous_integration/environment-3.10.yaml @@ -39,7 +39,7 @@ dependencies: - toolz - tornado - zict # overridden by git tip below - - zstandard + - zstandard >=0.9.0 - pip: - git+https://github.com/dask/dask - git+https://github.com/dask/s3fs diff --git a/continuous_integration/environment-3.8.yaml b/continuous_integration/environment-3.8.yaml index 0dbe6377f1b..7268d049531 100644 --- a/continuous_integration/environment-3.8.yaml +++ b/continuous_integration/environment-3.8.yaml @@ -41,7 +41,7 @@ dependencies: - toolz - tornado - zict - - zstandard + - zstandard >=0.9.0 - pip: - git+https://github.com/dask/dask - git+https://github.com/jcrist/crick # Only tested here diff --git a/continuous_integration/environment-3.9.yaml b/continuous_integration/environment-3.9.yaml index a79d010b278..7eabb56ac03 100644 --- a/continuous_integration/environment-3.9.yaml +++ b/continuous_integration/environment-3.9.yaml @@ -19,7 +19,7 @@ dependencies: - ipywidgets - jinja2 - locket >=1.0 - - lz4 # Only tested here + - lz4 >=0.23.1 # Only tested here - msgpack-python - netcdf4 - paramiko @@ -32,7 +32,7 @@ dependencies: - pytest-repeat - pytest-rerunfailures - pytest-timeout - - python-snappy # Only tested here + - python-snappy >=0.5.3 # Only tested here - pytorch # Only tested here - requests - s3fs @@ -44,7 +44,7 @@ dependencies: - torchvision # Only tested here - tornado - zict - - zstandard + - zstandard >=0.9.0 - pip: - git+https://github.com/dask/dask - keras diff --git a/distributed/protocol/compression.py b/distributed/protocol/compression.py index 72172e90dde..33e504137ff 100644 --- a/distributed/protocol/compression.py +++ b/distributed/protocol/compression.py @@ -11,6 +11,7 @@ from contextlib import suppress from typing import Literal +from packaging.version import parse as parse_version from tlz import identity import dask @@ -39,56 +40,36 @@ with suppress(ImportError): import snappy - def _fixed_snappy_decompress(data): - # snappy.decompress() doesn't accept memoryviews - if isinstance(data, (memoryview, bytearray)): - data = bytes(data) - return snappy.decompress(data) + # In python-snappy 0.5.3, support for the Python Buffer Protocol was added. + # This is needed to handle other objects (like `memoryview`s) without + # copying to `bytes` first. + # + # Note: `snappy.__version__` doesn't exist in a release yet. + # So do a little test that will fail if snappy is not 0.5.3 or later. + try: + snappy.compress(memoryview(b"")) + except TypeError: + raise ImportError("Need snappy >= 0.5.3") compressions["snappy"] = { "compress": snappy.compress, - "decompress": _fixed_snappy_decompress, + "decompress": snappy.decompress, } default_compression = "snappy" with suppress(ImportError): import lz4 - try: - # try using the new lz4 API - import lz4.block - - lz4_compress = lz4.block.compress - lz4_decompress = lz4.block.decompress - except ImportError: - # fall back to old one - lz4_compress = lz4.LZ4_compress - lz4_decompress = lz4.LZ4_uncompress - - # helper to bypass missing memoryview support in current lz4 - # (fixed in later versions) - - def _fixed_lz4_compress(data): - try: - return lz4_compress(data) - except TypeError: - if isinstance(data, (memoryview, bytearray)): - return lz4_compress(bytes(data)) - else: - raise - - def _fixed_lz4_decompress(data): - try: - return lz4_decompress(data) - except (ValueError, TypeError): - if isinstance(data, (memoryview, bytearray)): - return lz4_decompress(bytes(data)) - else: - raise + # Required to use `lz4.block` APIs and Python Buffer Protocol support. + if parse_version(lz4.__version__) < parse_version("0.23.1"): + raise ImportError("Need lz4 >= 0.23.1") + + from lz4.block import compress as lz4_compress + from lz4.block import decompress as lz4_decompress compressions["lz4"] = { - "compress": _fixed_lz4_compress, - "decompress": _fixed_lz4_decompress, + "compress": lz4_compress, + "decompress": lz4_decompress, } default_compression = "lz4" @@ -96,6 +77,10 @@ def _fixed_lz4_decompress(data): with suppress(ImportError): import zstandard + # Required for Python Buffer Protocol support. + if parse_version(zstandard.__version__) < parse_version("0.9.0"): + raise ImportError("Need zstandard >= 0.9.0") + zstd_compressor = zstandard.ZstdCompressor( level=dask.config.get("distributed.comm.zstd.level"), threads=dask.config.get("distributed.comm.zstd.threads"),