Callback collection through entry points

docs/source/extensions/callbacks.rst

@@ -173,6 +173,9 @@ Lightning would not be able to disambiguate the state for these two callbacks, a
 by default only defines the class name as the key, e.g., here ``Counter``.
 
 
+-----------
+
+
 Best Practices
 --------------
 The following are best practices when using/designing callbacks.
@@ -183,6 +186,53 @@ The following are best practices when using/designing callbacks.
 4. Directly calling methods (eg. `on_validation_end`) is strongly discouraged.
 5. Whenever possible, your callbacks should not depend on the order in which they are executed.
 
+
+-----------
+
+
+Entry Points
+------------
+
+Lightning supports registering Trainer callbacks directly through
+`Entry Points <https://setuptools.pypa.io/en/latest/userguide/entry_point.html>`_. Entry points allow an arbitrary
+package to *install* callbacks that the Lightning Trainer can automatically use, without you having to add them
+to the Trainer manually. This is useful in production environments where it is common to provide specialized monitoring
+and logging callbacks globally for every application.
+
+Here is a callback factory function that returns two special callbacks:
+
+.. code-block:: python
+    :caption: factories.py
+
+    def my_callbacks_factory():
+        return [MyCallback1(), MyCallback2()]
+
+If we make this `factories.py` file into an installable package, we can define an **entry point** for this factory function.
+Here is a minimal example of the `setup.py` file for the package `my-package`:
+
+.. code-block:: python
+    :caption: setup.py
+
+    from setuptools import setup
+
+    setup(
+        name="my-package",
+        version="0.0.1",
+        entry_points={
+            "pytorch_lightning.callbacks_factory": [
+                # The format here must be [any name]=[module path]:[function name]
+                "monitor_callbacks=factories:my_callbacks_factory"
+            ]
+        },
+    )
+
+The group name for the entry points is ``pytorch_lightning.callbacks_factory`` and it contains a list of strings that
+specify where to find the function within the package.
+
+Now, if you `pip install -e .` this package, it will register the ``my_callbacks_factory`` function and Lightning
+will automatically call it to collect the callbacks whenever you run the Trainer!
+
+
 -----------
 
 .. _callback_hooks:

dummy/factories.py

@@ -0,0 +1,13 @@
+from pytorch_lightning.callbacks import Callback
+
+
+class MyPrintingCallback(Callback):
+    def on_train_start(self, trainer, pl_module):
+        print("Training is starting")
+
+    def on_train_end(self, trainer, pl_module):
+        print("Training is ending")
+
+
+def my_callbacks_factory():
+    return [MyPrintingCallback()]

dummy/setup.py

@@ -0,0 +1,12 @@
+from setuptools import setup
+
+setup(
+    name="my-package",
+    version="0.0.1",
+    # entry_points={
+    #     "pytorch_lightning.callbacks_factory": [
+    #         # The format here must be [any name]=[module path]:[function name]
+    #         "monitor_callbacks=factories:my_callbacks_factory"
+    #     ]
+    # },
+)

pytorch_lightning/trainer/connectors/callback_connector.py

@@ -11,10 +11,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import importlib.metadata
+import logging
 import os
 from datetime import timedelta
 from typing import Dict, List, Optional, Sequence, Union
 
+import pkg_resources
+
 from pytorch_lightning.callbacks import (
     Callback,
     GradientAccumulationScheduler,
@@ -28,8 +32,11 @@
 from pytorch_lightning.callbacks.timer import Timer
 from pytorch_lightning.utilities.enums import ModelSummaryMode
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.imports import _PYTHON_GREATER_EQUAL_3_8_0
 from pytorch_lightning.utilities.rank_zero import rank_zero_deprecation, rank_zero_info
 
+_log = logging.getLogger(__name__)
+
 
 class CallbackConnector:
     def __init__(self, trainer):
@@ -91,6 +98,8 @@ def on_trainer_init(
         if self.trainer.state._fault_tolerant_mode.is_enabled:
             self._configure_fault_tolerance_callbacks()
 
+        self._configure_external_callbacks()
+
         # push all model checkpoint callbacks to the end
         # it is important that these are the last callbacks to run
         self.trainer.callbacks = self._reorder_callbacks(self.trainer.callbacks)
@@ -232,6 +241,35 @@ def _configure_fault_tolerance_callbacks(self):
         # don't use `log_dir` to minimize the chances of failure
         self.trainer.callbacks.append(_FaultToleranceCheckpoint(dirpath=self.trainer.default_root_dir))
 
+    def _configure_external_callbacks(self) -> None:
+        """Add external callbacks registered through entry points.
+
+        The entry points are expected to be functions returning a list of callbacks, which will be added to the Trainer
+        callback list.
+        """
+        if _PYTHON_GREATER_EQUAL_3_8_0:
+            factories = importlib.metadata.entry_points().get("pytorch_lightning.callbacks_factory", ())
+        else:
+            factories = pkg_resources.iter_entry_points("pytorch_lightning.callbacks_factory")
+
+        for factory in factories:
+            callbacks_list: List[Callback] = factory.load()()
+            if not isinstance(callbacks_list, list):
+                raise TypeError(
+                    f"The entry point '{factory.name}' returned a {type(callbacks_list)} but is expected to return"
+                    f" a list of `pytorch_lightning.callbacks.Callback`."
+                )
+            if not all(isinstance(cb, Callback) for cb in callbacks_list):
+                raise TypeError(
+                    f"The entry point '{factory.name}' is expected to return a list of callbacks, but at least one"
+                    " callack was not an instance of `pytorch_lightning.callbacks.Callback`."
+                )
+            _log.info(
+                f"Adding {len(callbacks_list)} callbacks from entry point '{factory.name}':"
+                f" {', '.join(type(cb).__name__ for cb in callbacks_list)}"
+            )
+            self.trainer.callbacks.extend(callbacks_list)
+
     def _attach_model_logging_functions(self):
         lightning_module = self.trainer.lightning_module
         for callback in self.trainer.callbacks:

pytorch_lightning/utilities/imports.py

@@ -64,6 +64,8 @@ def _compare_version(package: str, op: Callable, version: str, use_base_version:
 
     >>> _compare_version("torch", operator.ge, "0.1")
     True
+    >>> _compare_version("does_not_exist", operator.ge, "0.0")
+    False
     """
     try:
         pkg = importlib.import_module(package)
@@ -85,7 +87,7 @@ def _compare_version(package: str, op: Callable, version: str, use_base_version:
 
 _IS_WINDOWS = platform.system() == "Windows"
 _IS_INTERACTIVE = hasattr(sys, "ps1")  # https://stackoverflow.com/a/64523765
-_PYTHON_GREATER_EQUAL_3_8_0 = _compare_version("python", operator.ge, "3.8.0")
+_PYTHON_GREATER_EQUAL_3_8_0 = Version(platform.python_version()) >= Version("3.8.0")
 _TORCH_GREATER_EQUAL_1_8_1 = _compare_version("torch", operator.ge, "1.8.1")
 _TORCH_GREATER_EQUAL_1_9 = _compare_version("torch", operator.ge, "1.9.0")
 _TORCH_GREATER_EQUAL_1_9_1 = _compare_version("torch", operator.ge, "1.9.1")

tests/trainer/connectors/test_callback_connector.py

@@ -12,9 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
+from unittest.mock import Mock
 
+import pytest
 import torch
 
+import pytorch_lightning
 from pytorch_lightning import Callback, LightningModule, Trainer
 from pytorch_lightning.callbacks import (
     EarlyStopping,
@@ -26,6 +29,7 @@
     TQDMProgressBar,
 )
 from pytorch_lightning.trainer.connectors.callback_connector import CallbackConnector
+from pytorch_lightning.utilities.imports import _PYTHON_GREATER_EQUAL_3_8_0
 from tests.helpers import BoringModel
 
 
@@ -214,3 +218,71 @@ def test_attach_model_callbacks_override_info(caplog):
         cb_connector._attach_model_callbacks()
 
     assert "existing callbacks passed to Trainer: EarlyStopping, LearningRateMonitor" in caplog.text
+
+
+class ExternalCallback(Callback):
+    """A callback in another library that gets registered through entry points."""
+
+    pass
+
+
+def test_configure_external_callbacks_raises(monkeypatch):
+    """Test that the connector validates the return type of Callback factories registered through entry points."""
+
+    def factory_incorrect_return_type():
+        return "invalid"
+
+    def factory_incorrect_element_type():
+        return [ExternalCallback(), "invalid"]
+
+    _make_entry_point_query_mock(monkeypatch, factory_incorrect_return_type)
+    with pytest.raises(TypeError, match="The entry point 'mocked' returned a <class 'str'> but is expected to return"):
+        Trainer()
+
+    _make_entry_point_query_mock(monkeypatch, factory_incorrect_element_type)
+    with pytest.raises(TypeError, match="at least one callack was not an instance of"):
+        Trainer()
+
+
+def test_configure_external_callbacks(monkeypatch):
+    """Test that the connector collects Callback instances from factories registered through entry points."""
+
+    def factory_no_callback():
+        return []
+
+    def factory_one_callback():
+        return [ExternalCallback()]
+
+    def factory_multiple_callbacks():
+        return [ExternalCallback(), ExternalCallback()]
+
+    _make_entry_point_query_mock(monkeypatch, factory_no_callback)
+    trainer = Trainer(enable_checkpointing=False, enable_progress_bar=False, enable_model_summary=False)
+    assert trainer.callbacks == [trainer.accumulation_scheduler]  # this scheduler callback gets added by default
+
+    _make_entry_point_query_mock(monkeypatch, factory_one_callback)
+    trainer = Trainer(enable_checkpointing=False, enable_progress_bar=False, enable_model_summary=False)
+    assert isinstance(trainer.callbacks[1], ExternalCallback)
+
+    _make_entry_point_query_mock(monkeypatch, factory_multiple_callbacks)
+    trainer = Trainer(enable_checkpointing=False, enable_progress_bar=False, enable_model_summary=False)
+    assert isinstance(trainer.callbacks[1], ExternalCallback)
+    assert isinstance(trainer.callbacks[2], ExternalCallback)
+
+
+def _make_entry_point_query_mock(monkeypatch, callback_factory):
+    query_mock = Mock()
+    entry_point = Mock()
+    entry_point.name = "mocked"
+    entry_point.load.return_value = callback_factory
+    if _PYTHON_GREATER_EQUAL_3_8_0:
+        query_mock().get.return_value = [entry_point]
+        monkeypatch.setattr(
+            pytorch_lightning.trainer.connectors.callback_connector.importlib.metadata, "entry_points", query_mock
+        )
+        return query_mock.get
+    else:
+        query_mock.return_value = [entry_point]
+        monkeypatch.setattr(
+            pytorch_lightning.trainer.connectors.callback_connector.pkg_resources, "iter_entry_points", query_mock
+        )