talksik/mediapipe
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 43cd697ec87dcc5cab5051f27960bb77a057399d
This commit is contained in:
MediaPipe Team
2020-03-20 15:28:51 -07:00
committed by jqtang
parent 3b6d3c4058
commit 1722d4b8a2
71 changed files with 6114 additions and 626 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/stream_handler/default_input_stream_handler.cc
+15 -38
View File
@@ -17,16 +17,28 @@
#include <algorithm>
#include "absl/strings/substitute.h"
#include "mediapipe/framework/input_stream_handler.h"
namespace mediapipe {
REGISTER_INPUT_STREAM_HANDLER(DefaultInputStreamHandler);
// Returns all CollectionItemId's for a Collection TagMap.
std::vector<CollectionItemId> GetIds(
const std::shared_ptr<tool::TagMap>& tag_map) {
std::vector<CollectionItemId> result;
for (auto id = tag_map->BeginId(); id < tag_map->EndId(); ++id) {
result.push_back(id);
}
return result;
}
DefaultInputStreamHandler::DefaultInputStreamHandler(
std::shared_ptr<tool::TagMap> tag_map, CalculatorContextManager* cc_manager,
const MediaPipeOptions& options, bool calculator_run_in_parallel)
: InputStreamHandler(std::move(tag_map), cc_manager, options,
calculator_run_in_parallel) {
calculator_run_in_parallel),
sync_set_(this, GetIds(input_stream_managers_.TagMap())) {
if (options.HasExtension(DefaultInputStreamHandlerOptions::ext)) {
SetBatchSize(options.GetExtension(DefaultInputStreamHandlerOptions::ext)
.batch_size());
@@ -35,47 +47,12 @@ DefaultInputStreamHandler::DefaultInputStreamHandler(
NodeReadiness DefaultInputStreamHandler::GetNodeReadiness(
Timestamp* min_stream_timestamp) {
DCHECK(min_stream_timestamp);
*min_stream_timestamp = Timestamp::Done();
Timestamp min_bound = Timestamp::Done();
for (const auto& stream : input_stream_managers_) {
bool empty;
Timestamp stream_timestamp = stream->MinTimestampOrBound(&empty);
if (empty) {
min_bound = std::min(min_bound, stream_timestamp);
}
*min_stream_timestamp = std::min(*min_stream_timestamp, stream_timestamp);
}
if (*min_stream_timestamp == Timestamp::Done()) {
return NodeReadiness::kReadyForClose;
}
if (min_bound > *min_stream_timestamp) {
return NodeReadiness::kReadyForProcess;
}
CHECK_EQ(min_bound, *min_stream_timestamp);
return NodeReadiness::kNotReady;
return sync_set_.GetReadiness(min_stream_timestamp);
}
void DefaultInputStreamHandler::FillInputSet(Timestamp input_timestamp,
InputStreamShardSet* input_set) {
CHECK(input_timestamp.IsAllowedInStream());
CHECK(input_set);
for (CollectionItemId id = input_stream_managers_.BeginId();
id < input_stream_managers_.EndId(); ++id) {
auto& stream = input_stream_managers_.Get(id);
int num_packets_dropped = 0;
bool stream_is_done = false;
Packet current_packet = stream->PopPacketAtTimestamp(
input_timestamp, &num_packets_dropped, &stream_is_done);
CHECK_EQ(num_packets_dropped, 0)
<< absl::Substitute("Dropped $0 packet(s) on input stream \"$1\".",
num_packets_dropped, stream->Name());
AddPacketToShard(&input_set->Get(id), std::move(current_packet),
stream_is_done);
}
sync_set_.FillInputSet(input_timestamp, input_set);
}
} // namespace mediapipe
mediapipe/framework/stream_handler/default_input_stream_handler.h
+3
View File
@@ -45,6 +45,9 @@ class DefaultInputStreamHandler : public InputStreamHandler {
// Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
void FillInputSet(Timestamp input_timestamp,
InputStreamShardSet* input_set) override;
// The packet-set builder.
SyncSet sync_set_;
};
} // namespace mediapipe
mediapipe/framework/stream_handler/immediate_input_stream_handler.cc
+46 -33
View File
@@ -19,6 +19,8 @@
namespace mediapipe {
using SyncSet = InputStreamHandler::SyncSet;
// An input stream handler that delivers input packets to the Calculator
// immediately, with no dependency between input streams. It also invokes
// Calculator::Process when any input stream becomes done.
@@ -47,8 +49,11 @@ class ImmediateInputStreamHandler : public InputStreamHandler {
void FillInputSet(Timestamp input_timestamp,
InputStreamShardSet* input_set) override;
// Record of the last reported timestamp bound for each input stream.
mediapipe::internal::Collection<Timestamp> timestamp_bounds_;
absl::Mutex mutex_;
// The packet-set builder for each input stream.
std::vector<SyncSet> sync_sets_ ABSL_GUARDED_BY(mutex_);
// The input timestamp for each kReadyForProcess input stream.
std::vector<Timestamp> ready_timestamps_ ABSL_GUARDED_BY(mutex_);
};
REGISTER_INPUT_STREAM_HANDLER(ImmediateInputStreamHandler);
@@ -57,31 +62,47 @@ ImmediateInputStreamHandler::ImmediateInputStreamHandler(
CalculatorContextManager* calculator_context_manager,
const MediaPipeOptions& options, bool calculator_run_in_parallel)
: InputStreamHandler(tag_map, calculator_context_manager, options,
calculator_run_in_parallel),
timestamp_bounds_(std::move(tag_map)) {}
calculator_run_in_parallel) {
for (auto id = tag_map->BeginId(); id < tag_map->EndId(); ++id) {
sync_sets_.emplace_back(this, std::vector<CollectionItemId>{id});
ready_timestamps_.push_back(Timestamp::Unset());
}
}
NodeReadiness ImmediateInputStreamHandler::GetNodeReadiness(
Timestamp* min_stream_timestamp) {
*min_stream_timestamp = Timestamp::Done();
absl::MutexLock lock(&mutex_);
Timestamp input_timestamp = Timestamp::Done();
Timestamp min_bound = Timestamp::Done();
bool stream_became_done = false;
for (CollectionItemId i = input_stream_managers_.BeginId();
i < input_stream_managers_.EndId(); ++i) {
const auto& stream = input_stream_managers_.Get(i);
bool empty;
Timestamp stream_timestamp = stream->MinTimestampOrBound(&empty);
if (!empty) {
input_timestamp = std::min(input_timestamp, stream_timestamp);
for (int i = 0; i < sync_sets_.size(); ++i) {
if (ready_timestamps_[i] > Timestamp::Unset()) {
min_bound = std::min(min_bound, ready_timestamps_[i]);
input_timestamp = std::min(input_timestamp, ready_timestamps_[i]);
continue;
}
*min_stream_timestamp = std::min(*min_stream_timestamp, stream_timestamp);
if (stream_timestamp != timestamp_bounds_.Get(i)) {
if (stream_timestamp == Timestamp::Done()) {
Timestamp prev_ts = sync_sets_[i].LastProcessed();
Timestamp stream_ts;
NodeReadiness readiness = sync_sets_[i].GetReadiness(&stream_ts);
min_bound = std::min(min_bound, stream_ts);
if (readiness == NodeReadiness::kReadyForProcess) {
ready_timestamps_[i] = stream_ts;
input_timestamp = std::min(input_timestamp, stream_ts);
} else if (readiness == NodeReadiness::kReadyForClose) {
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]);
} else if (prev_ts < Timestamp::Done()) {
stream_became_done = true;
ready_timestamps_[i] = Timestamp::Done();
}
timestamp_bounds_.Get(i) = stream_timestamp;
}
}
*min_stream_timestamp = min_bound;
if (*min_stream_timestamp == Timestamp::Done()) {
return NodeReadiness::kReadyForClose;
@@ -94,6 +115,8 @@ NodeReadiness ImmediateInputStreamHandler::GetNodeReadiness(
}
if (stream_became_done) {
// The stream_became_done logic is kept for backward compatibility.
// Note that the minimum bound is returned in min_stream_timestamp.
return NodeReadiness::kReadyForProcess;
}
@@ -102,23 +125,13 @@ NodeReadiness ImmediateInputStreamHandler::GetNodeReadiness(
void ImmediateInputStreamHandler::FillInputSet(Timestamp input_timestamp,
InputStreamShardSet* input_set) {
CHECK(input_timestamp.IsAllowedInStream());
CHECK(input_set);
for (CollectionItemId id = input_stream_managers_.BeginId();
id < input_stream_managers_.EndId(); ++id) {
auto& stream = input_stream_managers_.Get(id);
if (stream->QueueHead().Timestamp() == input_timestamp) {
int num_packets_dropped = 0;
bool stream_is_done = false;
Packet current_packet = stream->PopPacketAtTimestamp(
input_timestamp, &num_packets_dropped, &stream_is_done);
AddPacketToShard(&input_set->Get(id), std::move(current_packet),
stream_is_done);
absl::MutexLock lock(&mutex_);
for (int i = 0; i < sync_sets_.size(); ++i) {
if (ready_timestamps_[i] == input_timestamp) {
sync_sets_[i].FillInputSet(input_timestamp, input_set);
ready_timestamps_[i] = Timestamp::Unset();
} else {
Timestamp bound = stream->MinTimestampOrBound(nullptr);
AddPacketToShard(&input_set->Get(id),
Packet().At(bound.PreviousAllowedInStream()),
bound == Timestamp::Done());
sync_sets_[i].FillInputBounds(input_set);
}
}
}
mediapipe/framework/stream_handler/sync_set_input_stream_handler.cc
+17 -53
View File
@@ -17,6 +17,7 @@
// TODO: Move protos in another CL after the C++ code migration.
#include "absl/strings/substitute.h"
#include "absl/synchronization/mutex.h"
#include "mediapipe/framework/collection_item_id.h"
#include "mediapipe/framework/input_stream_handler.h"
#include "mediapipe/framework/mediapipe_options.pb.h"
#include "mediapipe/framework/packet_set.h"
@@ -69,7 +70,7 @@ class SyncSetInputStreamHandler : public InputStreamHandler {
private:
absl::Mutex mutex_;
// The ids of each set of inputs.
std::vector<std::vector<CollectionItemId>> sync_sets_ ABSL_GUARDED_BY(mutex_);
std::vector<SyncSet> sync_sets_ ABSL_GUARDED_BY(mutex_);
// The index of the ready sync set. A value of -1 indicates that no
// sync sets are ready.
int ready_sync_set_index_ ABSL_GUARDED_BY(mutex_) = -1;
@@ -98,7 +99,7 @@ void SyncSetInputStreamHandler::PrepareForRun(
sync_sets_.clear();
std::set<CollectionItemId> used_ids;
for (const auto& sync_set : handler_options.sync_set()) {
sync_sets_.emplace_back();
std::vector<CollectionItemId> stream_ids;
CHECK_LT(0, sync_set.tag_index_size());
for (const auto& tag_index : sync_set.tag_index()) {
std::string tag;
@@ -109,8 +110,9 @@ void SyncSetInputStreamHandler::PrepareForRun(
CHECK(!::mediapipe::ContainsKey(used_ids, id))
<< "stream \"" << tag_index << "\" is in more than one sync set.";
used_ids.insert(id);
sync_sets_.back().push_back(id);
stream_ids.push_back(id);
}
sync_sets_.emplace_back(this, std::move(stream_ids));
}
std::vector<CollectionItemId> remaining_ids;
for (CollectionItemId id = input_stream_managers_.BeginId();
@@ -120,7 +122,7 @@ void SyncSetInputStreamHandler::PrepareForRun(
}
}
if (!remaining_ids.empty()) {
sync_sets_.push_back(std::move(remaining_ids));
sync_sets_.emplace_back(this, std::move(remaining_ids));
}
ready_sync_set_index_ = -1;
ready_timestamp_ = Timestamp::Done();
@@ -137,24 +139,14 @@ NodeReadiness SyncSetInputStreamHandler::GetNodeReadiness(
absl::MutexLock lock(&mutex_);
if (ready_sync_set_index_ >= 0) {
*min_stream_timestamp = ready_timestamp_;
// TODO: Return kNotReady unless a new ready syncset is found.
return NodeReadiness::kReadyForProcess;
}
for (int sync_set_index = 0; sync_set_index < sync_sets_.size();
++sync_set_index) {
const std::vector<CollectionItemId>& sync_set = sync_sets_[sync_set_index];
*min_stream_timestamp = Timestamp::Done();
Timestamp min_bound = Timestamp::Done();
for (CollectionItemId id : sync_set) {
const auto& stream = input_stream_managers_.Get(id);
bool empty;
Timestamp stream_timestamp = stream->MinTimestampOrBound(&empty);
if (empty) {
min_bound = std::min(min_bound, stream_timestamp);
}
*min_stream_timestamp = std::min(*min_stream_timestamp, stream_timestamp);
}
if (*min_stream_timestamp == Timestamp::Done()) {
NodeReadiness readiness =
sync_sets_[sync_set_index].GetReadiness(min_stream_timestamp);
if (readiness == NodeReadiness::kReadyForClose) {
// This sync set is done, remove it. Note that this invalidates
// sync set indexes higher than sync_set_index. However, we are
// guaranteed that we were not ready before entering the outer
@@ -165,15 +157,14 @@ NodeReadiness SyncSetInputStreamHandler::GetNodeReadiness(
continue;
}
if (min_bound > *min_stream_timestamp) {
if (readiness == NodeReadiness::kReadyForProcess) {
// TODO: Prioritize sync-sets to avoid starvation.
if (*min_stream_timestamp < ready_timestamp_) {
// Store the timestamp and corresponding sync set index for the
// sync set with the earliest arrival timestamp.
ready_timestamp_ = *min_stream_timestamp;
ready_sync_set_index_ = sync_set_index;
}
} else {
CHECK_EQ(min_bound, *min_stream_timestamp);
}
}
if (ready_sync_set_index_ >= 0) {
@@ -188,44 +179,17 @@ NodeReadiness SyncSetInputStreamHandler::GetNodeReadiness(
return NodeReadiness::kNotReady;
}
void SyncSetInputStreamHandler::FillInputBounds(
Timestamp input_timestamp, InputStreamShardSet* input_set) {
for (int i = 0; i < sync_sets_.size(); ++i) {
if (i != ready_sync_set_index_) {
// Set the input streams for the not-ready sync sets.
for (CollectionItemId id : sync_sets_[i]) {
const auto stream = input_stream_managers_.Get(id);
Timestamp bound = stream->MinTimestampOrBound(nullptr);
AddPacketToShard(&input_set->Get(id),
Packet().At(bound.PreviousAllowedInStream()),
bound == Timestamp::Done());
}
}
}
}
void SyncSetInputStreamHandler::FillInputSet(Timestamp input_timestamp,
InputStreamShardSet* input_set) {
// Assume that all current packets are already cleared.
CHECK(input_timestamp.IsAllowedInStream());
CHECK(input_set);
absl::MutexLock lock(&mutex_);
CHECK_LE(0, ready_sync_set_index_);
CHECK_EQ(input_timestamp, ready_timestamp_);
// Set the input streams for the ready sync set.
for (CollectionItemId id : sync_sets_[ready_sync_set_index_]) {
const auto& stream = input_stream_managers_.Get(id);
int num_packets_dropped = 0;
bool stream_is_done = false;
Packet current_packet = stream->PopPacketAtTimestamp(
input_timestamp, &num_packets_dropped, &stream_is_done);
CHECK_EQ(num_packets_dropped, 0)
<< absl::Substitute("Dropped $0 packet(s) on input stream \"$1\".",
num_packets_dropped, stream->Name());
AddPacketToShard(&input_set->Get(id), std::move(current_packet),
stream_is_done);
sync_sets_[ready_sync_set_index_].FillInputSet(input_timestamp, input_set);
for (int i = 0; i < sync_sets_.size(); ++i) {
if (i != ready_sync_set_index_) {
sync_sets_[i].FillInputBounds(input_set);
}
}
FillInputBounds(input_timestamp, input_set);
ready_sync_set_index_ = -1;
ready_timestamp_ = Timestamp::Done();
}