returnn 1.20240119.134529__tar.gz → 1.20240124.211635__tar.gz

{returnn-1.20240119.134529/returnn.egg-info → returnn-1.20240124.211635}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240119.134529
+Version: 1.20240124.211635
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20240124.211635/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240124.211635'
+long_version = '1.20240124.211635+git.780e425'

{returnn-1.20240119.134529 → returnn-1.20240124.211635}/returnn/datasets/hdf.py

@@ -438,7 +438,7 @@ class HDFDataset(CachedDataset):
         :param str key:
         :rtype: bool
         """
-        if self.get_data_dtype(key).startswith("int"):
+        if "int" in self.get_data_dtype(key):
             if key in self.num_outputs:
                 return self.num_outputs[key][1] <= 1
         return False

{returnn-1.20240119.134529 → returnn-1.20240124.211635}/returnn/frontend/__init__.py

@@ -49,4 +49,5 @@ from .types import *
 
 # Modules not in the main namespace but in sub namespaces.
 from . import audio
+from . import hooks
 from . import init

returnn-1.20240124.211635/returnn/frontend/hooks.py

@@ -0,0 +1,165 @@
+"""
+Hooks logic
+"""
+
+from __future__ import annotations
+from typing import Optional, Any, Dict, Callable, TypeVar, Tuple
+from types import MethodType
+from collections import OrderedDict
+import weakref
+
+
+__all__ = ["setup_post_hook_on_method", "RemovableHandle"]
+
+
+T = TypeVar("T")
+
+
+def setup_post_hook_on_method(
+    obj: Any,
+    attr: str,
+    hook: Callable[[T, Tuple[Any, ...], Dict[str, Any], Any], Optional[Any]],
+    *,
+    prepend: bool = False,
+) -> RemovableHandle:
+    """
+    :param obj: e.g. module
+    :param attr: attribute name to get the method, via ``getattr(obj, attr)``
+    :param hook:
+    :param prepend: if there are multiple hooks, this will be registered in front of all, otherwise at the end
+    """
+    method = MethodWithHooks.get(obj, attr)
+    return _register_new_hook(method, method.post_hooks, hook, prepend=prepend)
+
+
+class MethodWithHooks:
+    """
+    Method with hooks
+
+    (Not sure yet whether this should be exposed in the API...)
+    """
+
+    @classmethod
+    def get(cls, obj: Any, attr: str) -> MethodWithHooks:
+        """get existing or init new :class:`MethodWithHooks`"""
+        method = getattr(obj, attr)
+        if not isinstance(method, MethodWithHooks):
+            assert (
+                isinstance(method, MethodType)
+                and method.__self__ is obj
+                and method.__func__ is getattr(obj.__class__, attr)
+            ), (
+                f"{obj}.{attr} is {method}, {method.__self__} is not obj"
+                f" or {method.__func__} is not {getattr(obj.__class__, attr)}"
+            )  # not necessary, just what we expect here...
+            method = MethodWithHooks(obj, attr)
+            method.setup()
+        return method
+
+    def __init__(self, obj: Any, attr: str):
+        """
+        :param obj:
+        :param attr:
+        """
+        self.obj = obj
+        self.attr = attr
+        self.post_hooks: Dict[int, Callable[[T, Tuple[Any, ...], Dict[str, Any], Any], Optional[Any]]] = OrderedDict()
+
+    def setup(self):
+        """setup on object"""
+        setattr(self.obj, self.attr, self)
+        if self.attr == "__call__":
+            self.obj.__class__.__call__ = _CallWrapperClass(self.obj.__class__)
+
+    def uninstall(self):
+        """undo the setup"""
+        assert not self.post_hooks
+        delattr(self.obj, self.attr)
+        if self.attr == "__call__":
+            call_wrapper = self.obj.__class__.__call__
+            assert isinstance(call_wrapper, _CallWrapperClass)
+            self.obj.__class__.__call__ = call_wrapper.orig_call
+
+    def has_any_hook(self) -> bool:
+        """whether there is any hook"""
+        return bool(self.post_hooks)
+
+    def __call__(self, *args, **kwargs):
+        function = getattr(self.obj.__class__, self.attr)
+        if isinstance(function, _CallWrapperClass):
+            function = function.orig_call
+        result = function(self.obj, *args, **kwargs)
+        for hook in self.post_hooks.values():
+            result = hook(self.obj, args, kwargs, result)
+        return result
+
+
+class _CallWrapperClass:
+    """
+    We need special handling for ``obj(*args, **kwargs)`` - simply wrapping ``__call__`` is not enough.
+    An instance of this class here would overwrite the obj.__class__.__call__.
+    """
+
+    def __init__(self, cls):
+        self.cls = cls
+        self.orig_call = getattr(cls, "__call__")
+
+    def __call__(self, obj, *args, **kwargs):  # called on the class
+        call_func = getattr(obj, "__call__")
+        if isinstance(call_func, MethodWithHooks):  # obj has hooks installed
+            return call_func(*args, **kwargs)
+        # This is what we would expect currently, but the user might have dome some weird things,
+        # e.g. overwritten obj.__call__ or obj.__class__.__call__ in some other way,
+        # so maybe extend this then, if we encounter this.
+        # (Only do it when we really now the case properly, and add a test case for it.)
+        assert isinstance(call_func, MethodType)
+        assert call_func.__func__ is self
+        assert call_func.__self__ is obj
+        return self.orig_call(obj, *args, **kwargs)
+
+    def __get__(self, obj, objtype=None):
+        if obj is not None:
+            return MethodType(self, obj)
+        return self
+
+
+def _register_new_hook(
+    method: MethodWithHooks,
+    hooks_dict: Dict[int, Any],
+    hook: Callable[[T, Tuple[Any, ...], Dict[str, Any], Any], Optional[Any]],
+    *,
+    prepend: bool = False,
+) -> RemovableHandle:
+    """register"""
+    handle = RemovableHandle(method, hooks_dict)
+    hooks_dict[handle.id] = hook
+    if prepend:
+        assert isinstance(hooks_dict, OrderedDict)
+        hooks_dict.move_to_end(handle.id, last=False)
+    return handle
+
+
+class RemovableHandle:
+    """
+    A handle which provides the capability to remove a hook.
+
+    Code partly copied from PyTorch.
+    """
+
+    id: int
+    next_id: int = 0
+
+    def __init__(self, method: MethodWithHooks, hooks_dict: Dict[int, Any]):
+        self.method_ref = weakref.ref(method)
+        self.hooks_dict_ref = weakref.ref(hooks_dict)
+        self.id = RemovableHandle.next_id
+        RemovableHandle.next_id += 1
+
+    def remove(self) -> None:
+        """remove hook"""
+        method: MethodWithHooks = self.method_ref()
+        hooks_dict = self.hooks_dict_ref()
+        if hooks_dict is not None and self.id in hooks_dict:
+            del hooks_dict[self.id]
+            if method and not method.has_any_hook():  # we removed the last hook
+                method.uninstall()

{returnn-1.20240119.134529 → returnn-1.20240124.211635}/returnn/frontend/module.py

@@ -3,7 +3,7 @@ Base module class, :class:`Module`.
 """
 
 from __future__ import annotations
-from typing import Any, Optional, Sequence, List, Tuple, Union, Set, Iterator, Callable, TypeVar
+from typing import Any, Optional, Sequence, List, Tuple, Union, Set, Iterator, Callable, TypeVar, Dict
 from returnn.util.basic import OptionalNotImplementedError, RefIdEq
 from returnn.tensor import Tensor, Dim
 from .. import frontend as rf
@@ -247,6 +247,24 @@ class Module:
             param: rf.Parameter
             param.to(device=device, dtype=dtype)
 
+    def register_forward_hook(
+        self,
+        hook: Callable[[T, Tuple[Any, ...], Dict[str, Any], Any], Optional[Any]],
+        *,
+        prepend: bool = False,
+    ) -> rf.hooks.RemovableHandle:
+        """
+        Register forward hook. Very similar to PyTorch. (Code even partly copied from PyTorch.)
+        This will be called after __call__, following this logic (where module=self):
+
+            result = module(*args, **kwargs)
+            result = hook(args, kwargs, result)
+
+        :param hook:
+        :param prepend: if there are multiple hooks, this will be registered in front of all, otherwise at the end
+        """
+        return rf.hooks.setup_post_hook_on_method(self, "__call__", hook=hook, prepend=prepend)
+
 
 class Functional(Module):
     """

{returnn-1.20240119.134529 → returnn-1.20240124.211635}/returnn/util/multi_proc_non_daemonic_spawn.py

@@ -17,6 +17,8 @@ References:
 """
 
 from __future__ import annotations
+
+import time
 from typing import Optional, Any, Callable, Tuple
 import os
 import signal
@@ -66,7 +68,7 @@ class NonDaemonicSpawnProcess(SpawnProcess):
 
     def join(self, timeout=None):
         """join"""
-        if is_exiting():
+        if is_exiting() and self.is_alive():
             # Process._bootstrap (the subproc main logic) will do an early call to
             # multiprocessing.util._exit_function, which terminates all daemon procs,
             # and then joins all procs (daemon or non-daemon).
@@ -80,12 +82,21 @@ class NonDaemonicSpawnProcess(SpawnProcess):
             # but then it is too late, as we would hang here now in the join().
             # https://github.com/rwth-i6/returnn/issues/1497
             # https://github.com/python/cpython/issues/114220
-            try:
-                os.kill(self.ident, signal.SIGINT)
-            except ProcessLookupError:
-                pass
+            # First try SIGINT. See NonDaemonicSpawnProcess docstring.
+            # Then, if that fails, try the others.
+            signals = ["SIGINT"] * 3 + ["SIGTERM", "SIGKILL"]
+            for sig in signals:
+                try:
+                    print(f"Send signal {sig} to pid {self.ident}")
+                    os.kill(self.ident, getattr(signal, sig))
+                except ProcessLookupError:
+                    break
+                super().join(timeout=5)
+                if not self.is_alive():
+                    break
         super().join(timeout=timeout)
-        self._close_cleanup()
+        if not self.is_alive():
+            self._close_cleanup()
 
     def close(self):
         """close"""
@@ -155,15 +166,31 @@ class _AtExitCleanupProcess:
             return  # ignore
         if self.proc_pid is None:  # already cleaned
             return
-        # The proc might have been killed by some other code. That's ok.
-        try:
-            # Send SIGINT, not SIGTERM or SIGKILL. See NonDaemonicSpawnProcess docstring.
-            os.kill(self.proc_pid, signal.SIGINT)
-        except ProcessLookupError:
-            pass
-        else:
+        # First try SIGINT. See NonDaemonicSpawnProcess docstring.
+        # Then, if that fails, try the others.
+        signals = ["SIGINT"] * 3 + ["SIGTERM", "SIGKILL"]
+        for sig in signals:
+            # The proc might have been killed by some other code. That's ok.
             try:
-                os.waitpid(self.proc_pid, 0)
-            except ChildProcessError:
-                pass
+                print(f"Send signal {sig} to pid {self.proc_pid}")
+                os.kill(self.proc_pid, getattr(signal, sig))
+            except ProcessLookupError:
+                break
+            if self._waitpid():
+                break
+            time.sleep(1)
+            if self._waitpid():
+                break
+            time.sleep(4)
+            if self._waitpid():
+                break
         self.proc_pid = None
+
+    def _waitpid(self) -> bool:
+        try:
+            pid, _ = os.waitpid(self.proc_pid, os.WNOHANG)
+            if pid == self.proc_pid:
+                return True
+        except ChildProcessError:
+            return True
+        return False

{returnn-1.20240119.134529 → returnn-1.20240124.211635/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240119.134529
+Version: 1.20240124.211635
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

{returnn-1.20240119.134529 → returnn-1.20240124.211635}/returnn.egg-info/SOURCES.txt

@@ -172,6 +172,7 @@ returnn/frontend/dropout.py
 returnn/frontend/dtype.py
 returnn/frontend/gradient.py
 returnn/frontend/graph.py
+returnn/frontend/hooks.py
 returnn/frontend/init.py
 returnn/frontend/label_smoothing.py
 returnn/frontend/linear.py

{returnn-1.20240119.134529 → returnn-1.20240124.211635}/tests/test_torch_frontend.py

@@ -533,6 +533,54 @@ def test_Data_copy_tranpose_match_priority():
     numpy.testing.assert_equal(x_np, raw_np)
 
 
+def test_forward_hook():
+    from returnn.frontend import hooks
+    from types import MethodType, FunctionType
+
+    call_count = 0
+    hook_call_count = 0
+
+    class _Module(rf.Module):
+        def __call__(self, x: Tensor) -> Tensor:
+            nonlocal call_count
+            call_count += 1
+            return x * 0.5
+
+    def _hook(module, args, kwargs, result):
+        nonlocal hook_call_count
+        assert isinstance(module, _Module) and module is mod
+        assert isinstance(args, tuple) and len(args) == 1 and isinstance(args[0], Tensor)
+        assert isinstance(kwargs, dict) and len(kwargs) == 0
+        assert isinstance(result, Tensor)
+        hook_call_count += 1
+        return result * 2.0
+
+    mod = _Module()
+    time_dim = Dim(None, name="time")
+    x = Tensor("x", [time_dim], "float32", raw_tensor=torch.range(0.0, 5.0))
+    mod(x)
+    assert call_count == 1 and hook_call_count == 0
+
+    handle = mod.register_forward_hook(_hook)
+    assert isinstance(handle, hooks.RemovableHandle)
+    assert isinstance(mod.__call__, hooks.MethodWithHooks)
+    assert isinstance(mod.__class__.__call__, hooks._CallWrapperClass)
+    mod(x)
+    assert call_count == 2 and hook_call_count == 1
+
+    mod.__call__(x)
+    assert call_count == 3 and hook_call_count == 2
+
+    handle.remove()
+    assert isinstance(mod.__call__, MethodType)
+    assert isinstance(mod.__class__.__call__, FunctionType)
+    mod(x)
+    assert call_count == 4 and hook_call_count == 2
+
+    mod.__call__(x)
+    assert call_count == 5 and hook_call_count == 2
+
+
 if __name__ == "__main__":
     better_exchook.install()
     if len(sys.argv) <= 1:

returnn-1.20240119.134529/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240119.134529'
-long_version = '1.20240119.134529+git.43e68f4'