[export] Add schema version to serializer/deserializer (#107420)

Added a version number to the schema for BC issues. We will add this number to the serialized ExportedProgram and then when deserializing, if the number does not match up with the existing deserializer, we will error. We should update the number of there are any major changes to the schema.
Pull Request resolved: #107420
Approved by: https://github.com/zhxchen17

test/export/test_serialize.py

@@ -11,6 +11,7 @@
     ExportedProgramSerializer,
     deserialize,
     serialize,
+    SerializeError,
 )
 from torch._subclasses.fake_tensor import FakeTensor
 from torch.fx.experimental.symbolic_shapes import is_concrete_int
@@ -369,6 +370,19 @@ def test_exportdb_supported(self, name: str, case: ExportCase) -> None:
 
 instantiate_parametrized_tests(TestDeserialize)
 
+@unittest.skipIf(not torchdynamo.is_dynamo_supported(), "dynamo doesn't support")
+class TestSchemaVersioning(TestCase):
+    def test_error(self):
+        def f(x):
+            return x + x
+
+        ep = export(f, (torch.randn(1, 3),))
+
+        serialized_ep, serialized_state_dict = ExportedProgramSerializer().serialize(ep)
+        serialized_ep.schema_version = -1
+        with self.assertRaisesRegex(SerializeError, r"Serialized schema version -1 does not match our current"):
+            ExportedProgramDeserializer().deserialize(serialized_ep, serialized_state_dict)
+
 
 class TestOpVersioning(TestCase):
     """Test if serializer/deserializer behaves correctly if version mismatch."""

torch/_export/serde/schema.py

@@ -5,6 +5,10 @@
 from enum import IntEnum
 from typing import Dict, List, Optional, Tuple
 
+
+# NOTE: Please update this value if any modifications are made to the schema
+SCHEMA_VERSION = 0
+
 # TODO (zhxchen17) Move to a separate file.
 class _Union:
     @classmethod
@@ -247,3 +251,4 @@ class ExportedProgram:
     opset_version: Dict[str, int]
     range_constraints: Dict[str, RangeConstraint]
     equality_constraints: List[Tuple[Tuple[str, int], Tuple[str, int]]]
+    schema_version: int

torch/_export/serde/serialize.py

@@ -39,6 +39,7 @@
     OptionalTensorArgument,
     RangeConstraint,
     ScalarType,
+    SCHEMA_VERSION,
     SymBool,
     SymBoolArgument,
     SymExpr,
@@ -752,6 +753,7 @@ def serialize(self, exported_program: ep.ExportedProgram) -> Tuple[ExportedProgr
                 opset_version=self.opset_version,
                 range_constraints=serialized_range_constraints,
                 equality_constraints=serialized_equality_constraints,
+                schema_version=SCHEMA_VERSION,
             ),
             serialize_state_dict(exported_program.state_dict),
         )
@@ -1198,6 +1200,12 @@ def deserialize_range_constraints(
     def deserialize(
         self, serialized_exported_program: ExportedProgram, serialized_state_dict: bytes
     ) -> ep.ExportedProgram:
+        if serialized_exported_program.schema_version != SCHEMA_VERSION:
+            raise SerializeError(
+                f"Serialized schema version {serialized_exported_program.schema_version} "
+                f"does not match our current schema version {SCHEMA_VERSION}."
+            )
+
         symbol_name_to_range = {
             k: symbolic_shapes.ValueRanges(_int_to_sympy_int(v.min_val), _int_to_sympy_int(v.max_val))
             for k, v in serialized_exported_program.range_constraints.items()