Repository: pfnet/pytorch-pfn-extras

Issue:
Custom function trigger
Implement a helper Trigger class to be triggered by a user-implemented function.

```
class CustomTrigger(trigger.Trigger):
    class __init__(self, fn: Callable[..., bool], args: Sequential[Any], kwargs: Mapping[Any], check_trigger: TriggerLike) -> None
        ...
```


Discussion:

Resolving patch:
diff --git a/pytorch_pfn_extras/training/triggers/__init__.py b/pytorch_pfn_extras/training/triggers/__init__.py
--- a/pytorch_pfn_extras/training/triggers/__init__.py
+++ b/pytorch_pfn_extras/training/triggers/__init__.py
@@ -2,6 +2,9 @@
 from pytorch_pfn_extras.training.triggers.early_stopping_trigger import (  # NOQA
     EarlyStoppingTrigger,
 )
+from pytorch_pfn_extras.training.triggers.function_trigger import (  # NOQA
+    FunctionTrigger,
+)
 from pytorch_pfn_extras.training.triggers.interval_trigger import (  # NOQA
     IntervalTrigger,
 )
diff --git a/pytorch_pfn_extras/training/triggers/function_trigger.py b/pytorch_pfn_extras/training/triggers/function_trigger.py
new file mode 100644
--- /dev/null
+++ b/pytorch_pfn_extras/training/triggers/function_trigger.py
@@ -0,0 +1,54 @@
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Mapping,
+    Optional,
+    Sequence,
+)
+
+from pytorch_pfn_extras.training import trigger as trigger_module
+from pytorch_pfn_extras.training._manager_protocol import (
+    ExtensionsManagerProtocol,
+)
+
+if TYPE_CHECKING:
+    from pytorch_pfn_extras.training._trigger_util import TriggerLike
+
+
+class FunctionTrigger(trigger_module.Trigger):
+    def __init__(
+        self,
+        fn: Callable[..., bool],
+        args: Optional[Sequence[Any]] = None,
+        kwargs: Optional[Mapping[str, Any]] = None,
+        trigger: "TriggerLike" = (1, "iteration"),
+    ) -> None:
+        self._fn = fn
+        self._args = args or []
+        self._kwargs = kwargs or {}
+        self._interval_trigger = trigger_module.get_trigger(trigger)
+
+    def __call__(self, manager: ExtensionsManagerProtocol) -> bool:
+        if not self._interval_trigger(manager):
+            return False
+
+        return self._fn(*self._args, **self._kwargs)
+
+    def state_dict(self) -> Dict[str, Any]:
+        state = {
+            "interval_trigger": self._interval_trigger.state_dict(),
+        }
+        return state
+
+    def load_state_dict(self, to_load: Dict[str, Any]) -> None:
+        self._interval_trigger.load_state_dict(to_load["interval_trigger"])
+
+    def may_fire(self, iteration: int, epoch_len: int) -> bool:
+        if self._interval_trigger.may_fire(
+            iteration=iteration, epoch_len=epoch_len
+        ):
+            return self._fn(*self._args, **self._kwargs)
+        else:
+            return False


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/training_tests/triggers_tests/test_function_trigger.py b/tests/pytorch_pfn_extras_tests/training_tests/triggers_tests/test_function_trigger.py
new file mode 100644
--- /dev/null
+++ b/tests/pytorch_pfn_extras_tests/training_tests/triggers_tests/test_function_trigger.py
@@ -0,0 +1,53 @@
+from unittest.mock import MagicMock
+
+import pytest
+from pytorch_pfn_extras.training import ExtensionsManager
+from pytorch_pfn_extras.training import trigger as trigger_module
+from pytorch_pfn_extras.training._trigger_util import TriggerLike
+from pytorch_pfn_extras.training.triggers import FunctionTrigger
+
+
+def test_function_is_called() -> None:
+    fn = MagicMock()
+    args = [MagicMock()]
+    kwargs = {"a": MagicMock()}
+    trigger = FunctionTrigger(
+        fn=fn, args=args, kwargs=kwargs, trigger=(1, "iteration")
+    )
+    fn.assert_not_called()
+    manager = ExtensionsManager({}, {}, 1, iters_per_epoch=10)
+    with manager.run_iteration():
+        pass
+    trigger(manager)
+    fn.assert_called_once_with(*args, **kwargs)
+
+
+def test_trigger_with_value() -> None:
+    value = {"value": False}
+    args = [value]
+    trigger = FunctionTrigger(
+        fn=lambda x: x["value"], args=args, trigger=(1, "iteration")
+    )
+    manager = ExtensionsManager({}, {}, 1, iters_per_epoch=10)
+    with manager.run_iteration():
+        pass
+    assert not trigger(manager)
+    value["value"] = True
+    assert trigger(manager)
+
+
+@pytest.mark.parametrize(
+    "trigger, iters_per_epoch",
+    [((1, "iteration"), 10), ((1, "epoch"), 20), ((0.123, "epoch"), 17)],
+)
+def test_with_interval_trigger(
+    trigger: TriggerLike, iters_per_epoch: int
+) -> None:
+    trigger = trigger_module.get_trigger(trigger)
+    manager = ExtensionsManager({}, [], 10, iters_per_epoch=iters_per_epoch)
+    function_trigger = FunctionTrigger(fn=lambda: True, trigger=trigger)
+
+    while not manager.stop_trigger:
+        with manager.run_iteration():
+            pass
+        assert trigger(manager) == function_trigger(manager)