returnn 1.20240626.231046__tar.gz → 1.20240628.104058__tar.gz

{returnn-1.20240626.231046 → returnn-1.20240628.104058}/PKG-INFO

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

returnn-1.20240628.104058/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240628.104058'
+long_version = '1.20240628.104058+git.32c390d'

{returnn-1.20240626.231046 → returnn-1.20240628.104058}/returnn/datasets/distrib_files.py

@@ -138,6 +138,7 @@ class DistributeFilesDataset(CachedDataset2):
         buffer_size: int = 1,
         distrib_shard_files: bool = False,
         _meta_info_cache: Optional[Dict[str, Any]] = None,
+        _distrib_info: Optional[Dict[str, int]] = None,
         **kwargs,
     ):
         """
@@ -149,6 +150,7 @@ class DistributeFilesDataset(CachedDataset2):
         :param distrib_shard_files: set to true to shard the data across worker processes in
             distributed training scenaria
         :param _meta_info_cache: for internal use
+        :param _distrib_info: for internal use
         """
         super().__init__(**kwargs)
         self.files = files
@@ -159,12 +161,24 @@ class DistributeFilesDataset(CachedDataset2):
         self._file_sizes: Optional[Dict[str, int]] = None  # key -> size. for equal distribution across sub epochs
         self._data_keys: Optional[List[str]] = None
         self._num_seqs: Optional[int] = None
-        self.distrib_shard_files = distrib_shard_files
-        self._shard_index, self._num_shards = _get_rank_and_size() if distrib_shard_files else (0, 1)
 
         self._workers: Dict[int, _WorkerProcParent] = {}  # epoch -> worker
         self._files_order_cache: Dict[int, List[List[FileTree]]] = {}  # full epoch (0-indexed) -> files order
 
+        self.distrib_shard_files = distrib_shard_files
+        if distrib_shard_files:
+            if _distrib_info:
+                # If we're in a child process `_get_rank_and_size()` no longer works,
+                # so we pass the info about the shards via a pickled property.
+                # See also Dataset.__reduce__.
+                self._shard_index = _distrib_info["shard_index"]
+                self._num_shards = _distrib_info["num_shards"]
+            else:
+                self._shard_index, self._num_shards = _get_rank_and_size()
+        else:
+            self._shard_index = 0
+            self._num_shards = 1
+
         if _meta_info_cache:
             # This allows to skip the lazy init in self.initialize().
             # This can be used when pickling/deepcopying an instance of this dataset,
@@ -183,6 +197,10 @@ class DistributeFilesDataset(CachedDataset2):
         self._lazy_init_num_outputs()
         super().initialize()
 
+    @property
+    def _distrib_info(self):
+        return {"num_shards": self._num_shards, "shard_index": self._shard_index}
+
     @property
     def _meta_info_cache(self):
         if not self.num_outputs:
@@ -220,9 +238,10 @@ class DistributeFilesDataset(CachedDataset2):
         }
 
     def __del__(self):
-        for k, worker in self._workers.items():
-            try_run(worker.exit, kwargs={"join": False})
-        self._workers.clear()
+        if hasattr(self, "_workers"):  # might not be set after an early exception in __init__
+            for k, worker in self._workers.items():
+                try_run(worker.exit, kwargs={"join": False})
+            self._workers.clear()
 
     def init_seq_order(self, epoch: Optional[int] = None, seq_list=None, seq_order=None) -> bool:
         """

{returnn-1.20240626.231046 → returnn-1.20240628.104058}/returnn.egg-info/PKG-INFO

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

returnn-1.20240626.231046/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240626.231046'
-long_version = '1.20240626.231046+git.14f0203'