Repository navigation
Re-validate the root model only once per _update_fields call. - #2472
Open
copybara-service[bot] wants to merge 1 commit into
Open
copybara-service[bot] wants to merge 1 commit into
copybara-service[bot] wants to merge 1 commit into
Conversation
`BaseModelFrozen._update_fields` walked from every mutated submodel up to the root and called `m.__class__.from_dict(m.to_dict())` on each node along the way. Because `from_dict(to_dict())` already re-validates the whole subtree recursively, and because the ancestral paths of sibling submodels overlap, this re-validated the same models many times over. Updating the 18 parameters that a pulse optimization step writes triggered roughly 54 full serialization and validation passes, 18 of which were over the entire `ToraxConfig` tree. Deduplicate the traversal: clear the `functools.cached_property` caches once per unique ancestral node, then re-validate the root model exactly once. Validation coverage is unchanged, since validating the root recursively validates every mutated and ancestral submodel. This reduces the time to apply a 18-parameter update to a `ToraxConfig` from 3.88s to 0.24s (16x) in an optimized build, and from 12.18s to 0.75s in a fastbuild. For a TORAX simulation driven by an external optimizer this is pure startup overhead paid on every trial. PiperOrigin-RevId: 983683899
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Re-validate the root model only once per
_update_fieldscall.BaseModelFrozen._update_fieldswalked from every mutated submodel up to theroot and called
m.__class__.from_dict(m.to_dict())on each node along the way.Because
from_dict(to_dict())already re-validates the whole subtreerecursively, and because the ancestral paths of sibling submodels overlap, this
re-validated the same models many times over. Updating the 18 parameters that a
pulse optimization step writes triggered roughly 54 full serialization and
validation passes, 18 of which were over the entire
ToraxConfigtree.Deduplicate the traversal: clear the
functools.cached_propertycaches once perunique ancestral node, then re-validate the root model exactly once. Validation
coverage is unchanged, since validating the root recursively validates every
mutated and ancestral submodel.
This reduces the time to apply a 18-parameter update to a
ToraxConfigfrom3.88s to 0.24s (16x) in an optimized build, and from 12.18s to 0.75s in a
fastbuild. For a TORAX simulation driven by an external optimizer this is pure
startup overhead paid on every trial.