returnn 1.20241005.114831__tar.gz → 1.20241010.171521__tar.gz

{returnn-1.20241005.114831 → returnn-1.20241010.171521}/PKG-INFO

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

returnn-1.20241010.171521/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20241010.171521'
+long_version = '1.20241010.171521+git.4c3c84a'

{returnn-1.20241005.114831 → returnn-1.20241010.171521}/returnn/datasets/distrib_files.py

@@ -11,6 +11,7 @@ import os
 import sys
 import numpy
 from returnn.log import log
+from returnn.util import better_exchook
 from returnn.util.basic import override_env_var, try_run
 from returnn.util.multi_proc_non_daemonic_spawn import NonDaemonicSpawnContext
 from returnn.config import SubProcCopyGlobalConfigPreInitFunc
@@ -573,6 +574,7 @@ def _worker_proc_loop(
     if sys.platform == "linux":
         with open("/proc/self/comm", "w") as f:
             f.write(f"CFD worker {epoch}")
+    better_exchook.setup_all()
 
     assert isinstance(epoch, int) and isinstance(buffer_size, int)
     assert isinstance(dataset_dict, dict)

{returnn-1.20241005.114831 → returnn-1.20241010.171521}/returnn/datasets/meta.py

@@ -1832,25 +1832,35 @@ class VariableDataset(Dataset):
     based on a user-provided function.
     """
 
-    def __init__(self, *, get_dataset, **kwargs):
+    def __init__(self, *, get_dataset, dataset_lru_cache_size: int = 1, **kwargs):
         """
         :param get_dataset: function (*, epoch: int, **_) -> Dict[str,Any], will be called for every sub-epoch.
-            It will cache the dict from the prev call, and if the dict is the same, it will not recreate the dataset.
+            It will cache the dataset(s) from the prev call (dataset_lru_cache_size),
+            and if the dict is the same of those, it will not recreate the dataset.
+        :param dataset_lru_cache_size
         """
+        from functools import lru_cache
+
         super().__init__(**kwargs)
         self._get_dataset = get_dataset
         self._dataset_dict: Optional[Dict[str, Any]] = None
         self._dataset: Optional[Dataset] = None
+        self._dataset_lru_cache_size = dataset_lru_cache_size
+        self._make_dataset = lru_cache(maxsize=self._dataset_lru_cache_size)(
+            lambda dataset_dict: init_dataset(dataset_dict, parent_dataset=self)
+        )
         self._load_dataset(epoch=1)
         self.num_inputs = self._dataset.num_inputs
         self.num_outputs = self._dataset.num_outputs
         self.labels = self._dataset.labels
 
     def _load_dataset(self, epoch: int):
-        dataset_dict = self._get_dataset(epoch=epoch)
-        if dataset_dict != self._dataset_dict:
-            self._dataset_dict = dataset_dict
-            self._dataset = init_dataset(dataset_dict, parent_dataset=self)
+        from returnn.util.basic import get_fwd_compat_kwargs, make_hashable
+
+        dataset_dict = self._get_dataset(self=self, epoch=epoch, **get_fwd_compat_kwargs())
+        assert isinstance(dataset_dict, dict)
+        dataset_dict = make_hashable(dataset_dict)
+        self._dataset = self._make_dataset(dataset_dict)
 
     def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
         """init seq order"""

{returnn-1.20241005.114831 → returnn-1.20241010.171521}/returnn/datasets/multi_proc.py

@@ -7,6 +7,7 @@ from typing import Optional, Any, Dict, List
 import sys
 import gc
 import multiprocessing as mp
+from returnn.util import better_exchook
 from returnn.util.basic import try_run
 from returnn.config import SubProcCopyGlobalConfigPreInitFunc
 from returnn.util.multi_proc_non_daemonic_spawn import NonDaemonicSpawnContext
@@ -168,6 +169,7 @@ class MultiProcDataset(CachedDataset2):
         if sys.platform == "linux":
             with open("/proc/self/comm", "w") as f:
                 f.write(f"MPD worker {worker_index}")
+        better_exchook.setup_all()
 
         dataset: Optional[Dataset] = None
 

{returnn-1.20241005.114831 → returnn-1.20241010.171521}/returnn/util/basic.py

@@ -2469,7 +2469,12 @@ def make_hashable(obj):
 
         if isinstance(obj, tf.Tensor):
             return RefIdEq(obj)
-    assert False, "don't know how to make hashable: %r (%r)" % (obj, type(obj))
+    # Try if this is already hashable.
+    try:
+        hash(obj)
+    except Exception:
+        raise TypeError("don't know how to make hashable: %r (%r)" % (obj, type(obj)))
+    return obj
 
 
 class RefIdEq(Generic[T]):

{returnn-1.20241005.114831 → returnn-1.20241010.171521}/returnn/util/better_exchook.py

@@ -1258,6 +1258,8 @@ def format_tb(tb=None, limit=None, allLocals=None, allGlobals=None, withTitle=Fa
 
 def print_tb(tb, file=None, **kwargs):
     """
+    Replacement for traceback.print_tb.
+
     :param types.TracebackType|types.FrameType|StackSummary tb:
     :param io.TextIOBase|io.StringIO|typing.TextIO|None file: stderr by default
     :return: nothing, prints to ``file``
@@ -1269,8 +1271,43 @@ def print_tb(tb, file=None, **kwargs):
     file.flush()
 
 
+def print_exception(etype, value, tb, limit=None, file=None, chain=True):
+    """
+    Replacement for traceback.print_exception.
+
+    :param etype: exception type
+    :param value: exception value
+    :param tb: traceback
+    :param int|None limit:
+    :param io.TextIOBase|io.StringIO|typing.TextIO|None file: stderr by default
+    :param bool chain: whether to print the chain of exceptions
+    """
+    better_exchook(etype, value, tb, autodebugshell=False, file=file, limit=limit, chain=chain)
+
+
+def print_exc(limit=None, file=None, chain=True):
+    """
+    Replacement for traceback.print_exc.
+    Shorthand for 'print_exception(*sys.exc_info(), limit, file)'.
+
+    :param int|None limit:
+    :param io.TextIOBase|io.StringIO|typing.TextIO|None file: stderr by default
+    :param bool chain:
+    """
+    print_exception(*sys.exc_info(), limit=limit, file=file, chain=chain)
+
+
 def better_exchook(
-    etype, value, tb, debugshell=False, autodebugshell=True, file=None, with_color=None, with_preamble=True
+    etype,
+    value,
+    tb,
+    debugshell=False,
+    autodebugshell=True,
+    file=None,
+    with_color=None,
+    with_preamble=True,
+    limit=None,
+    chain=True,
 ):
     """
     Replacement for sys.excepthook.
@@ -1284,6 +1321,8 @@ def better_exchook(
         and exception information. stderr by default.
     :param bool|None with_color: whether to use ANSI escape codes for colored output
     :param bool with_preamble: print a short preamble for the exception
+    :param int|None limit:
+    :param bool chain: whether to print the chain of exceptions
     """
     if file is None:
         file = sys.stderr
@@ -1292,16 +1331,17 @@ def better_exchook(
     output = _OutputLinesCollector(color=color)
 
     rec_args = dict(autodebugshell=False, file=file, with_color=with_color, with_preamble=with_preamble)
-    if getattr(value, "__cause__", None):
-        better_exchook(type(value.__cause__), value.__cause__, value.__cause__.__traceback__, **rec_args)
-        output("")
-        output("The above exception was the direct cause of the following exception:")
-        output("")
-    elif getattr(value, "__context__", None):
-        better_exchook(type(value.__context__), value.__context__, value.__context__.__traceback__, **rec_args)
-        output("")
-        output("During handling of the above exception, another exception occurred:")
-        output("")
+    if chain:
+        if getattr(value, "__cause__", None):
+            better_exchook(type(value.__cause__), value.__cause__, value.__cause__.__traceback__, **rec_args)
+            output("")
+            output("The above exception was the direct cause of the following exception:")
+            output("")
+        elif getattr(value, "__context__", None):
+            better_exchook(type(value.__context__), value.__context__, value.__context__.__traceback__, **rec_args)
+            output("")
+            output("During handling of the above exception, another exception occurred:")
+            output("")
 
     def format_filename(s):
         """
@@ -1320,7 +1360,14 @@ def better_exchook(
     all_locals, all_globals = {}, {}
     if tb is not None:
         output.lines.extend(
-            format_tb(tb=tb, allLocals=all_locals, allGlobals=all_globals, withTitle=True, with_color=color.enable)
+            format_tb(
+                tb=tb,
+                limit=limit,
+                allLocals=all_locals,
+                allGlobals=all_globals,
+                withTitle=True,
+                with_color=color.enable,
+            )
         )
     else:
         output(color("better_exchook: traceback unknown", color.fg_colors[1]))
@@ -1710,3 +1757,34 @@ def replace_traceback_format_tb():
     if hasattr(traceback, "StackSummary"):
         traceback.StackSummary.format = format_tb
         traceback.StackSummary.extract = _StackSummary_extract
+
+
+def replace_traceback_print_tb():
+    """
+    Replaces these functions from the traceback module by our own:
+
+    - traceback.print_tb
+    - traceback.print_exception
+    - traceback.print_exc
+
+    Note that this kind of monkey patching might not be safe under all circumstances
+    and is not officially supported by Python.
+    """
+    import traceback
+
+    traceback.print_tb = print_tb
+    traceback.print_exception = print_exception
+    traceback.print_exc = print_exc
+
+
+def setup_all():
+    """
+    Calls:
+
+    - :func:`install`
+    - :func:`replace_traceback_format_tb`
+    - :func:`replace_traceback_print_tb`
+    """
+    install()
+    replace_traceback_format_tb()
+    replace_traceback_print_tb()

{returnn-1.20241005.114831 → returnn-1.20241010.171521}/returnn.egg-info/PKG-INFO

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

returnn-1.20241005.114831/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20241005.114831'
-long_version = '1.20241005.114831+git.c53ebb4'