Closed maffoo closed 7 months ago
In https://github.com/quantumlib/Cirq/pull/6115 we are considering migrating to orjson
to replace the std-library json
implementation. We should check that this proposal still works in that case, e.g. that orjson
preserves key order during deserialization. Note that the performance of contextual serialization was also called out as a problem in that PR, though the overhead was not as bad as what we've observed in internal benchmarks as noted above. This performance will depend heavily on the specifics of what is being serialized and how expensive it is to traverse.
All modified and coverable lines are covered by tests :white_check_mark:
Comparison is base (
a4ec796
) 97.82% compared to head (98ccf83
) 97.82%.
:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.
A "contextual serialization" mechanism was added in #3601 and #3673 that allows serializing certain objects once even when they occur multiple times within a larger structure. The only place we use this mechanism in cirq itself is for
FrozenCircuit
, since a given circuit can be included as a subcircuit multiple times within a larger circuit.On internal benchmarks, the existing contextual serialization mechanism is a major performance bottleneck, in particular the
has_serializable_by_key
function takes a significant amount of time even when no contextual serialization is needed because it traverses the entire object being serialized and makes lots ofisinstance
calls. The reason for these pre-checks is that the context serialization works by changing the outer-most serialized object to instead be a_ContextualSerialization
instance, which includes allSerializableByKey
instances in anobject_dag
list created from a pre-order traversal. The current implementation also has the downside that during serialization it finds already-seen instances by traversing theobject_dag
list itself, potentially incurring quadratic overhead if the list grows large; this is necessary becauseSerializableByKey
instances are not guaranteed to be hashable (though in practice the only existing case ofFrozenCircuit
is hashable).This mechanism can be simplified by simply keeping track of
SerializableByKey
instances as we encounter them and associating them with keys. The first time such an instance is encountered during serialization it is assigned a key and then serialized as a "value":{"cirq_type": "VAL": "key": key, "val": val}
wherekey
is the assigned key andval
is the normal serialization of the object. If the same instance is encountered later during serialization, it is serialized as a "reference":{"cirq_type": "REF", "key": key}
. We require thatSerializableByKey
instances be hashable so that the internal memoization can use a dict instead of a list, to avoid quadratic overheads from repeatedly searching for memoized instances. This would potentially allow us to use this mechanism on many more types without worrying about the runtime impact. Note that this requires that the traversal order when deserializing be consistent with that used when serializing, so that we never encounter aREF
before the correspondingVAL
. AFAICT this will always be the case, because thejson
library respects the order in which object keys appear when deserializing (functions likejson.load
can take anobject_pairs_hook
which takes key-value pairs as a list, and when these are collected in a dict and passed toobject_hook
the python dict implementation also preserves order).This change gives a significant improvement in serialization performance; on one internal benchmark the time spent in
json_serialization.to_json
when profiling dropped from 719s to 39s for an 18x speedup.There's still a bit of work to do to clean up json test cases since in a few places we already had existing
.json_inward
tests and I need to figure out how to preserve those as well as moving the existing.json
tests to be.json_inward
tests. But otherwise this is ready for review and would be good to discuss the proposed approach.