111 lines
3.9 KiB
C++
111 lines
3.9 KiB
C++
// Copyright 2019 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 "mediapipe/framework/calculator_context_manager.h"
|
|
|
|
#include <utility>
|
|
|
|
#include "absl/memory/memory.h"
|
|
#include "absl/synchronization/mutex.h"
|
|
#include "mediapipe/framework/port/logging.h"
|
|
|
|
namespace mediapipe {
|
|
|
|
void CalculatorContextManager::Initialize(
|
|
CalculatorState* calculator_state,
|
|
std::shared_ptr<tool::TagMap> input_tag_map,
|
|
std::shared_ptr<tool::TagMap> output_tag_map,
|
|
bool calculator_run_in_parallel) {
|
|
CHECK(calculator_state);
|
|
calculator_state_ = calculator_state;
|
|
input_tag_map_ = std::move(input_tag_map);
|
|
output_tag_map_ = std::move(output_tag_map);
|
|
calculator_run_in_parallel_ = calculator_run_in_parallel;
|
|
}
|
|
|
|
absl::Status CalculatorContextManager::PrepareForRun(
|
|
std::function<absl::Status(CalculatorContext*)> setup_shards_callback) {
|
|
setup_shards_callback_ = std::move(setup_shards_callback);
|
|
default_context_ = absl::make_unique<CalculatorContext>(
|
|
calculator_state_, input_tag_map_, output_tag_map_);
|
|
return setup_shards_callback_(default_context_.get());
|
|
}
|
|
|
|
void CalculatorContextManager::CleanupAfterRun() {
|
|
default_context_ = nullptr;
|
|
absl::MutexLock lock(&contexts_mutex_);
|
|
active_contexts_.clear();
|
|
idle_contexts_.clear();
|
|
}
|
|
|
|
CalculatorContext* CalculatorContextManager::GetDefaultCalculatorContext()
|
|
const {
|
|
CHECK(default_context_.get());
|
|
return default_context_.get();
|
|
}
|
|
|
|
CalculatorContext* CalculatorContextManager::GetFrontCalculatorContext(
|
|
Timestamp* context_input_timestamp) {
|
|
CHECK(calculator_run_in_parallel_);
|
|
absl::MutexLock lock(&contexts_mutex_);
|
|
CHECK(!active_contexts_.empty());
|
|
*context_input_timestamp = active_contexts_.begin()->first;
|
|
return active_contexts_.begin()->second.get();
|
|
}
|
|
|
|
CalculatorContext* CalculatorContextManager::PrepareCalculatorContext(
|
|
Timestamp input_timestamp) {
|
|
if (!calculator_run_in_parallel_) {
|
|
return GetDefaultCalculatorContext();
|
|
}
|
|
absl::MutexLock lock(&contexts_mutex_);
|
|
CHECK(!mediapipe::ContainsKey(active_contexts_, input_timestamp))
|
|
<< "Multiple invocations with the same timestamps are not allowed with "
|
|
"parallel execution, input_timestamp = "
|
|
<< input_timestamp;
|
|
CalculatorContext* calculator_context = nullptr;
|
|
if (idle_contexts_.empty()) {
|
|
auto new_context = absl::make_unique<CalculatorContext>(
|
|
calculator_state_, input_tag_map_, output_tag_map_);
|
|
MEDIAPIPE_CHECK_OK(setup_shards_callback_(new_context.get()));
|
|
calculator_context = new_context.get();
|
|
active_contexts_.emplace(input_timestamp, std::move(new_context));
|
|
} else {
|
|
// Retrieves an inactive calculator context from idle_contexts_.
|
|
calculator_context = idle_contexts_.front().get();
|
|
active_contexts_.emplace(input_timestamp,
|
|
std::move(idle_contexts_.front()));
|
|
idle_contexts_.pop_front();
|
|
}
|
|
return calculator_context;
|
|
}
|
|
|
|
void CalculatorContextManager::RecycleCalculatorContext() {
|
|
absl::MutexLock lock(&contexts_mutex_);
|
|
// The first element in active_contexts_ will be recycled.
|
|
auto iter = active_contexts_.begin();
|
|
idle_contexts_.push_back(std::move(iter->second));
|
|
active_contexts_.erase(iter);
|
|
}
|
|
|
|
bool CalculatorContextManager::HasActiveContexts() {
|
|
if (!calculator_run_in_parallel_) {
|
|
return false;
|
|
}
|
|
absl::MutexLock lock(&contexts_mutex_);
|
|
return !active_contexts_.empty();
|
|
}
|
|
|
|
} // namespace mediapipe
|