Add v2 RETURNN config serialization #601

NeoLegends · 2025-05-09T12:44:14Z

This PR adds an adapted version of the v2 serialization written by @albertz. I intend for this PR to house discussion on what features to support in v2 serialization and what not to support.

The adaptations I made mainly concern themselves around being compatible with existing RETURNN configs that use the v1 serialization mechanisms like the Import, PartialImport, etc.

serialization_v2/serialization.py

NeoLegends · 2025-05-14T09:55:46Z

I brought back Dim serialization in a way where RETURNN is not imported in the manager.

Stefanwuu · 2025-08-26T14:16:40Z

serialization_v2/serialization.py

+        known_modules: Collection[str] = (),
+        sis_path_handling: SisPathHandling = None,
+    ):
+        from returnn.tensor import Dim, batch_dim, single_step_dim


This import requires returnn to be available for i6_core, I think this is not the case in many setups, maybe we should make the import conditional.

Icemole

I didn't go through all of the code (until half of the _Serializer class), but here's my feedback. Feel free to fix what you think is appropriate :)

Icemole · 2025-09-24T11:09:51Z

serialization_v2/serialization.py

+
+We handle those objects specially:
+- primitive types (int, float, bool, str)
+- Sisyphus Path objects


Regarding Path serialization (1/2, also see below): as far as I understand, you also seem to be handling sisyphus variables by using _instanciate_delayed_copy (as seen here). Is this correct?

If it's indeed correct, are you sure that Path handling is triggered? I assume that you added this because you found some case in which it wasn't supported. But if it's not needed, maybe you could also drop explicit support for Path. Path inherits from AbstractPath, which in turn inherits from DelayedBase (source). That would make the code easier by also dropping e.g. SisPathHandling entirely.

Note: Path.get() calls Path.get_path(), as seen here, so either call is equivalent.

Icemole · 2025-09-24T11:15:48Z

serialization_v2/serialization.py

+            .strip()
+        )
+        assert isinstance(_base_sys_path_list, list) and all(isinstance(p, str) for p in _base_sys_path_list)
+    return _base_sys_path_list


you always use the output of this function along with the in operator. I'd recommend to convert to a set (all elements are strings anyway, so they're hashable).

Icemole · 2025-09-24T11:16:58Z

serialization_v2/serialization.py

+    return name.isidentifier() and not iskeyword(name)
+
+
+class ReturnnConfigWithNewSerialization(ReturnnConfig):


ReturnnConfigV2 is a shorter name, and matches the directory structure. Better or not, I leave it to you... I personally think it's better 🙂

Icemole · 2025-09-24T11:28:05Z

serialization_v2/serialization.py

+        except OSError:
+            # cannot get source, e.g. code is a lambda
+            assert name is not None
+            args = [


While it seems "obvious" that you're trying to serialize the function manually, the code seems extremely obscure. I'd recommend adding a source as a comment that we can look for in case of doubt about the functionality of this code.

Is it just an exhaustive enumeration of the second column in https://docs.python.org/3/library/inspect.html, type "code"? If so, here's a suggestion for a comment:

Suggested change

args = [

# Dump all elements that identify this piece of code as a function.

# The exhaustive enumeration below is taken from https://docs.python.org/3/library/inspect.html, type "code".

args = [

I see that this comes from https://github.com/rwth-i6/i6_core/blob/main/returnn/config.py#L293. The code is practically equivalent, so I strongly suggest to find ways of combining both chunks of code.

For a preliminary idea: you're only using _parse_python inside ReturnnConfigWithNewSerialization. Moreover, I see that the only difference is an assertion in the code (https://github.com/rwth-i6/i6_core/blob/main/returnn/config.py#L299) which you don't seem to have (is this intended?). Because of this two factors, I suggest to make the _parse_python function a method in ReturnnConfigWithNewSerialization with the following structure:

class ReturnnConfigWithNewSerialization(ReturnnConfig): # Nitpick: note the double underscore at the beginning to ideally match the definition of the superclass. def __parse_python(...): try: return super().__parse_python(...) except AssertionError: # Thrown by the superclass method in some circumstances: https://github.com/rwth-i6/i6_core/blob/main/returnn/config.py#L299. # We ignore the assertion on purpose because TODO. return self.__parse_python(code.get())

Icemole · 2025-09-24T11:39:58Z

serialization_v2/serialization.py

+
+class ReturnnConfigWithNewSerialization(ReturnnConfig):
+    """
+    Overwrites the serialization behavior of ReturnnConfig


Suggested change

Overwrites the serialization behavior of ReturnnConfig

V2 RETURNN config with simpler and better serialization. This class is completely retrocompatible with respect to V1.

Allows serializing lambda functions as well as sisyphus `Path` and `Variable`, among the supported data types in the base RETURNN config class.

Migration is recommended.

If possible add more details to motivate people to migrate to this better serialization config class.

Icemole · 2025-09-24T11:51:01Z

serialization_v2/serialization.py

+                python_prolog_code,
+                serialized.as_serialized_code(),
+                python_epilog_code,
+                "# -*- mode: python; tab-width: 4 -*-\n",


Isn't this supposed to be at the top of the file?

Icemole · 2025-09-24T11:54:44Z

serialization_v2/serialization.py

+            known_modules={"returnn"},
+            extra_sys_paths=extra_sys_paths,
+            # instanciate_delayed should already have handled it (e.g. Path),
+            # or if it has not, then we want to keep it as it is (e.g. PtCheckpoint),


Regarding Path serialization (2/2): from the comments here I see that you already thought of Path being handled by instanciate_delayed. Can you provide an example of an object (e.g. PtCheckpoint as per the comments) that might not be properly handled?

Icemole · 2025-09-24T13:01:26Z

serialization_v2/serialization.py

+        return "".join(code.py_code for code in self.code_list)
+
+
+class _Serializer:


Nitpick: I strongly prefer having basic code as much to the top of the file as possible (e.g. _Serializer), and then code which builds from such code (e.g. serialize_config).

This also includes methods from the classes, e.g. _handle_next_queue_item would precede work_queue.

But that's my preference, maybe yours differs.

Icemole · 2025-09-24T13:37:59Z

serialization_v2/serialization.py

+            _Ref(value): name for name, value in self._internal_reserved_names.items()
+        }
+
+    def work_queue(self):


How about renaming the function for clarity and adding a few more details as a TL;DR for the docstring?

Suggested change

def work_queue(self):

def resolve_serialization_values(self):

"""

Runs the serialization queue, making sure that all values can be properly serialized.

If a value depends on another value, the latter is serialized first.

Icemole · 2025-09-24T14:49:10Z

serialization_v2/serialization.py

+
+
+@dataclass
+class SerializedConfig:


I see that as_serialization_collection is not used (isn't it?). I propose dropping this class and using "".join(...) instead.

Add v2 serialization

574ce70

NeoLegends requested review from Atticus1806, DanEnergetics, Icemole, JackTemaki, Stefanwuu, albertz, curufinwe, hannah220, michelwi and robin-p-schmitt May 9, 2025 12:44

michelwi reviewed May 13, 2025

View reviewed changes

serialization_v2/serialization.py Outdated Show resolved Hide resolved

NeoLegends added 3 commits May 14, 2025 11:38

Reintroduce Dim support without importing it in manager

02c4769

Check for identifier using python functions

30fb57a

Add some docs

908c477

NeoLegends added 2 commits May 25, 2025 14:12

handle Checkpoints

3d67d86

Merge branch 'main' into serialization-v2

960e2b2

Stefanwuu requested changes Aug 26, 2025

View reviewed changes

Icemole reviewed Sep 30, 2025

View reviewed changes

		return name.isidentifier() and not iskeyword(name)


		class ReturnnConfigWithNewSerialization(ReturnnConfig):

-    Overwrites the serialization behavior of ReturnnConfig
+    V2 RETURNN config with simpler and better serialization. This class is completely retrocompatible with respect to V1.
+    Allows serializing lambda functions as well as sisyphus `Path` and `Variable`, among the supported data types in the base RETURNN config class.
+    Migration is recommended.

		return "".join(code.py_code for code in self.code_list)


		class _Serializer:

-    def work_queue(self):
+    def resolve_serialization_values(self):
+        """
+        Runs the serialization queue, making sure that all values can be properly serialized.
+        If a value depends on another value, the latter is serialized first.

Add v2 RETURNN config serialization #601

Are you sure you want to change the base?

Add v2 RETURNN config serialization #601

Uh oh!

Conversation

NeoLegends commented May 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

NeoLegends commented May 14, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Icemole left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

NeoLegends commented May 9, 2025 •

edited

Loading