Project import generated by Copybara.
PiperOrigin-RevId: 253489161
This commit is contained in:
MediaPipe Team
@@ -0,0 +1,330 @@
|
||||
// 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.
|
||||
|
||||
#ifndef MEDIAPIPE_FRAMEWORK_SCHEDULER_H_
|
||||
#define MEDIAPIPE_FRAMEWORK_SCHEDULER_H_
|
||||
|
||||
#include <atomic>
|
||||
#include <functional>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <queue>
|
||||
#include <set>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/base/macros.h"
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/calculator_node.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/status.h"
|
||||
#include "mediapipe/framework/scheduler_queue.h"
|
||||
#include "mediapipe/framework/scheduler_shared.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
class CalculatorGraph;
|
||||
class Executor;
|
||||
|
||||
namespace internal {
|
||||
|
||||
// The class scheduling a calculator graph.
|
||||
class Scheduler {
|
||||
public:
|
||||
Scheduler(const Scheduler&) = delete;
|
||||
Scheduler& operator=(const Scheduler&) = delete;
|
||||
|
||||
explicit Scheduler(CalculatorGraph* graph);
|
||||
|
||||
~Scheduler();
|
||||
|
||||
// Sets the executor that will run the nodes. Must be called before the
|
||||
// scheduler is started. This is the normal executor used for nodes that
|
||||
// do not use a special one.
|
||||
void SetExecutor(Executor* executor);
|
||||
|
||||
// Sets the executor that will run the nodes assigned to the executor
|
||||
// named |name|. Must be called before the scheduler is started.
|
||||
::mediapipe::Status SetNonDefaultExecutor(const std::string& name,
|
||||
Executor* executor);
|
||||
|
||||
// Resets the data members at the beginning of each graph run.
|
||||
void Reset();
|
||||
|
||||
// Starts scheduling nodes.
|
||||
void Start();
|
||||
|
||||
// Wait for the current run to finish (block the current thread until all
|
||||
// source calculators have returned StatusStop(), all graph input streams
|
||||
// have been closed, and no more calculators can be run).
|
||||
// This function can be called only after Start().
|
||||
// Runs application thread tasks while waiting.
|
||||
::mediapipe::Status WaitUntilDone() LOCKS_EXCLUDED(state_mutex_);
|
||||
|
||||
// Wait until the running graph is in the idle mode, which is when nothing can
|
||||
// be scheduled and nothing is running in the worker threads. This function
|
||||
// can be called only after Start().
|
||||
// Runs application thread tasks while waiting.
|
||||
::mediapipe::Status WaitUntilIdle() LOCKS_EXCLUDED(state_mutex_);
|
||||
|
||||
// Wait until any graph input stream has been unthrottled.
|
||||
// This is meant to be used by CalculatorGraph::AddPacketToInputStream, which
|
||||
// needs to check a status protected by its own mutex. That mutex, which
|
||||
// protects throttle changes on graph input streams, should be passed as the
|
||||
// secondary_mutex argument.
|
||||
// This function can be called by multiple threads concurrently.
|
||||
// Runs application thread tasks while waiting.
|
||||
void WaitUntilGraphInputStreamUnthrottled(absl::Mutex* secondary_mutex)
|
||||
LOCKS_EXCLUDED(state_mutex_) EXCLUSIVE_LOCKS_REQUIRED(secondary_mutex);
|
||||
|
||||
// Wait until any observed output emits a packet. Like a semaphore,
|
||||
// this function returns immediately if an observed packet has already been
|
||||
// emitted since the previous call. This relies on the fact that the calls are
|
||||
// in sequence. Runs application thread tasks while waiting.
|
||||
// Returns ::mediapipe::OutOfRangeError if the graph terminated.
|
||||
::mediapipe::Status WaitForObservedOutput() LOCKS_EXCLUDED(state_mutex_);
|
||||
|
||||
// Callback that is invoked by a node when it wants to be scheduled.
|
||||
// If the node is throttled, the call is ignored.
|
||||
// This method is thread-safe.
|
||||
void ScheduleNodeIfNotThrottled(CalculatorNode* node, CalculatorContext* cc);
|
||||
|
||||
// Schedules an OpenNode() call for |node|.
|
||||
void ScheduleNodeForOpen(CalculatorNode* node);
|
||||
|
||||
// Adds all the nodes in |nodes_to_schedule| to the scheduler queue, without
|
||||
// checking if they are ready. Called by the graph when unthrottling nodes.
|
||||
void ScheduleUnthrottledReadyNodes(
|
||||
const std::vector<CalculatorNode*>& nodes_to_schedule);
|
||||
|
||||
void QueueIdleStateChanged(bool idle);
|
||||
|
||||
// Called by DelegatingExecutor to add an application thread task.
|
||||
void AddApplicationThreadTask(std::function<void()> task);
|
||||
|
||||
// Adds |node| to |unopened_sources_|.
|
||||
// This can only be called before the scheduler is started.
|
||||
void AddUnopenedSourceNode(CalculatorNode* node);
|
||||
|
||||
// Adds |node| to |sources_queue_|.
|
||||
void AddNodeToSourcesQueue(CalculatorNode* node) LOCKS_EXCLUDED(state_mutex_);
|
||||
|
||||
// Assigns node to a scheduler queue.
|
||||
void AssignNodeToSchedulerQueue(CalculatorNode* node);
|
||||
|
||||
// Pauses the scheduler. Does nothing if Cancel has been called.
|
||||
void Pause() LOCKS_EXCLUDED(state_mutex_);
|
||||
|
||||
// Resumes the scheduler.
|
||||
void Resume() LOCKS_EXCLUDED(state_mutex_);
|
||||
|
||||
// Aborts the scheduler if the graph is started but is not terminated; no-op
|
||||
// otherwise. For the graph to properly be cancelled, graph_->HasError()
|
||||
// must also return true.
|
||||
void Cancel() LOCKS_EXCLUDED(state_mutex_);
|
||||
|
||||
// Returns true if scheduler is paused.
|
||||
bool IsPaused() LOCKS_EXCLUDED(state_mutex_);
|
||||
|
||||
// Returns true if scheduler is terminated.
|
||||
bool IsTerminated() LOCKS_EXCLUDED(state_mutex_);
|
||||
|
||||
// Cleanup any remaining state after the run.
|
||||
void CleanupAfterRun();
|
||||
|
||||
void SetHasError(bool error) { shared_.has_error = error; }
|
||||
|
||||
// Notifies the scheduler that a packet was added to a graph input stream.
|
||||
// The scheduler needs to check whether it is still deadlocked, and
|
||||
// unthrottle again if so.
|
||||
void AddedPacketToGraphInputStream() LOCKS_EXCLUDED(state_mutex_);
|
||||
|
||||
void ThrottledGraphInputStream() LOCKS_EXCLUDED(state_mutex_);
|
||||
void UnthrottledGraphInputStream() LOCKS_EXCLUDED(state_mutex_);
|
||||
void EmittedObservedOutput() LOCKS_EXCLUDED(state_mutex_);
|
||||
|
||||
// Closes all source nodes at the next scheduling opportunity.
|
||||
void CloseAllSourceNodes();
|
||||
|
||||
// Notifies the scheduler that all graph input streams have been closed.
|
||||
void ClosedAllGraphInputStreams();
|
||||
|
||||
// Returns the scheduler's runtime measures for overhead measurement.
|
||||
// Only meant for test purposes. See SchedulerTimer for details.
|
||||
internal::SchedulerTimes GetSchedulerTimes();
|
||||
|
||||
private:
|
||||
// State of the scheduler. The figure shows the allowed state transitons.
|
||||
//
|
||||
// NOT_STARTED
|
||||
// |
|
||||
// v
|
||||
// RUNNING--+
|
||||
// | | ^ |
|
||||
// | | \ |
|
||||
// | | \ v
|
||||
// | | PAUSED
|
||||
// | | |
|
||||
// | v v
|
||||
// | CANCELLING
|
||||
// | |
|
||||
// v v
|
||||
// TERMINATING
|
||||
// |
|
||||
// v
|
||||
// TERMINATED
|
||||
enum State {
|
||||
STATE_NOT_STARTED = 0, // The initial state.
|
||||
STATE_RUNNING = 1, // The scheduler is running and scheduling nodes.
|
||||
STATE_PAUSED = 2, // The scheduler is not scheduling nodes.
|
||||
STATE_CANCELLING = 3, // The scheduler is being cancelled. The scheduler
|
||||
// cannot be paused in this state so that
|
||||
// scheduler_queue_ can be drained.
|
||||
STATE_TERMINATED = 4, // The scheduler has terminated.
|
||||
};
|
||||
|
||||
// Sorts the source nodes in unopened_sources_ by source layer.
|
||||
struct SourceLayerCompare {
|
||||
bool operator()(CalculatorNode* lhs, CalculatorNode* rhs) const {
|
||||
if (lhs->source_layer() != rhs->source_layer()) {
|
||||
return lhs->source_layer() < rhs->source_layer();
|
||||
}
|
||||
return lhs->Id() < rhs->Id();
|
||||
}
|
||||
};
|
||||
|
||||
// Start (or resume) or stop all queues.
|
||||
void SetQueuesRunning(bool running);
|
||||
|
||||
// Submit waiting tasks on all queues after resuming.
|
||||
void SubmitWaitingTasksOnQueues();
|
||||
|
||||
// Returns true if nothing can be scheduled and no tasks are running or
|
||||
// scheduled to run on the Executor.
|
||||
bool IsIdle() EXCLUSIVE_LOCKS_REQUIRED(state_mutex_);
|
||||
|
||||
// Clean up active_sources_ by removing closed sources. If all the active
|
||||
// sources are closed, this will leave active_sources_ empty. If not, some
|
||||
// closed sources may be left in there.
|
||||
void CleanupActiveSources();
|
||||
|
||||
// Adds the next layer of sources to the scheduler queue if the previous layer
|
||||
// has finished running.
|
||||
// Returns true if it scheduled any sources.
|
||||
bool TryToScheduleNextSourceLayer() EXCLUSIVE_LOCKS_REQUIRED(state_mutex_);
|
||||
|
||||
// Takes care of three different operations, as needed:
|
||||
// - activating sources;
|
||||
// - unthrottling sources or graph input streams to resolve a deadlock.
|
||||
// - terminating the scheduler.
|
||||
// Thread-safe and reentrant.
|
||||
// TODO: analyze call sites, split it up further.
|
||||
void HandleIdle() EXCLUSIVE_LOCKS_REQUIRED(state_mutex_);
|
||||
|
||||
// Terminates the scheduler. Should only be called by HandleIdle.
|
||||
void Quit() EXCLUSIVE_LOCKS_REQUIRED(state_mutex_);
|
||||
|
||||
// Helper for the various Wait methods. Waits for the given condition,
|
||||
// running application thread tasks in the meantime.
|
||||
void ApplicationThreadAwait(const std::function<bool()>& stop_condition);
|
||||
|
||||
// The calculator graph to run.
|
||||
CalculatorGraph* graph_;
|
||||
|
||||
// Data accessed by all SchedulerQueues.
|
||||
SchedulerShared shared_;
|
||||
|
||||
// Queue of nodes that need to be run.
|
||||
SchedulerQueue default_queue_;
|
||||
|
||||
// Non-default scheduler queues, keyed by their executor names.
|
||||
std::map<std::string, std::unique_ptr<SchedulerQueue>> non_default_queues_;
|
||||
|
||||
// Holds pointers to all queues used by the scheduler, for convenience.
|
||||
std::vector<SchedulerQueue*> scheduler_queues_;
|
||||
|
||||
// Priority queue of source nodes ordered by layer and then source process
|
||||
// order. This stores the set of sources that are yet to be run.
|
||||
std::priority_queue<SchedulerQueue::Item> sources_queue_
|
||||
GUARDED_BY(state_mutex_);
|
||||
|
||||
// Source nodes with the smallest source layer are at the beginning of
|
||||
// unopened_sources_. Before the scheduler is started, all source nodes are
|
||||
// added to unopened_sources_. Once the scheduler starts running,
|
||||
// unopened_sources_ should only be accessed under the protection of
|
||||
// state_mutex_. A source node is removed from unopened_sources_ after it is
|
||||
// opened.
|
||||
std::set<CalculatorNode*, SourceLayerCompare> unopened_sources_;
|
||||
|
||||
// Keeps track of sources that can be considered for scheduling. Sources are
|
||||
// scheduled in layers, and those that are not currently active will not be
|
||||
// scheduled even if ready. Sources are removed once they are closed.
|
||||
std::vector<CalculatorNode*> active_sources_;
|
||||
|
||||
// Condition variable used to wait for some changes to the scheduler state.
|
||||
// These correspond to the Wait* methods in this class.
|
||||
// Not all state changes need to signal this, only those that enter one of
|
||||
// the waitable states.
|
||||
absl::CondVar state_cond_var_ GUARDED_BY(state_mutex_);
|
||||
|
||||
// Number of queues which are not idle.
|
||||
// Note: this indicates two slightly different things:
|
||||
// a. the number of queues which still have nodes running;
|
||||
// b. the number of queues whose executors may still access the scheduler.
|
||||
// When a queue becomes idle, it has stopped running nodes, and the scheduler
|
||||
// decrements the count. However, it is not done accessing the scheduler
|
||||
// until HandleIdle returns. Therefore, a and b are briefly out of sync.
|
||||
// This is ok, because it happens within a single critical section, which is
|
||||
// guarded by state_mutex_. If we wanted to split this critical section, we
|
||||
// would have to separate a and b into two variables.
|
||||
int non_idle_queue_count_ GUARDED_BY(state_mutex_) = 0;
|
||||
|
||||
// Tasks to be executed on the application thread.
|
||||
std::deque<std::function<void()>> app_thread_tasks_ GUARDED_BY(state_mutex_);
|
||||
|
||||
// Used by HandleIdle to avoid multiple concurrent executions.
|
||||
// We cannot simply hold a mutex throughout it, for two reasons:
|
||||
// - We need it to be reentrant, which Mutex does not support.
|
||||
// - We want simultaneous calls to return immediately instead of waiting,
|
||||
// and Mutex's TryLock is not guaranteed to work.
|
||||
bool handling_idle_ GUARDED_BY(state_mutex_) = false;
|
||||
|
||||
// Mutex for the scheduler state and related things.
|
||||
// Note: state_ is declared as atomic so that its getter methods don't need
|
||||
// to acquire state_mutex_.
|
||||
absl::Mutex state_mutex_;
|
||||
|
||||
// Current state of the scheduler.
|
||||
std::atomic<State> state_ = ATOMIC_VAR_INIT(STATE_NOT_STARTED);
|
||||
|
||||
// True if all graph input streams are closed.
|
||||
bool graph_input_streams_closed_ GUARDED_BY(state_mutex_) = false;
|
||||
|
||||
// Number of throttled graph input streams.
|
||||
int throttled_graph_input_stream_count_ GUARDED_BY(state_mutex_) = 0;
|
||||
|
||||
// Used to stop WaitUntilGraphInputStreamUnthrottled.
|
||||
int unthrottle_seq_num_ GUARDED_BY(state_mutex_) = 0;
|
||||
|
||||
// Used to stop WaitForObservedOutput.
|
||||
bool observed_output_signal_ GUARDED_BY(state_mutex_) = false;
|
||||
|
||||
// True if an application thread is waiting in WaitForObservedOutput.
|
||||
bool waiting_for_observed_output_ GUARDED_BY(state_mutex_) = false;
|
||||
};
|
||||
|
||||
} // namespace internal
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_FRAMEWORK_SCHEDULER_H_
|
||||
Reference in New Issue
Block a user