Skip to content

Re-validate the root model only once per _update_fields call. - #2472

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_983683899
Open

copybara-service[bot] wants to merge 1 commit into
mainfrom
test_983683899

Conversation

@copybara-service

Copy link
Copy Markdown

Re-validate the root model only once per _update_fields call.

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.

`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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant