returnn 1.20230823.135336__tar.gz → 1.20230824.648__tar.gz

{returnn-1.20230823.135336/returnn.egg-info → returnn-1.20230824.648}/PKG-INFO

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

returnn-1.20230824.648/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20230824.000648'
+long_version = '1.20230824.000648+git.e565223'

{returnn-1.20230823.135336 → returnn-1.20230824.648}/returnn/datasets/basic.py

@@ -209,6 +209,7 @@ class Dataset(object):
         """
         assert issubclass(cls, Dataset)
         ds = cls(**kwargs)
+        assert isinstance(ds, Dataset)
         for attr, value in state.items():
             setattr(ds, attr, value)
         return ds

{returnn-1.20230823.135336 → returnn-1.20230824.648}/returnn/datasets/map.py

@@ -84,7 +84,7 @@ class MapDatasetWrapper(CachedDataset2):
         super(MapDatasetWrapper, self).__init__(**kwargs)
 
         if callable(map_dataset):
-            map_dataset = map_dataset(**kwargs)
+            map_dataset = map_dataset()
         assert isinstance(map_dataset, MapDatasetBase)
 
         self._dataset = map_dataset
@@ -92,6 +92,13 @@ class MapDatasetWrapper(CachedDataset2):
         assert map_dataset.data_types, f"{self}: map_dataset {map_dataset} needs to provide data_types"
         self.num_outputs = {key: _get_num_outputs_entry(key, opts) for key, opts in map_dataset.data_types.items()}
 
+    @property
+    def map_dataset(self) -> MapDatasetBase:  # also important for Dataset.__reduce__
+        """
+        :return: the wrapped MapDataset
+        """
+        return self._dataset
+
     @property
     def num_seqs(self):
         """
@@ -118,22 +125,32 @@ class MapDatasetWrapper(CachedDataset2):
         """
         super(MapDatasetWrapper, self).init_seq_order(epoch=epoch, seq_list=seq_list, seq_order=seq_order)
 
-        if seq_list is not None or seq_order is not None:
-            raise NotImplementedError
-
-        try:
-            self._seq_order = self._dataset.get_seq_order(epoch=epoch)
-        except OptionalNotImplementedError:
+        if seq_list is not None:
+            assert seq_order is None
+            tag_to_idx = {
+                self._dataset.get_seq_tag(corpus_seq_idx): corpus_seq_idx
+                for corpus_seq_idx in range(len(self._dataset))
+            }
+            seq_order = [tag_to_idx[tag] for tag in seq_list]
+        if seq_order is not None:
+            self._seq_order = seq_order
+        else:
             try:
-                self._seq_order = self.get_seq_order_for_epoch(
-                    epoch=epoch, num_seqs=len(self._dataset), get_seq_len=self._dataset.get_seq_len
-                )
+                self._seq_order = self._dataset.get_seq_order(epoch=epoch)
             except OptionalNotImplementedError:
-                # only support seq_ordering that need no length here
-                assert self.seq_ordering in ["default", "reverse", "random"]
-                self._seq_order = self.get_seq_order_for_epoch(
-                    epoch=epoch, num_seqs=len(self._dataset), get_seq_len=None
-                )
+                try:
+                    self._seq_order = self.get_seq_order_for_epoch(
+                        epoch=epoch, num_seqs=len(self._dataset), get_seq_len=self._dataset.get_seq_len
+                    )
+                except OptionalNotImplementedError:
+                    # only support seq_ordering that need no length here
+                    assert self.seq_ordering in ["default", "reverse", "random"], (
+                        f"{self}: dataset {self._dataset}.get_seq_len is not implemented,"
+                        f" seq_ordering {self.seq_ordering!r} is not supported"
+                    )
+                    self._seq_order = self.get_seq_order_for_epoch(
+                        epoch=epoch, num_seqs=len(self._dataset), get_seq_len=None
+                    )
 
         return True
 

{returnn-1.20230823.135336 → returnn-1.20230824.648}/returnn/datasets/multi_proc.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 from typing import Optional, Any, Dict, List
 from .basic import init_dataset, DatasetSeq
 from .cached2 import CachedDataset2
+from returnn.config import Config, get_global_config, set_global_config
 import multiprocessing as mp
 
 # noinspection PyProtectedMember
@@ -51,7 +52,18 @@ class MultiProcDataset(CachedDataset2):
 
     def initialize(self):
         """init"""
+        self._lazy_init()
+        super().initialize()
+
+    def _lazy_init(self):
         if not self._worker_procs:
+            # We send the global config to the subprocesses,
+            # because some datasets might use custom functions inside the config,
+            # and pickling them would fail otherwise.
+            # https://github.com/rwth-i6/returnn/issues/1384
+            # Pickling the config works, as the config has special pickling support.
+            global_config = get_global_config(raise_exception=False)
+
             # Seq order proc directly sends the seq order to each worker.
             seq_order_to_worker = []  # type: List[mpConnection]
             worker_from_seq_order = []  # type: List[mpConnection]
@@ -71,7 +83,12 @@ class MultiProcDataset(CachedDataset2):
             seq_order_proc = _mp.Process(
                 name=f"{self.name} seq order proc",
                 target=self._seq_order_proc_loop,
-                args=(self.dataset, seq_order_proc_child_conn, seq_order_to_worker),
+                # We deliberately put global_config first, because pickling of dataset might require
+                # that the global config is already loaded, which will have the side effect
+                # that it is registered as global config, and the Python module will be registered,
+                # so references to it will work.
+                args=(global_config, self.dataset, seq_order_proc_child_conn, seq_order_to_worker),
+                daemon=True,
             )
             seq_order_proc.start()
             # Make sure the child connection is closed here.
@@ -86,11 +103,18 @@ class MultiProcDataset(CachedDataset2):
                 worker_proc = _mp.Process(
                     name=f"{self.name} worker proc {i + 1}/{self.num_workers}",
                     target=self._worker_proc_loop,
-                    args=(self.dataset, self.buffer_size, worker_child_conns[i], worker_from_seq_order[i]),
+                    args=(
+                        global_config,  # see above
+                        self.dataset,
+                        self.buffer_size,
+                        worker_child_conns[i],
+                        worker_from_seq_order[i],
+                    ),
+                    daemon=True,
                 )
                 worker_proc.start()
                 worker_procs.append(worker_proc)
-                worker_child_conns[i].close()
+                worker_child_conns[i].close()  # see above
 
             self._seq_order_proc_parent_conn = seq_order_proc_parent_conn  # type: mpConnection
             self._seq_order_proc = seq_order_proc
@@ -107,8 +131,6 @@ class MultiProcDataset(CachedDataset2):
             msg, self.labels = self._seq_order_proc_parent_conn.recv()
             assert msg == "labels"
 
-        super().initialize()
-
     def __del__(self):
         if self._seq_order_proc:
             self._seq_order_proc_parent_conn.send(("exit", {}))
@@ -120,7 +142,11 @@ class MultiProcDataset(CachedDataset2):
                 worker_proc.join()
 
     @staticmethod
-    def _seq_order_proc_loop(dataset_dict: Dict[str, Any], parent: mpConnection, workers: List[mpConnection]):
+    def _seq_order_proc_loop(
+        global_config: Optional[Config], dataset_dict: Dict[str, Any], parent: mpConnection, workers: List[mpConnection]
+    ):
+        if global_config:
+            set_global_config(global_config)
         num_workers = len(workers)
         dataset = init_dataset(dataset_dict)
         try:
@@ -150,8 +176,14 @@ class MultiProcDataset(CachedDataset2):
 
     @staticmethod
     def _worker_proc_loop(
-        dataset_dict: Dict[str, Any], buffer_size: int, parent: mpConnection, seq_order: mpConnection
+        global_config: Optional[Config],
+        dataset_dict: Dict[str, Any],
+        buffer_size: int,
+        parent: mpConnection,
+        seq_order: mpConnection,
     ):
+        if global_config:
+            set_global_config(global_config)
         dataset = init_dataset(dataset_dict)
 
         got_init_seq_order = False
@@ -244,8 +276,8 @@ class MultiProcDataset(CachedDataset2):
         :returns whether the order changed (True is always safe to return)
         """
         super().init_seq_order(epoch=epoch, seq_list=seq_list, seq_order=seq_order)
-
         if epoch is not None or seq_list is not None or seq_order is not None:
+            self._lazy_init()
             self._seq_order_proc_parent_conn.send(
                 ("init_seq_order", {"epoch": epoch, "seq_list": seq_list, "seq_order": seq_order})
             )

{returnn-1.20230823.135336 → returnn-1.20230824.648/returnn.egg-info}/PKG-INFO

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

{returnn-1.20230823.135336 → returnn-1.20230824.648}/tests/test_Config.py

@@ -296,6 +296,42 @@ def test_config_pickle_function():
         assert f_() == 42
 
 
+def _config_pickle_proc_main(config, f):
+    assert isinstance(config, Config)
+    assert get_global_config() is config
+    assert callable(f)
+    f()
+
+
+def test_config_pickle_function_multi_proc():
+    # Same as test_config_pickle_function but via multiprocessing,
+    # so across process boundaries.
+    import multiprocessing
+
+    _mp = multiprocessing.get_context("spawn")
+
+    config = Config()
+    config.load_file(
+        StringIO(
+            textwrap.dedent(
+                """\
+                #!returnn.py
+
+                def my_custom_func():
+                    import sys
+                    sys.exit(42)
+                """
+            )
+        )
+    )
+    with global_config_ctx(config):
+        f = config.typed_dict["my_custom_func"]
+        proc = _mp.Process(target=_config_pickle_proc_main, args=(config, f))
+        proc.start()
+        proc.join()
+        assert proc.exitcode == 42
+
+
 if __name__ == "__main__":
     better_exchook.install()
     if len(sys.argv) <= 1:

{returnn-1.20230823.135336 → returnn-1.20230824.648}/tests/test_MultiProcDataset.py

@@ -82,6 +82,47 @@ def test_MultiProcDataset_meta():
         compare_dataset_seqs(meta_dataset_seqs, mp_dataset_seqs)
 
 
+def test_MultiProcDataset_via_config():
+    # https://github.com/rwth-i6/returnn/issues/1384
+    from io import StringIO
+    import textwrap
+    from returnn.config import Config, global_config_ctx
+
+    config = Config()
+    config.load_file(
+        StringIO(
+            textwrap.dedent(
+                """\
+                #!returnn.py
+
+                import numpy
+                from returnn.datasets.map import MapDatasetBase
+
+                class MyCustomMapDatasetInConfig(MapDatasetBase):
+                    def __init__(self):
+                        super().__init__(data_types={"data": {"shape": (None, 3)}})
+
+                    def __len__(self):
+                        return 2
+
+                    def __getitem__(self, item):
+                        return {"data": numpy.zeros((5, 3))}
+                """
+            )
+        )
+    )
+
+    with timeout(), global_config_ctx(config):
+        mp_dataset = MultiProcDataset(
+            dataset={"class": "MapDatasetWrapper", "map_dataset": config.typed_dict["MyCustomMapDatasetInConfig"]},
+            num_workers=1,
+            buffer_size=1,
+        )
+        mp_dataset.initialize()
+        items = dummy_iter_dataset(mp_dataset)
+        assert len(items) == 2
+
+
 class _MyCustomMapDatasetException(Exception):
     pass
 
@@ -138,6 +179,37 @@ def test_MultiProcDataset_exception_at_item():
             raise Exception("Expected exception")
 
 
+class _MyCustomDummyMapDataset(MapDatasetBase):
+    def __init__(self):
+        super().__init__(data_types={"data": {"shape": (None, 3)}})
+
+    def __len__(self):
+        return 2
+
+    def __getitem__(self, item):
+        return {"data": numpy.zeros((item * 2 + 5, 3))}
+
+
+def test_MultiProcDataset_pickle():
+    import pickle
+
+    with timeout():
+        mp_dataset = MultiProcDataset(
+            dataset={"class": "MapDatasetWrapper", "map_dataset": _MyCustomDummyMapDataset},
+            num_workers=1,
+            buffer_size=1,
+        )
+        mp_dataset.initialize()
+        mp_dataset_seqs = dummy_iter_dataset(mp_dataset)
+        mp_dataset_serialized = pickle.dumps(mp_dataset)
+
+    with timeout():
+        mp_dataset_ = pickle.loads(mp_dataset_serialized)
+        mp_dataset_seqs_ = dummy_iter_dataset(mp_dataset_)
+
+    compare_dataset_seqs(mp_dataset_seqs, mp_dataset_seqs_)
+
+
 if __name__ == "__main__":
     better_exchook.install()
     if len(sys.argv) <= 1:

{returnn-1.20230823.135336 → returnn-1.20230824.648}/tests/test_PTDataset.py

@@ -85,6 +85,33 @@ def test_HDFDataset():
     assert c == n
 
 
+def test_MultiProcDataset_HDFDataset():
+    from test_HDFDataset import generate_hdf_from_other
+    from test_MultiProcDataset import timeout
+    from returnn.datasets.multi_proc import MultiProcDataset
+
+    hdf_fn = generate_hdf_from_other({"class": "Task12AXDataset", "num_seqs": 23})
+    with timeout(10):
+        mp_dataset = MultiProcDataset(
+            dataset={"class": "HDFDataset", "files": [hdf_fn], "cache_byte_size": 0},
+            num_workers=1,
+            buffer_size=1,
+        )
+        mp_dataset.initialize()
+
+        mp_manager = torch.multiprocessing.Manager()
+        loader = get_loader_from_returnn_dataset(mp_dataset, mp_manager)
+        c = 0
+        n = 3
+        for batch in loader:
+            print(batch)
+            c += 1
+            if c >= n:
+                break
+
+        assert c == n
+
+
 if __name__ == "__main__":
     better_exchook.install()
     if len(sys.argv) <= 1:

returnn-1.20230823.135336/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20230823.135336'
-long_version = '1.20230823.135336+git.48dd63e'