wehub-resource-sync
172 lines
7.2 KiB
C++
172 lines
7.2 KiB
C++
/* 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.
|
|
==============================================================================*/
|
|
#ifndef TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_UTIL_H_
|
|
#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_UTIL_H_
|
|
|
|
#include <cstdint>
|
|
#include <functional>
|
|
#include <memory>
|
|
#include <optional>
|
|
#include <string>
|
|
#include <utility>
|
|
#include <variant>
|
|
#include <vector>
|
|
|
|
#include "absl/status/status.h"
|
|
#include "absl/status/statusor.h"
|
|
#include "absl/strings/string_view.h"
|
|
#include "riegeli/bytes/fd_reader.h" // from @riegeli
|
|
#include "riegeli/bytes/string_reader.h" // from @riegeli
|
|
#include "riegeli/records/record_reader.h" // from @riegeli
|
|
#include "tensorflow/tools/proto_splitter/chunk.pb.h"
|
|
#include "tsl/platform/protobuf.h"
|
|
|
|
namespace tensorflow {
|
|
namespace tools::proto_splitter {
|
|
|
|
using MessageBytes = std::variant<std::shared_ptr<tsl::protobuf::Message>,
|
|
tsl::protobuf::Message*, std::string>;
|
|
|
|
struct ChunkedProto {
|
|
std::vector<MessageBytes>* chunks = nullptr;
|
|
::tensorflow::proto_splitter::ChunkedMessage* chunked_message = nullptr;
|
|
};
|
|
|
|
// TODO(b/282796592): Consider switching to `tsl::protobuf::FieldPath` in the
|
|
// future.
|
|
|
|
// Fields can be represented using their name (string) or number (int). Map key
|
|
// and repeated field index types are also included in this variant type.
|
|
using FieldType = std::variant<std::string, int, bool>;
|
|
using Field = std::pair<FieldType, std::optional<FieldType>>;
|
|
|
|
// Convert a sequence of field tags to a vector of fields. A single field is a
|
|
// std::vector<FieldType>, since multiple field tags may correspond to a single
|
|
// field when the field is repeated or a map.
|
|
absl::StatusOr<const std::vector<Field>> GetFieldTypes(
|
|
const tsl::protobuf::RepeatedPtrField<
|
|
::tensorflow::proto_splitter::FieldIndex>& field_tags);
|
|
|
|
// Sets message.field_desc[field_index] to the data contained in chunk,
|
|
// according to the (cpp) type described by field_desc. Uses message_callback
|
|
// (instead of simply assigning) when field_desc describes a message.
|
|
absl::Status SetRepeatedFieldElement(
|
|
tsl::protobuf::Message* message,
|
|
const tsl::protobuf::FieldDescriptor* field_desc, uint64_t field_index,
|
|
const std::string& chunk,
|
|
std::function<absl::Status(void)> message_callback);
|
|
|
|
// Sets message.field_desc to the data contained in chunk, according to the
|
|
// (cpp) type described by field_desc. Uses message_callback (instead of simply
|
|
// assigning) when field_desc describes a message.
|
|
absl::Status SetFieldElement(
|
|
tsl::protobuf::Message* message,
|
|
const tsl::protobuf::FieldDescriptor* field_desc, const std::string& chunk,
|
|
std::function<absl::Status(void)> message_callback);
|
|
|
|
// Adds a new map entry (repeated message element with key/value fields) to
|
|
// message.field_desc (a map). The new map entry's key is set using map_key,
|
|
// according to its type.
|
|
absl::Status AddMapEntry(tsl::protobuf::Message* message,
|
|
const tsl::protobuf::FieldDescriptor* field_desc,
|
|
FieldType map_key);
|
|
|
|
// Struct returned by `GetMutableField`. The field can be retrieved by using the
|
|
// Reflection API on the parent.
|
|
struct MutableFieldResult {
|
|
tsl::protobuf::Message* parent;
|
|
const tsl::protobuf::FieldDescriptor* field;
|
|
// If field is repeated or map, `index` is set to the list index or the
|
|
// position at which the map key appears. If the field is not repeated,
|
|
// `index` is -1.
|
|
int index;
|
|
};
|
|
|
|
// Returns a mutable parent, field descriptor, and int index in the case of a
|
|
// repeated or map value field (or -1 if a non-repeated/map field).
|
|
absl::StatusOr<MutableFieldResult> GetMutableField(
|
|
tsl::protobuf::Message* message, const std::vector<FieldType>& fields);
|
|
|
|
absl::StatusOr<MutableFieldResult> GetMutableField(
|
|
tsl::protobuf::Message* message, const Field& field);
|
|
|
|
// Gets info about the mutable field that's directly attached to message.
|
|
absl::StatusOr<MutableFieldResult> GetMutableField(
|
|
tsl::protobuf::Message* message, const FieldType& field_type);
|
|
|
|
// Struct returned by `GetField`. The field can be retrieved by using the
|
|
// Reflection API on the parent.
|
|
struct FieldResult {
|
|
const tsl::protobuf::Message* parent;
|
|
const tsl::protobuf::FieldDescriptor* field;
|
|
// If field is repeated or map, `index` is set to the list index or the
|
|
// position at which the map key appears. If the field is not repeated,
|
|
// `index` is -1.
|
|
int index;
|
|
};
|
|
|
|
// Returns the parent message, field descriptor, and int index from following
|
|
// the provided message and fields.
|
|
absl::StatusOr<FieldResult> GetField(const tsl::protobuf::Message& message,
|
|
const std::vector<FieldType>& fields);
|
|
|
|
// Updates `field_tag` in the ChunkedField proto.
|
|
absl::Status AddFieldTag(
|
|
const tsl::protobuf::Descriptor& desc, const std::vector<FieldType>& fields,
|
|
::tensorflow::proto_splitter::ChunkedField& chunked_field);
|
|
|
|
absl::Status AddFieldTag(
|
|
const tsl::protobuf::Descriptor& desc, const Field& field,
|
|
::tensorflow::proto_splitter::ChunkedField& chunked_field);
|
|
|
|
// Returns the index of the map key in the map field. If the key is not found,
|
|
// returns -1.
|
|
absl::StatusOr<int> FindMapKey(const tsl::protobuf::Message& parent,
|
|
const tsl::protobuf::FieldDescriptor& map_field,
|
|
const tsl::protobuf::FieldDescriptor* key_field,
|
|
FieldType map_key);
|
|
|
|
// Formats bytes into something more readable. (e.g. 52428800 -> "50.0MiB")
|
|
std::string HumanReadableBytes(int64_t byte_count);
|
|
// Formats microseconds into a more readable string.
|
|
std::string HumanReadableDuration(int64_t microseconds);
|
|
|
|
// Construct a reader object to read in records from the .cpb file.
|
|
absl::StatusOr<riegeli::RecordReader<riegeli::FdReader<>>> GetRiegeliReader(
|
|
absl::string_view cpb_file);
|
|
|
|
// Construct a reader object to read in records from a string.
|
|
absl::StatusOr<riegeli::RecordReader<riegeli::StringReader<>>>
|
|
GetRiegeliStringReader(absl::string_view cpb_data);
|
|
|
|
// Read the last chunk, which contains metadata necessary for reading the
|
|
// remaining chunks.
|
|
absl::StatusOr<::tensorflow::proto_splitter::ChunkMetadata> GetChunkMetadata(
|
|
riegeli::RecordReaderBase& reader);
|
|
|
|
// Use the `reader` to read in the chunk specified by `chunk_info`.
|
|
absl::StatusOr<std::string> ReadChunk(
|
|
riegeli::RecordReaderBase& reader,
|
|
const ::tensorflow::proto_splitter::ChunkInfo& chunk_info);
|
|
|
|
// Returns true if prefix can only be found as a .pb file, and false if a .cpb
|
|
// file exists. Returns an error if neither .pb nor .cpb exist.
|
|
absl::StatusOr<bool> OnlyContainsPb(absl::string_view prefix);
|
|
|
|
} // namespace tools::proto_splitter
|
|
} // namespace tensorflow
|
|
|
|
#endif // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_UTIL_H_
|