Project import generated by Copybara.

GitOrigin-RevId: e3a43e4e5e519cd14df7095749059e2613bdcf76
This commit is contained in:
MediaPipe Team
2020-07-08 18:24:46 -07:00
committed by jqtang
parent 67bd8a2bf0
commit e9fbe868e5
96 changed files with 2546 additions and 1175 deletions
+19 -2
View File
@@ -36,7 +36,7 @@ exports_files([
mediapipe_proto_library(
name = "calculator_proto",
srcs = ["calculator.proto"],
visibility = [":mediapipe_internal"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:calculator_options_proto",
"//mediapipe/framework:mediapipe_options_proto",
@@ -68,7 +68,7 @@ mediapipe_proto_library(
mediapipe_proto_library(
name = "calculator_profile_proto",
srcs = ["calculator_profile.proto"],
visibility = [":mediapipe_internal"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:calculator_options_proto",
"//mediapipe/framework:calculator_proto",
@@ -830,6 +830,8 @@ cc_library(
":port",
":timestamp",
":type_map",
"//mediapipe/framework/deps:no_destructor",
"//mediapipe/framework/deps:registration",
"//mediapipe/framework/port:core_proto",
"//mediapipe/framework/port:integral_types",
"//mediapipe/framework/port:logging",
@@ -1524,6 +1526,21 @@ cc_test(
],
)
cc_test(
name = "packet_registration_test",
size = "small",
srcs = ["packet_registration_test.cc"],
deps = [
":calculator_framework",
":packet",
":packet_test_cc_proto",
":type_map",
"//mediapipe/framework/port:core_proto",
"//mediapipe/framework/port:gtest_main",
"@com_google_absl//absl/strings",
],
)
cc_test(
name = "packet_generator_test",
size = "small",
+4 -1
View File
@@ -113,8 +113,11 @@ class CalculatorContract {
// calculations should use SetProcessTimestampBounds.
// When true, Process is called for every new timestamp bound, with or without
// new packets. A call to Process with only an input timestamp bound is
// new packets. A call to Process with only an input timestamp bound is
// normally used to compute a new output timestamp bound.
// NOTE: Also, when true, Process is called when input streams become done,
// which means, Process needs to handle input streams in "done" state.
// (Usually, by closing calculators' outputs where and when appropriate.)
void SetProcessTimestampBounds(bool process_timestamps) {
process_timestamps_ = process_timestamps;
}
+13 -9
View File
@@ -91,6 +91,9 @@ typedef ::mediapipe::StatusOr<OutputStreamPoller> StatusOrPoller;
// {{"video_id", mediapipe::MakePacket<std::string>("Ex-uGhDzue4")}}));
// // See mediapipe/framework/graph_runner.h for an interface
// // to insert and extract packets from a graph as it runs.
// // Once it is done using the graph, close its streams and wait till done.
// MP_RETURN_IF_ERROR(graph->CloseAllInputStreams());
// MP_RETURN_IF_ERROR(graph->WaitUntilDone());
class CalculatorGraph {
public:
// Defines possible modes for adding a packet to a graph input stream.
@@ -157,8 +160,9 @@ class CalculatorGraph {
std::function<::mediapipe::Status(const Packet&)> packet_callback);
// Adds an OutputStreamPoller for a stream. This provides a synchronous,
// polling API for accessing a stream's output. For asynchronous output, use
// ObserveOutputStream. See also the helpers in tool/sink.h.
// polling API for accessing a stream's output. Should only be called before
// Run() or StartRun(). For asynchronous output, use ObserveOutputStream. See
// also the helpers in tool/sink.h.
StatusOrPoller AddOutputStreamPoller(const std::string& stream_name);
// Gets output side packet by name after the graph is done. However, base
@@ -300,6 +304,13 @@ class CalculatorGraph {
void RecordError(const ::mediapipe::Status& error)
ABSL_LOCKS_EXCLUDED(error_mutex_);
// Combines errors into a status. Returns true if the vector of errors is
// non-empty.
bool GetCombinedErrors(const std::string& error_prefix,
::mediapipe::Status* error_status);
// Convenience overload which specifies a default error prefix.
bool GetCombinedErrors(::mediapipe::Status* error_status);
// Returns the maximum input stream queue size.
int GetMaxInputStreamQueueSize();
@@ -501,13 +512,6 @@ class CalculatorGraph {
void CleanupAfterRun(::mediapipe::Status* status)
ABSL_LOCKS_EXCLUDED(error_mutex_);
// Combines errors into a status. Returns true if the vector of errors is
// non-empty.
bool GetCombinedErrors(const std::string& error_prefix,
::mediapipe::Status* error_status);
// Convenience overload which specifies a default error prefix.
bool GetCombinedErrors(::mediapipe::Status* error_status);
// Calls HandlePreRunStatus or HandleStatus on the StatusHandlers. Which one
// is called depends on the GraphRunState parameter (PRE_RUN or POST_RUN).
// current_run_side_packets_ must be set before this function is called.
+3 -2
View File
@@ -458,8 +458,9 @@ class Vector3
// return the index of the largest component (fabs)
int LargestAbsComponent() const {
Vector3 temp = Abs();
return temp[0] > temp[1] ? temp[0] > temp[2] ? 0 : 2
: temp[1] > temp[2] ? 1 : 2;
return temp[0] > temp[1] ? temp[0] > temp[2] ? 0 : 2
: temp[1] > temp[2] ? 1
: 2;
}
// return the index of the smallest, median ,largest component of the vector
+1 -1
View File
@@ -155,7 +155,7 @@ class InputStreamHandler {
// max number of invocations that are allowed to be scheduled is reached.
// Returns true if at least one invocation has been scheduled.
// The latest minimum timestamp bound of the input streams is returned in
// *input_bound iff the latest readiness of the node is kNotReady when the
// *input_bound if the latest readiness of the node is kNotReady when the
// function returns. During batching, this value will be equal to the
// timestamp of the first set of inputs in the batch. In other cases,
// Timestamp::Unset() is returned.
@@ -61,11 +61,25 @@ class LegacyCalculatorSupport {
// platforms.
#ifndef __APPLE__
ABSL_CONST_INIT
#endif // !__APPLE__
#endif // !__APPLE__
static thread_local C* current_; // NOLINT
};
};
// We only declare this variable for two specializations of the template because
// it is only meant to be used for these two types.
// Note that, since these variables are members of specific template
// _specializations_, they are not themselves templates, and therefore their
// definitions must be in the .cc file. However, a declaration still needs to be
// included in the header, or some compilers will assume they have no
// definition.
template <>
thread_local CalculatorContext*
LegacyCalculatorSupport::Scoped<CalculatorContext>::current_;
template <>
thread_local CalculatorContract*
LegacyCalculatorSupport::Scoped<CalculatorContract>::current_;
} // namespace mediapipe
#endif // MEDIAPIPE_FRAMEWORK_LEGACY_CALCULATOR_SUPPORT_H_
+12
View File
@@ -51,6 +51,18 @@ const HolderBase* GetHolder(const Packet& packet) {
return packet.holder_.get();
}
::mediapipe::StatusOr<Packet> PacketFromDynamicProto(
const std::string& type_name, const std::string& serialized) {
ASSIGN_OR_RETURN(
auto message_holder,
packet_internal::MessageHolderRegistry::CreateByName(type_name));
auto* message =
const_cast<proto_ns::MessageLite*>(message_holder->GetProtoMessageLite());
RET_CHECK_NE(message, nullptr);
RET_CHECK(message->ParseFromString(serialized));
return packet_internal::Create(message_holder.release());
}
} // namespace packet_internal
Packet Packet::At(class Timestamp timestamp) const& {
+82 -105
View File
@@ -27,6 +27,8 @@
#include "absl/memory/memory.h"
#include "absl/strings/str_cat.h"
#include "absl/synchronization/mutex.h"
#include "mediapipe/framework/deps/no_destructor.h"
#include "mediapipe/framework/deps/registration.h"
#include "mediapipe/framework/port.h"
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/logging.h"
@@ -51,6 +53,8 @@ Packet Create(HolderBase* holder, Timestamp timestamp);
Packet Create(std::shared_ptr<HolderBase> holder, Timestamp timestamp);
const HolderBase* GetHolder(const Packet& packet);
const std::shared_ptr<HolderBase>& GetHolderShared(const Packet& packet);
::mediapipe::StatusOr<Packet> PacketFromDynamicProto(
const std::string& type_name, const std::string& serialized);
} // namespace packet_internal
// A generic container class which can hold data of any type. The type of
@@ -355,112 +359,11 @@ class HolderBase {
// Downcasts this to Holder<T>. Returns nullptr if deserialization
// failed or if the requested type is not what is stored.
template <typename T>
inline Holder<T>* As(
typename std::enable_if<
(!std::is_base_of<proto_ns::MessageLite, T>::value &&
!std::is_base_of<proto_ns::Message, T>::value) ||
(std::is_same<proto_ns::MessageLite, T>::value ||
std::is_same<proto_ns::Message, T>::value)>::type* = 0) {
if (HolderIsOfType<Holder<T>>() || HolderIsOfType<ForeignHolder<T>>()) {
return static_cast<Holder<T>*>(this);
}
// Does not hold a T.
return nullptr;
}
// For proto Message/MessageLite subclasses.
// When holder data is a concrete proto, the method downcasts this to
// Holder<T> if the requested type is what is stored.
// When holder data is a generic proto Message/MessageLite and a concrete
// proto type T is requested, the method will downcast the HolderBase to
// Holder<T> if the proto data is an instance of T.
template <typename T>
inline Holder<T>* As(
typename std::enable_if<
(std::is_base_of<proto_ns::MessageLite, T>::value ||
std::is_base_of<proto_ns::Message, T>::value) &&
(!std::is_same<proto_ns::MessageLite, T>::value &&
!std::is_same<proto_ns::Message, T>::value)>::type* = 0) {
// Holder data is an instance of subclass type T.
if (HolderIsOfType<Holder<T>>() || HolderIsOfType<ForeignHolder<T>>()) {
return static_cast<Holder<T>*>(this);
}
// Holder data is a generic proto Message/MessageLite and a subclass type T
// is requested.
if (HolderIsOfType<Holder<proto_ns::Message>>() ||
HolderIsOfType<ForeignHolder<proto_ns::Message>>() ||
HolderIsOfType<Holder<proto_ns::MessageLite>>() ||
HolderIsOfType<ForeignHolder<proto_ns::MessageLite>>()) {
// TODO: Holder<proto_ns::Message/MessageLite> cannot be
// legally downcast to Holder<T>, even though that downcast works in
// practice. Need to propose a better way to do the downcast.
Holder<T>* holder = static_cast<Holder<T>*>(this);
T tmp;
VLOG(2) << "Holder proto data type: " << holder->data().GetTypeName()
<< " vs requested proto type: " << tmp.GetTypeName();
if (tmp.GetTypeName() == holder->data().GetTypeName()) {
return holder;
}
}
// Does not hold a T.
return nullptr;
}
Holder<T>* As();
// Same as non-const As() function.
template <typename T>
inline const Holder<T>* As(
typename std::enable_if<
(!std::is_base_of<proto_ns::MessageLite, T>::value &&
!std::is_base_of<proto_ns::Message, T>::value) ||
(std::is_same<proto_ns::MessageLite, T>::value ||
std::is_same<proto_ns::Message, T>::value)>::type* = 0) const {
if (HolderIsOfType<Holder<T>>() || HolderIsOfType<ForeignHolder<T>>()) {
return static_cast<const Holder<T>*>(this);
}
// Does not hold a T.
return nullptr;
}
// For proto Message/MessageLite subclasses.
// When holder data is a concrete proto, the method downcasts this to
// Holder<T> if the requested type is what is stored.
// When holder data is a generic proto Message/MessageLite and a concrete
// proto type T is requested, the method will downcast the HolderBase to
// Holder<T> if the proto data is an instance of T.
template <typename T>
inline const Holder<T>* As(
typename std::enable_if<
(std::is_base_of<proto_ns::MessageLite, T>::value ||
std::is_base_of<proto_ns::Message, T>::value) &&
(!std::is_same<proto_ns::MessageLite, T>::value &&
!std::is_same<proto_ns::Message, T>::value)>::type* = 0) const {
if (HolderIsOfType<Holder<T>>() || HolderIsOfType<ForeignHolder<T>>()) {
return static_cast<const Holder<T>*>(this);
}
// Holder data is a generic proto Message/MessageLite and a subclass type T
// is requested.
if (HolderIsOfType<Holder<proto_ns::Message>>() ||
HolderIsOfType<ForeignHolder<proto_ns::Message>>() ||
HolderIsOfType<Holder<proto_ns::MessageLite>>() ||
HolderIsOfType<ForeignHolder<proto_ns::MessageLite>>()) {
// TODO: Holder<proto_ns::Message/MessageLite> cannot be
// legally downcast to Holder<T>, even though that downcast works in
// practice. Need to propose a better way to do the downcast.
Holder<T>* holder = static_cast<const Holder<T>*>(this);
T tmp;
VLOG(2) << "Holder proto data type: " << holder->data().GetTypeName()
<< " vs requested proto type: " << tmp.GetTypeName();
if (tmp.GetTypeName() == holder->data().GetTypeName()) {
return holder;
}
}
// Does not hold a T.
return nullptr;
}
const Holder<T>* As() const;
// Returns the pointer to MessageLite type for the data in holder, if
// underlying object is protocol buffer type, otherwise, nullptr is returned.
@@ -520,12 +423,68 @@ ConvertToVectorOfProtoMessageLitePtrs(const T* data,
return result;
}
// This registry is used to create Holders of the right concrete C++ type given
// a proto type std::string (which is used as the registration key).
class MessageHolderRegistry
: public GlobalFactoryRegistry<std::unique_ptr<HolderBase>> {};
template <typename T>
struct is_concrete_proto_t
: public std::integral_constant<
bool, std::is_base_of<proto_ns::MessageLite, T>{} &&
!std::is_same<proto_ns::MessageLite, T>{} &&
!std::is_same<proto_ns::Message, T>{}> {};
// Registers a message type. T must be a non-cv-qualified concrete proto type.
template <typename T>
struct MessageRegistrationImpl {
static NoDestructor<mediapipe::RegistrationToken> registration;
};
// Static members of template classes can be defined in the header.
template <typename T>
NoDestructor<mediapipe::RegistrationToken>
MessageRegistrationImpl<T>::registration(MessageHolderRegistry::Register(
T{}.GetTypeName(), [] { return absl::make_unique<Holder<T>>(new T); }));
// For non-Message payloads, this does nothing.
template <typename T, typename Enable = void>
struct HolderSupport {
static void EnsureStaticInit() {}
};
// This template ensures that, for each concrete MessageLite subclass that is
// stored in a Packet, we register a function that allows us to create a
// Holder with the correct payload type from the proto's type name.
template <typename T>
struct HolderSupport<T,
typename std::enable_if<is_concrete_proto_t<T>{}>::type> {
// We must use std::remove_cv to ensure we don't try to register Foo twice if
// there are Holder<Foo> and Holder<const Foo>. TODO: lift this
// up to Holder?
using R = MessageRegistrationImpl<typename std::remove_cv<T>::type>;
// For the registration static member to be instantiated, it needs to be
// referenced in a context that requires the definition to exist (see ISO/IEC
// C++ 2003 standard, 14.7.1). Calling this ensures that's the case.
// We need two different call-sites to cover proto types for which packets
// are only ever created (i.e. the protos are only produced by calculators)
// and proto types for which packets are only ever consumed (i.e. the protos
// are only consumed by calculators).
static void EnsureStaticInit() { CHECK(R::registration.get() != nullptr); }
};
template <typename T>
class Holder : public HolderBase {
public:
explicit Holder(const T* ptr) : ptr_(ptr) { SetHolderTypeId<Holder>(); }
explicit Holder(const T* ptr) : ptr_(ptr) {
HolderSupport<T>::EnsureStaticInit();
SetHolderTypeId<Holder>();
}
~Holder() override { delete_helper(); }
const T& data() const { return *ptr_; }
const T& data() const {
HolderSupport<T>::EnsureStaticInit();
return *ptr_;
}
size_t GetTypeId() const final { return tool::GetTypeHash<T>(); }
// Releases the underlying data pointer and transfers the ownership to a
// unique pointer.
@@ -622,6 +581,24 @@ class ForeignHolder : public Holder<T> {
}
};
template <typename T>
Holder<T>* HolderBase::As() {
if (HolderIsOfType<Holder<T>>() || HolderIsOfType<ForeignHolder<T>>()) {
return static_cast<Holder<T>*>(this);
}
// Does not hold a T.
return nullptr;
}
template <typename T>
const Holder<T>* HolderBase::As() const {
if (HolderIsOfType<Holder<T>>() || HolderIsOfType<ForeignHolder<T>>()) {
return static_cast<const Holder<T>*>(this);
}
// Does not hold a T.
return nullptr;
}
} // namespace packet_internal
inline Packet::Packet(const Packet& packet)
@@ -0,0 +1,57 @@
// Copyright 2020 The MediaPipe Authors.
//
// 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 "absl/strings/str_cat.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/packet.h"
#include "mediapipe/framework/packet_test.pb.h"
#include "mediapipe/framework/port/core_proto_inc.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/status_matchers.h"
namespace mediapipe {
namespace {
namespace test_ns {
class TestSinkCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Tag("IN").Set<mediapipe::InputOnlyProto>();
cc->Outputs().Tag("OUT").Set<int>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) override {
int x = cc->Inputs().Tag("IN").Get<mediapipe::InputOnlyProto>().x();
cc->Outputs().Tag("OUT").AddPacket(
MakePacket<int>(x).At(cc->InputTimestamp()));
return ::mediapipe::OkStatus();
}
};
REGISTER_CALCULATOR(::mediapipe::test_ns::TestSinkCalculator);
} // namespace test_ns
TEST(PacketTest, InputTypeRegistration) {
using testing::Contains;
ASSERT_EQ(mediapipe::InputOnlyProto{}.GetTypeName(),
"mediapipe.InputOnlyProto");
EXPECT_THAT(packet_internal::MessageHolderRegistry::GetRegisteredNames(),
Contains("mediapipe.InputOnlyProto"));
}
} // namespace
} // namespace mediapipe
+21 -41
View File
@@ -174,54 +174,13 @@ TEST(PacketTest, ReturnGenericProtobufMessage) {
.x(0));
}
TEST(PacketTest, ReturnProtobufMessageSubType) {
std::unique_ptr<::mediapipe::PacketTestProto> proto_ptr(
new ::mediapipe::PacketTestProto);
proto_ptr->add_x(123);
Packet packet = Adopt(static_cast<proto_ns::Message*>(proto_ptr.release()));
EXPECT_EQ(123, packet.Get<::mediapipe::PacketTestProto>().x(0));
EXPECT_EQ(123, packet.Get<const ::mediapipe::PacketTestProto>().x(0));
}
TEST(PacketTest, TryWrongProtobufMessageSubType) {
// Packet of PacketTestProto.
std::unique_ptr<::mediapipe::PacketTestProto> proto_ptr(
new ::mediapipe::PacketTestProto);
proto_ptr->add_x(123);
Packet packet = Adopt(proto_ptr.release());
EXPECT_FALSE(packet.ValidateAsType<::mediapipe::SimpleProto>().ok());
EXPECT_TRUE(packet.ValidateAsType<::mediapipe::PacketTestProto>().ok());
// Packet of proto_ns::Message.
proto_ptr.reset(new ::mediapipe::PacketTestProto);
proto_ptr->add_x(456);
Packet packet2 = Adopt(static_cast<proto_ns::Message*>(proto_ptr.release()));
EXPECT_FALSE(packet2.ValidateAsType<::mediapipe::SimpleProto>().ok());
EXPECT_TRUE(packet2.ValidateAsType<::mediapipe::PacketTestProto>().ok());
EXPECT_EQ(123, packet.Get<::mediapipe::PacketTestProto>().x(0));
}
TEST(PacketTest, ReturnProtobufMessageLiteSubType) {
std::unique_ptr<::mediapipe::PacketTestProto> proto_ptr(
new ::mediapipe::PacketTestProto);
proto_ptr->add_x(123);
Packet packet =
Adopt(static_cast<proto_ns::MessageLite*>(proto_ptr.release()));
EXPECT_EQ(123, packet.Get<::mediapipe::PacketTestProto>().x(0));
EXPECT_EQ(123, packet.Get<const ::mediapipe::PacketTestProto>().x(0));
}
TEST(PacketTest, TryWrongProtobufMessageLiteSubType) {
// Packet of PacketTestProto.
std::unique_ptr<::mediapipe::PacketTestProto> proto_ptr(
new ::mediapipe::PacketTestProto);
// Packet of proto_ns::MessageLite.
proto_ptr->add_x(456);
Packet packet =
Adopt(static_cast<proto_ns::MessageLite*>(proto_ptr.release()));
EXPECT_FALSE(packet.ValidateAsType<::mediapipe::SimpleProto>().ok());
EXPECT_TRUE(packet.ValidateAsType<::mediapipe::PacketTestProto>().ok());
EXPECT_EQ(456, packet.Get<::mediapipe::PacketTestProto>().x(0));
}
TEST(PacketTest, GetProtoBase) {
@@ -505,5 +464,26 @@ TEST(PacketTest, TestConsumeOrCopyBoundedArray) {
EXPECT_TRUE(packet2.IsEmpty());
}
TEST(PacketTest, MessageHolderRegistration) {
using testing::Contains;
Packet packet = MakePacket<mediapipe::SimpleProto>();
ASSERT_EQ(mediapipe::SimpleProto{}.GetTypeName(), "mediapipe.SimpleProto");
EXPECT_THAT(packet_internal::MessageHolderRegistry::GetRegisteredNames(),
Contains("mediapipe.SimpleProto"));
}
TEST(PacketTest, PacketFromSerializedProto) {
mediapipe::SimpleProto original;
original.add_value("foo");
std::string serialized = original.SerializeAsString();
StatusOr<Packet> maybe_packet = packet_internal::PacketFromDynamicProto(
"mediapipe.SimpleProto", serialized);
MP_ASSERT_OK(maybe_packet);
Packet packet = maybe_packet.ValueOrDie();
MP_EXPECT_OK(packet.ValidateAsType<::mediapipe::SimpleProto>());
EXPECT_FALSE(packet.ValidateAsType<::mediapipe::PacketTestProto>().ok());
}
} // namespace
} // namespace mediapipe
+6
View File
@@ -39,3 +39,9 @@ message SerializationProxyProto {
repeated float float_value = 2;
repeated string string_value = 3;
}
// This proto should be used only as an input to a calculator, to verify that
// that case is covered.
message InputOnlyProto {
optional int32 x = 1;
}
+1 -1
View File
@@ -46,7 +46,7 @@
// but may or may not still be able to run other OpenGL code.
#if !defined(MEDIAPIPE_DISABLE_GL_COMPUTE) && \
(defined(__APPLE__) || defined(__EMSCRIPTEN__) || \
defined(MEDIAPIPE_DISABLE_GPU))
defined(MEDIAPIPE_DISABLE_GPU) || MEDIAPIPE_USING_SWIFTSHADER)
#define MEDIAPIPE_DISABLE_GL_COMPUTE
#endif
+144 -133
View File
@@ -143,29 +143,29 @@ TEST_F(GraphTracerTest, CalculatorTrace) {
{{MakePacket<std::string>("goodbye").At(start_timestamp_)}});
// Validate the GraphTrace data.
EXPECT_THAT(GetTrace(),
EqualsProto(::mediapipe::ParseTextProtoOrDie<GraphTrace>(R"(
base_time: 1608911100000000
base_timestamp: 1608911100000000
stream_name: ""
stream_name: "input_stream"
stream_name: "output_stream"
calculator_trace {
node_id: 0
input_timestamp: 0
event_type: PROCESS
start_time: 0
finish_time: 10000
thread_id: 0
input_trace {
finish_time: 0
packet_timestamp: 0
stream_id: 1
event_data: 1
}
output_trace { packet_timestamp: 0 stream_id: 2 }
}
)")));
EXPECT_THAT(
GetTrace(), EqualsProto(::mediapipe::ParseTextProtoOrDie<GraphTrace>(R"(
base_time: 1608911100000000
base_timestamp: 1608911100000000
stream_name: ""
stream_name: "input_stream"
stream_name: "output_stream"
calculator_trace {
node_id: 0
input_timestamp: 0
event_type: PROCESS
start_time: 0
finish_time: 10000
thread_id: 0
input_trace {
finish_time: 0
packet_timestamp: 0
stream_id: 1
event_data: 1
}
output_trace { packet_timestamp: 0 stream_id: 2 event_data: 2 }
}
)")));
}
TEST_F(GraphTracerTest, GraphTrace) {
@@ -205,92 +205,101 @@ TEST_F(GraphTracerTest, GraphTrace) {
LogOutputPackets("PCalculator_3", GraphTrace::PROCESS, curr_time,
{{MakePacket<std::string>("out").At(start_timestamp_)}});
curr_time += absl::Microseconds(2000);
ClearCalculatorContext("PCalculator_3");
LogInputPackets("PCalculator_3", GraphTrace::PROCESS, curr_time,
// Note: the packet data ID is based on the packet's payload address, which
// means the same ID can be reused if data is allocated in the same location
// as a previously expired packet (b/160212191). This means the generated
// trace can change depending on the allocator. To keep results stable, we
// must keep the packets used in this test alive until the end. Each
// TestContextBuilder happens to keep a reference to all packets for the last
// context, so for now we just create a separate TestContextBuilder instead of
// clearing it. TODO: revise this test.
SetUpCalculatorContext("PCalculator_3a", /*node_id=*/2, {"up_2"}, {"down_2"});
LogInputPackets("PCalculator_3a", GraphTrace::PROCESS, curr_time,
{MakePacket<std::string>("pup").At(start_timestamp_ + 5)});
curr_time += absl::Microseconds(20000);
LogOutputPackets(
"PCalculator_3", GraphTrace::PROCESS, curr_time,
"PCalculator_3a", GraphTrace::PROCESS, curr_time,
{{MakePacket<std::string>("pout").At(start_timestamp_ + 5)}});
curr_time += absl::Microseconds(1000);
// Validate the GraphTrace data.
EXPECT_THAT(GetTrace(),
EqualsProto(::mediapipe::ParseTextProtoOrDie<GraphTrace>(R"(
base_time: 1608911100000000
base_timestamp: 1608911100000000
stream_name: ""
stream_name: "input_stream"
stream_name: "up_1"
stream_name: "up_2"
stream_name: "down_1"
stream_name: "down_2"
calculator_trace {
node_id: 0
input_timestamp: 0
event_type: PROCESS
start_time: 0
finish_time: 10000
thread_id: 0
input_trace {
finish_time: 0
packet_timestamp: 0
stream_id: 1
event_data: 1
}
output_trace { packet_timestamp: 0 stream_id: 2 }
output_trace { packet_timestamp: 0 stream_id: 3 }
output_trace { packet_timestamp: 5 stream_id: 3 }
}
calculator_trace {
node_id: 1
input_timestamp: 0
event_type: PROCESS
start_time: 11000
finish_time: 21000
thread_id: 0
input_trace {
start_time: 10000
finish_time: 11000
packet_timestamp: 0
stream_id: 2
event_data: 2
}
output_trace { packet_timestamp: 0 stream_id: 4 }
}
calculator_trace {
node_id: 2
input_timestamp: 0
event_type: PROCESS
start_time: 16000
finish_time: 36000
thread_id: 0
input_trace {
start_time: 10000
finish_time: 16000
packet_timestamp: 0
stream_id: 3
event_data: 3
}
output_trace { packet_timestamp: 0 stream_id: 5 }
}
calculator_trace {
node_id: 2
input_timestamp: 5
event_type: PROCESS
start_time: 38000
finish_time: 58000
thread_id: 0
input_trace {
start_time: 10000
finish_time: 38000
packet_timestamp: 5
stream_id: 3
event_data: 4
}
output_trace { packet_timestamp: 5 stream_id: 5 }
}
)")));
EXPECT_THAT(
GetTrace(), EqualsProto(::mediapipe::ParseTextProtoOrDie<GraphTrace>(R"(
base_time: 1608911100000000
base_timestamp: 1608911100000000
stream_name: ""
stream_name: "input_stream"
stream_name: "up_1"
stream_name: "up_2"
stream_name: "down_1"
stream_name: "down_2"
calculator_trace {
node_id: 0
input_timestamp: 0
event_type: PROCESS
start_time: 0
finish_time: 10000
thread_id: 0
input_trace {
finish_time: 0
packet_timestamp: 0
stream_id: 1
event_data: 1
}
output_trace { packet_timestamp: 0 stream_id: 2 event_data: 2 }
output_trace { packet_timestamp: 0 stream_id: 3 event_data: 3 }
output_trace { packet_timestamp: 5 stream_id: 3 event_data: 4 }
}
calculator_trace {
node_id: 1
input_timestamp: 0
event_type: PROCESS
start_time: 11000
finish_time: 21000
thread_id: 0
input_trace {
start_time: 10000
finish_time: 11000
packet_timestamp: 0
stream_id: 2
event_data: 5
}
output_trace { packet_timestamp: 0 stream_id: 4 event_data: 6 }
}
calculator_trace {
node_id: 2
input_timestamp: 0
event_type: PROCESS
start_time: 16000
finish_time: 36000
thread_id: 0
input_trace {
start_time: 10000
finish_time: 16000
packet_timestamp: 0
stream_id: 3
event_data: 7
}
output_trace { packet_timestamp: 0 stream_id: 5 event_data: 8 }
}
calculator_trace {
node_id: 2
input_timestamp: 5
event_type: PROCESS
start_time: 38000
finish_time: 58000
thread_id: 0
input_trace {
start_time: 10000
finish_time: 38000
packet_timestamp: 5
stream_id: 3
event_data: 9
}
output_trace { packet_timestamp: 5 stream_id: 5 event_data: 10 }
}
)")));
// No timestamps are completed before start_time_.
// One timestamp is completed before start_time_ + 10ms.
@@ -1275,37 +1284,39 @@ TEST_F(GraphTracerE2ETest, GpuTaskTrace) {
GraphTrace trace_1;
builder.CreateTrace(buffer, absl::InfinitePast(), absl::InfiniteFuture(),
&trace_1);
EXPECT_THAT(trace_1, EqualsProto(::mediapipe::ParseTextProtoOrDie<GraphTrace>(
R"(
base_time: 1100
base_timestamp: 1000
stream_name: ""
stream_name: "stream_1"
stream_name: "stream_2"
calculator_trace {
node_id: 333
input_timestamp: 0
event_type: PROCESS
start_time: 0
finish_time: 1000
input_trace {
finish_time: 0
packet_timestamp: 0
stream_id: 1
event_data: 0
}
output_trace { packet_timestamp: 0 stream_id: 2 }
thread_id: 0
}
calculator_trace {
node_id: 333
input_timestamp: 0
event_type: GPU_TASK
start_time: 100
finish_time: 2100
thread_id: 0
}
)")));
EXPECT_THAT(
trace_1,
EqualsProto(::mediapipe::ParseTextProtoOrDie<GraphTrace>(
R"(
base_time: 1100
base_timestamp: 1000
stream_name: ""
stream_name: "stream_1"
stream_name: "stream_2"
calculator_trace {
node_id: 333
input_timestamp: 0
event_type: PROCESS
start_time: 0
finish_time: 1000
input_trace {
finish_time: 0
packet_timestamp: 0
stream_id: 1
event_data: 0
}
output_trace { packet_timestamp: 0 stream_id: 2 event_data: 0 }
thread_id: 0
}
calculator_trace {
node_id: 333
input_timestamp: 0
event_type: GPU_TASK
start_time: 100
finish_time: 2100
thread_id: 0
}
)")));
GraphTrace trace_2;
builder.CreateLog(buffer, absl::InfinitePast(), absl::InfiniteFuture(),
@@ -330,13 +330,12 @@ class TraceBuilder::Impl {
if (trace_event_registry_[event->event_type].is_stream_event()) {
auto stream_trace = event->is_finish ? result->add_output_trace()
: result->add_input_trace();
if (event->is_finish) {
// Log only the packet id for each output event.
stream_trace->set_stream_id(stream_id_map_[event->stream_id]);
stream_trace->set_packet_timestamp(LogTimestamp(event->packet_ts));
} else {
// Log the full stream trace for each input event.
BuildStreamTrace(*event, stream_trace);
BuildStreamTrace(*event, stream_trace);
if (!event->is_finish) {
// Note: is_finish is true for output events, false for input events.
// For input events, we log some additional timing information. The
// finish_time is the start_time of this Process call, the start_time
// is the finish_time of the Process call that output the packet.
stream_trace->set_finish_time(LogTime(event->event_time));
const TraceEvent* output_event = FindOutputEvent(*event);
if (output_event) {
@@ -116,10 +116,19 @@ NodeReadiness ImmediateInputStreamHandler::GetNodeReadiness(
CHECK_EQ(stream_ts, Timestamp::Done());
if (ProcessTimestampBounds()) {
// With kReadyForClose, the timestamp-bound Done is returned.
// This bound is processed using the preceding input-timestamp.
// TODO: Make all InputStreamHandlers process Done() like this.
ready_timestamps_[i] = stream_ts.PreviousAllowedInStream();
input_timestamp = std::min(input_timestamp, ready_timestamps_[i]);
static const Timestamp kDonePrecedingTimestamp =
Timestamp::Done().PreviousAllowedInStream();
if (prev_ts < kDonePrecedingTimestamp) {
// When kReadyForClose is received for the first time for a sync set,
// it is processed using the timestamp preceding Done() to indicate
// input stream is done, but still needs to be processed.
min_bound = std::min(min_bound, kDonePrecedingTimestamp);
input_timestamp = std::min(input_timestamp, kDonePrecedingTimestamp);
ready_timestamps_[i] = kDonePrecedingTimestamp;
} else {
ready_timestamps_[i] = Timestamp::Done();
}
} else if (prev_ts < Timestamp::Done()) {
stream_became_done = true;
ready_timestamps_[i] = Timestamp::Done();
@@ -133,6 +133,11 @@ class ImmediateInputStreamHandlerTest : public ::testing::Test {
}
}
const InputStream& Input(const CollectionItemId& id) {
CHECK(cc_);
return cc_->Inputs().Get(id);
}
PacketType packet_type_;
std::function<void()> headers_ready_callback_;
std::function<void()> notification_callback_;
@@ -262,6 +267,344 @@ TEST_F(ImmediateInputStreamHandlerTest, ReadyForClose) {
EXPECT_TRUE(errors_.empty());
}
TEST_F(ImmediateInputStreamHandlerTest, ProcessTimestampBounds) {
input_stream_handler_->SetProcessTimestampBounds(true);
Timestamp min_stream_timestamp;
ASSERT_FALSE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::PreStream());
const auto& input_a_id = name_to_id_["input_a"];
const auto& input_b_id = name_to_id_["input_b"];
const auto& input_c_id = name_to_id_["input_c"];
std::list<Packet> packets;
packets.push_back(Adopt(new std::string("packet 1")).At(Timestamp(1)));
input_stream_handler_->AddPackets(input_b_id, packets);
input_stream_handler_->SetNextTimestampBound(input_b_id, Timestamp::Done());
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::Unset());
EXPECT_EQ(cc_->InputTimestamp(), Timestamp(1));
ExpectPackets(cc_->Inputs(), {{"input_b", "packet 1"}});
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp(1));
EXPECT_TRUE(Input(input_a_id).Value().IsEmpty());
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp::Unstarted());
EXPECT_TRUE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp::Unstarted());
// FinalizeInputSet() is a no-op.
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
&cc_->Inputs());
input_stream_handler_->ClearCurrentInputs(cc_);
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::Unset());
EXPECT_EQ(cc_->InputTimestamp(), Timestamp::Max());
EXPECT_TRUE(Input(input_b_id).Value().IsEmpty());
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp::Max());
EXPECT_TRUE(Input(input_a_id).Value().IsEmpty());
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp::Unstarted());
EXPECT_TRUE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp::Unstarted());
// FinalizeInputSet() is a no-op.
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
&cc_->Inputs());
input_stream_handler_->ClearCurrentInputs(cc_);
EXPECT_TRUE(
input_stream_handler_->GetInputStreamManager(input_b_id)->IsEmpty());
input_stream_handler_->SetNextTimestampBound(input_a_id, Timestamp::Done());
input_stream_handler_->SetNextTimestampBound(input_c_id, Timestamp::Done());
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::Unset());
EXPECT_EQ(cc_->InputTimestamp(), Timestamp::Max());
EXPECT_TRUE(Input(input_b_id).Value().IsEmpty());
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp::Max());
EXPECT_TRUE(Input(input_a_id).Value().IsEmpty());
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp::Max());
EXPECT_TRUE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp::Max());
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
&cc_->Inputs());
input_stream_handler_->ClearCurrentInputs(cc_);
EXPECT_TRUE(errors_.empty());
// Schedule invocation for Close.
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::Unset());
EXPECT_EQ(cc_->InputTimestamp(), Timestamp::Done());
EXPECT_TRUE(Input(input_b_id).Value().IsEmpty());
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp::Unset());
EXPECT_TRUE(Input(input_a_id).Value().IsEmpty());
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp::Unset());
EXPECT_TRUE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp::Unset());
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
&cc_->Inputs());
input_stream_handler_->ClearCurrentInputs(cc_);
EXPECT_TRUE(errors_.empty());
}
TEST_F(ImmediateInputStreamHandlerTest,
ProcessTimestampBoundsNoOpScheduleInvocations) {
input_stream_handler_->SetProcessTimestampBounds(true);
const auto& input_a_id = name_to_id_["input_a"];
const auto& input_b_id = name_to_id_["input_b"];
const auto& input_c_id = name_to_id_["input_c"];
Timestamp min_stream_timestamp;
std::list<Packet> packets;
packets.push_back(Adopt(new std::string("packet 1")).At(Timestamp(1)));
input_stream_handler_->AddPackets(input_b_id, packets);
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::Unset());
EXPECT_EQ(cc_->InputTimestamp(), Timestamp(1));
ExpectPackets(cc_->Inputs(), {{"input_b", "packet 1"}});
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp(1));
EXPECT_TRUE(Input(input_a_id).Value().IsEmpty());
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp::Unstarted());
EXPECT_TRUE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp::Unstarted());
// FinalizeInputSet() is a no-op.
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
&cc_->Inputs());
input_stream_handler_->ClearCurrentInputs(cc_);
input_stream_handler_->SetNextTimestampBound(input_a_id, Timestamp::Done());
input_stream_handler_->SetNextTimestampBound(input_c_id, Timestamp::Done());
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::Unset());
EXPECT_EQ(cc_->InputTimestamp(), Timestamp::Max());
EXPECT_TRUE(Input(input_b_id).Value().IsEmpty());
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp(1));
EXPECT_TRUE(Input(input_a_id).Value().IsEmpty());
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp::Max());
EXPECT_TRUE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp::Max());
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
&cc_->Inputs());
input_stream_handler_->ClearCurrentInputs(cc_);
EXPECT_TRUE(errors_.empty());
// Try to schedule invocations several times again. Considering nothing
// changed since last invocation nothing should be scheduled.
for (int i = 0; i < 3; ++i) {
ASSERT_FALSE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp(2));
EXPECT_TRUE(Input(input_b_id).Value().IsEmpty());
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp::Unset());
EXPECT_TRUE(Input(input_a_id).Value().IsEmpty());
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp::Unset());
EXPECT_TRUE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp::Unset());
}
input_stream_handler_->SetNextTimestampBound(input_b_id, Timestamp::Done());
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::Unset());
EXPECT_EQ(cc_->InputTimestamp(), Timestamp::Max());
EXPECT_TRUE(Input(input_b_id).Value().IsEmpty());
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp::Max());
EXPECT_TRUE(Input(input_a_id).Value().IsEmpty());
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp::Max());
EXPECT_TRUE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp::Max());
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
&cc_->Inputs());
input_stream_handler_->ClearCurrentInputs(cc_);
EXPECT_TRUE(errors_.empty());
// Schedule invocation for Close.
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::Unset());
EXPECT_EQ(cc_->InputTimestamp(), Timestamp::Done());
EXPECT_TRUE(Input(input_b_id).Value().IsEmpty());
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp::Unset());
EXPECT_TRUE(Input(input_a_id).Value().IsEmpty());
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp::Unset());
EXPECT_TRUE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp::Unset());
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
&cc_->Inputs());
input_stream_handler_->ClearCurrentInputs(cc_);
EXPECT_TRUE(errors_.empty());
// Try to schedule invocations several times again. Considering nothing
// changed since last invocation nothing should be scheduled.
for (int i = 0; i < 3; ++i) {
ASSERT_FALSE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::Unset());
EXPECT_TRUE(Input(input_b_id).Value().IsEmpty());
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp::Unset());
EXPECT_TRUE(Input(input_a_id).Value().IsEmpty());
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp::Unset());
EXPECT_TRUE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp::Unset());
}
}
// Due to some temporary changes in ImmediateInputStreamHandler some packets
// - were queued but never released
// - were released in incorrect order
// As other test cases were passing, this test case is designed to ensure that.
TEST_F(ImmediateInputStreamHandlerTest, VerifyPacketsReleaseOrder) {
input_stream_handler_->SetProcessTimestampBounds(true);
const auto& input_a_id = name_to_id_["input_a"];
const auto& input_b_id = name_to_id_["input_b"];
const auto& input_c_id = name_to_id_["input_c"];
Packet packet_a = Adopt(new std::string("packet a"));
Packet packet_b = Adopt(new std::string("packet b"));
Packet packet_c = Adopt(new std::string("packet c"));
input_stream_handler_->AddPackets(input_a_id, {packet_a.At(Timestamp(1))});
input_stream_handler_->AddPackets(input_b_id, {packet_b.At(Timestamp(2))});
input_stream_handler_->AddPackets(input_c_id, {packet_c.At(Timestamp(3))});
Timestamp min_stream_timestamp;
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::Unset());
EXPECT_EQ(cc_->InputTimestamp(), Timestamp(1));
ASSERT_FALSE(Input(input_a_id).IsEmpty());
EXPECT_EQ(Input(input_a_id).Get<std::string>(), "packet a");
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp(1));
EXPECT_TRUE(Input(input_b_id).Value().IsEmpty());
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp(1));
EXPECT_TRUE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp(2));
// FinalizeInputSet() is a no-op.
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
&cc_->Inputs());
input_stream_handler_->ClearCurrentInputs(cc_);
input_stream_handler_->AddPackets(input_a_id, {packet_a.At(Timestamp(5))});
input_stream_handler_->AddPackets(input_b_id, {packet_b.At(Timestamp(5))});
input_stream_handler_->AddPackets(input_c_id, {packet_c.At(Timestamp(5))});
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::Unset());
EXPECT_EQ(cc_->InputTimestamp(), Timestamp(2));
EXPECT_TRUE(Input(input_a_id).Value().IsEmpty());
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp(4));
ASSERT_FALSE(Input(input_b_id).IsEmpty());
EXPECT_EQ(Input(input_b_id).Get<std::string>(), "packet b");
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp(2));
EXPECT_TRUE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp(2));
// FinalizeInputSet() is a no-op.
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
&cc_->Inputs());
input_stream_handler_->ClearCurrentInputs(cc_);
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::Unset());
EXPECT_EQ(cc_->InputTimestamp(), Timestamp(3));
EXPECT_TRUE(Input(input_a_id).Value().IsEmpty());
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp(4));
EXPECT_TRUE(Input(input_b_id).Value().IsEmpty());
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp(4));
ASSERT_FALSE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Get<std::string>(), "packet c");
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp(3));
// FinalizeInputSet() is a no-op.
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
&cc_->Inputs());
input_stream_handler_->ClearCurrentInputs(cc_);
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::Unset());
EXPECT_EQ(cc_->InputTimestamp(), Timestamp(5));
ASSERT_FALSE(Input(input_a_id).Value().IsEmpty());
EXPECT_EQ(Input(input_a_id).Get<std::string>(), "packet a");
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp(5));
ASSERT_FALSE(Input(input_b_id).Value().IsEmpty());
EXPECT_EQ(Input(input_b_id).Get<std::string>(), "packet b");
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp(5));
ASSERT_FALSE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Get<std::string>(), "packet c");
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp(5));
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
&cc_->Inputs());
input_stream_handler_->ClearCurrentInputs(cc_);
input_stream_handler_->SetNextTimestampBound(input_a_id, Timestamp::Done());
input_stream_handler_->SetNextTimestampBound(input_b_id, Timestamp::Done());
input_stream_handler_->SetNextTimestampBound(input_c_id, Timestamp::Done());
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::Unset());
EXPECT_EQ(cc_->InputTimestamp(), Timestamp::Max());
EXPECT_TRUE(Input(input_b_id).Value().IsEmpty());
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp::Max());
EXPECT_TRUE(Input(input_a_id).Value().IsEmpty());
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp::Max());
EXPECT_TRUE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp::Max());
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
&cc_->Inputs());
input_stream_handler_->ClearCurrentInputs(cc_);
// Schedule invocation for Close.
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::Unset());
EXPECT_EQ(cc_->InputTimestamp(), Timestamp::Done());
EXPECT_TRUE(Input(input_b_id).Value().IsEmpty());
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp::Unset());
EXPECT_TRUE(Input(input_a_id).Value().IsEmpty());
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp::Unset());
EXPECT_TRUE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp::Unset());
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
&cc_->Inputs());
input_stream_handler_->ClearCurrentInputs(cc_);
ASSERT_FALSE(input_stream_handler_->ScheduleInvocations(
/*max_allowance=*/1, &min_stream_timestamp));
EXPECT_EQ(min_stream_timestamp, Timestamp::Unset());
EXPECT_TRUE(Input(input_b_id).Value().IsEmpty());
EXPECT_EQ(Input(input_b_id).Value().Timestamp(), Timestamp::Unset());
EXPECT_TRUE(Input(input_a_id).Value().IsEmpty());
EXPECT_EQ(Input(input_a_id).Value().Timestamp(), Timestamp::Unset());
EXPECT_TRUE(Input(input_c_id).Value().IsEmpty());
EXPECT_EQ(Input(input_c_id).Value().Timestamp(), Timestamp::Unset());
}
// This test simulates how CalculatorNode::ProcessNode() uses an input
// stream handler and the associated input streams.
TEST_F(ImmediateInputStreamHandlerTest, SimulateProcessNode) {
+57
View File
@@ -641,4 +641,61 @@ class DummyTestCalculator : public CalculatorBase {
};
REGISTER_CALCULATOR(DummyTestCalculator);
// A Calculator that passes the input value to the output after sleeping for
// a set number of microseconds.
class PassThroughWithSleepCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).Set<int>();
cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
cc->InputSidePackets().Tag("SLEEP_MICROS").Set<int>();
cc->InputSidePackets().Tag("CLOCK").Set<std::shared_ptr<Clock>>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status Open(CalculatorContext* cc) final {
cc->SetOffset(TimestampDiff(0));
sleep_micros_ = cc->InputSidePackets().Tag("SLEEP_MICROS").Get<int>();
if (sleep_micros_ < 0) {
return ::mediapipe::InternalError("SLEEP_MICROS should be >= 0");
}
clock_ = cc->InputSidePackets().Tag("CLOCK").Get<std::shared_ptr<Clock>>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
clock_->Sleep(absl::Microseconds(sleep_micros_));
int value = cc->Inputs().Index(0).Value().Get<int>();
cc->Outputs().Index(0).Add(new int(value), cc->InputTimestamp());
return ::mediapipe::OkStatus();
}
private:
int sleep_micros_ = 0;
std::shared_ptr<Clock> clock_;
};
REGISTER_CALCULATOR(PassThroughWithSleepCalculator);
// A Calculator that multiples two input values.
class MultiplyIntCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).Set<int>();
cc->Inputs().Index(1).SetSameAs(&cc->Inputs().Index(0));
// cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
RET_CHECK(cc->Outputs().HasTag("OUT"));
cc->Outputs().Tag("OUT").SetSameAs(&cc->Inputs().Index(0));
return ::mediapipe::OkStatus();
}
::mediapipe::Status Open(CalculatorContext* cc) final {
cc->SetOffset(TimestampDiff(0));
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
int x = cc->Inputs().Index(0).Value().Get<int>();
int y = cc->Inputs().Index(1).Value().Get<int>();
cc->Outputs().Tag("OUT").Add(new int(x * y), cc->InputTimestamp());
return ::mediapipe::OkStatus();
}
};
REGISTER_CALCULATOR(MultiplyIntCalculator);
} // namespace mediapipe
+7
View File
@@ -101,6 +101,13 @@ std::string ParseNameFromStream(const std::string& stream) {
return name;
}
std::pair<std::string, int> ParseTagIndex(const std::string& tag_index) {
std::string tag;
int index;
MEDIAPIPE_CHECK_OK(tool::ParseTagIndex(tag_index, &tag, &index));
return {tag, index};
}
std::pair<std::string, int> ParseTagIndexFromStream(const std::string& stream) {
std::string tag, name;
int index;
+3
View File
@@ -76,6 +76,9 @@ std::string CanonicalNodeName(const CalculatorGraphConfig& graph_config,
// Parses the name from a "tag:index:name".
std::string ParseNameFromStream(const std::string& stream);
// Parses the TagIndex from a "tag:index".
std::pair<std::string, int> ParseTagIndex(const std::string& tag_index);
// Parses the TagIndex from a "tag:index:name".
std::pair<std::string, int> ParseTagIndexFromStream(const std::string& stream);
+6 -4
View File
@@ -13,15 +13,15 @@
# limitations under the License.
#
licenses(["notice"]) # Apache 2.0
package(default_visibility = ["//mediapipe:__subpackages__"])
load(
"//mediapipe/framework/tool:mediapipe_graph.bzl",
"mediapipe_simple_subgraph",
)
licenses(["notice"]) # Apache 2.0
package(default_visibility = ["//mediapipe:__subpackages__"])
filegroup(
name = "test_graph",
srcs = ["test.pbtxt"],
@@ -31,6 +31,8 @@ exports_files([
"test.pbtxt",
"dub_quad_test_subgraph.pbtxt",
"nested_test_subgraph.pbtxt",
"single_flow_container_test.pbtxt",
"dual_flow_container_test.pbtxt",
])
mediapipe_simple_subgraph(