Feat: read entrypoints in dask.sizeof (#7688)

agoose77 · web-flow · commit 0a7c41dd1ffd · 2022-06-29T11:28:43.000-04:00
Entry-points are only used when there is no dispatch for a given type which eliminates the cost for applications which do not need entry-point based implementations.
diff --git a/dask/sizeof.py b/dask/sizeof.py
@@ -1,4 +1,6 @@
+import importlib.metadata
 import itertools
+import logging
 import random
 import sys
 from array import array
@@ -7,6 +9,8 @@
 
 sizeof = Dispatch(name="sizeof")
 
+logger = logging.getLogger(__name__)
+
 
 @sizeof.register(object)
 def sizeof_default(o):
@@ -208,3 +212,23 @@ def sizeof_pyarrow_table(table):
     @sizeof.register(pa.ChunkedArray)
     def sizeof_pyarrow_chunked_array(data):
         return int(_get_col_size(data)) + 1000
+
+
+def _register_entry_point_plugins():
+    """Register sizeof implementations exposed by the entry_point mechanism."""
+    if sys.version_info >= (3, 10):
+        sizeof_entry_points = importlib.metadata.entry_points(group="dask.sizeof")
+    else:
+        sizeof_entry_points = importlib.metadata.entry_points().get("dask.sizeof", [])
+
+    for entry_point in sizeof_entry_points:
+        registrar = entry_point.load()
+        try:
+            registrar(sizeof)
+        except Exception:
+            logger.exception(
+                f"Failed to register sizeof entry point {entry_point.name}"
+            )
+
+
+_register_entry_point_plugins()
diff --git a/dask/tests/test_sizeof.py b/dask/tests/test_sizeof.py
@@ -1,8 +1,10 @@
+import os
 import sys
 from array import array
 
 import pytest
 
+from dask.multiprocessing import get_context
 from dask.sizeof import sizeof
 from dask.utils import funcname
 
@@ -147,3 +149,47 @@ def test_dict():
     d = {i: x for i in range(100)}
     assert sizeof(d) > x.nbytes * 100
     assert isinstance(sizeof(d), int)
+
+
+def _get_sizeof_on_path(path, size):
+    sys.path.append(os.fsdecode(path))
+
+    # Dask will have already called _register_entry_point_plugins
+    # before we can modify sys.path, so we re-register here.
+    import dask.sizeof
+
+    dask.sizeof._register_entry_point_plugins()
+
+    import class_impl
+
+    cls = class_impl.Impl(size)
+    return sizeof(cls)
+
+
+def test_register_backend_entrypoint(tmp_path):
+    # Create special sizeof implementation for a dummy class
+    (tmp_path / "impl_sizeof.py").write_bytes(
+        b"def sizeof_plugin(sizeof):\n"
+        b'    print("REG")\n'
+        b'    @sizeof.register_lazy("class_impl")\n'
+        b"    def register_impl():\n"
+        b"        import class_impl\n"
+        b"        @sizeof.register(class_impl.Impl)\n"
+        b"        def sizeof_impl(obj):\n"
+        b"            return obj.size \n"
+    )
+    # Define dummy class that possesses a size attribute
+    (tmp_path / "class_impl.py").write_bytes(
+        b"class Impl:\n    def __init__(self, size):\n        self.size = size"
+    )
+    dist_info = tmp_path / "impl_sizeof-0.0.0.dist-info"
+    dist_info.mkdir()
+    (dist_info / "entry_points.txt").write_bytes(
+        b"[dask.sizeof]\nimpl = impl_sizeof:sizeof_plugin\n"
+    )
+
+    with get_context().Pool(1) as pool:
+        assert (
+            pool.apply(_get_sizeof_on_path, args=(tmp_path, 3_14159265)) == 3_14159265
+        )
+    pool.join()
diff --git a/docs/source/how-to/extend-sizeof.rst b/docs/source/how-to/extend-sizeof.rst
@@ -0,0 +1,36 @@
+Extend `sizeof`
+============
+
+When Dask needs to compute the size of an object in bytes, e.g. to determine which objects to spill to disk, it uses the ``dask.sizeof.sizeof`` registration mechanism. Users who need to define a ``sizeof`` implementation for their own objects can use ``sizeof.register``:
+
+.. code-block:: python
+
+   >>> import numpy as np
+   >>> from dask.sizeof import sizeof
+   >>> @sizeof.register(np.ndarray)
+   >>> def sizeof_numpy_like(array):
+   ...     return array.nbytes
+
+This code can be executed in order to register the implementation with Dask by placing it in one of the library's modules e.g. ``__init__.py``. However, this introduces a maintenance burden on the developers of these libraries, and must be manually imported on all workers in the event that these libraries do not accept the patch. 
+
+Therefore, Dask also exposes an `entrypoint <https://packaging.python.org/specifications/entry-points/>`_ under the group ``dask.sizeof`` to enable third-party libraries to develop and maintain these ``sizeof`` implementations. 
+
+For a fictitious library ``numpy_sizeof_dask.py``, the necessary ``setup.cfg`` configuration would be as follows:
+
+.. code-block:: ini
+
+   [options.entry_points]
+   dask.sizeof = 
+      numpy = numpy_sizeof_dask:sizeof_plugin
+
+whilst ``numpy_sizeof_dask.py`` would contain
+
+.. code-block:: python
+
+   >>> import numpy as np
+   >>> def sizeof_plugin(sizeof):
+   ...    @sizeof.register(np.ndarray)
+   ...    def sizeof_numpy_like(array):
+   ...        return array.nbytes 
+
+Upon the first import of `dask.sizeof`, Dask calls the entrypoint (``sizeof_plugin``) with the ``dask.sizeof.sizeof`` object, which can then be used to register a sizeof implementation.
