Skip to content

Preserve table signature when constructing tables from checkpoint protos - #158

Draft
rootkiller6788 wants to merge 1 commit into
google-deepmind:masterfrom
rootkiller6788:fix-tables-from-proto-signature
Draft

Preserve table signature when constructing tables from checkpoint protos#158
rootkiller6788 wants to merge 1 commit into
google-deepmind:masterfrom
rootkiller6788:fix-tables-from-proto-signature

Conversation

@rootkiller6788

Copy link
Copy Markdown

Problem

server_from_proto.tables_from_proto reconstructs reverb.Table objects from PriorityTableCheckpoint protos but never passed the signature field to reverb.Table. PriorityTableCheckpoint.signature is a tensorflow.StructuredValue (field 9 in reverb/cc/checkpointing/checkpoint.proto) that the C++ Table::Checkpoint() writes and Table::InitializeFromCheckpoint() restores. The Python path dropped it, so any table rebuilt from a checkpoint proto lost its signature.

A signature-less table means TrajectoryDataset.from_table_signature and Client.sample(..., unpack_as_table_signature=True) can no longer reconstruct the nested data structure, and Table.info().signature returns None.

Fix

Decode config.signature with tensorflow.python.saved_model.nested_structure_coder.decode_proto (the same codec TableInfo.from_serialized_proto uses) and pass the result to reverb.Table(signature=...). When the field is unset, signature=None preserves the previous behavior.

Test

Added ServerFromProtoTest.test_table_from_proto_preserves_signature, which encodes a tf.TensorSpec signature into a PriorityTableCheckpoint, converts it via tables_from_proto, and asserts tables[0].info.signature equals the original signature.

server_from_proto.tables_from_proto ignored the signature field of
PriorityTableCheckpoint, so a table rebuilt from a checkpoint proto lost
its signature. Decode the signature with nested_structure_coder and pass
it to reverb.Table, matching how the C++ checkpoint path restores it.

Add a regression test verifying the signature round-trips.
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