wehub-resource-sync
204 lines
6.6 KiB
Python
204 lines
6.6 KiB
Python
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Utilities for Proto Splitter modules."""
|
|
|
|
from collections.abc import Sequence
|
|
from typing import Any, Optional, Union
|
|
|
|
from google.protobuf import descriptor
|
|
from google.protobuf import message
|
|
from tensorflow.tools.proto_splitter import chunk_pb2
|
|
|
|
_BYTE_UNITS = [(1, "B"), (1 << 10, "KiB"), (1 << 20, "MiB"), (1 << 30, "GiB")]
|
|
|
|
|
|
def format_bytes(b: int) -> str:
|
|
"""Formats bytes into a human-readable string."""
|
|
for i in range(1, len(_BYTE_UNITS)):
|
|
if b < _BYTE_UNITS[i][0]:
|
|
n = f"{b / _BYTE_UNITS[i-1][0]:.2f}"
|
|
units = _BYTE_UNITS[i - 1][1]
|
|
break
|
|
else:
|
|
n = f"{b / _BYTE_UNITS[-1][0]:.2f}"
|
|
units = _BYTE_UNITS[-1][1]
|
|
n = n.rstrip("0").rstrip(".")
|
|
return f"{n}{units}"
|
|
|
|
|
|
FieldTypes = Union[str, int, bool, Sequence[Union[str, int, bool]]]
|
|
|
|
|
|
def get_field(
|
|
proto: message.Message, fields: FieldTypes
|
|
) -> tuple[Any, Optional[descriptor.FieldDescriptor]]:
|
|
"""Returns the field and field descriptor from the proto.
|
|
|
|
Args:
|
|
proto: Parent proto of any message type.
|
|
fields: List of string/int/map key fields, e.g. ["nodes", "attr", "value"]
|
|
can represent `proto.nodes.attr["value"]`.
|
|
|
|
Returns:
|
|
Tuple of (
|
|
Field in the proto or `None` if none are found,
|
|
Field descriptor
|
|
)
|
|
"""
|
|
field_proto = proto
|
|
field_desc = None
|
|
for field_proto, field_desc, _, _ in _walk_fields(proto, fields):
|
|
pass
|
|
return field_proto, field_desc
|
|
|
|
|
|
def get_field_tag(
|
|
proto: message.Message, fields: FieldTypes
|
|
) -> Sequence[chunk_pb2.FieldIndex]:
|
|
"""Generates FieldIndex proto for a nested field within a proto.
|
|
|
|
Args:
|
|
proto: Parent proto of any message type.
|
|
fields: List of string/int/map key fields, e.g. ["nodes", "attr", "value"]
|
|
can represent `proto.nodes.attr["value"]`.
|
|
|
|
Returns:
|
|
A list of FieldIndex protos with the same length as `fields`.
|
|
"""
|
|
field_tags = []
|
|
for _, field_desc, map_key, list_index in _walk_fields(proto, fields):
|
|
field_tags.append(chunk_pb2.FieldIndex(field=field_desc.number))
|
|
if map_key is not None:
|
|
key_type = field_desc.message_type.fields_by_name["key"].type
|
|
field_tags.append(
|
|
chunk_pb2.FieldIndex(map_key=_map_key_proto(key_type, map_key))
|
|
)
|
|
elif list_index is not None:
|
|
field_tags.append(chunk_pb2.FieldIndex(index=list_index))
|
|
return field_tags
|
|
|
|
|
|
def _walk_fields(proto: message.Message, fields: FieldTypes):
|
|
"""Yields fields in a proto.
|
|
|
|
Args:
|
|
proto: Parent proto of any message type.
|
|
fields: List of string/int/map key fields, e.g. ["nodes", "attr", "value"]
|
|
can represent `proto.nodes.attr["value"]`.
|
|
|
|
Yields:
|
|
Tuple of (
|
|
Field in the proto or `None` if none are found,
|
|
Field descriptor,
|
|
Key into this map field (or None),
|
|
Index into this repeated field (or None))
|
|
"""
|
|
if not isinstance(fields, list):
|
|
fields = [fields]
|
|
|
|
field_proto = proto
|
|
parent_desc = proto.DESCRIPTOR
|
|
i = 0
|
|
while i < len(fields):
|
|
field = fields[i]
|
|
field_desc = None
|
|
map_key = None
|
|
index = None
|
|
|
|
if parent_desc is None:
|
|
raise ValueError(
|
|
f"Unable to find fields: {fields} in proto of type {type(proto)}."
|
|
)
|
|
|
|
if isinstance(field, int):
|
|
try:
|
|
field_desc = parent_desc.fields_by_number[field]
|
|
except KeyError:
|
|
raise KeyError( # pylint:disable=raise-missing-from
|
|
f"Unable to find field number {field} in {parent_desc.full_name}. "
|
|
f"Valid field numbers: {parent_desc.fields_by_number.keys()}"
|
|
)
|
|
elif isinstance(field, str):
|
|
try:
|
|
field_desc = parent_desc.fields_by_name[field]
|
|
except KeyError:
|
|
raise KeyError( # pylint:disable=raise-missing-from
|
|
f"Unable to find field '{field}' in {parent_desc.full_name}. "
|
|
f"Valid field names: {parent_desc.fields_by_name.keys()}"
|
|
)
|
|
else: # bool (only expected as map key)
|
|
raise TypeError("Unexpected bool found in field list.")
|
|
|
|
i += 1
|
|
parent_desc = field_desc.message_type
|
|
if field_proto is not None:
|
|
field_proto = getattr(field_proto, field_desc.name)
|
|
|
|
# Handle special fields types (map key and list index).
|
|
if _is_map(parent_desc) and i < len(fields):
|
|
# Next field is the map key.
|
|
map_key = fields[i]
|
|
|
|
try:
|
|
field_proto = field_proto[map_key] if field_proto is not None else None
|
|
except KeyError:
|
|
field_proto = None
|
|
i += 1
|
|
|
|
if i < len(fields):
|
|
# The next field must be from the Value Message.
|
|
value_desc = parent_desc.fields_by_name["value"]
|
|
assert value_desc.message_type is not None
|
|
parent_desc = value_desc.message_type
|
|
|
|
elif is_repeated(field_desc) and i < len(fields):
|
|
# The next field is the index within the list.
|
|
index = fields[i]
|
|
try:
|
|
field_proto = field_proto[index] if field_proto is not None else None
|
|
except IndexError:
|
|
field_proto = None
|
|
i += 1
|
|
|
|
yield field_proto, field_desc, map_key, index
|
|
|
|
|
|
def _is_map(desc: descriptor.Descriptor) -> bool:
|
|
return desc.GetOptions().map_entry if desc is not None else False
|
|
|
|
|
|
def is_repeated(field_desc: descriptor.FieldDescriptor) -> bool:
|
|
# copybara:uncomment_begin(google-only)
|
|
# return field_desc.is_repeated
|
|
# copybara:uncomment_end_and_comment_begin
|
|
return field_desc.label == descriptor.FieldDescriptor.LABEL_REPEATED
|
|
# copybara:comment_end
|
|
|
|
|
|
_FIELD_DESC = descriptor.FieldDescriptor
|
|
_MAP_KEY = {
|
|
_FIELD_DESC.TYPE_STRING: lambda key: chunk_pb2.FieldIndex.MapKey(s=key),
|
|
_FIELD_DESC.TYPE_BOOL: lambda key: chunk_pb2.FieldIndex.MapKey(boolean=key),
|
|
_FIELD_DESC.TYPE_UINT32: lambda key: chunk_pb2.FieldIndex.MapKey(ui32=key),
|
|
_FIELD_DESC.TYPE_UINT64: lambda key: chunk_pb2.FieldIndex.MapKey(ui64=key),
|
|
_FIELD_DESC.TYPE_INT32: lambda key: chunk_pb2.FieldIndex.MapKey(i32=key),
|
|
_FIELD_DESC.TYPE_INT64: lambda key: chunk_pb2.FieldIndex.MapKey(i64=key),
|
|
}
|
|
|
|
|
|
def _map_key_proto(key_type, key):
|
|
"""Returns MapKey proto for a key of key_type."""
|
|
return _MAP_KEY[key_type](key)
|