Repository: rwth-i6/returnn

Issue:
`ConcatFilesDataset` combines poorly with `MetaDataset`
I'm working on running a test config using `ConcatFilesDataset` (#1521, #1519).

The original setup I have uses `MetaDataset` to load the features and the alignment targets from two distinct sets of HDFs. This makes integrating the `ConcatFilesDataset` pretty difficult because it cannot deal with heterogenous data. Take the following example:

```python
train = {
    "class": "MetaDataset",
    "data_map": {"classes": ("alignments", "data"), "data": ("features", "data")},
    "datasets": {
        "alignments": {
            "class": "HDFDataset",
            "files": [
                "/alignment/files.hdf",
                "/..."
            ],
            "partition_epoch": 250,
            "seq_ordering": "random",
        },
        "features": {
            "class": "HDFDataset",
            "files": [
                "/feature/files.hdf",
                "/..."
            ],
        },
    },
    "seq_order_control_dataset": "alignments",
}
```

Where do we integrate the `ConcatFilesDataset`?

Thoughts:
1. We can place `ConcatFilesDataset` around the `MetaDataset`. This is problematic because we a) cannot give both the alignment caches and the feature caches to the `ConcatFilesDataset`'s `files` list, as they are fundamentally different in size and heterogenous in contents. We b) could give the `ConcatFilesDataset` only the features and always include all alignment caches in the `MetaDataset` emitted by `get_sub_epoch_dataset`. This would work but would be slow as we'd have to reload all alignment caches for every subepoch. We could also c) try and compute the relevant alignment caches given a list of feature caches in `get_sub_epoch_dataset`, but I'm not sure we can assume this always works. It might be worth precomputing this info though.
2. We can place `ConcatFilesDataset` inside the `MetaDataset` where the `HDFDataset`s currently are. This works as long as only one of the datasets is a `ConcatFilesDataset`, because it is the one that needs to be the `seq_order_control_dataset` (it does not accept a `seq_order` from outside https://github.com/rwth-i6/returnn/commit/02fa44b986042c7dbb6baa6a42a06205c2acc932#diff-6f16d6edae7b1113b8c292acded8b6ab78875f54dca08cd58edd11110709372fR189-R190). This means we can only use `ConcatFilesDataset` for one of the data streams, but the advantage over option 1b) is that we don't need to reload the HDFs for the other data stream after every subepoch. This might be enough if only one stream is data-heavy (features) and the other ones are small enough they can be loaded the normal way (targets). If we want both datasets to be `ConcatFilesDataset` to leverage the advanced prefetching and caching behavior we can precompute a mapping between sequence ID <-> containing file path. Then one of the datasets can be the `seq_order_control_dataset` and the others use that mapping to load up the data files and set up the sub dataset for precisely the relevant segments on-demand.
3. There might be issues with `partition_epoch` being specified on both the parent `MetaDataset` and on the `ConcatFilesDataset` but I haven't put much thought into that yet.

WDYT? Is there anything I'm missing?


Discussion:
`ConcatFilesDataset` is intended to be around everything else if you have a hierarchy of multiple datasets (like here with `MetaDataset`). It's the only possible way because otherwise it must be able to operate on the whole dataset at once, which is not really possible with `ConcatFilesDataset`, which is exactly the point. E.g. `init_seq_order` will get `seq_list` or `seq_order` otherwise, which will fail.

> This is problematic because we a) cannot give both the alignment caches and the feature caches to the `ConcatFilesDataset`'s `files` list, as they are fundamentally different in size and heterogenous in contents.

Oh yea I did not think about this before. But fortunately this is simple to extend: `ConcatFilesDataset`'s `files` could be a list of an arbitrarily nested structure (e.g. tuple or dict), where all the leaves should be files (i.e. just `str`). So in this specific case, you could give it `list[tuple[str,str]]`, or maybe use `list[dict[str,str]]`, where all the `dict`s would have `"alignments": file` and `"features": file`. Then `get_sub_epoch_dataset` would get a subset of that list. Calculating the size would just iterate over all leaves. This should be all fairly straightforward with the `tree` package (which we use in other places, e.g. also in `FileCache`, see `import tree`).

Do you want to make a PR for that?
Or, another possibility which should already work right now: In `ConcatFilesDataset`'s `files`, you only specify one of them, either features or alignments. I assume you can infer the filename of the other one, i.e. given an features filename, I assume you can infer the filename of the alignments. Or otherwise you can also have a dict in your config, sth like `alignment_file_per_feature_file`, which maps one to the other. Then in `get_sub_epoch_dataset`, you will get one of it, but you can infer the other.

But I think the extension as discussed in my previous post makes sense anyway. It's probably cleaner.

I like the tree idea. I'm going to think about that for a bit and run a PR. I think it still requires consistently prepared HDFs or a mapping that's known beforehand.

> Or, another possibility which should already work right now: In ConcatFilesDataset's files, you only specify one of them, either features or alignments. I assume you can infer the filename of the other one, i.e. given an features filename, I assume you can infer the filename of the alignments. Or otherwise you can also have a dict in your config, sth like alignment_file_per_feature_file, which maps one to the other. Then in get_sub_epoch_dataset, you will get one of it, but you can infer the other.

Yes I wrote this in my large comment, this would work if both types of caches are prepared in a consistent way or if the mapping is known from somewhere in advance. But it's very cumbersome and probably also brittle from setup to setup. I think in that case I'd rather re-dump all the data into a new, consistent set of HDFs and drop the meta dataset altogether.
> this would work if both types of caches are prepared in a consistent way or if the mapping is known from somewhere in advance

But also the approach with `tree` assumes there is such a mapping, which you also know in advance, or not?

> this would work but would be slow as we'd have to reload all alignment caches for every subepoch. 

I don't understand this. Why do you need to load them all?

* Because you actually don't have individual alignment files but only one single alignment file? But the tree-approach would have the same problem then. I think the best solution is that you split up the HDFs, consistent to the features.
* Or if you already have individual alignments, I don't understand your comment. Why do you need to load them all?

> But also the approach with tree assumes there is such a mapping, which you also know in advance, or not?

Yes I edited my comment, did not have this in before.

> I don't understand this. Why do you need to load them all?

I think I was unclear here, I did not mean loading the data, but opening the HDFs, reading out which sequences are stored in which HDFs, etc. All the prep-work.

If you only had a single alignment HDF this would apply to the tree approach as well, yes. I think in that case it might really be helpful to have the concat dataset not operating as primary dataset, but as a sub dataset of the meta dataset, because this avoids redoing the initial prep work for the alignments.

I think it all depends on how large the alignment dataset is. It's probably not very large. Anyways, we should support the tree use case. I believe we often have consistently prepared data. I'm PRing this.
> I think I was unclear here, I did not mean loading the data, but opening the HDFs, reading out which sequences are stored in which HDFs, etc. All the prep-work.

I also don't understand this. Why? You mean under the assumption that you don't know the mapping in advance? But if this is the case, there you don't know the mapping, or that there is no clear mapping, then also the `tree` approach would not work.

Or if you have the mapping, you don't need to open all the HDFs but just the right one. But then the tree approach or doing the mapping logic inside `get_sub_epoch_dataset` would basically lead to the same solution, only implemented slightly differently.

I think the only reasonable case is that you have a clear mapping, and you know it in advance. If this is not the case yet, it should be easy to prepare it that way.


Resolving patch:
diff --git a/returnn/datasets/concat_files.py b/returnn/datasets/concat_files.py
--- a/returnn/datasets/concat_files.py
+++ b/returnn/datasets/concat_files.py
@@ -6,7 +6,7 @@
 
 from __future__ import annotations
 
-from typing import Optional, Any, Dict, Tuple, List, Sequence, Collection, Callable
+from typing import Optional, Any, Dict, Tuple, List, Sequence, Collection, Callable, Union
 import os
 import sys
 import numpy
@@ -24,6 +24,9 @@
 
 __all__ = ["ConcatFilesDataset"]
 
+Filename = str
+FileTree = Union[Filename, Tuple["FileTree", ...], Dict[Any, "FileTree"], List["FileTree"]]
+
 
 class ConcatFilesDataset(CachedDataset2):
     """
@@ -36,7 +39,7 @@ class ConcatFilesDataset(CachedDataset2):
     This scheme allows to shuffle over the files,
     which makes shuffling much more efficient over a large dataset
     at the cost of no longer shuffling over the full dataset in every subepoch.
-    Instead the quality of the shuffle depends on the number of files the dataset is
+    Instead, the quality of the shuffle depends on the number of files the dataset is
     split into -- the more files per subepoch, the better.
 
     Additionally, this scheme allows to prefetch and cache the upcoming needed files,
@@ -75,14 +78,53 @@ def get_sub_epoch_dataset(files_subepoch: List[str]) -> Dict[str, Any]:
           "partition_epoch": P,  # P << M
         }
 
-    For some discussion, see https://github.com/rwth-i6/returnn/issues/1519.
+    Instead of a plain list of strings, you can also provide a list of any nested structures to keep
+    multimodal data together::
+
+        def get_sub_epoch_dataset(files_subepoch: List[Tuple[str, str]]) -> Dict[str, Any]:
+          from returnn.util.file_cache import CachedFile
+
+          alignments, features = tuple(zip(*files_subepoch)) # transpose
+
+          return {
+            "class": "MetaDataset",
+            "data_map": {"classes": ("alignments", "data"), "data": ("features", "data")},
+            "datasets": {
+              "alignments": {
+                "class": "HDFDataset",
+                "files": alignments,
+                "seq_ordering": "random",
+              },
+              "features": {
+                "class": "HDFDataset",
+                "files": features,
+              },
+            },
+            "seq_order_control_dataset": "alignments",
+          }
+
+        train = {
+          "class": "ConcatFilesDataset",
+          "files": [
+            ("/nfs/alignment_1.hdf", "/nfs/features_1.hdf"),
+            ...
+          ],  # M entries
+          "get_sub_epoch_dataset": get_sub_epoch_dataset,
+          "partition_epoch": P,  # P << M
+        }
+
+    In this case the file sizes for sub epoch distribution are summed up per list entry
+    by iterating over the structure leaves.
+
+    For some discussion, see https://github.com/rwth-i6/returnn/issues/1519 and
+    https://github.com/rwth-i6/returnn/issues/1524.
     """
 
     def __init__(
         self,
         *,
-        files: List[str],
-        get_sub_epoch_dataset: Callable[[List[str]], Dict[str, Any]],
+        files: List[FileTree],
+        get_sub_epoch_dataset: Callable[[List[FileTree]], Dict[str, Any]],
         preload_next_n_sub_epochs: int = 1,
         buffer_size: int = 1,
         file_cache_opts: Optional[Dict[str, Any]] = None,
@@ -90,7 +132,8 @@ def __init__(
         **kwargs,
     ):
         """
-        :param files: the files to shuffle over
+        :param files: the files to shuffle over, can also be a list of arbitrarily nested python objects
+            to keep associated heterogeneous data together
         :param get_sub_epoch_dataset: callable which returns a dataset dict for a given subset of files
         :param preload_next_n_sub_epochs: how many sub epoch datasets to preload
         :param buffer_size: buffer size for each worker, amount of seqs to prefetch
@@ -103,13 +146,13 @@ def __init__(
         self.preload_next_n_sub_epochs = preload_next_n_sub_epochs
         self.buffer_size = buffer_size
         self.file_cache_opts = file_cache_opts or {}
-        self._file_sizes: Optional[Dict[str, int]] = None  # file -> size. for equal distribution across sub epochs
+        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._file_cache: Optional[_FileCacheProc] = None
         self._workers: Dict[int, _WorkerProcParent] = {}  # epoch -> worker
-        self._files_order_cache: Dict[int, List[List[str]]] = {}  # full epoch (0-indexed) -> files order
+        self._files_order_cache: Dict[int, List[List[FileTree]]] = {}  # full epoch (0-indexed) -> files order
 
         if _meta_info_cache:
             # This allows to skip the lazy init in self.initialize().
@@ -157,9 +200,13 @@ def _lazy_init_num_outputs(self):
             exit_hook()
 
     def _lazy_init_file_sizes(self):
+        import tree
+
         if self._file_sizes:
             return
-        self._file_sizes = {fn: os.path.getsize(fn) for fn in self.files}
+        self._file_sizes = {
+            _get_key_for_file_tree(t): sum((os.path.getsize(fn) for fn in tree.flatten(t)), 0) for t in self.files
+        }
 
     def _lazy_init_file_cache_proc(self):
         if self._file_cache:
@@ -226,7 +273,7 @@ def init_seq_order(self, epoch: Optional[int] = None, seq_list=None, seq_order=N
             if ep_ in self._workers:
                 continue
             full_epoch_0idx_ = (ep_ - 1) // self.partition_epoch
-            files_order: List[List[str]] = self._files_order_cache[full_epoch_0idx_]
+            files_order: List[List[FileTree]] = self._files_order_cache[full_epoch_0idx_]
             dataset_dict, exit_hook = self._get_sub_dataset_dict(files=files_order[(ep_ - 1) % self.partition_epoch])
             worker = _WorkerProcParent(
                 name=f"{self.__class__.__name__} {self.name} ep {epoch}",
@@ -241,7 +288,7 @@ def init_seq_order(self, epoch: Optional[int] = None, seq_list=None, seq_order=N
         self._num_seqs = self._workers[epoch].get_num_seqs()
         return True
 
-    def _get_sub_dataset_dict(self, files: List[str]) -> Tuple[Dict[str, Any], _FileCacheExitHook]:
+    def _get_sub_dataset_dict(self, files: List[FileTree]) -> Tuple[Dict[str, Any], _FileCacheExitHook]:
         dataset_dict = self.get_sub_epoch_dataset(files)
         if "random_seed_offset" not in dataset_dict:
             dataset_dict["random_seed_offset"] = self.random_seed_offset
@@ -255,8 +302,8 @@ def _get_sub_dataset_dict(self, files: List[str]) -> Tuple[Dict[str, Any], _File
 
     @staticmethod
     def _get_files_per_sub_epochs(
-        *, partition_epoch: int, file_sizes: Dict[str, int], files_order: Sequence[str]
-    ) -> List[List[str]]:
+        *, partition_epoch: int, file_sizes: Dict[str, int], files_order: Sequence[FileTree]
+    ) -> List[List[FileTree]]:
         total_size = sum(file_sizes.values())
         avg_size_per_sub_epoch = total_size / partition_epoch
         # Now evenly distribute the files over the sub epochs.
@@ -275,24 +322,24 @@ def _get_files_per_sub_epochs(
         assert len(files_per_sub_epochs) == partition_epoch
         sub_epoch_idx = 0
         size_taken = 0
-        for i, fn in enumerate(files_order):
+        for i, f_tree in enumerate(files_order):
+            size = file_sizes[_get_key_for_file_tree(f_tree)]
             num_remaining = len(files_order) - i
-            size = file_sizes[fn]
             if num_remaining <= partition_epoch - sub_epoch_idx - 1:
                 # All remaining sub epochs must be filled.
                 assert files_per_sub_epochs[sub_epoch_idx]
                 sub_epoch_idx += 1
-                files_per_sub_epochs[sub_epoch_idx].append(fn)
+                files_per_sub_epochs[sub_epoch_idx].append(f_tree)
                 size_taken = size
                 continue
             if sub_epoch_idx == partition_epoch - 1:
                 # We are done. Just add the rest to the last sub epoch.
-                files_per_sub_epochs[sub_epoch_idx].append(fn)
+                files_per_sub_epochs[sub_epoch_idx].append(f_tree)
                 size_taken += size
                 continue
             assert size_taken <= avg_size_per_sub_epoch
             if size_taken + size <= avg_size_per_sub_epoch:
-                files_per_sub_epochs[sub_epoch_idx].append(fn)
+                files_per_sub_epochs[sub_epoch_idx].append(f_tree)
                 size_taken += size
                 continue
             # We should increase the sub epoch index.
@@ -302,10 +349,10 @@ def _get_files_per_sub_epochs(
                 abs((size_taken + size) - avg_size_per_sub_epoch)
                 <= abs(size_taken - avg_size_per_sub_epoch)
             ):
-                files_per_sub_epochs[sub_epoch_idx].append(fn)
+                files_per_sub_epochs[sub_epoch_idx].append(f_tree)
                 size_taken = 0
             else:
-                files_per_sub_epochs[sub_epoch_idx + 1].append(fn)
+                files_per_sub_epochs[sub_epoch_idx + 1].append(f_tree)
                 size_taken = size
             sub_epoch_idx += 1
         assert all(files_per_sub_epochs)
@@ -348,6 +395,13 @@ def get_data_keys(self) -> List[str]:
         return self._data_keys
 
 
+def _get_key_for_file_tree(t: FileTree) -> str:
+    """generates a deterministic key given a file tree"""
+    import tree
+
+    return ":".join(tree.flatten(t))
+
+
 class _WorkerProcParent:
     def __init__(
         self,


Test patch: