returnn 1.20251023.135024__tar.gz → 1.20251106.185107__tar.gz

{returnn-1.20251023.135024/returnn.egg-info → returnn-1.20251106.185107}/PKG-INFO

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

returnn-1.20251106.185107/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20251106.185107'
+long_version = '1.20251106.185107+git.fd8d2d8'

{returnn-1.20251023.135024 → returnn-1.20251106.185107}/returnn/config.py

@@ -801,7 +801,7 @@ class SubProcCopyGlobalConfigPreInitFunc:
         from returnn.log import log
         from returnn import __old_mod_loader__
 
-        better_exchook.install()
+        better_exchook.setup_all()
         __old_mod_loader__.disable_lazy_mod_loads()
 
         if self.global_config:

{returnn-1.20251023.135024 → returnn-1.20251106.185107}/returnn/datasets/generating.py

@@ -1164,11 +1164,9 @@ class StaticDataset(CachedDataset2):
         """supports sorting"""
         return True
 
-    def _collect_single_seq(self, seq_idx):
-        """
-        :param int seq_idx:
-        :rtype: DatasetSeq
-        """
+    def _collect_single_seq(self, seq_idx: int) -> Optional[DatasetSeq]:
+        if seq_idx >= len(self._seq_order):
+            return None
         corpus_seq_idx = self._seq_order[seq_idx]
         data = self.data[corpus_seq_idx]
         return DatasetSeq(

{returnn-1.20251023.135024 → returnn-1.20251106.185107}/returnn/datasets/lm.py

@@ -694,6 +694,26 @@ class LmDataset(CachedDataset2):
             self.next_seq_idx = seq_idx + 1
             return DatasetSeq(seq_idx=seq_idx, features=data, targets=targets, seq_tag=seq_tag)
 
+    def finish_epoch(self, *, free_resources: bool = False):
+        """finish epoch"""
+        super().finish_epoch(free_resources=free_resources)
+
+        if free_resources:
+            self._orths_offsets_and_lens = None
+            if self._orth_mmaps is not None:
+                for m in self._orth_mmaps:
+                    if m is not None:
+                        m.close()
+                self._orth_mmaps = None
+            if self._orth_files is not None:
+                for f in self._orth_files:
+                    if f is not None:
+                        f.close()
+                self._orth_files = None
+
+            self._seq_list = None
+            self._seq_index_by_tag = None
+
 
 def _is_bliss(filename):
     """

{returnn-1.20251023.135024 → returnn-1.20251106.185107}/returnn/datasets/meta.py

@@ -964,7 +964,6 @@ class CombinedDataset(CachedDataset2):
         self.dataset_keys = set([m[0] for m in data_map.keys()])  # type: typing.Set[str]
         self.dataset_idx2key_map = dict(enumerate(sorted(self.dataset_keys)))  # idx -> dataset-key
         self.data_keys = set(data_map.values())  # type: typing.Set[str]
-        assert "data" in self.data_keys
         self.target_list = sorted(self.data_keys - {"data"})
 
         # Build target lookup table that maps from dataset_key and data_key (data key used by CombinedDataset)
@@ -994,8 +993,7 @@ class CombinedDataset(CachedDataset2):
         if data_dims:
             data_dims = convert_data_dims(data_dims)
             self.data_dims = data_dims
-            assert "data" in data_dims
-            for key in self.target_list:
+            for key in self.data_keys:
                 assert key in data_dims
         else:
             self.data_dims = {}
@@ -1009,7 +1007,7 @@ class CombinedDataset(CachedDataset2):
             if dataset_data_key in dataset.labels:
                 self.labels[data_key] = dataset.labels[dataset_data_key]
 
-        self.num_inputs = self.data_dims["data"][0]
+        self.num_inputs = self.data_dims["data"][0] if "data" in self.data_dims else 0
         self.num_outputs = self.data_dims
 
         self.data_dtypes = {
@@ -1019,6 +1017,9 @@ class CombinedDataset(CachedDataset2):
 
         self.dataset_seq_idx_boundaries: Optional[List[int]] = None
         self.dataset_sorted_seq_idx_list: Optional[List[Tuple[int, int]]] = None
+        self._sub_dataset_cur_loaded_seq_range: Optional[List[Tuple[int, int]]] = None
+        # The usage is about the seqs already covered in dataset_sorted_seq_idx_list,
+        # in case we dynamically build up this list.
         self.used_num_seqs_per_subset: Optional[List[int]] = None
 
     def init_seq_order(self, epoch=None, seq_list=None, seq_order=None):
@@ -1030,7 +1031,7 @@ class CombinedDataset(CachedDataset2):
         """
 
         assert seq_list is None and seq_order is None, "seq_list and seq_order not supported for %s" % self.__class__
-        need_reinit = self.epoch is None or self.epoch != epoch
+        need_reinit = self.epoch is None or self.epoch != epoch or self.expected_load_seq_start > 0
         num_seqs_saved = self._num_seqs
         super(CombinedDataset, self).init_seq_order(
             epoch=epoch, seq_list=seq_list, seq_order=seq_order
@@ -1047,13 +1048,15 @@ class CombinedDataset(CachedDataset2):
         for dataset in self.datasets.values():
             dataset.init_seq_order(epoch=epoch)
 
+        self._sub_dataset_cur_loaded_seq_range = [(0, 0)] * len(self.datasets)
+
         # noinspection PyBroadException
         try:
             total_num_seqs = sum([self.datasets[k].num_seqs for k in sorted(self.datasets.keys())])
         except Exception:
             total_num_seqs = None
 
-        if total_num_seqs is not None:
+        if total_num_seqs is not None and self.seq_ordering != "interleave":
             self.dataset_seq_idx_boundaries = self._create_dataset_seq_idx_boundaries()
 
             if self.sampling_sizes:
@@ -1090,7 +1093,7 @@ class CombinedDataset(CachedDataset2):
 
             # Re-initialize sequence orders of sub-datasets with created sequence list.
             self.used_num_seqs_per_subset = []
-            for dataset_idx, dataset_key in self.dataset_idx2key_map.items():
+            for dataset_idx, dataset_key in sorted(self.dataset_idx2key_map.items()):
                 assert self.datasets[dataset_key].have_corpus_seq_idx()
                 self.datasets[dataset_key].init_seq_order(epoch=epoch, seq_order=seq_order_subdatasets[dataset_idx])
                 self.used_num_seqs_per_subset.append(len(seq_order_subdatasets[dataset_idx]))
@@ -1098,6 +1101,11 @@ class CombinedDataset(CachedDataset2):
         else:
             self.dataset_sorted_seq_idx_list = []  # We will fill this as we go
             self.used_num_seqs_per_subset = [0] * len(self.datasets)
+            self._num_seqs = total_num_seqs
+
+            # These are currently not supported/implemented.
+            # All of these should just be done in the sub-datasets directly.
+            assert self.partition_epoch == 1 and self.repeat_epoch == 1 and self._num_shards == 1
 
         return True
 
@@ -1236,13 +1244,30 @@ class CombinedDataset(CachedDataset2):
 
         return dataset.get_estimated_seq_length(dataset_seq_idx)
 
-    def _expand_dataset_sec_idxs(self, num_values):
+    def _sub_dataset_make_cur_loaded(self, dataset_idx: int) -> bool:
+        # Cur meaning for the next sequence to be added to dataset_sorted_seq_idx_list.
+        seq_idx = self.used_num_seqs_per_subset[dataset_idx]
+        cur_start, cur_end = self._sub_dataset_cur_loaded_seq_range[dataset_idx]
+        if seq_idx >= cur_end:
+            self._sub_dataset_load_seqs(dataset_idx, cur_start, seq_idx + 1)
+            return True
+        elif seq_idx < cur_start:
+            return False
+        else:
+            return True
+
+    def _expand_dataset_seq_idxs(self, num_values: int) -> bool:
         """
-        :param int num_values: Add num_values entries to the dataset-segment-idx mapping table
-        :return: something?
-        :rtype: bool
+        Try to extend dataset_sorted_seq_idx_list.
+        We expect that we have reached the end of it.
+
+        :param num_values: Add num_values entries to the dataset-segment-idx mapping table
+        :return: whether we added num_values entries
         """
-        for i in range(num_values):
+        for _ in range(num_values):
+            for j in range(len(self.datasets)):
+                self._sub_dataset_make_cur_loaded(j)
+
             if self.seq_ordering == "default":  # i.e. in order
                 dataset_idx = 0
                 while dataset_idx < len(self.datasets):
@@ -1265,6 +1290,32 @@ class CombinedDataset(CachedDataset2):
                 else:
                     return False  # No dataset has remaining data
 
+            elif self.seq_ordering == "interleave":
+                complete_fracs_and_ds_idx = [
+                    (
+                        self.datasets[self.dataset_idx2key_map[j]].get_complete_frac(
+                            self.used_num_seqs_per_subset[j] - 1, allow_only_lr_suitable=True
+                        )
+                        if self.used_num_seqs_per_subset[j] > 0
+                        else 0.0,
+                        j,
+                    )
+                    for j in range(len(self.datasets))
+                ]
+                assert all(frac is not None for frac, _ in complete_fracs_and_ds_idx), (
+                    f"{self}: Datasets must provide complete frac for interleave,"
+                    f" got {complete_fracs_and_ds_idx}, dataset idx2key map {self.dataset_idx2key_map}"
+                )
+                # Sort by complete frac, i.e. datasets with the lowest complete frac first.
+                complete_fracs_and_ds_idx.sort()
+                for complete_frac, dataset_idx in complete_fracs_and_ds_idx:
+                    if self.datasets[self.dataset_idx2key_map[dataset_idx]].is_less_than_num_seqs(
+                        self.used_num_seqs_per_subset[dataset_idx]
+                    ):
+                        break
+                else:
+                    return False  # No dataset has remaining data
+
             elif self.seq_ordering == "random_dataset":
                 while True:
                     # Build probability table
@@ -1323,19 +1374,23 @@ class CombinedDataset(CachedDataset2):
     def _load_seqs(self, start, end):
         # If the segment order is not yet known, fix the next few segments
         if end > len(self.dataset_sorted_seq_idx_list):
-            self._expand_dataset_sec_idxs(end - len(self.dataset_sorted_seq_idx_list))
+            self._expand_dataset_seq_idxs(end - len(self.dataset_sorted_seq_idx_list))
 
         requested_seqs = self.dataset_sorted_seq_idx_list[start:end]
 
         for dataset_idx in range(len(self.datasets)):
-            dataset = self.datasets[self.dataset_idx2key_map[dataset_idx]]
             sub_requested_seqs = [s[1] for s in requested_seqs if s[0] == dataset_idx]
             if not sub_requested_seqs:
                 continue
             sub_start, sub_end = min(sub_requested_seqs), max(sub_requested_seqs)
-            dataset.load_seqs(sub_start, sub_end + 1)
+            self._sub_dataset_load_seqs(dataset_idx, sub_start, sub_end + 1)
         super(CombinedDataset, self)._load_seqs(start=start, end=end)
 
+    def _sub_dataset_load_seqs(self, dataset_idx: int, start: int, end: int):
+        self._sub_dataset_cur_loaded_seq_range[dataset_idx] = (start, end)
+        dataset = self.datasets[self.dataset_idx2key_map[dataset_idx]]
+        dataset.load_seqs(start, end)
+
     def _get_data(self, dataset_key, dataset_seq_idx, data_key):
         """
         :type dataset_seq_idx: int
@@ -1348,7 +1403,10 @@ class CombinedDataset(CachedDataset2):
         if dataset_data_key is not None:
             return dataset.get_data(dataset_seq_idx, dataset_data_key)
         else:
-            return numpy.array([], self.data_dtypes[data_key])
+            shape: List[int] = [0] * self.num_outputs[data_key][1]
+            if shape and not self.is_data_sparse(data_key):
+                shape[-1] = self.get_data_dim(data_key)
+            return numpy.zeros(shape, dtype=self.data_dtypes[data_key])
 
     def _collect_single_seq(self, seq_idx):
         """
@@ -1362,19 +1420,30 @@ class CombinedDataset(CachedDataset2):
         dataset = self.datasets[dataset_key]
 
         seq_tag = dataset.get_tag(dataset_seq_idx)
-        features = self._get_data(dataset_key, dataset_seq_idx, "data")
-        targets = {target: self._get_data(dataset_key, dataset_seq_idx, target) for target in self.target_list}
-        return DatasetSeq(seq_idx=seq_idx, seq_tag=seq_tag, features=features, targets=targets)
+        features = {key: self._get_data(dataset_key, dataset_seq_idx, key) for key in self.data_keys}
+        complete_frac = None
+        if self.seq_ordering == "interleave":
+            # In the interleave case, by design, this should be monotonically increasing,
+            # as per how we select the next seq in _expand_dataset_seq_idxs.
+            complete_frac = dataset.get_complete_frac(dataset_seq_idx, allow_only_lr_suitable=True)
+        # In other cases, complete_frac is not so straightforward.
+        # In the case that the total num seqs is known, then it's anyway not necessary.
+        return DatasetSeq(seq_idx=seq_idx, complete_frac=complete_frac, seq_tag=seq_tag, features=features)
 
-    def is_less_than_num_seqs(self, n):
+    def is_less_than_num_seqs(self, n: int) -> bool:
         """
-        :param int n:
-        :rtype: bool
+        :param n:
         """
         if n < len(self.dataset_sorted_seq_idx_list):
             return True
         else:
-            return self._expand_dataset_sec_idxs(n - len(self.dataset_sorted_seq_idx_list) + 1)
+            return self._expand_dataset_seq_idxs(n - len(self.dataset_sorted_seq_idx_list) + 1)
+
+    def get_data_keys(self) -> List[str]:
+        """data keys"""
+        if "data" in self.data_keys:
+            return ["data"] + sorted(self.data_keys - {"data"})
+        return sorted(self.data_keys)
 
     def get_target_list(self):
         """

{returnn-1.20251023.135024 → returnn-1.20251106.185107}/returnn/datasets/util/vocabulary.py

@@ -11,6 +11,7 @@ __all__ = [
     "SentencePieces",
     "CharacterTargets",
     "Utf8ByteTargets",
+    "HuggingFaceTokenizer",
 ]
 
 from typing import Optional, Union, Type, Callable, List, Dict
@@ -691,3 +692,92 @@ class Utf8ByteTargets(Vocabulary):
             assert ((seq >= 0) & (seq < 256)).all(), f"invalid byte value, must be within 0-255: {seq}"
             seq = seq.astype(numpy.uint8)
         return bytearray(seq).decode(encoding="utf8")
+
+
+class HuggingFaceTokenizer(Vocabulary):
+    """
+    Uses the `AutoTokenizer` class from the `transformers` package.
+    """
+
+    def __init__(self, *, huggingface_repo_dir: str):
+        """
+        :param str huggingface_repo_dir: the directory containing the `tokenizer_config.json` file.
+        """
+        import transformers  # noqa
+
+        # Make sure it is a string. (Could be e.g. Sis Path.)
+        huggingface_repo_dir = str(huggingface_repo_dir)
+        self._opts = {"huggingface_repo_dir": huggingface_repo_dir}
+        self._cache_key = huggingface_repo_dir
+        self.tokenizer = transformers.AutoTokenizer.from_pretrained(huggingface_repo_dir, trust_remote_code=True)
+        super().__init__(
+            vocab_file=None,
+            seq_postfix=None,
+            unknown_label=self.tokenizer.unk_token_id,
+            eos_label=self.tokenizer.eos_token_id,
+            bos_label=self.tokenizer.bos_token_id,
+            pad_label=self.tokenizer.pad_token_id,
+        )
+
+    def __repr__(self):
+        return "%s(%r)" % (self.__class__.__name__, self._opts)
+
+    def _parse_vocab(self):
+        self.num_labels = len(self.tokenizer)
+        # Do not load labels/vocab here. This is not really needed.
+
+    @property
+    def labels(self) -> List[str]:
+        """list of labels"""
+        if self._cache_key and self._cache_key in self._cache:
+            self._vocab, self._labels = self._cache[self._cache_key]
+            assert self.num_labels == len(self._vocab) == len(self._labels)
+        else:
+            self._labels = [self.tokenizer._convert_id_to_token(i) for i in range(self.num_labels)]  # noqa
+            self._vocab = {label: i for (i, label) in enumerate(self._labels)}
+            if self._cache_key:
+                self._cache[self._cache_key] = (self._vocab, self._labels)
+        return self._labels
+
+    def is_id_valid(self, idx: int) -> bool:
+        """
+        :param idx:
+        """
+        return 0 <= idx < len(self.tokenizer)
+
+    def id_to_label(self, idx: int, default: Union[str, Type[KeyError], None] = KeyError) -> Optional[str]:
+        """
+        :param idx:
+        :param default:
+        """
+        if default is not KeyError and not self.is_id_valid(idx):
+            return default
+        return self.tokenizer.convert_ids_to_tokens(idx)
+
+    def label_to_id(self, label: str, default: Union[int, Type[KeyError], None] = KeyError) -> Optional[int]:
+        """
+        :param label:
+        :param default:
+        """
+        res = self.tokenizer.convert_token_to_id(label)
+        if res == self.unknown_label_id or res < 0 or res is None:
+            # It could be that the label really is the unknown-label, or it could be that the label is unknown.
+            if label == self.id_to_label(self.unknown_label_id):
+                return self.unknown_label_id
+            if default is KeyError:
+                raise KeyError("label %r not found" % label)
+            return default
+        return res
+
+    def get_seq(self, sentence: str) -> List[int]:
+        """
+        :param sentence: assumed to be seq of vocab entries separated by whitespace
+        """
+        return self.tokenizer(sentence)["input_ids"]
+
+    def get_seq_labels(self, seq):
+        """
+        :param list[int]|numpy.ndarray seq: 1D sequence
+        :rtype: str
+        """
+        return self.tokenizer.decode(seq, skip_special_tokens=True)

{returnn-1.20251023.135024 → returnn-1.20251106.185107}/returnn/frontend/attention.py

@@ -483,6 +483,7 @@ class RelPosSelfAttention(SelfAttentionBase):
         matrix_bd = _rel_pos_enc_shift(matrix_bd, axis, pos_emb_spatial_dim, hist_dim)
 
         scores = matrix_ac + matrix_bd  # (batch, head, time1, time2)
+        del matrix_ac, matrix_bd
         scores *= self.key_dim_per_head.dimension**-0.5
         att_weights = rf.softmax(scores, axis=hist_dim)
         att_weights = rf.dropout(att_weights, self.att_dropout, axis=self.att_dropout_broadcast and hist_dim)

{returnn-1.20251023.135024 → returnn-1.20251106.185107}/returnn/frontend/encoder/conformer.py

@@ -273,6 +273,7 @@ class ConformerEncoderLayer(rf.Module):
         x_mhsa = self.self_att(x_mhsa_ln, axis=spatial_dim)
         x_mhsa = rf.dropout(x_mhsa, self.dropout, axis=self.dropout_broadcast and self.out_dim)
         x_mhsa_out = x_mhsa + x_ffn1_out
+        del x_mhsa
 
         # Conv
         x_conv_ln = self.conv_layer_norm(x_mhsa_out)

{returnn-1.20251023.135024 → returnn-1.20251106.185107}/returnn/frontend/module.py

@@ -274,10 +274,17 @@ class Functional(Module):
     (This is often not necessary, but sometimes useful.)
     """
 
-    def __init__(self, func):
+    def __init__(self, func, *, attribs: Optional[Dict[str, Any]] = None):
+        """
+        :param func: callable. you might want to use functools.partial if you want to fix some arguments.
+        :param attribs: optional dict of attributes to set on this module. e.g. ``out_dim``.
+        """
         super().__init__()
         assert callable(func)
         self.func = func
+        if attribs:
+            for k, v in attribs.items():
+                setattr(self, k, v)
 
     def __repr__(self):
         return f"{self.__class__.__name__}({self.func.__qualname__})"

{returnn-1.20251023.135024 → returnn-1.20251106.185107}/returnn/torch/engine.py

@@ -532,7 +532,7 @@ class Engine(EngineBase):
                     for key, val in eval_info.items():
                         self._tensorboard_writer.add_scalar(f"train/{key}", val, global_step=self.global_train_step)
                     self._tensorboard_writer.add_scalar(
-                        f"train/learning_rate",
+                        "train/learning_rate",
                         self._updater.get_effective_learning_rate(),
                         global_step=self.global_train_step,
                     )
@@ -930,7 +930,7 @@ class Engine(EngineBase):
             if not os.path.exists(filename) and os.path.exists(model_epoch_filename):
                 filename = model_epoch_filename
             print("Load model %s" % (filename,), file=log.v4)
-            checkpoint_state = torch.load(filename, map_location=self._device)
+            checkpoint_state = _torch_load(filename, device=self._device)
             if epoch is None:
                 epoch = checkpoint_state.get("epoch", self._start_epoch or 1)
             step = checkpoint_state.get("step", 1)
@@ -1030,7 +1030,7 @@ class Engine(EngineBase):
                         print("(No relevant parameters matching.)", file=log.v3)
                     continue
                 print(f"Pre-load weights for key '{preload_key}' from {opts['filename']}", file=log.v3)
-                preload_model_state = torch.load(opts["filename"], map_location=self._device)
+                preload_model_state = _torch_load(opts["filename"], device=self._device)
                 if opts.get("checkpoint_key", "model") is not None:
                     # This can be used if an external checkpoint saves a checkpoint a different structure that just the
                     # model state dict. E.g., if a checkpoint is created using
@@ -1063,6 +1063,28 @@ class Engine(EngineBase):
                 preload_model_state_keys = set(preload_model_state.keys())
                 loaded_state_keys.update(preload_model_state.keys())
                 missing_keys.difference_update(preload_model_state.keys())
+
+                custom_missing_load_func = opts.get("custom_missing_load_func")
+                if custom_missing_load_func:
+                    custom_missing_vars_map = {}
+                    for var_name in missing_keys_preload:
+                        var_shape = self._pt_model.state_dict()[var_name].shape
+                        var_val = custom_missing_load_func(
+                            name=var_name,
+                            shape=var_shape,
+                            preload_model_state=preload_model_state,
+                            **util.get_fwd_compat_kwargs(),
+                        )
+                        if var_val is not None:
+                            assert var_val.shape == var_shape
+                            custom_missing_vars_map[var_name] = var_val
+                    preload_model_state.update(custom_missing_vars_map)
+                    missing_keys_preload, unexpected_keys_preload = self._pt_model.load_state_dict(
+                        preload_model_state, strict=False
+                    )
+                    loaded_state_keys.update(preload_model_state.keys())
+                    missing_keys.difference_update(preload_model_state.keys())
+
                 del preload_model_state
                 gc.collect()
 
@@ -1700,3 +1722,15 @@ def _get_total_grad_norm(model: torch.nn.Module, p: float) -> float:
             p=p,
         ).item()
     )
+
+
+def _torch_load(filename: Union[str, os.PathLike], *, device: str) -> Dict[str, Any]:
+    # Might resolve PtCheckpoint or Sisyphus Path objects or so.
+    filename = os.fspath(filename)
+
+    if filename.endswith(".safetensors"):
+        from safetensors.torch import load_file as safetensors_load
+
+        return safetensors_load(filename, device=device)
+
+    return torch.load(filename, map_location=device)

returnn-1.20251106.185107/returnn/util/collect_outputs_dict.py

@@ -0,0 +1,79 @@
+"""
+Customized (derived) dict to pass as ``collected_outputs`` to some of the RF modules,
+or potential other use cases.
+
+You can predefine (by pattern) what kind of outputs you want to collect and store in this dict.
+"""
+
+from typing import Optional, Union, Sequence
+import fnmatch
+
+
+class CollectOutputsDict(dict):
+    """
+    Customized (derived) dict, where you can predefine (by key pattern)
+    what kind of keys you want to collect and store in this dict.
+    Other keys will be ignored.
+    """
+
+    def __init__(self, *args, allowed_key_patterns: Optional[Sequence[str]] = None, **kwargs):
+        """
+        Initialize the CollectOutputsDict.
+
+        :param allowed_key_patterns:
+            List of key patterns (with wildcards) that are allowed to be stored in the dict.
+            If None, all keys are allowed.
+        """
+        super().__init__(*args, **kwargs)
+        self.allowed_key_patterns = allowed_key_patterns
+
+    def __setitem__(self, key, value):
+        """
+        Set an item in the dict if the key matches allowed patterns.
+        """
+        if self.is_key_allowed(key):
+            super().__setitem__(key, value)
+
+    def setdefault(self, key, default=None):
+        """
+        Set default value for a key if it matches allowed patterns.
+        """
+        if self.is_key_allowed(key):
+            return super().setdefault(key, default)
+        return None
+
+    def update(self, mapping, **kwargs):
+        """
+        Update the dict with another mapping, only adding allowed keys.
+        """
+        assert not kwargs
+        for key, value in mapping.items():
+            if self.is_key_allowed(key):
+                super().__setitem__(key, value)
+
+    def is_key_allowed(self, key: str) -> bool:
+        """
+        Check if the key matches any of the allowed patterns.
+
+        :param key:
+        :return: True if the key is allowed, False otherwise.
+        """
+        if self.allowed_key_patterns is None:
+            return True  # If no patterns defined, allow all keys
+        for pattern in self.allowed_key_patterns:
+            if fnmatch.fnmatch(key, pattern):
+                return True
+        return False
+
+
+def is_key_allowed_in_collect_outputs_dict(collect_outputs: Union[CollectOutputsDict, dict], key: str) -> bool:
+    """
+    Check if a key is allowed in the given CollectOutputsDict.
+
+    :param collect_outputs:
+    :param key:
+    :return: True if the key is allowed, False otherwise.
+    """
+    if isinstance(collect_outputs, CollectOutputsDict):
+        return collect_outputs.is_key_allowed(key)
+    return True  # If it's a regular dict, all keys are allowed

{returnn-1.20251023.135024 → returnn-1.20251106.185107}/returnn/util/file_cache.py

@@ -426,7 +426,21 @@ class FileCache:
                 orig_mtime_ns = os.stat(src_filename).st_mtime_ns
                 FileInfo(mtime_ns=orig_mtime_ns).save(info_file_name)
 
-                _copy_with_prealloc(src_filename, dst_tmp_filename)
+                try:
+                    _copy_with_prealloc(src_filename, dst_tmp_filename)
+                except Exception:
+                    # Cleanup if it was created already.
+                    # That avoids some of the ambiguity of the existence of the .copy file.
+                    # https://github.com/rwth-i6/returnn/issues/1785
+                    try:
+                        os.remove(dst_tmp_filename)
+                    except FileNotFoundError:
+                        pass
+                    try:
+                        os.remove(info_file_name)
+                    except FileNotFoundError:  # not really expected here, but safe to ignore
+                        pass
+                    raise
                 os.rename(dst_tmp_filename, dst_filename)
 
     @staticmethod

{returnn-1.20251023.135024 → returnn-1.20251106.185107/returnn.egg-info}/PKG-INFO

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

{returnn-1.20251023.135024 → returnn-1.20251106.185107}/returnn.egg-info/SOURCES.txt

@@ -320,6 +320,7 @@ returnn/util/__init__.py
 returnn/util/basic.py
 returnn/util/better_exchook.py
 returnn/util/bpe.py
+returnn/util/collect_outputs_dict.py
 returnn/util/debug.py
 returnn/util/debug_helpers.py
 returnn/util/file_cache.py