wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 6 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/tools/proto_splitter/cc/BUILD
+335
View File
@@ -0,0 +1,335 @@
load(
"//tensorflow:tensorflow.bzl",
"if_oss",
"tf_cc_test",
)
# Description:
# Utilities for splitting and joining large protos > 2GB.
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//learning/brain/tfrt/aot:__subpackages__",
"//learning/brain/tfrt/tf_tpu:__subpackages__",
"//tensorflow/cc/saved_model/image_format:__subpackages__",
"//tensorflow/tools/proto_splitter:__subpackages__",
],
licenses = ["notice"],
)
cc_library(
name = "split",
hdrs = ["split.h"],
deps = [
":util",
"//tensorflow/tools/proto_splitter:chunk_proto_cc",
"//tensorflow/tools/proto_splitter:versions_proto_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@tsl//tsl/platform:protobuf",
] + if_oss([
"//tensorflow/tools/proto_splitter:protos_impl",
]),
)
cc_library(
name = "composable_splitter_base",
srcs = ["composable_splitter_base.cc"],
hdrs = ["composable_splitter_base.h"],
deps = [
":max_size",
":split",
":util",
"//tensorflow/core:lib",
"//tensorflow/tools/proto_splitter:chunk_proto_cc",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:cord",
"@riegeli//riegeli/base:initializer",
"@riegeli//riegeli/bytes:cord_writer",
"@riegeli//riegeli/bytes:fd_writer",
"@riegeli//riegeli/bytes:string_writer",
"@riegeli//riegeli/records:record_writer",
"@tsl//tsl/platform:errors",
"@tsl//tsl/platform:protobuf",
] + if_oss([
"//tensorflow/tools/proto_splitter:protos_impl",
]),
)
cc_library(
name = "composable_splitter",
hdrs = ["composable_splitter.h"],
deps = [
":composable_splitter_base",
":util",
"//tensorflow/tools/proto_splitter:chunk_proto_cc",
"@tsl//tsl/platform:protobuf",
] + if_oss([
"//tensorflow/tools/proto_splitter:protos_impl",
]),
)
tf_cc_test(
name = "composable_splitter_test",
srcs = ["composable_splitter_test.cc"],
data = [],
tags = [
"no_mac", # b/291933687
"no_windows", # b/291001524
],
deps = [
":composable_splitter",
":test_util",
":util",
"//tensorflow/core:lib",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:path",
"//tensorflow/tools/proto_splitter:chunk_proto_cc",
"//tensorflow/tools/proto_splitter/testdata:test_message_proto_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:cord",
"@com_google_googletest//:gtest_main",
"@riegeli//riegeli/base:initializer",
"@riegeli//riegeli/bytes:cord_reader",
"@riegeli//riegeli/bytes:fd_reader",
"@riegeli//riegeli/bytes:string_reader",
"@riegeli//riegeli/records:record_reader",
"@tsl//tsl/platform:protobuf",
"@tsl//tsl/platform:status_matchers",
"@xla//xla/tsl/lib/core:status_test_util",
] + if_oss([
"//tensorflow/tools/proto_splitter:protos_impl",
]),
)
cc_library(
name = "util",
srcs = ["util.cc"],
hdrs = ["util.h"],
visibility = [
"//learning/brain/tfrt/aot:__subpackages__",
"//tensorflow/cc/saved_model:__subpackages__",
"//tensorflow/cc/saved_model/image_format:__subpackages__",
"//tensorflow/tools/proto_splitter:__subpackages__",
],
deps = [
"//tensorflow/core/platform",
"//tensorflow/core/platform:env",
"//tensorflow/tools/proto_splitter:chunk_proto_cc",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:cord",
"@com_google_absl//absl/strings:str_format",
"@riegeli//riegeli/base:initializer",
"@riegeli//riegeli/base:types",
"@riegeli//riegeli/bytes:fd_reader",
"@riegeli//riegeli/bytes:string_reader",
"@riegeli//riegeli/records:record_reader",
"@tsl//tsl/platform:errors",
"@tsl//tsl/platform:protobuf",
"@tsl//tsl/platform:statusor",
] + if_oss([
"//tensorflow/tools/proto_splitter:protos_impl",
]),
)
tf_cc_test(
name = "util_test",
srcs = ["util_test.cc"],
data = [
"//tensorflow/tools/proto_splitter/testdata:split-standard.cpb",
],
tags = [
"no_mac", # b/291933687
"no_windows", # b/291001524
],
deps = [
":test_util",
":util",
"//tensorflow/core/platform:macros",
"//tensorflow/core/platform:path",
"//tensorflow/core/platform:test",
"//tensorflow/tools/proto_splitter:chunk_proto_cc",
"//tensorflow/tools/proto_splitter/testdata:test_message_proto_cc",
"@com_google_absl//absl/status",
"@com_google_googletest//:gtest_main",
"@tsl//tsl/platform:protobuf",
"@tsl//tsl/platform:status",
"@tsl//tsl/platform:status_matchers",
"@tsl//tsl/platform:statusor",
"@xla//xla/tsl/lib/core:status_test_util",
] + if_oss([
"//tensorflow/tools/proto_splitter:protos_impl",
]),
)
cc_library(
name = "graph_def_splitter",
srcs = ["graph_def_splitter.cc"],
hdrs = ["graph_def_splitter.h"],
deps = [
":composable_splitter",
":composable_splitter_base",
":large_node_splitter",
":max_size",
":repeated_field_splitter",
":size_splitter",
":util",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@tsl//tsl/platform:errors",
"@tsl//tsl/platform:protobuf",
"@tsl//tsl/platform:statusor",
],
)
cc_library(
name = "size_splitter",
hdrs = ["size_splitter.h"],
deps = [
":composable_splitter",
":util",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/types:optional",
],
)
cc_library(
name = "repeated_field_splitter",
hdrs = ["repeated_field_splitter.h"],
deps = [
":composable_splitter",
":max_size",
":size_splitter",
":util",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@tsl//tsl/platform:protobuf",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "graph_def_splitter_test",
srcs = ["graph_def_splitter_test.cc"],
data = [
"//tensorflow/tools/proto_splitter/testdata:function-large-nodes.pb",
"//tensorflow/tools/proto_splitter/testdata:function-lots-of-nodes.pb",
"//tensorflow/tools/proto_splitter/testdata:graph-def-and-function.pb",
"//tensorflow/tools/proto_splitter/testdata:split-large-constant.pb",
"//tensorflow/tools/proto_splitter/testdata:split-large-nodes.pb",
"//tensorflow/tools/proto_splitter/testdata:split-lots-nodes.pb",
],
tags = [
"no_mac", # b/291933687
"no_windows", # b/291001524
],
deps = [
":graph_def_splitter",
":max_size",
":test_util",
":util",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core/platform:path",
"//tensorflow/tools/proto_splitter/testdata:test_message_proto_cc",
"@com_google_absl//absl/strings:cord",
"@com_google_googletest//:gtest_main",
"@tsl//tsl/platform:protobuf",
] + if_oss([
"//tensorflow/tools/proto_splitter:protos_impl",
]),
)
cc_library(
name = "max_size",
srcs = ["max_size.cc"],
hdrs = ["max_size.h"],
deps = ["@com_google_absl//absl/synchronization"],
)
cc_library(
name = "saved_model_splitter",
srcs = ["saved_model_splitter.cc"],
hdrs = ["saved_model_splitter.h"],
deps = [
":composable_splitter",
":graph_def_splitter",
":large_node_splitter",
":max_size",
":util",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/status",
"@tsl//tsl/platform:errors",
"@tsl//tsl/platform:protobuf",
],
)
tf_cc_test(
name = "saved_model_splitter_test",
srcs = ["saved_model_splitter_test.cc"],
data = [
"//tensorflow/cc/saved_model:saved_model_test_files",
],
tags = ["no_windows"], # b/291001524
deps = [
":max_size",
":saved_model_splitter",
":util",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core/platform:path",
"//tensorflow/tools/proto_splitter/testdata:test_message_proto_cc",
"@com_google_googletest//:gtest_main",
"@tsl//tsl/platform:protobuf",
] + if_oss([
"//tensorflow/tools/proto_splitter:protos_impl",
]),
)
cc_library(
name = "large_node_splitter",
hdrs = ["large_node_splitter.h"],
deps = [
":composable_splitter",
":composable_splitter_base",
":max_size",
":size_splitter",
":util",
"@com_google_absl//absl/status:statusor",
"@tsl//tsl/platform:errors",
],
)
cc_library(
name = "test_util",
testonly = True,
hdrs = ["test_util.h"],
deps = [
"//tensorflow/core:test",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@tsl//tsl/platform:errors",
"@tsl//tsl/platform:protobuf",
"@tsl//tsl/platform:statusor",
],
)
tensorflow/tools/proto_splitter/cc/composable_splitter.h
+52
View File
@@ -0,0 +1,52 @@
/* 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_COMPOSABLE_SPLITTER_H_
#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_COMPOSABLE_SPLITTER_H_
#include <vector>
#include "tensorflow/tools/proto_splitter/cc/composable_splitter_base.h"
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include "tensorflow/tools/proto_splitter/chunk.pb.h"
#include "tsl/platform/protobuf.h"
namespace tensorflow {
namespace tools::proto_splitter {
// A Splitter that can be composed with other splitters.
class ComposableSplitter : public ComposableSplitterBase {
public:
// Initializer.
explicit ComposableSplitter(tsl::protobuf::Message* message)
: ComposableSplitterBase(message), message_(message) {}
// Initialize a child ComposableSplitter.
explicit ComposableSplitter(tsl::protobuf::Message* message,
ComposableSplitterBase* parent_splitter,
std::vector<FieldType>* fields_in_parent)
: ComposableSplitterBase(message, parent_splitter, fields_in_parent),
message_(message) {}
protected:
tsl::protobuf::Message* message() { return message_; }
private:
tsl::protobuf::Message* message_;
};
} // namespace tools::proto_splitter
} // namespace tensorflow
#endif // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_COMPOSABLE_SPLITTER_H_
tensorflow/tools/proto_splitter/cc/composable_splitter_base.cc
+321
View File
@@ -0,0 +1,321 @@
/* 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.
==============================================================================*/
#include "tensorflow/tools/proto_splitter/cc/composable_splitter_base.h"
#include <unistd.h>
#include <cstddef>
#include <cstdint>
#include <deque>
#include <iterator>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include <variant>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/cord.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "riegeli/base/maker.h" // from @riegeli
#include "riegeli/bytes/cord_writer.h" // from @riegeli
#include "riegeli/bytes/fd_writer.h" // from @riegeli
#include "riegeli/bytes/string_writer.h" // from @riegeli
#include "riegeli/records/record_writer.h" // from @riegeli
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/tools/proto_splitter/cc/max_size.h"
#include "tensorflow/tools/proto_splitter/cc/split.h"
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include "tensorflow/tools/proto_splitter/chunk.pb.h"
#include "tsl/platform/env.h"
#include "tsl/platform/errors.h"
#define IS_OSS true
namespace tensorflow {
namespace tools::proto_splitter {
using ::tensorflow::proto_splitter::ChunkMetadata;
VersionDef ComposableSplitterBase::Version() {
VersionDef version;
version.set_splitter_version(1);
version.set_join_version(0);
return version;
}
void ComposableSplitterBase::SetInitialSize(size_t size) { size_ = size; }
size_t ComposableSplitterBase::GetInitialSize() {
if (size_ == 0) {
size_ = message_->ByteSizeLong();
}
return size_;
}
absl::StatusOr<ChunkedProto> ComposableSplitterBase::Split() {
if (parent_splitter_ != nullptr) {
return absl::UnimplementedError(
"The `Split` function behavior for children ComposableSplitter has not "
"been defined. Please call `parent_splitter.Split()` instead.");
}
if (!built_) {
LOG(INFO) << "Splitting message '" << message_->GetDescriptor()->full_name()
<< "' (size " << HumanReadableBytes(GetInitialSize()) << ") "
<< "into chunks of size " << HumanReadableBytes(GetMaxSize());
uint64_t start_time = Env::Default()->NowMicros();
TF_RETURN_IF_ERROR(BuildChunks());
TF_RETURN_IF_ERROR(FixChunks());
uint64_t end_time = Env::Default()->NowMicros();
std::string chunk_msg;
if (chunked_message_.chunked_fields().empty()) {
chunk_msg = "No chunks were generated.";
} else {
chunk_msg = absl::StrCat(
"Generated ", chunked_message_.chunked_fields_size(), " chunks.");
}
LOG(INFO) << "Finished chunking '" << message_->GetDescriptor()->full_name()
<< "', took " << HumanReadableDuration(end_time - start_time)
<< ". " << chunk_msg;
built_ = true;
}
return (ChunkedProto){.chunks = &chunks_,
.chunked_message = &chunked_message_};
}
static absl::Status WriteToRecordWriter(
riegeli::RecordWriterBase& writer, const std::vector<MessageBytes>& chunks,
ChunkedMessage& chunked_message,
const ::tensorflow::proto_splitter::VersionDef& version) {
// Export Riegeli / chunked file.
ChunkMetadata metadata;
*metadata.mutable_message() = chunked_message;
*metadata.mutable_version() = version;
auto* metadata_chunks = metadata.mutable_chunks();
for (const auto& chunk : chunks) {
auto* chunk_metadata = metadata_chunks->Add();
if (std::holds_alternative<std::shared_ptr<tsl::protobuf::Message>>(
chunk)) {
const auto& msg_chunk =
std::get<std::shared_ptr<tsl::protobuf::Message>>(chunk);
LOG(INFO) << "Writing chunk of size " << msg_chunk->ByteSizeLong();
writer.WriteRecord(*msg_chunk);
chunk_metadata->set_size(msg_chunk->ByteSizeLong());
chunk_metadata->set_type(
::tensorflow::proto_splitter::ChunkInfo::MESSAGE);
} else if (std::holds_alternative<tsl::protobuf::Message*>(chunk)) {
auto* msg_chunk = std::get<tsl::protobuf::Message*>(chunk);
writer.WriteRecord(*msg_chunk);
chunk_metadata->set_size(msg_chunk->ByteSizeLong());
chunk_metadata->set_type(
::tensorflow::proto_splitter::ChunkInfo::MESSAGE);
} else {
const auto& str_chunk = std::get<std::string>(chunk);
writer.WriteRecord(str_chunk);
chunk_metadata->set_size(str_chunk.size());
chunk_metadata->set_type(::tensorflow::proto_splitter::ChunkInfo::BYTES);
}
chunk_metadata->set_offset(writer.LastPos().get().numeric());
}
writer.WriteRecord(metadata);
return absl::OkStatus();
}
absl::Status ComposableSplitterBase::CheckIfWriteImplemented() {
if (parent_splitter_ != nullptr) {
return absl::UnimplementedError(
"The `Write` function behavior for children ComposableSplitter has not "
"been defined. Please call `parent_splitter.Write()` instead.");
}
return absl::OkStatus();
}
absl::StatusOr<std::string> ComposableSplitterBase::Write(
std::string file_prefix) {
TF_RETURN_IF_ERROR(CheckIfWriteImplemented());
auto split_results = Split();
if (!split_results.ok()) return split_results.status();
std::vector<MessageBytes>* chunks = split_results.value().chunks;
ChunkedMessage* chunked_message = split_results.value().chunked_message;
tsl::Env* env = tsl::Env::Default();
TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(
std::string{tensorflow::io::Dirname(file_prefix)}));
std::string output_path;
if (chunked_message->chunked_fields().empty()) {
// Export regular pb.
output_path = absl::StrCat(file_prefix, ".pb");
TF_RETURN_IF_ERROR(
tensorflow::WriteBinaryProto(env, output_path, *message_));
} else {
// Export Riegeli / chunked file.
output_path = absl::StrCat(file_prefix, ".cpb");
riegeli::RecordWriter writer(
riegeli::Maker<riegeli::FdWriter>(output_path));
if (!writer.is_open()) return writer.status();
TF_RETURN_IF_ERROR(
WriteToRecordWriter(writer, *chunks, *chunked_message, Version()));
if (!writer.Close()) return writer.status();
}
LOG(INFO) << "Splitter output written to " << output_path;
return output_path;
}
absl::StatusOr<std::tuple<std::string, bool>>
ComposableSplitterBase::WriteToString() {
TF_RETURN_IF_ERROR(CheckIfWriteImplemented());
auto split_results = Split();
if (!split_results.ok()) return split_results.status();
std::vector<MessageBytes>* chunks = split_results.value().chunks;
ChunkedMessage* chunked_message = split_results.value().chunked_message;
std::string output;
if (chunked_message->chunked_fields().empty()) {
// Export regular pb.
if (!message_->SerializeToString(&output))
return absl::InvalidArgumentError("Serialization to string failed");
LOG(INFO) << "Splitter output written to string";
return std::make_tuple(output, false);
} else {
// Export Riegeli / chunked file.
riegeli::RecordWriter writer(
riegeli::Maker<riegeli::StringWriter>(&output));
if (!writer.is_open()) return writer.status();
TF_RETURN_IF_ERROR(
WriteToRecordWriter(writer, *chunks, *chunked_message, Version()));
if (!writer.Close()) return writer.status();
LOG(INFO) << "Splitter output written to string";
return std::make_tuple(output, true);
}
}
#if !IS_OSS
absl::StatusOr<std::tuple<absl::Cord, bool>>
ComposableSplitterBase::WriteToCord() {
TF_RETURN_IF_ERROR(CheckIfWriteImplemented());
auto split_results = Split();
if (!split_results.ok()) return split_results.status();
std::vector<MessageBytes>* chunks = split_results.value().chunks;
ChunkedMessage* chunked_message = split_results.value().chunked_message;
absl::Cord output;
if (chunked_message->chunked_fields().empty()) {
// Export regular pb.
if (!message_->SerializeToString(&output))
return absl::InvalidArgumentError("Serialization to absl::Cord failed");
LOG(INFO) << "Splitter output written to absl::Cord";
return std::make_tuple(output, false);
} else {
// Export Riegeli / chunked file.
riegeli::RecordWriter writer(riegeli::Maker<riegeli::CordWriter>(&output));
if (!writer.is_open()) return writer.status();
TF_RETURN_IF_ERROR(
WriteToRecordWriter(writer, *chunks, *chunked_message, Version()));
if (!writer.Close()) return writer.status();
LOG(INFO) << "Splitter output written to absl::Cord";
return std::make_tuple(output, true);
}
}
#endif
absl::Status ComposableSplitterBase::SetMessageAsBaseChunk() {
if (!chunks_.empty()) {
return absl::FailedPreconditionError(
"Cannot set `message` as the base chunk since there are already "
"created chunks.");
}
chunks_.push_back(message_);
chunked_message_.set_chunk_index(0);
chunks_order_.push_back(0);
add_chunk_order_.push_back(0);
return absl::OkStatus();
}
absl::Status ComposableSplitterBase::AddChunk(
std::unique_ptr<MessageBytes> chunk, std::vector<FieldType>* fields,
int* index) {
if (parent_splitter_ != nullptr) {
std::vector<FieldType> all_fields(fields_in_parent_->begin(),
fields_in_parent_->end());
all_fields.insert(all_fields.end(), fields->begin(), fields->end());
return parent_splitter_->AddChunk(std::move(chunk), &all_fields, index);
}
auto new_chunk_index = chunks_.size();
auto new_field = chunked_message_.add_chunked_fields();
new_field->mutable_message()->set_chunk_index(new_chunk_index);
TF_RETURN_IF_ERROR(
AddFieldTag(*message_->GetDescriptor(), *fields, *new_field));
// Add chunk at the end or insert at the index position.
if (index == nullptr) {
chunks_.push_back(*chunk);
chunks_order_.push_back(new_chunk_index);
} else {
auto it = chunks_.begin();
std::advance(it, *index);
chunks_.insert(it, *chunk);
fix_chunk_order_ = true;
auto it2 = chunks_order_.begin();
std::advance(it2, *index);
chunks_order_.insert(it2, new_chunk_index);
}
add_chunk_order_.push_back(new_chunk_index);
return absl::OkStatus();
}
absl::Status ComposableSplitterBase::FixChunks() {
if (!fix_chunk_order_) return absl::OkStatus();
// Use `add_chunk_order_` and `chunks_order_` to update the chunk indices.
absl::flat_hash_map<int, int> chunk_indices;
for (int i = 0; i < chunks_order_.size(); ++i) {
chunk_indices[chunks_order_[i]] = i;
}
std::deque<ChunkedMessage*> to_fix = {&chunked_message_};
while (!to_fix.empty()) {
auto msg = to_fix.front();
to_fix.pop_front();
for (int i = 0; i < msg->chunked_fields_size(); ++i) {
to_fix.push_back(msg->mutable_chunked_fields(i)->mutable_message());
}
if (!msg->has_chunk_index()) continue;
int current_chunk_idx = msg->chunk_index();
int new_chunk_index = chunk_indices[add_chunk_order_[current_chunk_idx]];
msg->set_chunk_index(new_chunk_index);
}
fix_chunk_order_ = false;
return absl::OkStatus();
}
} // namespace tools::proto_splitter
} // namespace tensorflow
tensorflow/tools/proto_splitter/cc/composable_splitter_base.h
+127
View File
@@ -0,0 +1,127 @@
/* 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_COMPOSABLE_SPLITTER_BASE_H_
#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_COMPOSABLE_SPLITTER_BASE_H_
#include <cstddef>
#include <memory>
#include <string>
#include <tuple>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/cord.h"
#include "tensorflow/tools/proto_splitter/cc/split.h"
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include "tensorflow/tools/proto_splitter/chunk.pb.h"
#include "tsl/platform/protobuf.h"
#define IS_OSS true
namespace tensorflow {
namespace tools::proto_splitter {
// Parts of the ComposableSplitter that are independent of the template (to
// avoid linking issues).
class ComposableSplitterBase : public Splitter {
public:
explicit ComposableSplitterBase(tsl::protobuf::Message* message)
: built_(false), message_(message), parent_splitter_(nullptr) {}
explicit ComposableSplitterBase(tsl::protobuf::Message* message,
ComposableSplitterBase* parent_splitter,
std::vector<FieldType>* fields_in_parent)
: built_(false),
message_(message),
parent_splitter_(parent_splitter),
fields_in_parent_(fields_in_parent) {}
// Splits a proto message.
// Returns a pair of (
// chunks: List of messages/bytes,
// ChunkedMessage: Metadata about the chunked fields.)
// If the message is not split, `chunks` should only contain the original
// message.
absl::StatusOr<ChunkedProto> Split() override;
// Serializes a proto to disk.
// The writer writes all chunks into a Riegeli file. The chunk metadata
// (ChunkMetadata) is written at the very end.
// file_prefix: string prefix of the filepath. The writer will automatically
// attach a `.pb` or `.cpb` (chunked pb) suffix depending on whether the
// proto is split.
// Returns the full output path.
absl::StatusOr<std::string> Write(std::string file_prefix) override;
// The bool field record whether it's saved as a chunked protobuf (true) or
// regular protobuf (false).
absl::StatusOr<std::tuple<std::string, bool>> WriteToString();
#if !IS_OSS
absl::StatusOr<std::tuple<absl::Cord, bool>> WriteToCord();
#endif
VersionDef Version() override;
// Builds the Splitter object by generating chunks from the proto.
// Subclasses of `ComposableChunks` should only need to override this method.
// This method should be called once per Splitter to create the chunks.
// Users should call the methods `Split` or `Write` instead.
virtual absl::Status BuildChunks() = 0;
// Set or get the size of the initial user-provided proto.
void SetInitialSize(size_t size);
size_t GetInitialSize();
protected:
// Initializes `chunks_` and `chunked_message_`, with the user-provided
// message as the initial chunk. This step is optional.
absl::Status SetMessageAsBaseChunk();
// Adds a new chunk and updates the ChunkedMessage proto. If set, the index
// indicates where to insert the chunk.
absl::Status AddChunk(std::unique_ptr<MessageBytes> chunk,
std::vector<FieldType>* fields, int* index = nullptr);
std::vector<FieldType>* fields_in_parent() { return fields_in_parent_; }
private:
// After BuildChunks(), the chunk_index of each nested ChunkedMessage is set
// to the length of the list when the chunk was added. This would be fine if
// the chunks were always added to the end of the list. However, this is not
// always the case the indices must be updated.
absl::Status FixChunks();
absl::Status CheckIfWriteImplemented();
bool built_;
tsl::protobuf::Message* message_;
std::vector<MessageBytes> chunks_;
::tensorflow::proto_splitter::ChunkedMessage chunked_message_;
ComposableSplitterBase* parent_splitter_;
std::vector<FieldType>* fields_in_parent_;
size_t size_ = 0;
// Keep list of chunk keys in the order in which they were added to
// the list and the actual order of chunks. If the user inserts a
// chunk out of order, then these arrays will be used to fix the message chunk
// indices.
std::vector<int> add_chunk_order_;
std::vector<int> chunks_order_;
bool fix_chunk_order_ = false;
};
} // namespace tools::proto_splitter
} // namespace tensorflow
#endif // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_COMPOSABLE_SPLITTER_BASE_H_
tensorflow/tools/proto_splitter/cc/composable_splitter_test.cc
+361
View File
@@ -0,0 +1,361 @@
/* 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.
==============================================================================*/
#include "tensorflow/tools/proto_splitter/cc/composable_splitter.h"
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/strings/cord.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "riegeli/base/maker.h" // from @riegeli
#include "riegeli/bytes/cord_reader.h" // from @riegeli
#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 "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/file_system_helper.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/tools/proto_splitter/cc/test_util.h"
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include "tensorflow/tools/proto_splitter/chunk.pb.h"
#include "tensorflow/tools/proto_splitter/testdata/test_message.pb.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/status_matchers.h"
#include "tsl/platform/statusor.h"
#define IS_OSS true
namespace tensorflow {
namespace tools::proto_splitter {
namespace {
using ::tensorflow::proto_splitter::ChunkedMessage;
using ::tensorflow::proto_splitter::ChunkMetadata;
using ::tensorflow::proto_splitter_testdata::RepeatedRepeatedString;
using ::tensorflow::proto_splitter_testdata::RepeatedString;
using ::testing::HasSubstr;
using ::testing::SizeIs;
using tsl::testing::StatusIs;
// Required in OSS to prevent string to bool conversion in FieldType variant.
using namespace std::string_literals; // NOLINT
// Splits each string in a RepeatedString into separate chunks.
class RepeatedStringSplitter : public ComposableSplitter {
friend class ComposableSplitter;
public:
using ComposableSplitter::ComposableSplitter;
absl::Status BuildChunks() override {
RepeatedString* repeated_string =
google::protobuf::DynamicCastMessage<RepeatedString>(message());
auto strings = repeated_string->strings();
if (strings.empty()) {
TF_RETURN_IF_ERROR(SetMessageAsBaseChunk());
return absl::OkStatus();
}
for (int i = 0; i < strings.size(); i++) {
auto s = std::make_unique<MessageBytes>(strings[i]);
std::vector<FieldType> fields = {"strings"s, i};
TF_RETURN_IF_ERROR(AddChunk(std::move(s), &fields));
}
return absl::OkStatus();
}
};
RepeatedString SetUpRepeatedString(std::vector<std::string> strings) {
RepeatedString message;
*message.mutable_strings() = {strings.begin(), strings.end()};
return message;
}
TEST(RepeatedStringSplitterTest, TestSplitChunks) {
std::vector<std::string> strings = {"piece-1", "piece-2", "piece-3"};
auto message = SetUpRepeatedString(strings);
RepeatedStringSplitter splitter = RepeatedStringSplitter(&message);
TF_ASSERT_OK_AND_ASSIGN(auto ret, splitter.Split());
std::vector<MessageBytes>* chunks = ret.chunks;
ASSERT_NE(chunks, nullptr);
ChunkedMessage* chunked_message = ret.chunked_message;
ASSERT_NE(chunked_message, nullptr);
for (int i = 0; i < chunks->size(); i++) {
MessageBytes chunk = (*chunks)[i];
EXPECT_THAT(chunk, ::testing::VariantWith<std::string>(strings[i]));
}
EXPECT_THAT(*chunked_message, EqualsProto(R"pb(chunked_fields {
field_tag { field: 1 }
field_tag { index: 0 }
message { chunk_index: 0 }
}
chunked_fields {
field_tag { field: 1 }
field_tag { index: 1 }
message { chunk_index: 1 }
}
chunked_fields {
field_tag { field: 1 }
field_tag { index: 2 }
message { chunk_index: 2 }
})pb"));
// Calling split again should return the same chunks/ChunkedMessage.
TF_ASSERT_OK_AND_ASSIGN(auto ret2, splitter.Split());
std::vector<MessageBytes>* chunks2 = ret2.chunks;
ChunkedMessage* chunked_message2 = ret2.chunked_message;
EXPECT_EQ(chunks2, chunks);
EXPECT_EQ(chunked_message2, chunked_message);
}
static void CheckChunks(riegeli::RecordReaderBase& reader,
std::vector<std::string>& strings) {
ChunkMetadata chunk_metadata;
reader.Seek(reader.Size().value());
reader.SeekBack();
reader.ReadRecord(chunk_metadata);
auto& chunk_info = chunk_metadata.chunks();
EXPECT_EQ(chunk_info.size(), strings.size());
for (int i = 0; i < chunk_info.size(); i++) {
reader.Seek(chunk_info[i].offset());
absl::string_view chunk;
reader.ReadRecord(chunk);
EXPECT_EQ(strings[i], std::string(chunk));
}
EXPECT_THAT(chunk_metadata.message(),
EqualsProto(R"pb(chunked_fields {
field_tag { field: 1 }
field_tag { index: 0 }
message { chunk_index: 0 }
}
chunked_fields {
field_tag { field: 1 }
field_tag { index: 1 }
message { chunk_index: 1 }
}
chunked_fields {
field_tag { field: 1 }
field_tag { index: 2 }
message { chunk_index: 2 }
})pb"));
}
TEST(RepeatedStringSplitterTest, TestWrite) {
std::vector<std::string> strings = {"piece-1", "piece-2", "piece-3"};
auto message = SetUpRepeatedString(strings);
RepeatedStringSplitter splitter = RepeatedStringSplitter(&message);
std::string output_prefix = tensorflow::io::GetTempFilename("");
TF_ASSERT_OK_AND_ASSIGN(std::string expected_file,
splitter.Write(output_prefix));
TF_ASSERT_OK_AND_ASSIGN(auto exists,
internal::FileExists(Env::Default(), expected_file));
EXPECT_TRUE(exists);
// Look for the last chunk, which should contain a ChunkMetadata proto.
riegeli::RecordReader file_reader(
riegeli::Maker<riegeli::FdReader>(std::move(expected_file)));
CheckChunks(file_reader, strings);
}
TEST(RepeatedStringSplitterTest, TestWriteToString) {
std::vector<std::string> strings = {"piece-1", "piece-2", "piece-3"};
auto message = SetUpRepeatedString(strings);
RepeatedStringSplitter splitter = RepeatedStringSplitter(&message);
auto string_output_results = splitter.WriteToString();
TF_EXPECT_OK(string_output_results.status());
std::string string_output = std::get<0>(string_output_results.value());
bool is_chunked = std::get<1>(string_output_results.value());
EXPECT_TRUE(is_chunked);
// Look for the last chunk, which should contain a ChunkMetadata proto.
riegeli::RecordReader string_reader(
riegeli::Maker<riegeli::StringReader>(string_output));
CheckChunks(string_reader, strings);
}
#if !IS_OSS
TEST(RepeatedStringSplitterTest, TestWriteToCord) {
std::vector<std::string> strings = {"piece-1", "piece-2", "piece-3"};
auto message = SetUpRepeatedString(strings);
RepeatedStringSplitter splitter = RepeatedStringSplitter(&message);
auto cord_output_results = splitter.WriteToCord();
TF_EXPECT_OK(cord_output_results.status());
absl::Cord cord_output = std::get<0>(cord_output_results.value());
bool is_chunked = std::get<1>(cord_output_results.value());
EXPECT_TRUE(is_chunked);
// Look for the last chunk, which should contain a ChunkMetadata proto.
riegeli::RecordReader cord_reader(
riegeli::Maker<riegeli::CordReader>(&cord_output));
CheckChunks(cord_reader, strings);
}
#endif
TEST(RepeatedStringSplitterTest, TestNoSplit) {
RepeatedString message; // No strings
RepeatedStringSplitter splitter = RepeatedStringSplitter(&message);
TF_ASSERT_OK_AND_ASSIGN(auto ret, splitter.Split());
std::vector<MessageBytes>* chunks = ret.chunks;
ASSERT_NE(chunks, nullptr);
ChunkedMessage* chunked_message = ret.chunked_message;
ASSERT_NE(chunked_message, nullptr);
EXPECT_THAT(*chunks, SizeIs(1));
EXPECT_THAT(*std::get<tsl::protobuf::Message*>((*chunks)[0]),
EqualsProto(""));
EXPECT_THAT(*chunked_message, EqualsProto(R"pb(chunk_index: 0)pb"));
}
// Splits each string in a RepeatedString into separate chunks.
class RepeatedRepeatedStringSplitter : public ComposableSplitter {
public:
using ComposableSplitter::ComposableSplitter;
absl::Status BuildChunks() override {
TF_RETURN_IF_ERROR(SetMessageAsBaseChunk());
RepeatedRepeatedString* msg =
google::protobuf::DynamicCastMessage<RepeatedRepeatedString>(message());
auto repeated_strings = msg->rs();
for (int i = 0; i < repeated_strings.size(); i++) {
std::vector<FieldType> fields = {"rs"s, i};
auto splitter =
RepeatedStringSplitter(&repeated_strings[i], this, &fields);
TF_RETURN_IF_ERROR(splitter.BuildChunks());
}
return absl::OkStatus();
}
};
TEST(ComposableTest, RepeatedRepeatedStringTest) {
std::vector<std::string> strings1 = {"piece-1", "piece-2", "piece-3"};
auto rs1 = SetUpRepeatedString(strings1);
std::vector<std::string> strings2 = {"new-strings-1"};
auto rs2 = SetUpRepeatedString(strings2);
std::vector<std::string> strings3 = {"foo-1", "foo-2"};
auto rs3 = SetUpRepeatedString(strings3);
std::vector<RepeatedString> rs = {rs1, rs2, rs3};
RepeatedRepeatedString message;
message.mutable_rs()->Add(rs.begin(), rs.end());
RepeatedRepeatedStringSplitter splitter =
RepeatedRepeatedStringSplitter(&message);
TF_ASSERT_OK_AND_ASSIGN(auto ret, splitter.Split());
std::vector<MessageBytes>* chunks = ret.chunks;
ASSERT_NE(chunks, nullptr);
ChunkedMessage* chunked_message = ret.chunked_message;
ASSERT_NE(chunked_message, nullptr);
std::vector<std::string> expected_chunks = {
"piece-1", "piece-2", "piece-3", "new-strings-1", "foo-1", "foo-2"};
// RepeatedRepeatedStringSplitter sets the first chunk as the user-provided
// message, so the expected size is 7.
EXPECT_THAT(*chunks, SizeIs(7));
EXPECT_THAT(*std::get<tsl::protobuf::Message*>((*chunks)[0]),
EqualsProto(message));
for (int i = 1; i < chunks->size(); i++) {
MessageBytes chunk = (*chunks)[i];
EXPECT_THAT(chunk,
::testing::VariantWith<std::string>(expected_chunks[i - 1]));
}
// message.rs[2].strings[0] (value = "foo-1") should be the chunk at index 5.
EXPECT_THAT(chunked_message->chunked_fields()[4],
EqualsProto(R"pb(field_tag { field: 2 }
field_tag { index: 2 }
field_tag { field: 1 }
field_tag { index: 0 }
message { chunk_index: 5 })pb"));
}
TEST(ComposableTest, ChildSplitterTest) {
std::vector<std::string> strings1 = {"piece-1", "piece-2", "piece-3"};
auto message1 = SetUpRepeatedString(strings1);
RepeatedStringSplitter splitter(&message1);
std::vector<FieldType> fields = {};
std::vector<std::string> strings2 = {"s1", "s2"};
auto message2 = SetUpRepeatedString(strings2);
RepeatedStringSplitter child(&message2, &splitter, &fields);
TF_EXPECT_OK(child.BuildChunks());
TF_ASSERT_OK_AND_ASSIGN(auto ret, splitter.Split());
std::vector<MessageBytes>* chunks = ret.chunks;
ASSERT_NE(chunks, nullptr);
EXPECT_THAT(*chunks, SizeIs(5)); // Total 5 chunks should be generated.
}
TEST(ComposableTest, ChildSplitterUnimplementedTest) {
RepeatedString message;
RepeatedStringSplitter splitter(&message);
std::vector<FieldType> fields = {};
RepeatedStringSplitter child(&message, &splitter, &fields);
EXPECT_THAT(child.Split(),
absl_testing::StatusIs(absl::StatusCode::kUnimplemented,
HasSubstr("`Split` function behavior")));
EXPECT_THAT(child.Write("str"),
absl_testing::StatusIs(absl::StatusCode::kUnimplemented,
HasSubstr("`Write` function behavior")));
}
class NoOpSplitter : public ComposableSplitter {
public:
using ComposableSplitter::ComposableSplitter;
absl::Status BuildChunks() override { return absl::OkStatus(); }
};
TEST(NoOpSplitterTest, TestWrite) {
std::vector<std::string> strings = {"piece-1", "piece-2", "piece-3"};
auto message = SetUpRepeatedString(strings);
NoOpSplitter splitter(&message);
std::string output_prefix = tensorflow::io::GetTempFilename("");
TF_ASSERT_OK_AND_ASSIGN(std::string expected_file,
splitter.Write(output_prefix));
TF_ASSERT_OK_AND_ASSIGN(auto exists,
internal::FileExists(Env::Default(), expected_file));
EXPECT_TRUE(exists);
RepeatedString read_message;
auto status = tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
expected_file, &read_message);
EXPECT_THAT(read_message, EqualsProto(message));
}
} // namespace
} // namespace tools::proto_splitter
} // namespace tensorflow
tensorflow/tools/proto_splitter/cc/graph_def_splitter.cc
+241
View File
@@ -0,0 +1,241 @@
/* 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.
==============================================================================*/
#include "tensorflow/tools/proto_splitter/cc/graph_def_splitter.h"
#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/tools/proto_splitter/cc/composable_splitter.h"
#include "tensorflow/tools/proto_splitter/cc/composable_splitter_base.h"
#include "tensorflow/tools/proto_splitter/cc/large_node_splitter.h"
#include "tensorflow/tools/proto_splitter/cc/max_size.h"
#include "tensorflow/tools/proto_splitter/cc/repeated_field_splitter.h"
#include "tensorflow/tools/proto_splitter/cc/size_splitter.h"
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/protobuf.h"
#include "tsl/platform/statusor.h"
namespace tensorflow::tools::proto_splitter {
namespace {
// Required in OSS to prevent string to bool conversion in FieldType variant.
using namespace std::string_literals; // NOLINT
// Extracts the value of large constants into a new chunk.
// Only works on NodeDef messages.
class ConstantSplitter : public SizeSplitter {
public:
using SizeSplitter::SizeSplitter;
absl::StatusOr<int> BuildChunksReturnSize() override {
NodeDef* node = google::protobuf::DynamicCastMessage<NodeDef>(message());
std::vector<FieldType> tensor_field = {"attr"s, "value"s, "tensor"s};
std::vector<FieldType> content_field = {"attr"s, "value"s, "tensor"s,
"tensor_content"s};
TF_ASSIGN_OR_RETURN(auto ret, GetMutableField(node, tensor_field));
auto tensor_msg =
ret.parent->GetReflection()->MutableMessage(ret.parent, ret.field);
TensorProto* tensor_proto =
google::protobuf::DynamicCastMessage<TensorProto>(tensor_msg);
int size_diff;
if (tensor_proto->tensor_content().empty()) {
Tensor t;
if (!t.FromProto(*tensor_proto)) {
return absl::InvalidArgumentError(
"Invalid Const NodeDef.attr[\"value\"].tensor value.");
}
TensorProto container;
t.AsProtoTensorContent(&container);
size_diff = container.tensor_content().size();
auto x = std::make_unique<std::string>(
std::move(*container.mutable_tensor_content()));
auto y = std::make_unique<MessageBytes>(std::move(*x));
TF_RETURN_IF_ERROR(AddChunk(std::move(y), &content_field));
} else {
size_diff = tensor_proto->tensor_content().size();
auto x = std::make_unique<std::string>(
std::move(*tensor_proto->mutable_tensor_content()));
auto y = std::make_unique<MessageBytes>(std::move(*x));
TF_RETURN_IF_ERROR(AddChunk(std::move(y), &content_field));
}
// Keep the TensorProto's dtype, tensor_shape, and version_number fields,
// but clear the raw tensor content / "xxx_val" attributes.
auto dtype = tensor_proto->dtype();
auto tensor_shape = tensor_proto->tensor_shape();
auto version_number = tensor_proto->version_number();
tensor_proto->Clear();
tensor_proto->set_dtype(dtype);
*tensor_proto->mutable_tensor_shape() = tensor_shape;
tensor_proto->set_version_number(version_number);
return size_diff;
}
};
class ConstantSplitterFactory : public SizeSplitterFactory {
public:
using SizeSplitterFactory::SizeSplitterFactory;
absl::StatusOr<std::unique_ptr<SizeSplitter>> CreateSplitter(
tsl::protobuf::Message* message, ComposableSplitterBase* parent_splitter,
std::vector<FieldType>* fields_in_parent, int size) override {
if (size < GetMaxSize()) return nullptr;
NodeDef* node = google::protobuf::DynamicCastMessage<NodeDef>(message);
if (node->op() != "Const")
return absl::UnimplementedError(absl::StrCat(
"Currently only able to split 'Const' nodes that are larger than the "
"2GB maximum proto size. Got node of type '",
node->op(), "' with size: ", size, "."));
ConstantSplitter* splitter =
new ConstantSplitter(message, parent_splitter, fields_in_parent);
return absl::WrapUnique(splitter);
}
};
class FunctionDefSplitter : public SizeSplitter {
public:
using SizeSplitter::SizeSplitter;
absl::StatusOr<int> BuildChunksReturnSize() override {
size_t current_size = GetInitialSize();
uint64_t max_size = GetMaxSize();
std::vector<FieldType> fields = {};
// First check if the entire FunctionDef can be split into a separate chunk.
// We do this before the `RepeatedMessageSplitter`, which is costly because
// it iterates through every `node_def`.
if (LARGE_SIZE_CHECK(current_size, max_size) && current_size < max_size) {
auto splitter = LargeNodeSplitter<FunctionDef>(message(), this, &fields);
splitter.SetInitialSize(current_size);
return splitter.BuildChunksReturnSize();
} else if (current_size > max_size) {
ConstantSplitterFactory constant_splitter_factory;
LargeNodeSplitterFactory<NodeDef> large_node_splitter_factory;
std::vector<SizeSplitterFactory*> factories = {
&constant_splitter_factory, &large_node_splitter_factory};
auto ret = RepeatedFieldSplitter<FunctionDef, NodeDef>::Create(
message(), this, &fields, "node_def"s, &factories);
if (!ret.ok()) return ret.status();
auto splitter = ret.value();
return splitter.BuildChunksReturnSize();
}
return 0;
}
};
class FunctionDefSplitterFactory : public SizeSplitterFactory {
public:
using SizeSplitterFactory::SizeSplitterFactory;
absl::StatusOr<std::unique_ptr<SizeSplitter>> CreateSplitter(
tsl::protobuf::Message* message, ComposableSplitterBase* parent_splitter,
std::vector<FieldType>* fields_in_parent, int size) override {
FunctionDefSplitter* splitter =
new FunctionDefSplitter(message, parent_splitter, fields_in_parent);
return absl::WrapUnique(splitter);
}
};
} // namespace
absl::Status GraphDefSplitter::BuildChunks() {
TF_RETURN_IF_ERROR(SetMessageAsBaseChunk());
GraphDef* g = google::protobuf::DynamicCastMessage<GraphDef>(message());
uint64_t max_size = GetMaxSize();
size_t graph_size = GetInitialSize();
if (graph_size < max_size) return absl::OkStatus();
// Set up GraphDef.node and GraphDef.library.function splitters.
std::vector<FieldType> field_in_parent = {};
ConstantSplitterFactory constant_splitter_factory;
LargeNodeSplitterFactory<NodeDef> large_node_splitter_factory;
std::vector<SizeSplitterFactory*> factories = {&constant_splitter_factory,
&large_node_splitter_factory};
auto node_splitter_ret = RepeatedFieldSplitter<GraphDef, NodeDef>::Create(
g, this, &field_in_parent, "node"s, &factories);
if (!node_splitter_ret.ok()) return node_splitter_ret.status();
auto node_splitter = node_splitter_ret.value();
FunctionDefSplitterFactory function_splitter_factory;
std::vector<FieldType> library_field = {"library"s};
std::vector<SizeSplitterFactory*> fn_factories = {&function_splitter_factory};
auto library_splitter_ret =
RepeatedFieldSplitter<FunctionDefLibrary, FunctionDef>::Create(
g->mutable_library(), this, &library_field, "function"s,
&fn_factories);
if (!library_splitter_ret.ok()) return library_splitter_ret.status();
auto library_splitter = library_splitter_ret.value();
size_t library_size = g->library().ByteSizeLong();
library_splitter.SetInitialSize(library_size);
size_t approx_node_size = graph_size - library_size;
node_splitter.SetInitialSize(approx_node_size);
// Call node and library splitters.
if (library_size > approx_node_size) {
TF_ASSIGN_OR_RETURN(int size_diff,
library_splitter.BuildChunksReturnSize());
library_size -= size_diff;
if (approx_node_size + library_size > max_size) {
TF_ASSIGN_OR_RETURN(int size_diff, node_splitter.BuildChunksReturnSize());
approx_node_size -= size_diff;
}
} else {
TF_ASSIGN_OR_RETURN(int size_diff, node_splitter.BuildChunksReturnSize());
approx_node_size -= size_diff;
if (approx_node_size + library_size > max_size) {
TF_ASSIGN_OR_RETURN(int size_diff,
library_splitter.BuildChunksReturnSize());
library_size -= size_diff;
}
}
// Recompute graph size. If the graph size is still greater than the max size,
// separate the entire library into a separate chunk.
if (g->ByteSizeLong() > max_size) {
LargeNodeSplitter<FunctionDefLibrary> entire_library_splitter(
g->mutable_library(), this, &library_field);
// The library chunk must be inserted before the function chunks generated
// by `library_splitter` which may have the same field tags.
int index = 1;
entire_library_splitter.SetChunkIndex(&index);
TF_RETURN_IF_ERROR(entire_library_splitter.BuildChunks());
}
return absl::OkStatus();
}
} // namespace tensorflow::tools::proto_splitter
tensorflow/tools/proto_splitter/cc/graph_def_splitter.h
+35
View File
@@ -0,0 +1,35 @@
/* 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_GRAPH_DEF_SPLITTER_H_
#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_GRAPH_DEF_SPLITTER_H_
#include "tensorflow/tools/proto_splitter/cc/composable_splitter.h"
namespace tensorflow {
namespace tools::proto_splitter {
// GraphDef Splitter.
// Modifies the user-provided message when building chunks.
class GraphDefSplitter : public ComposableSplitter {
public:
// Initializer.
using ComposableSplitter::ComposableSplitter;
absl::Status BuildChunks() override;
};
} // namespace tools::proto_splitter
} // namespace tensorflow
#endif // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_GRAPH_DEF_SPLITTER_H_
tensorflow/tools/proto_splitter/cc/graph_def_splitter_test.cc
+299
View File
@@ -0,0 +1,299 @@
/* 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.
==============================================================================*/
#include "tensorflow/tools/proto_splitter/cc/graph_def_splitter.h"
#include <cstdint>
#include <memory>
#include <string>
#include <variant>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/strings/cord.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/tools/proto_splitter/cc/max_size.h"
#include "tensorflow/tools/proto_splitter/cc/test_util.h"
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include "tensorflow/tools/proto_splitter/testdata/test_message.pb.h"
#include "tsl/platform/protobuf.h"
#include "tsl/platform/statusor.h"
namespace tensorflow {
namespace tools::proto_splitter {
namespace {
using ::tensorflow::proto_splitter::ChunkedMessage;
TEST(GraphDefSplitterTest, TestLargeConstant) {
GraphDef proto;
const std::string graph_def_path =
io::JoinPath(testing::TensorFlowSrcRoot(),
"tools/proto_splitter/testdata", "split-large-constant.pb");
int64_t max_size = 500;
DebugSetMaxSize(max_size);
TF_EXPECT_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
graph_def_path, &proto));
EXPECT_GE(proto.ByteSizeLong(), GetMaxSize());
std::string large_constant_1, large_constant_2;
const std::variant<std::string, absl::Cord>& tensor_constant_1 =
proto.node(2).attr().at("value").tensor().tensor_content();
const std::variant<std::string, absl::Cord>& tensor_constant_2 =
proto.node(4).attr().at("value").tensor().tensor_content();
if (std::holds_alternative<std::string>(tensor_constant_1)) {
large_constant_1 = std::get<std::string>(tensor_constant_1);
} else {
absl::CopyCordToString(std::get<absl::Cord>(tensor_constant_1),
&large_constant_1);
}
if (std::holds_alternative<std::string>(tensor_constant_2)) {
large_constant_2 = std::get<std::string>(tensor_constant_2);
} else {
absl::CopyCordToString(std::get<absl::Cord>(tensor_constant_2),
&large_constant_2);
}
GraphDefSplitter splitter(&proto);
TF_ASSERT_OK_AND_ASSIGN(auto x, splitter.Split());
ChunkedMessage* chunked_message = x.chunked_message;
ASSERT_NE(chunked_message, nullptr);
EXPECT_THAT(*chunked_message,
EqualsProto(R"pb(chunk_index: 0
chunked_fields {
field_tag { field: 1 }
field_tag { index: 2 }
field_tag { field: 5 }
field_tag { map_key { s: "value" } }
field_tag { field: 8 }
field_tag { field: 4 }
message { chunk_index: 1 }
}
chunked_fields {
field_tag { field: 1 }
field_tag { index: 4 }
field_tag { field: 5 }
field_tag { map_key { s: "value" } }
field_tag { field: 8 }
field_tag { field: 4 }
message { chunk_index: 2 }
})pb"));
std::vector<MessageBytes>* chunks = x.chunks;
ASSERT_NE(chunks, nullptr);
EXPECT_CHUNK_SIZES(chunks, max_size);
EXPECT_THAT((*chunks)[1],
::testing::VariantWith<std::string>(large_constant_1));
EXPECT_THAT((*chunks)[2],
::testing::VariantWith<std::string>(large_constant_2));
}
TEST(GraphDefSplitterTest, TestLargeNodes) {
GraphDef proto;
const std::string graph_def_path =
io::JoinPath(testing::TensorFlowSrcRoot(),
"tools/proto_splitter/testdata", "split-large-nodes.pb");
int64_t max_size = 200;
DebugSetMaxSize(max_size);
TF_EXPECT_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
graph_def_path, &proto));
EXPECT_GE(proto.ByteSize(), GetMaxSize());
// Get the large nodes that should be extracted for later comparison.
NodeDef node_1 = proto.node(1);
NodeDef node_2 = proto.node(2);
NodeDef node_3 = proto.node(3);
NodeDef node_5 = proto.node(5);
GraphDefSplitter splitter(&proto);
TF_ASSERT_OK_AND_ASSIGN(auto x, splitter.Split());
ChunkedMessage* chunked_message = x.chunked_message;
ASSERT_NE(chunked_message, nullptr);
EXPECT_THAT(*chunked_message, EqualsProto(R"pb(chunk_index: 0
chunked_fields {
field_tag { field: 1 }
field_tag { index: 1 }
message { chunk_index: 1 }
}
chunked_fields {
field_tag { field: 1 }
field_tag { index: 2 }
message { chunk_index: 2 }
}
chunked_fields {
field_tag { field: 1 }
field_tag { index: 3 }
message { chunk_index: 3 }
}
chunked_fields {
field_tag { field: 1 }
field_tag { index: 5 }
message { chunk_index: 4 }
})pb"));
std::vector<MessageBytes>* chunks = x.chunks;
ASSERT_NE(chunks, nullptr);
EXPECT_CHUNK_SIZES(chunks, max_size);
EXPECT_TRUE(std::holds_alternative<std::shared_ptr<tsl::protobuf::Message>>(
(*chunks)[1]));
EXPECT_TRUE(std::holds_alternative<std::shared_ptr<tsl::protobuf::Message>>(
(*chunks)[2]));
EXPECT_TRUE(std::holds_alternative<std::shared_ptr<tsl::protobuf::Message>>(
(*chunks)[3]));
EXPECT_TRUE(std::holds_alternative<std::shared_ptr<tsl::protobuf::Message>>(
(*chunks)[4]));
EXPECT_THAT(
*std::get<std::shared_ptr<tsl::protobuf::Message>>((*chunks)[1]).get(),
EqualsProto(node_1));
EXPECT_THAT(
*std::get<std::shared_ptr<tsl::protobuf::Message>>((*chunks)[2]).get(),
EqualsProto(node_2));
EXPECT_THAT(
*std::get<std::shared_ptr<tsl::protobuf::Message>>((*chunks)[3]).get(),
EqualsProto(node_3));
EXPECT_THAT(
*std::get<std::shared_ptr<tsl::protobuf::Message>>((*chunks)[4]).get(),
EqualsProto(node_5));
}
TEST(GraphDefSplitterTest, TestLotsNodes) {
GraphDef proto;
const std::string graph_def_path =
io::JoinPath(testing::TensorFlowSrcRoot(),
"tools/proto_splitter/testdata", "split-lots-nodes.pb");
// split-lots-nodes.pb has 15 nodes that are 95 or 96 bytes each. The max size
// is set to "exactly" the size of 5 nodes, but with the extra encoding bytes,
// only 4 nodes should fit in each chunk. Thus, there should be exactly 4
// chunks created for all 15 nodes.
int64_t max_size = 96 * 5;
DebugSetMaxSize(max_size);
TF_EXPECT_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
graph_def_path, &proto));
EXPECT_GE(proto.ByteSize(), GetMaxSize());
int expected_node_size = proto.node_size();
GraphDefSplitter splitter(&proto);
TF_ASSERT_OK_AND_ASSIGN(auto x, splitter.Split());
ChunkedMessage* chunked_message = x.chunked_message;
ASSERT_NE(chunked_message, nullptr);
EXPECT_THAT(
*chunked_message,
EqualsProto(R"pb(chunk_index: 0
chunked_fields { message { chunk_index: 1 } }
chunked_fields { message { chunk_index: 2 } }
chunked_fields { message { chunk_index: 3 } }
chunked_fields { message { chunk_index: 4 } })pb"));
std::vector<MessageBytes>* chunks = x.chunks;
ASSERT_NE(chunks, nullptr);
EXPECT_CHUNK_SIZES(chunks, max_size);
int actual_node_size = 0;
for (MessageBytes& chunk : *chunks) {
GraphDef* message = nullptr;
if (std::holds_alternative<std::shared_ptr<tsl::protobuf::Message>>(
chunk)) {
message = google::protobuf::DynamicCastMessage<GraphDef>(
std::get<std::shared_ptr<tsl::protobuf::Message>>(chunk).get());
} else if (std::holds_alternative<tsl::protobuf::Message*>(chunk)) {
message = google::protobuf::DynamicCastMessage<GraphDef>(
std::get<tsl::protobuf::Message*>(chunk));
} else {
EXPECT_FALSE(std::holds_alternative<std::string>(chunk));
}
actual_node_size += message->node_size();
}
EXPECT_EQ(actual_node_size, expected_node_size);
}
TEST(GraphDefSplitterTest, TestFunctionLotsOfNodes) {
GraphDef proto;
const std::string graph_def_path = io::JoinPath(
testing::TensorFlowSrcRoot(), "tools/proto_splitter/testdata",
"function-lots-of-nodes.pb");
int64_t max_size = 500;
DebugSetMaxSize(max_size);
TF_EXPECT_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
graph_def_path, &proto));
EXPECT_GE(proto.ByteSize(), GetMaxSize());
GraphDefSplitter splitter(&proto);
TF_ASSERT_OK_AND_ASSIGN(auto x, splitter.Split());
std::vector<MessageBytes>* chunks = x.chunks;
ASSERT_NE(chunks, nullptr);
EXPECT_CHUNK_SIZES(chunks, max_size);
}
TEST(GraphDefSplitterTest, TestFunctionLargeNodes) {
GraphDef proto;
const std::string graph_def_path =
io::JoinPath(testing::TensorFlowSrcRoot(),
"tools/proto_splitter/testdata", "function-large-nodes.pb");
int64_t max_size = 200;
DebugSetMaxSize(max_size);
TF_EXPECT_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
graph_def_path, &proto));
EXPECT_GE(proto.ByteSize(), GetMaxSize());
GraphDefSplitter splitter(&proto);
TF_ASSERT_OK_AND_ASSIGN(auto x, splitter.Split());
std::vector<MessageBytes>* chunks = x.chunks;
ASSERT_NE(chunks, nullptr);
EXPECT_CHUNK_SIZES(chunks, max_size);
}
TEST(GraphDefSplitterTest, TestGraphAndFunction) {
GraphDef proto;
const std::string graph_def_path = io::JoinPath(
testing::TensorFlowSrcRoot(), "tools/proto_splitter/testdata",
"graph-def-and-function.pb");
int64_t max_size = 200;
DebugSetMaxSize(max_size);
TF_EXPECT_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
graph_def_path, &proto));
EXPECT_GE(proto.ByteSize(), GetMaxSize());
GraphDefSplitter splitter(&proto);
TF_ASSERT_OK_AND_ASSIGN(auto x, splitter.Split());
std::vector<MessageBytes>* chunks = x.chunks;
ASSERT_NE(chunks, nullptr);
EXPECT_CHUNK_SIZES(chunks, max_size);
TF_ASSERT_OK(splitter.Write("/tmp/hoi"));
}
} // namespace
} // namespace tools::proto_splitter
} // namespace tensorflow
tensorflow/tools/proto_splitter/cc/large_node_splitter.h
+84
View File
@@ -0,0 +1,84 @@
/* 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_LARGE_NODE_SPLITTER_H_
#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_LARGE_NODE_SPLITTER_H_
#include <memory>
#include <vector>
#include "absl/status/statusor.h"
#include "tensorflow/tools/proto_splitter/cc/composable_splitter.h"
#include "tensorflow/tools/proto_splitter/cc/composable_splitter_base.h"
#include "tensorflow/tools/proto_splitter/cc/max_size.h"
#include "tensorflow/tools/proto_splitter/cc/size_splitter.h"
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include "tsl/platform/errors.h"
namespace tensorflow {
namespace tools::proto_splitter {
// Extracts messages that are large but not over the limit into a separate
// message chunk.
template <typename MessageType>
class LargeNodeSplitter : public SizeSplitter {
public:
using SizeSplitter::SizeSplitter;
void SetChunkIndex(int* index) { index_ = index; }
absl::StatusOr<int> BuildChunksReturnSize() override;
private:
int* index_ = nullptr;
};
template <typename MessageType>
absl::StatusOr<int> LargeNodeSplitter<MessageType>::BuildChunksReturnSize() {
MessageType* msg =
tsl::protobuf::DynamicCastToGenerated<MessageType>(message());
int initial_size = GetInitialSize();
std::shared_ptr<MessageType> new_msg = std::make_shared<MessageType>();
msg->Swap(new_msg.get());
std::vector<FieldType> fields = {};
auto x = std::make_unique<MessageBytes>(new_msg);
TF_RETURN_IF_ERROR(AddChunk(std::move(x), &fields, index_));
return initial_size;
}
template <typename MessageType>
class LargeNodeSplitterFactory : public SizeSplitterFactory {
public:
using SizeSplitterFactory::SizeSplitterFactory;
absl::StatusOr<std::unique_ptr<SizeSplitter>> CreateSplitter(
tsl::protobuf::Message* message, ComposableSplitterBase* parent_splitter,
std::vector<FieldType>* fields_in_parent, int size) override;
};
template <typename MessageType>
absl::StatusOr<std::unique_ptr<SizeSplitter>>
LargeNodeSplitterFactory<MessageType>::CreateSplitter(
tsl::protobuf::Message* message, ComposableSplitterBase* parent_splitter,
std::vector<FieldType>* fields_in_parent, int size) {
if (!(LARGE_SIZE_CHECK(size, GetMaxSize()))) return nullptr;
LargeNodeSplitter<MessageType>* splitter = new LargeNodeSplitter<MessageType>(
message, parent_splitter, fields_in_parent);
return absl::WrapUnique(splitter);
}
} // namespace tools::proto_splitter
} // namespace tensorflow
#endif // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_LARGE_NODE_SPLITTER_H_
tensorflow/tools/proto_splitter/cc/max_size.cc
+32
View File
@@ -0,0 +1,32 @@
/* 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.
==============================================================================*/
#include "tensorflow/tools/proto_splitter/cc/max_size.h"
#include "absl/synchronization/mutex.h"
namespace tensorflow {
namespace tools::proto_splitter {
ABSL_CONST_INIT absl::Mutex global_mutex(absl::kConstInit);
// The default max size is set to a bit less than 2GB, since the proto splitter
// isn't extremely precise.
uint64_t ProtoMaxSize = ((uint64_t)1 << 31) - 500;
uint64_t GetMaxSize() { return ProtoMaxSize; }
void DebugSetMaxSize(uint64_t size) { ProtoMaxSize = size; }
} // namespace tools::proto_splitter
} // namespace tensorflow
tensorflow/tools/proto_splitter/cc/max_size.h
+34
View File
@@ -0,0 +1,34 @@
/* 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_MAX_SIZE_H_
#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_MAX_SIZE_H_
#include <cstdint>
namespace tensorflow {
namespace tools::proto_splitter {
// Get max size allowed by each chunk in the proto splitter.
uint64_t GetMaxSize();
// Set the max size. Should only be used for testing purposes.
void DebugSetMaxSize(uint64_t size);
// Heuristic for determine whether to greedily generate a chunk.
#define LARGE_SIZE_CHECK(size, max_size) size > max_size / 3
} // namespace tools::proto_splitter
} // namespace tensorflow
#endif // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_MAX_SIZE_H_
tensorflow/tools/proto_splitter/cc/repeated_field_splitter.h
+169
View File
@@ -0,0 +1,169 @@
/* 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_REPEATED_FIELD_SPLITTER_H_
#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_REPEATED_FIELD_SPLITTER_H_
#include <cstddef>
#include <cstdint>
#include <memory>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/tools/proto_splitter/cc/composable_splitter.h"
#include "tensorflow/tools/proto_splitter/cc/max_size.h"
#include "tensorflow/tools/proto_splitter/cc/size_splitter.h"
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include "tsl/platform/protobuf.h"
namespace tensorflow::tools::proto_splitter {
// Splitter that works on repeated message fields.
template <typename ParentMessage, typename RepeatedMessage>
class RepeatedFieldSplitter : public SizeSplitter {
public:
static absl::StatusOr<RepeatedFieldSplitter> Create(
tsl::protobuf::Message* message, ComposableSplitter* parent_splitter,
std::vector<FieldType>* fields_in_parent, const FieldType& repeated_field,
std::vector<SizeSplitterFactory*>* splitter_factories);
absl::StatusOr<int> BuildChunksReturnSize() override;
private:
explicit RepeatedFieldSplitter(
tsl::protobuf::Message* message, ComposableSplitter* parent_splitter,
std::vector<FieldType>* fields_in_parent, const FieldType& repeated_field,
std::vector<SizeSplitterFactory*>* splitter_factories)
: SizeSplitter(message, parent_splitter, fields_in_parent),
repeated_field_(repeated_field),
splitter_factories_(splitter_factories) {}
FieldType repeated_field_;
std::vector<SizeSplitterFactory*>* splitter_factories_;
};
// Additional bytes added to each node to account for the extra info needed to
// encode the field key (realistically 3 but making it 5 for some wiggle room).
constexpr int kExtraBytes = 5;
template <typename ParentMessage, typename RepeatedMessage>
absl::StatusOr<RepeatedFieldSplitter<ParentMessage, RepeatedMessage>>
RepeatedFieldSplitter<ParentMessage, RepeatedMessage>::Create(
tsl::protobuf::Message* message, ComposableSplitter* parent_splitter,
std::vector<FieldType>* fields_in_parent, const FieldType& repeated_field,
std::vector<SizeSplitterFactory*>* splitter_factories) {
TF_ASSIGN_OR_RETURN(auto field_ret, GetField(*message, {repeated_field}));
if (!field_ret.field->is_repeated()) {
return absl::FailedPreconditionError("Unable to split non-repeated field.");
}
auto ret = RepeatedFieldSplitter<ParentMessage, RepeatedMessage>(
message, parent_splitter, fields_in_parent, repeated_field,
splitter_factories);
return ret;
}
template <typename ParentMessage, typename RepeatedMessage>
absl::StatusOr<int>
RepeatedFieldSplitter<ParentMessage, RepeatedMessage>::BuildChunksReturnSize() {
TF_ASSIGN_OR_RETURN(MutableFieldResult mfr,
GetMutableField(message(), {repeated_field_}));
tsl::protobuf::Message* parent = mfr.parent;
const tsl::protobuf::FieldDescriptor* repeated_field = mfr.field;
uint64_t max_size = GetMaxSize();
size_t initial_size = GetInitialSize();
// List of indices at which to split the repeated field. For example, [3, 5]
// means that the field list is split into: [:3], [3:5], [5:]
std::vector<int> repeated_msg_split;
// Track the total byte size of the current node split.
uint64_t total_size = 0;
// Linearly iterate through all nodes. It may be possible to optimize this
// further by making best guesses as to where to split the nodes, since
// most nodes (aside from constants) are relatively small.
int repeated_field_length =
parent->GetReflection()->FieldSize(*parent, repeated_field);
for (int i = 0; i < repeated_field_length; ++i) {
tsl::protobuf::Message* node =
parent->GetReflection()->MutableRepeatedMessage(parent, repeated_field,
i);
auto node_size = node->ByteSizeLong();
std::vector<FieldType> new_fields = {repeated_field_, i};
for (auto factory : *splitter_factories_) {
TF_ASSIGN_OR_RETURN(
std::unique_ptr<SizeSplitter> new_splitter,
factory->CreateSplitter(node, this, &new_fields, node_size));
if (new_splitter != nullptr) {
TF_ASSIGN_OR_RETURN(auto size_diff,
new_splitter->BuildChunksReturnSize());
node_size -= size_diff;
}
}
if (total_size + node_size > max_size) {
repeated_msg_split.push_back(i);
total_size = 0;
}
total_size += node_size + kExtraBytes;
}
if (!repeated_msg_split.empty()) {
auto repeated_nodes_ptrs =
parent->GetReflection()
->template MutableRepeatedPtrField<RepeatedMessage>(parent,
repeated_field);
std::vector<RepeatedMessage*> extracted_nodes(repeated_field_length);
repeated_nodes_ptrs->ExtractSubrange(0, repeated_field_length,
&extracted_nodes.at(0));
// Last range end is the size of the repeated field.
repeated_msg_split.push_back(repeated_field_length);
int range_start = 0;
for (int range_end : repeated_msg_split) {
auto new_msg = std::make_shared<ParentMessage>();
std::vector<FieldType> empty_fields;
auto x = std::make_unique<MessageBytes>(new_msg);
TF_RETURN_IF_ERROR(AddChunk(std::move(x), &empty_fields));
// Move nodes into new_msg.
TF_ASSIGN_OR_RETURN(auto new_ret,
GetMutableField(new_msg.get(), repeated_field_));
for (int j = range_start; j < range_end; ++j) {
new_msg->GetReflection()->AddAllocatedMessage(
new_msg.get(), new_ret.field, extracted_nodes[j]);
}
range_start = range_end;
}
}
// Estimate the size diff by subtracting the first computed chunk size from
// the initial size of the repeated field.
return initial_size - message()->ByteSizeLong();
}
} // namespace tensorflow::tools::proto_splitter
#endif // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_REPEATED_FIELD_SPLITTER_H_
tensorflow/tools/proto_splitter/cc/saved_model_splitter.cc
+62
View File
@@ -0,0 +1,62 @@
/* 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.
==============================================================================*/
#include "tensorflow/tools/proto_splitter/cc/saved_model_splitter.h"
#include <vector>
#include "absl/status/status.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"
#include "tensorflow/tools/proto_splitter/cc/graph_def_splitter.h"
#include "tensorflow/tools/proto_splitter/cc/large_node_splitter.h"
#include "tensorflow/tools/proto_splitter/cc/max_size.h"
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/protobuf.h"
namespace tensorflow {
namespace tools::proto_splitter {
// Required in OSS to prevent string to bool conversion in FieldType variant.
using namespace std::string_literals; // NOLINT
absl::Status SavedModelSplitter::BuildChunks() {
TF_RETURN_IF_ERROR(SetMessageAsBaseChunk());
SavedModel* sm = google::protobuf::DynamicCastMessage<SavedModel>(message());
int max_size = GetMaxSize();
if (GetInitialSize() < max_size) return absl::OkStatus();
std::vector<FieldType> fields_to_graph_def = {"meta_graphs"s, 0,
"graph_def"s};
GraphDefSplitter graph_def_splitter(
sm->mutable_meta_graphs(0)->mutable_graph_def(), this,
&fields_to_graph_def);
TF_RETURN_IF_ERROR(graph_def_splitter.BuildChunks());
if (sm->ByteSizeLong() < max_size) return absl::OkStatus();
LargeNodeSplitter<GraphDef> entire_graph_splitter(
sm->mutable_meta_graphs(0)->mutable_graph_def(), this,
&fields_to_graph_def);
int index = 1;
entire_graph_splitter.SetChunkIndex(&index);
TF_RETURN_IF_ERROR(entire_graph_splitter.BuildChunks());
return absl::OkStatus();
}
} // namespace tools::proto_splitter
} // namespace tensorflow
tensorflow/tools/proto_splitter/cc/saved_model_splitter.h
+35
View File
@@ -0,0 +1,35 @@
/* 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_SAVED_MODEL_SPLITTER_H_
#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_SAVED_MODEL_SPLITTER_H_
#include "tensorflow/tools/proto_splitter/cc/composable_splitter.h"
namespace tensorflow {
namespace tools::proto_splitter {
// SavedModel Splitter.
// Modifies the user-provided message when building chunks.
class SavedModelSplitter : public ComposableSplitter {
public:
// Initializer.
using ComposableSplitter::ComposableSplitter;
absl::Status BuildChunks() override;
};
} // namespace tools::proto_splitter
} // namespace tensorflow
#endif // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_SAVED_MODEL_SPLITTER_H_
tensorflow/tools/proto_splitter/cc/saved_model_splitter_test.cc
+89
View File
@@ -0,0 +1,89 @@
/* 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.
==============================================================================*/
#include "tensorflow/tools/proto_splitter/cc/saved_model_splitter.h"
#include <cstdint>
#include <memory>
#include <string>
#include <variant>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"
#include "tensorflow/tools/proto_splitter/cc/max_size.h"
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include "tensorflow/tools/proto_splitter/testdata/test_message.pb.h"
#include "tsl/platform/protobuf.h"
#include "tsl/platform/statusor.h"
namespace tensorflow {
namespace tools::proto_splitter {
namespace {
// Ensures that all Messages are less than the max size. std::string chunks are
// not limited by the max size, so they are ignored in this check.
#define EXPECT_CHUNK_SIZES(chunks, max_size) \
do { \
for (auto chunk : *chunks) { \
if (std::holds_alternative<std::shared_ptr<tsl::protobuf::Message>>( \
chunk)) { \
EXPECT_LE(std::get<std::shared_ptr<tsl::protobuf::Message>>(chunk) \
->ByteSizeLong(), \
max_size); \
} else if (std::holds_alternative<tsl::protobuf::Message*>(chunk)) { \
EXPECT_LE(std::get<tsl::protobuf::Message*>(chunk)->ByteSizeLong(), \
max_size); \
} \
} \
} while (0)
std::string NonChunkedSavedModel() {
return io::JoinPath(testing::TensorFlowSrcRoot(), "cc", "saved_model",
"testdata", "chunked_saved_model", "non_chunked_model",
"saved_model.pb");
}
TEST(SavedModelSplitterTest, TestSplit) {
SavedModel proto;
int64_t max_size = 80000;
DebugSetMaxSize(max_size);
TF_EXPECT_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(),
NonChunkedSavedModel(), &proto));
EXPECT_GE(proto.ByteSizeLong(), GetMaxSize());
SavedModelSplitter splitter(&proto);
TF_ASSERT_OK_AND_ASSIGN(auto x, splitter.Split());
std::vector<MessageBytes>* chunks = x.chunks;
ASSERT_NE(chunks, nullptr);
// Should create a new chunk with the single large constant.
EXPECT_EQ(2, chunks->size());
EXPECT_CHUNK_SIZES(chunks, max_size);
}
} // namespace
} // namespace tools::proto_splitter
} // namespace tensorflow
tensorflow/tools/proto_splitter/cc/size_splitter.h
+59
View File
@@ -0,0 +1,59 @@
/* 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_SIZE_SPLITTER_H_
#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_SIZE_SPLITTER_H_
#include <memory>
#include <vector>
#include "absl/status/statusor.h"
#include "tensorflow/tools/proto_splitter/cc/composable_splitter.h"
#include "tensorflow/tools/proto_splitter/cc/util.h"
namespace tensorflow {
namespace tools::proto_splitter {
// Abstract composable Splitter that returns the size of chunks created.
// Modifies the user-provided message when building chunks.
class SizeSplitter : public ComposableSplitter {
public:
// Initializer.
using ComposableSplitter::ComposableSplitter;
absl::Status BuildChunks() override {
return BuildChunksReturnSize().status();
}
// Chunks the message and returns the int size diff.
virtual absl::StatusOr<int> BuildChunksReturnSize() = 0;
};
class SizeSplitterFactory {
public:
explicit SizeSplitterFactory() = default;
// Creates a new SizeSplitter object based on the input parameters. May
// return nullptr if no split is necessary.
virtual absl::StatusOr<std::unique_ptr<SizeSplitter>> CreateSplitter(
tsl::protobuf::Message* message, ComposableSplitterBase* parent_splitter,
std::vector<FieldType>* fields_in_parent, int size) = 0;
virtual ~SizeSplitterFactory() = default;
};
} // namespace tools::proto_splitter
} // namespace tensorflow
#endif // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_SIZE_SPLITTER_H_
tensorflow/tools/proto_splitter/cc/split.h
+60
View File
@@ -0,0 +1,60 @@
/* 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_SPLIT_H_
#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_SPLIT_H_
#include <string>
#include <utility>
#include <variant>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include "tensorflow/tools/proto_splitter/chunk.pb.h"
#include "tensorflow/tools/proto_splitter/versions.pb.h"
#include "tsl/platform/protobuf.h"
namespace tensorflow {
namespace tools::proto_splitter {
using ::tensorflow::proto_splitter::ChunkedMessage;
using ::tensorflow::proto_splitter::VersionDef;
// Interface for proto message splitters.
class Splitter {
public:
virtual ~Splitter() = default;
// Split message into chunks.
virtual absl::StatusOr<ChunkedProto> Split() = 0;
// Write message to disk and returns the full output path.
virtual absl::StatusOr<std::string> Write(std::string file_prefix) = 0;
// Version info about the Splitter and required Merger versions.
virtual VersionDef Version() = 0;
protected:
Splitter() = default;
Splitter(const Splitter&) = default;
Splitter(Splitter&&) = default;
};
} // namespace tools::proto_splitter
} // namespace tensorflow
#endif // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_SPLIT_H_
tensorflow/tools/proto_splitter/cc/test_util.h
+136
View File
@@ -0,0 +1,136 @@
/* 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_TEST_UTIL_H_
#define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_TEST_UTIL_H_
#include <ostream>
#include <string>
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "tensorflow/core/platform/test.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/protobuf.h"
#include "tsl/platform/statusor.h"
namespace tensorflow {
namespace tools::proto_splitter {
// Ensures that all Messages are less than the max size. std::string chunks are
// not limited by the max size, so they are ignored in this check.
#define EXPECT_CHUNK_SIZES(chunks, max_size) \
do { \
for (auto chunk : *chunks) { \
if (std::holds_alternative<std::shared_ptr<tsl::protobuf::Message>>( \
chunk)) { \
EXPECT_LE(std::get<std::shared_ptr<tsl::protobuf::Message>>(chunk) \
->ByteSizeLong(), \
max_size); \
} else if (std::holds_alternative<tsl::protobuf::Message*>(chunk)) { \
EXPECT_LE(std::get<tsl::protobuf::Message*>(chunk)->ByteSizeLong(), \
max_size); \
} \
} \
} while (0)
inline std::string SerializeAsString(const tsl::protobuf::Message& message) {
std::string result;
{
// Use a nested block to guarantee triggering coded_stream's destructor
// before `result` is used. Due to copy elision, this code works without
// the nested block, but small, innocent looking changes can break it.
tsl::protobuf::io::StringOutputStream string_stream(&result);
tsl::protobuf::io::CodedOutputStream coded_stream(&string_stream);
coded_stream.SetSerializationDeterministic(true);
message.SerializeToCodedStream(&coded_stream);
}
return result;
}
template <typename MessageType>
tsl::StatusOr<MessageType> ParseTextProto(const char* text_proto) {
tsl::protobuf::TextFormat::Parser parser;
MessageType parsed_proto;
bool success =
tsl::protobuf::TextFormat::ParseFromString(text_proto, &parsed_proto);
if (!success) {
return absl::InvalidArgumentError(
"Input text could not be parsed to proto.");
}
return parsed_proto;
}
// TODO(b/229726259): EqualsProto is not available in OSS.
// Simple implementation of a proto matcher comparing string representations.
// This will fail if the string representations are not deterministic.
class ProtoStringMatcher {
public:
explicit ProtoStringMatcher(const tsl::protobuf::Message& expected)
: expected_(SerializeAsString(expected)) {}
explicit ProtoStringMatcher(const std::string& expected)
: text_format_(expected), use_text_format_(true) {}
template <typename Message>
bool MatchAndExplain(const Message& p,
::testing::MatchResultListener*) const {
if (use_text_format_) {
Message* m = p.New();
CHECK(tsl::protobuf::TextFormat::ParseFromString(text_format_, m))
<< "Failed to parse proto text: " << text_format_;
std::string expected = SerializeAsString(*m);
delete m;
return SerializeAsString(p) == expected;
} else {
return SerializeAsString(p) == expected_;
}
}
void DescribeTo(::std::ostream* os) const {
if (use_text_format_)
*os << text_format_;
else
*os << expected_;
}
void DescribeNegationTo(::std::ostream* os) const {
*os << "not equal to expected message: ";
auto out_message = use_text_format_ ? &text_format_ : &expected_;
if (out_message->size() < 1e5)
*os << *out_message;
else
*os << "(too large to print) \n";
}
private:
const std::string expected_;
const std::string text_format_;
const bool use_text_format_ = false;
};
inline ::testing::PolymorphicMatcher<ProtoStringMatcher> EqualsProto(
const tsl::protobuf::Message& x) {
return ::testing::MakePolymorphicMatcher(ProtoStringMatcher(x));
}
inline ::testing::PolymorphicMatcher<ProtoStringMatcher> EqualsProto(
const std::string& x) {
return ::testing::MakePolymorphicMatcher(ProtoStringMatcher(x));
}
} // namespace tools::proto_splitter
} // namespace tensorflow
#endif // TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_TEST_UTIL_H_
tensorflow/tools/proto_splitter/cc/util.cc
+805
View File
@@ -0,0 +1,805 @@
/* 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.
==============================================================================*/
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include <cstdint>
#include <functional>
#include <ios>
#include <optional>
#include <sstream>
#include <string>
#include <utility>
#include <variant>
#include <vector>
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/numbers.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "riegeli/base/maker.h" // from @riegeli
#include "riegeli/base/types.h" // from @riegeli
#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/core/platform/env.h"
#include "tensorflow/core/platform/file_system_helper.h"
#include "tensorflow/tools/proto_splitter/chunk.pb.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/protobuf.h"
#include "tsl/platform/statusor.h"
namespace tensorflow {
namespace tools::proto_splitter {
using ::tensorflow::proto_splitter::ChunkedField;
namespace {
absl::StatusOr<int> FieldInt(const FieldType& field) {
switch (field.index()) {
case 0: { // std::string
int value;
if (absl::SimpleAtoi(std::get<std::string>(field), &value)) return value;
return absl::InvalidArgumentError(absl::StrCat(
"Unable to convert '", std::get<std::string>(field), "' to int."));
break;
}
case 1: // int
return std::get<int>(field);
case 2: // bool
return std::get<bool>(field) ? 1 : 0;
default:
// This should not happen.
return absl::InvalidArgumentError("Invalid field type.");
}
}
absl::StatusOr<bool> FieldBool(const FieldType& field) {
switch (field.index()) {
case 0: { // std::string
std::string s = std::get<std::string>(field);
if (s == "true") return true;
if (s == "false") return false;
return absl::InvalidArgumentError("Unable to convert '" + s +
"' to bool.");
}
case 1: // int
return std::get<int>(field) != 0;
case 2: // bool
return std::get<bool>(field);
default:
// This should not happen.
return absl::InvalidArgumentError("Invalid field type.");
}
}
std::string FieldString(const FieldType& field) {
switch (field.index()) {
case 0: // std::string
return std::get<std::string>(field);
case 1: // int
return std::to_string(std::get<int>(field));
case 2: // bool
return std::get<bool>(field) ? "true" : "false";
default:
// This should not happen.
CHECK(false) << "Invalid field type: " << field.index(); // Crash OK.
return "";
}
}
template <typename FieldCallback, typename MapKeyCallback,
typename ListIndexCallback>
absl::Status WalkFields(const tsl::protobuf::Descriptor& desc,
const std::vector<FieldType>& fields,
FieldCallback field_callback,
MapKeyCallback map_key_callback,
ListIndexCallback list_index_callback) {
const tsl::protobuf::Descriptor* message_desc = &desc;
std::vector<FieldType>::const_iterator it = fields.begin();
while (it != fields.end()) {
const tsl::protobuf::FieldDescriptor* field_desc;
if (std::holds_alternative<std::string>(*it)) {
field_desc = message_desc->FindFieldByName(std::get<std::string>(*it));
} else if (std::holds_alternative<int>(*it)) {
field_desc = message_desc->FindFieldByNumber(std::get<int>(*it));
} else {
return absl::InvalidArgumentError(absl::StrCat(
"Invalid field, expected int or str: ", FieldString(*it)));
}
if (field_desc == nullptr) {
return absl::NotFoundError(
absl::StrCat("Field not found: ", FieldString(*it)));
}
TF_RETURN_IF_ERROR(field_callback(field_desc));
message_desc = field_desc->message_type();
it++;
// Handle special field types (map key and list index).
if (it != fields.end()) {
if (field_desc->is_map()) {
auto key_field = message_desc->FindFieldByNumber(1);
auto status = map_key_callback(*key_field, *it);
if (!status.ok()) {
return status;
}
// The next element in `fields` must be a field in the value message.
message_desc = message_desc->FindFieldByName("value")->message_type();
it++;
} else if (field_desc->is_repeated()) {
auto status = list_index_callback(*it);
if (!status.ok()) {
return status;
}
it++;
}
}
}
return absl::OkStatus();
}
absl::Status AddMapKey(const tsl::protobuf::FieldDescriptor& key_field,
const FieldType& map_key, ChunkedField& chunked_field) {
auto msg = chunked_field.add_field_tag()->mutable_map_key();
switch (key_field.type()) {
case tsl::protobuf::FieldDescriptor::TYPE_BOOL: {
TF_ASSIGN_OR_RETURN(auto key, FieldBool(map_key));
msg->set_boolean(key);
break;
}
case tsl::protobuf::FieldDescriptor::TYPE_STRING: {
msg->set_s(FieldString(map_key));
break;
}
case tsl::protobuf::FieldDescriptor::TYPE_INT32: {
TF_ASSIGN_OR_RETURN(auto key, FieldInt(map_key));
msg->set_i32(key);
break;
}
case tsl::protobuf::FieldDescriptor::TYPE_INT64: {
TF_ASSIGN_OR_RETURN(auto key, FieldInt(map_key));
msg->set_i64(key);
break;
}
case tsl::protobuf::FieldDescriptor::TYPE_UINT32: {
TF_ASSIGN_OR_RETURN(auto key, FieldInt(map_key));
msg->set_ui32(key);
break;
}
case tsl::protobuf::FieldDescriptor::TYPE_UINT64: {
TF_ASSIGN_OR_RETURN(auto key, FieldInt(map_key));
msg->set_ui64(key);
break;
}
case tsl::protobuf::FieldDescriptor::TYPE_FIXED64:
case tsl::protobuf::FieldDescriptor::TYPE_FIXED32:
case tsl::protobuf::FieldDescriptor::TYPE_SINT32:
case tsl::protobuf::FieldDescriptor::TYPE_SINT64:
case tsl::protobuf::FieldDescriptor::TYPE_SFIXED32:
case tsl::protobuf::FieldDescriptor::TYPE_SFIXED64:
return absl::UnimplementedError(
"Fixed map key types not implemented yet.");
case tsl::protobuf::FieldDescriptor::TYPE_DOUBLE:
case tsl::protobuf::FieldDescriptor::TYPE_FLOAT:
case tsl::protobuf::FieldDescriptor::TYPE_ENUM:
case tsl::protobuf::FieldDescriptor::TYPE_BYTES:
case tsl::protobuf::FieldDescriptor::TYPE_MESSAGE:
case tsl::protobuf::FieldDescriptor::TYPE_GROUP:
return absl::FailedPreconditionError(absl::StrCat(
"Encountered type that is not supported for MapKey.type: ",
key_field.type()));
default:
return absl::FailedPreconditionError("Proto field type not recognized.");
}
return absl::OkStatus();
}
absl::StatusOr<FieldType> GetMapKeyFromFieldIndex(
::tensorflow::proto_splitter::FieldIndex field_index) {
if (!field_index.has_map_key())
return absl::FailedPreconditionError(
"Field index doesn't contain a map key.");
switch (field_index.map_key().type_case()) {
case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::kBoolean:
return field_index.map_key().boolean();
break;
case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::kS:
return field_index.map_key().s();
break;
case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::kI32:
return field_index.map_key().i32();
break;
case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::kI64:
// Cast to int type, which may be lossy. We'll deal with it when it
// becomes an issue.
return static_cast<int>(field_index.map_key().i64());
break;
case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::kUi32:
return static_cast<int>(field_index.map_key().ui32());
break;
case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::kUi64:
return static_cast<int>(field_index.map_key().ui64());
break;
case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::
TYPE_NOT_SET:
default:
return absl::FailedPreconditionError(
absl::StrCat("Unknown map key type: ", field_index.DebugString()));
}
}
} // namespace
absl::StatusOr<const std::vector<Field>> GetFieldTypes(
const tsl::protobuf::RepeatedPtrField<
::tensorflow::proto_splitter::FieldIndex>& field_tags) {
std::vector<Field> fields;
for (int fti = 0; fti < field_tags.size();) {
switch (field_tags[fti].kind_case()) {
case ::tensorflow::proto_splitter::FieldIndex::KindCase::kField:
fields.push_back(
Field(static_cast<int>(field_tags[fti].field()), std::nullopt));
fti++;
if (fti == field_tags.size()) break;
// Multiple field tags may correspond to a single field when the field
// is repeated or a map.
if (field_tags[fti].has_index()) {
fields.back().second = static_cast<int>(field_tags[fti++].index());
} else if (field_tags[fti].has_map_key()) {
TF_ASSIGN_OR_RETURN(const FieldType& map_key,
GetMapKeyFromFieldIndex(field_tags[fti++]));
fields.back().second = map_key;
}
break;
case ::tensorflow::proto_splitter::FieldIndex::KindCase::kIndex:
return absl::FailedPreconditionError(
"Index doesn't belong to any field.");
break;
case ::tensorflow::proto_splitter::FieldIndex::KindCase::kMapKey:
return absl::FailedPreconditionError(
"Map key doesn't belong to any field.");
break;
case ::tensorflow::proto_splitter::FieldIndex::KindCase::KIND_NOT_SET:
default:
return absl::FailedPreconditionError(absl::StrCat(
"Unknown field kind: ", field_tags[fti].DebugString()));
}
}
return fields;
}
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) {
if (field_desc->is_map())
return absl::FailedPreconditionError("Field is a map.");
const tsl::protobuf::Reflection* reflection = message->GetReflection();
if (field_index >= reflection->FieldSize(*message, field_desc))
return absl::OutOfRangeError(
absl::StrCat("Field index out of range: ", field_index));
switch (field_desc->cpp_type()) {
case tsl::protobuf::FieldDescriptor::CPPTYPE_INT32:
reflection->SetRepeatedInt32(message, field_desc, std::stoi(chunk),
field_index);
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_INT64:
reflection->SetRepeatedInt64(message, field_desc, std::stoll(chunk),
field_index);
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_UINT32:
reflection->SetRepeatedUInt32(message, field_desc, std::stoul(chunk),
field_index);
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_UINT64:
reflection->SetRepeatedUInt64(message, field_desc, std::stoull(chunk),
field_index);
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_DOUBLE:
reflection->SetRepeatedDouble(message, field_desc, std::stod(chunk),
field_index);
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_FLOAT:
reflection->SetRepeatedFloat(message, field_desc, std::stof(chunk),
field_index);
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_BOOL:
bool b;
std::istringstream(chunk) >> std::boolalpha >> b;
reflection->SetRepeatedBool(message, field_desc, b, field_index);
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_ENUM:
reflection->SetRepeatedEnum(
message, field_desc, field_index,
field_desc->enum_type()->FindValueByName(chunk));
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_STRING:
reflection->SetRepeatedString(message, field_desc, field_index, chunk);
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_MESSAGE:
return message_callback();
break;
default:
return absl::FailedPreconditionError(absl::StrCat(
"Proto field type not recognized: ", field_desc->cpp_type_name()));
break;
}
return absl::OkStatus();
}
absl::Status SetFieldElement(
tsl::protobuf::Message* message,
const tsl::protobuf::FieldDescriptor* field_desc, const std::string& chunk,
std::function<absl::Status(void)> message_callback) {
const tsl::protobuf::Reflection* reflection = message->GetReflection();
switch (field_desc->cpp_type()) {
case tsl::protobuf::FieldDescriptor::CPPTYPE_INT32:
reflection->SetInt32(message, field_desc, std::stoi(chunk));
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_INT64:
reflection->SetInt64(message, field_desc, std::stoll(chunk));
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_UINT32:
reflection->SetUInt32(message, field_desc, std::stoul(chunk));
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_UINT64:
reflection->SetUInt64(message, field_desc, std::stoull(chunk));
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_DOUBLE:
reflection->SetDouble(message, field_desc, std::stod(chunk));
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_FLOAT:
reflection->SetFloat(message, field_desc, std::stof(chunk));
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_BOOL:
bool b;
std::istringstream(chunk) >> std::boolalpha >> b;
reflection->SetBool(message, field_desc, b);
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_ENUM:
reflection->SetEnum(message, field_desc,
field_desc->enum_type()->FindValueByName(chunk));
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_STRING:
reflection->SetString(message, field_desc, chunk);
break;
case tsl::protobuf::FieldDescriptor::CPPTYPE_MESSAGE:
return message_callback();
break;
default:
return absl::FailedPreconditionError("Proto field type not recognized.");
break;
}
return absl::OkStatus();
}
absl::Status AddMapEntry(tsl::protobuf::Message* message,
const tsl::protobuf::FieldDescriptor* field_desc,
FieldType map_key) {
tsl::protobuf::Message* map_entry =
message->GetReflection()->AddMessage(message, field_desc);
const tsl::protobuf::Reflection* reflection = map_entry->GetReflection();
const tsl::protobuf::FieldDescriptor* key =
field_desc->message_type()->FindFieldByNumber(1);
switch (key->cpp_type()) {
case tsl::protobuf::FieldDescriptor::CPPTYPE_INT32: {
TF_ASSIGN_OR_RETURN(auto map_key_int32, FieldInt(map_key));
reflection->SetInt32(map_entry, key, map_key_int32);
break;
}
case tsl::protobuf::FieldDescriptor::CPPTYPE_INT64: {
TF_ASSIGN_OR_RETURN(auto map_key_int64, FieldInt(map_key));
reflection->SetInt64(map_entry, key, map_key_int64);
break;
}
case tsl::protobuf::FieldDescriptor::CPPTYPE_UINT32: {
TF_ASSIGN_OR_RETURN(auto map_key_uint32, FieldInt(map_key));
reflection->SetUInt32(map_entry, key, map_key_uint32);
break;
}
case tsl::protobuf::FieldDescriptor::CPPTYPE_UINT64: {
TF_ASSIGN_OR_RETURN(auto map_key_uint64, FieldInt(map_key));
reflection->SetUInt64(map_entry, key, map_key_uint64);
break;
}
case tsl::protobuf::FieldDescriptor::CPPTYPE_BOOL: {
TF_ASSIGN_OR_RETURN(auto map_key_bool, FieldBool(map_key));
reflection->SetBool(map_entry, key, map_key_bool);
break;
}
case tsl::protobuf::FieldDescriptor::CPPTYPE_STRING: {
reflection->SetString(map_entry, key, FieldString(map_key));
break;
}
default:
return absl::FailedPreconditionError(absl::StrCat(
"Proto field type not recognized.", key->cpp_type_name()));
}
return absl::OkStatus();
}
absl::StatusOr<MutableFieldResult> GetMutableField(
tsl::protobuf::Message* message, const std::vector<FieldType>& fields) {
tsl::protobuf::Message* parent = message;
const tsl::protobuf::FieldDescriptor* field = nullptr;
int index = -1;
auto field_callback =
[&parent, &field, &index](
const tsl::protobuf::FieldDescriptor* field_desc) -> absl::Status {
const tsl::protobuf::Reflection* reflection = parent->GetReflection();
if (field != nullptr) {
// Update the parent proto using the previous FieldDescriptor.
if (field->is_map()) {
parent = reflection->MutableRepeatedMessage(parent, field, index);
parent = parent->GetReflection()->MutableMessage(
parent, parent->GetDescriptor()->FindFieldByNumber(2));
index = -1;
} else if (field->is_repeated()) {
parent = reflection->MutableRepeatedMessage(parent, field, index);
index = -1;
} else {
parent = reflection->MutableMessage(parent, field);
}
}
field = field_desc;
return absl::OkStatus();
};
auto map_key_callback = [&parent, &field, &index](
const tsl::protobuf::FieldDescriptor& key_field,
const FieldType& map_key) -> absl::Status {
// Finds the index of the map field w/ the corresponding map_key.
const tsl::protobuf::Reflection* reflection = parent->GetReflection();
TF_ASSIGN_OR_RETURN(int found_index,
FindMapKey(*parent, *field, &key_field, map_key));
if (found_index == -1) {
return absl::NotFoundError(
absl::StrCat("Error when getting map field, couldn't find key: ",
FieldString(map_key)));
}
if (reflection->FieldSize(*parent, field) <= found_index) {
return absl::NotFoundError(absl::StrCat(
"Can't access index ", index, " in field '", field->name(),
"' (size = ", reflection->FieldSize(*parent, field), ")."));
}
index = found_index;
return absl::OkStatus();
};
auto list_index_callback = [&parent, &field,
&index](FieldType list_index) -> absl::Status {
TF_ASSIGN_OR_RETURN(index, FieldInt(list_index));
const tsl::protobuf::Reflection* reflection = parent->GetReflection();
if (reflection->FieldSize(*parent, field) <= index) {
return absl::NotFoundError(absl::StrCat(
"Can't access index ", index, " in field '", field->name(),
"' (size = ", reflection->FieldSize(*parent, field), ")."));
}
return absl::OkStatus();
};
TF_RETURN_IF_ERROR(WalkFields(*message->GetDescriptor(), fields,
field_callback, map_key_callback,
list_index_callback));
MutableFieldResult result{.parent = parent, .field = field, .index = index};
return result;
}
absl::StatusOr<MutableFieldResult> GetMutableField(
tsl::protobuf::Message* message, const Field& field) {
std::vector<FieldType> fields = {field.first};
if (field.second != std::nullopt) fields.push_back(field.second.value());
return GetMutableField(message, fields);
}
absl::StatusOr<MutableFieldResult> GetMutableField(
tsl::protobuf::Message* message, const FieldType& field_type) {
std::vector<FieldType> fields = {field_type};
return GetMutableField(message, fields);
}
absl::StatusOr<FieldResult> GetField(const tsl::protobuf::Message& message,
const std::vector<FieldType>& fields) {
const tsl::protobuf::Message* parent = &message;
const tsl::protobuf::FieldDescriptor* field = nullptr;
int index = -1;
auto field_callback =
[&parent, &field, &index](
const tsl::protobuf::FieldDescriptor* field_desc) -> absl::Status {
const tsl::protobuf::Reflection* reflection = parent->GetReflection();
if (field != nullptr) {
// Update the parent proto using the previous FieldDescriptor.
if (field->is_map()) {
parent = &reflection->GetRepeatedMessage(*parent, field, index);
parent = &parent->GetReflection()->GetMessage(
*parent, parent->GetDescriptor()->FindFieldByNumber(2));
index = -1;
} else if (field->is_repeated()) {
parent = &reflection->GetRepeatedMessage(*parent, field, index);
index = -1;
} else {
parent = &reflection->GetMessage(*parent, field);
}
}
field = field_desc;
return absl::OkStatus();
};
auto map_key_callback = [&parent, &field, &index](
const tsl::protobuf::FieldDescriptor& key_field,
const FieldType& map_key) -> absl::Status {
// Finds the index of the map field w/ the corresponding map_key.
const tsl::protobuf::Reflection* reflection = parent->GetReflection();
TF_ASSIGN_OR_RETURN(int found_index,
FindMapKey(*parent, *field, &key_field, map_key));
if (found_index == -1) {
return absl::NotFoundError(
absl::StrCat("Error when getting map field, couldn't find key: ",
FieldString(map_key)));
}
if (reflection->FieldSize(*parent, field) <= found_index) {
return absl::NotFoundError(absl::StrCat(
"Can't access index ", index, " in field '", field->name(),
"' (size = ", reflection->FieldSize(*parent, field), ")."));
}
index = found_index;
return absl::OkStatus();
};
auto list_index_callback = [&parent, &field,
&index](FieldType list_index) -> absl::Status {
TF_ASSIGN_OR_RETURN(index, FieldInt(list_index));
const tsl::protobuf::Reflection* reflection = parent->GetReflection();
if (reflection->FieldSize(*parent, field) <= index) {
return absl::NotFoundError(absl::StrCat(
"Can't access index ", index, " in field '", field->name(),
"' (size = ", reflection->FieldSize(*parent, field), ")."));
}
return absl::OkStatus();
};
TF_RETURN_IF_ERROR(WalkFields(*message.GetDescriptor(), fields,
field_callback, map_key_callback,
list_index_callback));
FieldResult result{
.parent = parent,
.field = field,
.index = index,
};
return result;
}
absl::Status AddFieldTag(const tsl::protobuf::Descriptor& desc,
const std::vector<FieldType>& fields,
ChunkedField& chunked_field) {
auto field_callback =
[&chunked_field](
const tsl::protobuf::FieldDescriptor* field_desc) -> absl::Status {
chunked_field.add_field_tag()->set_field(field_desc->number());
return absl::OkStatus();
};
auto map_key_callback = [&chunked_field](
const tsl::protobuf::FieldDescriptor& key_field,
FieldType map_key) -> absl::Status {
return AddMapKey(key_field, map_key, chunked_field);
};
auto list_index_callback =
[&chunked_field](FieldType list_index) -> absl::Status {
TF_ASSIGN_OR_RETURN(auto index, FieldInt(list_index));
chunked_field.add_field_tag()->set_index(index);
return absl::OkStatus();
};
return WalkFields(desc, fields, field_callback, map_key_callback,
list_index_callback);
}
absl::Status AddFieldTag(const tsl::protobuf::Descriptor& desc,
const Field& field, ChunkedField& chunked_field) {
std::vector<FieldType> fields = {field.first};
if (field.second != std::nullopt) fields.push_back(field.second.value());
return AddFieldTag(desc, fields, chunked_field);
}
absl::StatusOr<int> FindMapKey(const tsl::protobuf::Message& parent,
const tsl::protobuf::FieldDescriptor& map_field,
const tsl::protobuf::FieldDescriptor* key_field,
FieldType map_key) {
const tsl::protobuf::Reflection* reflection = parent.GetReflection();
if (!map_field.is_map()) {
return absl::InvalidArgumentError(absl::StrCat(
"FindMapKey() was given a non map field: ", map_field.full_name()));
}
if (key_field == nullptr) {
key_field = map_field.message_type()->FindFieldByNumber(1);
}
bool found = false;
for (int i = 0; i < reflection->FieldSize(parent, &map_field) && !found;
i++) {
auto kv_pair = &reflection->GetRepeatedMessage(parent, &map_field, i);
auto kv_reflection = kv_pair->GetReflection();
switch (key_field->cpp_type()) {
case tsl::protobuf::FieldDescriptor::CPPTYPE_INT32: {
TF_ASSIGN_OR_RETURN(auto key, FieldInt(map_key));
if (kv_reflection->GetInt32(*kv_pair, key_field) == key) {
found = true;
}
break;
}
case tsl::protobuf::FieldDescriptor::CPPTYPE_INT64: {
TF_ASSIGN_OR_RETURN(auto key, FieldInt(map_key));
if (kv_reflection->GetInt64(*kv_pair, key_field) == key) {
found = true;
}
break;
}
case tsl::protobuf::FieldDescriptor::CPPTYPE_UINT32: {
TF_ASSIGN_OR_RETURN(auto key, FieldInt(map_key));
if (kv_reflection->GetUInt32(*kv_pair, key_field) == key) {
found = true;
}
break;
}
case tsl::protobuf::FieldDescriptor::CPPTYPE_UINT64: {
TF_ASSIGN_OR_RETURN(auto key, FieldInt(map_key));
if (kv_reflection->GetUInt64(*kv_pair, key_field) == key) {
found = true;
}
break;
}
case tsl::protobuf::FieldDescriptor::CPPTYPE_BOOL: {
TF_ASSIGN_OR_RETURN(auto key, FieldBool(map_key));
if (kv_reflection->GetBool(*kv_pair, key_field) == key) {
found = true;
}
break;
}
case tsl::protobuf::FieldDescriptor::CPPTYPE_STRING: {
if (kv_reflection->GetString(*kv_pair, key_field) ==
FieldString(map_key)) {
found = true;
}
break;
}
case tsl::protobuf::FieldDescriptor::CPPTYPE_DOUBLE:
case tsl::protobuf::FieldDescriptor::CPPTYPE_FLOAT:
case tsl::protobuf::FieldDescriptor::CPPTYPE_ENUM:
case tsl::protobuf::FieldDescriptor::CPPTYPE_MESSAGE:
return absl::FailedPreconditionError(absl::StrCat(
"Type is not supported for MapKey.type: ", key_field->cpp_type()));
default:
return absl::FailedPreconditionError(
"Proto field type not recognized.");
}
if (found) return i;
}
return -1;
}
namespace {
const std::vector<std::pair<double, std::string>>& ByteUnits() {
static const auto& byte_units =
*(new std::vector<std::pair<double, std::string>>(
{{1, "B"}, {1 << 10, "KiB"}, {1 << 20, "MiB"}, {1 << 30, "GiB"}}));
return byte_units;
}
} // namespace
std::string HumanReadableBytes(int64_t byte_count) {
for (int i = 1; i < ByteUnits().size(); i++) {
if (byte_count < ByteUnits()[i].first) {
return absl::StrFormat("%.1f%s", byte_count / ByteUnits()[i - 1].first,
ByteUnits()[i - 1].second);
}
}
int i = ByteUnits().size() - 1;
return absl::StrFormat("%.1f%s", byte_count / ByteUnits()[i].first,
ByteUnits()[i].second);
}
std::string HumanReadableDuration(int64_t microseconds) {
if (microseconds < 1000) {
return absl::StrFormat("%d microseconds", microseconds);
} else if (microseconds < 1e6) {
return absl::StrFormat("%.2f ms", microseconds / 1e3);
} else {
return absl::StrFormat("%.2f s", microseconds / 1e6);
}
}
absl::StatusOr<riegeli::RecordReader<riegeli::FdReader<>>> GetRiegeliReader(
absl::string_view cpb_file) {
riegeli::RecordReader reader(riegeli::Maker<riegeli::FdReader>(cpb_file));
if (!reader.ok()) {
return reader.status();
}
return reader;
}
absl::StatusOr<riegeli::RecordReader<riegeli::StringReader<>>>
GetRiegeliStringReader(absl::string_view cpb_data) {
riegeli::RecordReader reader(
riegeli::Maker<riegeli::StringReader<>>(cpb_data));
if (!reader.ok()) {
return reader.status();
}
return reader;
}
absl::StatusOr<::tensorflow::proto_splitter::ChunkMetadata> GetChunkMetadata(
riegeli::RecordReaderBase& reader) {
::tensorflow::proto_splitter::ChunkMetadata chunk_metadata;
bool read_metadata_success = reader.Seek(reader.Size().value()) &&
reader.SeekBack() &&
reader.ReadRecord(chunk_metadata);
if (read_metadata_success) return chunk_metadata;
return reader.status();
}
absl::StatusOr<std::string> ReadChunk(
riegeli::RecordReaderBase& reader,
const ::tensorflow::proto_splitter::ChunkInfo& chunk_info) {
riegeli::Position pos = chunk_info.offset();
std::string chunk(chunk_info.size(), '\0');
if (reader.Seek(pos) && reader.ReadRecord(chunk)) return chunk;
return absl::NotFoundError(
absl::StrCat("Chunk could not be found at position: ", pos, ".\n",
reader.status().ToString()));
}
absl::StatusOr<bool> OnlyContainsPb(absl::string_view prefix) {
const std::string pb_file = absl::StrCat(prefix, ".pb");
const std::string cpb_file = absl::StrCat(prefix, ".cpb");
TF_ASSIGN_OR_RETURN(bool is_pb,
internal::FileExists(Env::Default(), pb_file));
TF_ASSIGN_OR_RETURN(bool is_cpb,
internal::FileExists(Env::Default(), cpb_file));
if (is_pb && !is_cpb) {
return true;
} else if (!is_pb && !is_cpb) {
return absl::NotFoundError(
absl::StrCat("Could not find SavedModel .pb or .cpb at supplied "
"export directory path with prefix: ",
prefix,
". Check that "
"the directory exists and that you have the right "
"permissions for accessing it."));
}
LOG(INFO) << "Reading chunked proto from " << cpb_file;
return false;
}
} // namespace tools::proto_splitter
} // namespace tensorflow
tensorflow/tools/proto_splitter/cc/util.h
+171
View File
@@ -0,0 +1,171 @@
/* 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_
tensorflow/tools/proto_splitter/cc/util_test.cc
+550
View File
@@ -0,0 +1,550 @@
/* 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.
==============================================================================*/
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include <cstdint>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/tools/proto_splitter/cc/test_util.h"
#include "tensorflow/tools/proto_splitter/chunk.pb.h"
#include "tensorflow/tools/proto_splitter/testdata/test_message.pb.h"
#include "tsl/platform/protobuf.h"
#include "tsl/platform/status.h"
#include "tsl/platform/status_matchers.h"
#include "tsl/platform/statusor.h"
namespace tensorflow {
namespace tools::proto_splitter {
namespace {
using ::tensorflow::proto_splitter::ChunkedField;
using ::tensorflow::proto_splitter_testdata::ManyFields;
using ::testing::HasSubstr;
using tsl::testing::IsOkAndHolds;
using tsl::testing::StatusIs;
// Required in OSS to prevent string to bool conversion in FieldType variant.
using namespace std::string_literals; // NOLINT
absl::StatusOr<ManyFields> MakeManyFields() {
return ParseTextProto<ManyFields>(
R"pb(field_one {
repeated_field {}
repeated_field {
string_field: "inner_inner_string"
map_field_uint32 { key: 324 value: "map_value_324" }
map_field_uint32 { key: 543 value: "map_value_543" }
}
}
map_field_int64 { key: -1345 value: "map_value_-1345" }
nested_map_bool {
key: false
value { string_field: "string_false" }
}
nested_map_bool {
key: true
value { string_field: "string_true" }
})pb");
}
absl::StatusOr<
tsl::protobuf::RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>>
MakeFieldTags() {
TF_ASSIGN_OR_RETURN(auto ret, ParseTextProto<ChunkedField>(R"pb(
field_tag { field: 2 }
field_tag { index: 1505 }
field_tag { field: 5 }
field_tag { map_key { ui32: 123 } }
)pb"));
return ret.field_tag();
}
absl::StatusOr<
tsl::protobuf::RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>>
MakeFieldTagsTooManyIndices() {
TF_ASSIGN_OR_RETURN(auto ret, ParseTextProto<ChunkedField>(R"pb(
field_tag { field: 2 }
field_tag { index: 1505 }
field_tag { index: 1506 }
field_tag { field: 5 }
field_tag { map_key { ui32: 123 } }
)pb"));
return ret.field_tag();
}
absl::StatusOr<
tsl::protobuf::RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>>
MakeFieldTagsTooManyMapKeys() {
TF_ASSIGN_OR_RETURN(auto ret, ParseTextProto<ChunkedField>(R"pb(
field_tag { field: 2 }
field_tag { index: 1505 }
field_tag { field: 5 }
field_tag { map_key { ui32: 123 } }
field_tag { map_key: { ui32: 124 } }
)pb"));
return ret.field_tag();
}
absl::StatusOr<
tsl::protobuf::RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>>
MakeFieldTagsMisplacedIndex() {
TF_ASSIGN_OR_RETURN(auto ret, ParseTextProto<ChunkedField>(R"pb(
field_tag { field: 2 }
field_tag { index: 1505 }
field_tag { field: 5 }
field_tag { map_key { ui32: 123 } }
field_tag { index: 1504 }
)pb"));
return ret.field_tag();
}
absl::StatusOr<
tsl::protobuf::RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>>
MakeFieldTagsMisplacedMapKey() {
TF_ASSIGN_OR_RETURN(auto ret, ParseTextProto<ChunkedField>(R"pb(
field_tag { field: 2 }
field_tag { index: 1505 }
field_tag { map_key: { ui32: 321 } }
field_tag { field: 5 }
field_tag { map_key { ui32: 123 } }
)pb"));
return ret.field_tag();
}
TEST(UtilTest, TestFieldTag) {
ManyFields message;
ChunkedField field;
std::vector<FieldType> fields = {"nested_map_bool"s, true, 2, 50, 7, false};
TF_ASSERT_OK(AddFieldTag(*message.descriptor(), fields, field));
EXPECT_THAT(field,
EqualsProto(R"pb(field_tag { field: 7 }
field_tag { map_key { boolean: true } }
field_tag { field: 2 }
field_tag { index: 50 }
field_tag { field: 7 }
field_tag { map_key { boolean: false } })pb"));
}
TEST(UtilTest, TestFieldTagWithInt64MapKey) {
ManyFields message;
ChunkedField field;
std::vector<FieldType> fields = {"map_field_int64"s, -4234};
TF_ASSERT_OK(AddFieldTag(*message.descriptor(), fields, field));
EXPECT_THAT(field, EqualsProto(R"pb(field_tag { field: 6 }
field_tag { map_key { i64: -4234 } }
)pb"));
}
TEST(UtilTest, TestFieldTagWithUInt32MapKey) {
ManyFields message;
ChunkedField field;
std::vector<FieldType> fields = {"repeated_field"s, 1505, "map_field_uint32"s,
123};
TF_ASSERT_OK(AddFieldTag(*message.descriptor(), fields, field));
EXPECT_THAT(field, EqualsProto(R"pb(field_tag { field: 2 }
field_tag { index: 1505 }
field_tag { field: 5 }
field_tag { map_key { ui32: 123 } }
)pb"));
}
TEST(UtilTest, TestFieldTagConversion) {
ManyFields message;
ChunkedField field;
std::vector<FieldType> fields = {"repeated_field"s, "1505"s,
"map_field_uint32"s, 123};
TF_ASSERT_OK(AddFieldTag(*message.descriptor(), fields, field));
EXPECT_THAT(field, EqualsProto(R"pb(field_tag { field: 2 }
field_tag { index: 1505 }
field_tag { field: 5 }
field_tag { map_key { ui32: 123 } }
)pb"));
}
TEST(UtilTest, TestGetFieldTypes) {
TF_ASSERT_OK_AND_ASSIGN(auto tags, MakeFieldTags());
EXPECT_THAT(
GetFieldTypes(tags),
absl_testing::IsOkAndHolds(std::vector<Field>{{2, 1505}, {5, 123}}));
}
TEST(UtilTest, TestGetFieldTypesThenAddFieldTags) {
TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
ChunkedField chunked_field;
TF_ASSERT_OK_AND_ASSIGN(auto tags, MakeFieldTags());
TF_ASSERT_OK_AND_ASSIGN(std::vector<Field> fields, GetFieldTypes(tags));
for (const auto& field : fields) {
TF_ASSERT_OK(AddFieldTag(*message.descriptor(), field, chunked_field));
}
EXPECT_THAT(chunked_field,
EqualsProto(R"pb(field_tag { field: 2 }
field_tag { index: 1505 }
field_tag { field: 5 }
field_tag { map_key { ui32: 123 } }
)pb"));
}
TEST(UtilTest, TestGetFieldTypesThenAddFieldTagsTooManyIndices) {
TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
ChunkedField chunked_field;
TF_ASSERT_OK_AND_ASSIGN(auto tags, MakeFieldTagsTooManyIndices());
EXPECT_THAT(
GetFieldTypes(tags),
absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition,
HasSubstr("Index doesn't belong to any field")));
}
TEST(UtilTest, TestGetFieldTypesThenAddFieldTagsTooManyMapKeys) {
TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
ChunkedField chunked_field;
TF_ASSERT_OK_AND_ASSIGN(auto tags, MakeFieldTagsTooManyMapKeys());
EXPECT_THAT(
GetFieldTypes(tags),
absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition,
HasSubstr("Map key doesn't belong to any field")));
}
TEST(UtilTest, TestGetFieldTypesThenAddFieldTagsMisplacedIndex) {
TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
ChunkedField chunked_field;
TF_ASSERT_OK_AND_ASSIGN(auto tags, MakeFieldTagsMisplacedIndex());
EXPECT_THAT(
GetFieldTypes(tags),
absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition,
HasSubstr("Index doesn't belong to any field")));
}
TEST(UtilTest, TestGetFieldTypesThenAddFieldTagsMisplacedMapKey) {
TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
ChunkedField chunked_field;
TF_ASSERT_OK_AND_ASSIGN(auto tags, MakeFieldTagsMisplacedMapKey());
EXPECT_THAT(
GetFieldTypes(tags),
absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition,
HasSubstr("Map key doesn't belong to any field")));
}
TEST(UtilTest, TestSetRepeatedFieldElement) {
TF_ASSERT_OK_AND_ASSIGN(
auto message,
ParseTextProto<ManyFields>(R"pb(repeated_string_field: ""
repeated_string_field: "")pb"));
const tsl::protobuf::FieldDescriptor* field_desc =
message.GetDescriptor()->FindFieldByName("repeated_string_field");
const std::vector<std::string> chunks = {"repeated_string_one",
"repeated_string_two"};
TF_ASSERT_OK(SetRepeatedFieldElement(
&message, field_desc, 0, chunks[0],
[]() -> absl::Status { return absl::OkStatus(); }));
TF_ASSERT_OK(SetRepeatedFieldElement(
&message, field_desc, 1, chunks[1],
[]() -> absl::Status { return absl::OkStatus(); }));
EXPECT_FALSE(message.repeated_string_field().empty());
EXPECT_EQ(message.repeated_string_field().at(0), "repeated_string_one");
EXPECT_EQ(message.repeated_string_field().at(1), "repeated_string_two");
}
TEST(UtilTest, TestSetRepeatedFieldElementInvalidIndex) {
TF_ASSERT_OK_AND_ASSIGN(
auto message,
ParseTextProto<ManyFields>(R"pb(repeated_string_field: "")pb"));
const tsl::protobuf::FieldDescriptor* field_desc =
message.GetDescriptor()->FindFieldByName("repeated_string_field");
EXPECT_THAT(SetRepeatedFieldElement(
&message, field_desc, 1, "",
[]() -> absl::Status { return absl::OkStatus(); }),
absl_testing::StatusIs(absl::StatusCode::kOutOfRange,
HasSubstr("Field index out of range")));
}
TEST(UtilTest, TestSetRepeatedFieldElementAlreadyExists) {
TF_ASSERT_OK_AND_ASSIGN(
auto message, ParseTextProto<ManyFields>(
R"pb(repeated_string_field: "existing_string")pb"));
const tsl::protobuf::FieldDescriptor* field_desc =
message.GetDescriptor()->FindFieldByName("repeated_string_field");
const std::string chunk = "inner_inner_string_v2";
TF_EXPECT_OK(SetRepeatedFieldElement(
&message, field_desc, 0, chunk,
[]() -> absl::Status { return absl::OkStatus(); }));
EXPECT_FALSE(message.repeated_string_field().empty());
EXPECT_NE(message.repeated_string_field().at(0), "existing_string");
EXPECT_EQ(message.repeated_string_field().at(0), "inner_inner_string_v2");
}
TEST(UtilTest, TestSetRepeatedFieldElementBadFieldDesc) {
TF_ASSERT_OK_AND_ASSIGN(
auto message, ParseTextProto<ManyFields>(
R"pb(repeated_string_field: "existing_string")pb"));
const tsl::protobuf::FieldDescriptor* singular_field_desc =
message.GetDescriptor()->FindFieldByName("string_field");
EXPECT_DEATH(
auto status = SetRepeatedFieldElement(
&message, singular_field_desc, 0, "",
[]() -> absl::Status { return absl::OkStatus(); }),
HasSubstr("Field is singular; the method requires a repeated field"));
const tsl::protobuf::FieldDescriptor* map_field_desc =
message.GetDescriptor()->FindFieldByName("map_field_uint32");
EXPECT_THAT(SetRepeatedFieldElement(
&message, map_field_desc, 0, "",
[]() -> absl::Status { return absl::OkStatus(); }),
absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition,
HasSubstr("Field is a map")));
}
TEST(UtilTest, TestSetRepeatedFieldElementMessage) {
TF_ASSERT_OK_AND_ASSIGN(
auto message, ParseTextProto<ManyFields>(R"pb(repeated_field {})pb"));
const tsl::protobuf::FieldDescriptor* field_desc =
message.GetDescriptor()->FindFieldByName("repeated_field");
TF_ASSERT_OK_AND_ASSIGN(
auto chunk_proto, ParseTextProto<ManyFields>(R"pb(string_field: "")pb"));
const std::string chunk = chunk_proto.SerializeAsString();
TF_ASSERT_OK(SetRepeatedFieldElement(
&message, field_desc, 0, chunk,
[&message, &field_desc]() -> absl::Status {
tsl::protobuf::Message* inner_message =
message.GetReflection()->MutableRepeatedMessage(&message,
field_desc, 0);
inner_message->GetReflection()->SetString(
inner_message,
field_desc->message_type()->FindFieldByName("string_field"),
"inner_string");
return absl::OkStatus();
}));
EXPECT_FALSE(message.repeated_field().empty());
EXPECT_EQ(message.repeated_field().at(0).string_field(), "inner_string");
}
TEST(UtilTest, TestSetFieldElement) {
TF_ASSERT_OK_AND_ASSIGN(
auto message, ParseTextProto<ManyFields>(R"pb(string_field: "")pb"));
const tsl::protobuf::FieldDescriptor* field_desc =
message.GetDescriptor()->FindFieldByName("string_field");
const std::string chunk = "string_field_v2";
TF_ASSERT_OK(
SetFieldElement(&message, field_desc, chunk,
[]() -> absl::Status { return absl::OkStatus(); }));
EXPECT_EQ(message.string_field(), "string_field_v2");
}
TEST(UtilTest, TestSetFieldElementInvalidRepeated) {
TF_ASSERT_OK_AND_ASSIGN(
auto message,
ParseTextProto<ManyFields>(R"pb(repeated_string_field: "")pb"));
const tsl::protobuf::FieldDescriptor* field_desc =
message.GetDescriptor()->FindFieldByName("repeated_string_field");
EXPECT_DEATH(
auto status =
SetFieldElement(&message, field_desc, "",
[]() -> absl::Status { return absl::OkStatus(); }),
HasSubstr("Field is repeated; the method requires a singular field"));
}
TEST(UtilTest, TestSetFieldElementAlreadyExists) {
TF_ASSERT_OK_AND_ASSIGN(
auto message,
ParseTextProto<ManyFields>(R"pb(string_field: "existing_string")pb"));
const tsl::protobuf::FieldDescriptor* field_desc =
message.GetDescriptor()->FindFieldByName("string_field");
const std::string chunk = "inner_inner_string_v2";
TF_EXPECT_OK(
SetFieldElement(&message, field_desc, chunk,
[]() -> absl::Status { return absl::OkStatus(); }));
EXPECT_EQ(message.string_field(), "inner_inner_string_v2");
}
TEST(UtilTest, TestSetFieldElementMessage) {
TF_ASSERT_OK_AND_ASSIGN(auto message,
ParseTextProto<ManyFields>(R"pb(field_one {})pb"));
const tsl::protobuf::FieldDescriptor* field_desc =
message.GetDescriptor()->FindFieldByName("field_one");
TF_ASSERT_OK_AND_ASSIGN(
auto chunk_proto, ParseTextProto<ManyFields>(R"pb(string_field: "")pb"));
const std::string chunk = chunk_proto.SerializeAsString();
TF_ASSERT_OK(SetFieldElement(
&message, field_desc, chunk, [&message, &field_desc]() -> absl::Status {
tsl::protobuf::Message* inner_message =
message.GetReflection()->MutableMessage(&message, field_desc);
inner_message->GetReflection()->SetString(
inner_message,
field_desc->message_type()->FindFieldByName("string_field"),
"inner_string");
return absl::OkStatus();
}));
EXPECT_TRUE(message.has_field_one());
EXPECT_EQ(message.field_one().string_field(), "inner_string");
}
TEST(UtilTest, TestAddMapEntry) {
TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
const tsl::protobuf::FieldDescriptor* field_desc =
message.GetDescriptor()->FindFieldByName("map_field_int64");
FieldType map_key = -4234;
TF_ASSERT_OK(AddMapEntry(&message, field_desc, map_key));
TF_ASSERT_OK_AND_ASSIGN(int map_entry_index,
FindMapKey(message, *field_desc, nullptr, map_key));
tsl::protobuf::Message* map_entry =
message.GetReflection()->MutableRepeatedMessage(&message, field_desc,
map_entry_index);
map_entry->GetReflection()->SetString(
map_entry, field_desc->message_type()->FindFieldByNumber(2),
"map_value_-4234");
EXPECT_EQ(message.map_field_int64().at(-4234), "map_value_-4234");
}
TEST(UtilTest, GetFieldInvalidIndex) {
std::vector<FieldType> fields = {"field_one"s, "repeated_field"s, 100};
TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
EXPECT_THAT(GetField(message, fields),
absl_testing::StatusIs(absl::StatusCode::kNotFound,
HasSubstr("Can't access index 100")));
}
TEST(UtilTest, GetFieldInvalidField) {
std::vector<FieldType> fields = {"field_one"s, "INVALID"s};
TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
EXPECT_THAT(GetField(message, fields),
absl_testing::StatusIs(absl::StatusCode::kNotFound,
HasSubstr("Field not found: INVALID")));
}
TEST(UtilTest, GetFieldInvalidMapKey) {
std::vector<FieldType> fields = {"map_field_int64"s, 10000};
TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
EXPECT_THAT(GetField(message, fields),
absl_testing::StatusIs(absl::StatusCode::kNotFound,
HasSubstr("couldn't find key: 10000")));
}
TEST(UtilTest, GetField) {
TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
std::vector<FieldType> fields = {"field_one"s, "repeated_field"s};
TF_ASSERT_OK_AND_ASSIGN(auto ret, GetField(message, fields));
EXPECT_EQ(-1, ret.index);
EXPECT_EQ(2, ret.parent->GetReflection()->FieldSize(*ret.parent, ret.field));
fields = {"nested_map_bool"s, "true"s, "string_field"s};
TF_ASSERT_OK_AND_ASSIGN(auto ret2, GetField(message, fields));
EXPECT_EQ(-1, ret2.index);
EXPECT_EQ("string_true",
ret2.parent->GetReflection()->GetString(*ret2.parent, ret2.field));
}
TEST(UtilTest, FindMapKey) {
TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
const tsl::protobuf::FieldDescriptor* map_field =
message.GetDescriptor()->FindFieldByName("nested_map_bool");
TF_ASSERT_OK_AND_ASSIGN(int i,
FindMapKey(message, *map_field, nullptr, true));
EXPECT_EQ(1, i);
TF_ASSERT_OK_AND_ASSIGN(i, FindMapKey(message, *map_field, nullptr, false));
EXPECT_EQ(0, i);
map_field = message.GetDescriptor()->FindFieldByName("map_field_int64");
TF_ASSERT_OK_AND_ASSIGN(i, FindMapKey(message, *map_field, nullptr, -1345));
EXPECT_EQ(0, i);
TF_ASSERT_OK_AND_ASSIGN(i, FindMapKey(message, *map_field, nullptr, 1345));
EXPECT_EQ(-1, i);
}
TEST(UtilTest, FindMapKeyInvalid) {
TF_ASSERT_OK_AND_ASSIGN(auto message, MakeManyFields());
const tsl::protobuf::FieldDescriptor* not_map_field =
message.GetDescriptor()->FindFieldByName("string_field");
EXPECT_THAT(FindMapKey(message, *not_map_field, nullptr, -1345),
absl_testing::StatusIs(absl::StatusCode::kInvalidArgument,
HasSubstr("was given a non map field")));
}
TEST(UtilTest, TestHumanReadableBytes) {
EXPECT_EQ("1.0KiB", HumanReadableBytes(1024));
EXPECT_EQ("5.5KiB", HumanReadableBytes(5632));
EXPECT_EQ("52.2KiB", HumanReadableBytes(53432));
EXPECT_EQ("72.9MiB", HumanReadableBytes(76493281));
EXPECT_EQ("57.0MiB", HumanReadableBytes(5.977e7));
EXPECT_EQ("1.0GiB", HumanReadableBytes(1.074e9));
EXPECT_EQ("15.4GiB", HumanReadableBytes(16493342281));
}
TEST(UtilTest, TestHumanReadableDuration) {
EXPECT_EQ("534 microseconds", HumanReadableDuration(534));
EXPECT_EQ("1.00 ms", HumanReadableDuration(1000));
EXPECT_EQ("14.33 ms", HumanReadableDuration(14328));
EXPECT_EQ("95.83 s", HumanReadableDuration(95825433));
}
TEST(UtilTest, TestReadChunk) {
std::string cpb_file =
io::JoinPath(tensorflow::testing::TensorFlowSrcRoot(),
"tools/proto_splitter/testdata", "split-standard.cpb");
TF_ASSERT_OK_AND_ASSIGN(auto reader, GetRiegeliReader(cpb_file));
auto read_metadata = GetChunkMetadata(reader);
if (!read_metadata.ok()) {
reader.Close();
TF_ASSERT_OK(read_metadata.status());
}
::tensorflow::proto_splitter::ChunkMetadata metadata = read_metadata.value();
std::vector<::tensorflow::proto_splitter::ChunkInfo> chunks_info(
metadata.chunks().begin(), metadata.chunks().end());
for (const auto& chunk_info : chunks_info) {
TF_ASSERT_OK_AND_ASSIGN(std::string chunk, ReadChunk(reader, chunk_info));
ASSERT_EQ(chunk.size(), chunk_info.size());
}
}
} // namespace
} // namespace tools::proto_splitter
} // namespace tensorflow