servable_state_monitor.h

/* Copyright 2016 Google Inc. 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_SERVING_CORE_SERVABLE_STATE_MONITOR_H_
#define TENSORFLOW_SERVING_CORE_SERVABLE_STATE_MONITOR_H_
 
#include <deque>
#include <functional>
#include <map>
 
#include "absl/time/time.h"
#include "absl/types/optional.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow_serving/core/manager.h"
#include "tensorflow_serving/core/servable_id.h"
#include "tensorflow_serving/core/servable_state.h"
#include "tensorflow_serving/util/event_bus.h"
 
namespace tensorflow {
namespace serving {
 
class ServableStateMonitor {
 public:
  // START_SKIP_DOXYGEN
  struct ServableStateAndTime {
    ServableStateAndTime() = default;
    ServableStateAndTime(ServableState servable_state,
                         const uint64_t event_time)
        : state(std::move(servable_state)), event_time_micros(event_time) {}
 
    ServableState state;
 
    uint64_t event_time_micros;
 
    string DebugString() const;
  };
 
  using ServableName = string;
  using Version = int64;
  using VersionMap =
      std::map<Version, ServableStateAndTime, std::greater<Version>>;
  using ServableMap = std::map<ServableName, VersionMap>;
  using ServableSet = std::set<ServableName>;
 
  struct Options {
    Options() {}
 
    uint64_t max_count_log_events = 0;
  };
  // END_SKIP_DOXYGEN
 
  using BoundedLog = std::deque<ServableStateAndTime>;
 
  explicit ServableStateMonitor(EventBus<ServableState>* bus,
                                const Options& options = Options());
  virtual ~ServableStateMonitor();
 
  absl::optional<ServableState> GetState(const ServableId& servable_id) const
      TF_LOCKS_EXCLUDED(mu_);
 
  absl::optional<ServableStateAndTime> GetStateAndTime(
      const ServableId& servable_id) const TF_LOCKS_EXCLUDED(mu_);
 
  VersionMap GetVersionStates(const string& servable_name) const
      TF_LOCKS_EXCLUDED(mu_);
 
  ServableMap GetAllServableStates() const TF_LOCKS_EXCLUDED(mu_);
 
  ServableMap GetLiveServableStates() const TF_LOCKS_EXCLUDED(mu_);
 
  void ForgetUnloadedServableStates() TF_LOCKS_EXCLUDED(mu_);
 
  // Returns all servables that are in state
  // ServableState::ManagerState::kAvailable.
  // Note that as opposed to GetAllServableStates() and GetLiveServableStates(),
  // this method loops over all the tracked servables.
  ServableSet GetAvailableServableStates() const TF_LOCKS_EXCLUDED(mu_);
 
  BoundedLog GetBoundedLog() const TF_LOCKS_EXCLUDED(mu_);
 
  using ServableStateNotifierFn = std::function<void(
      bool reached_goal_state,
      const std::map<ServableId, ServableState::ManagerState>& states_reached)>;
  void NotifyWhenServablesReachState(
      const std::vector<ServableRequest>& servables,
      ServableState::ManagerState goal_state,
      const ServableStateNotifierFn& notifier_fn) TF_LOCKS_EXCLUDED(mu_);
 
  bool WaitUntilServablesReachStateWithTimeout(
      const std::vector<ServableRequest>& servables,
      ServableState::ManagerState goal_state, absl::Duration timeout,
      std::map<ServableId, ServableState::ManagerState>* states_reached =
          nullptr) TF_LOCKS_EXCLUDED(mu_) TF_MUST_USE_RESULT;
  bool WaitUntilServablesReachState(
      const std::vector<ServableRequest>& servables,
      ServableState::ManagerState goal_state,
      std::map<ServableId, ServableState::ManagerState>* states_reached =
          nullptr) TF_MUST_USE_RESULT;
 
  // Subscribes to all servable state changes hitting this monitor. This is
  // called after the monitor updates its own state based on the event.
  using NotifyFn = std::function<void(const ServableState&)>;
  void Notify(const NotifyFn& notify_fn) TF_LOCKS_EXCLUDED(notify_mu_);
 
 private:
  absl::optional<ServableStateMonitor::ServableStateAndTime>
  GetStateAndTimeInternal(const ServableId& servable_id) const
      TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
  // Request to send notification, setup using
  // NotifyWhenServablesReachState(...).
  struct ServableStateNotificationRequest {
    std::vector<ServableRequest> servables;
    ServableState::ManagerState goal_state;
    ServableStateNotifierFn notifier_fn;
  };
 
  // Checks whether the notification request is satisfied and we cand send it.
  // If so, returns the 'reached_goal_state' bool and the 'states_reached' by
  // each servable.  Oterwise returns nullopt.
  absl::optional<
      std::pair<bool, std::map<ServableId, ServableState::ManagerState>>>
  ShouldSendStateReachedNotification(
      const ServableStateNotificationRequest& notification_request)
      TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
  // Goes through the notification requests and tries to see if any of them can
  // be sent. If a notification is sent, the corresponding request is removed.
  void MaybeSendStateReachedNotifications() TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
  // Goes through the notify_fns list and calls each one with the currently
  // received ServableState.
  void SendNotifications(const ServableState& servable_state)
      TF_LOCKS_EXCLUDED(notify_mu_);
 
  // This method is called when an event comes in, but before we update our
  // state with the contents of the event. Subclasses may override this method
  // to do custom prepreocessing based on the event and the previous state of
  // the monitor, like calculate load-time, etc.
  virtual void PreHandleEvent(
      const EventBus<ServableState>::EventAndTime& state_and_time);
 
  // Handles a bus event.
  void HandleEvent(const EventBus<ServableState>::EventAndTime& event_and_time)
      TF_LOCKS_EXCLUDED(mu_, notify_mu_);
 
  const Options options_;
 
  std::unique_ptr<EventBus<ServableState>::Subscription> bus_subscription_;
 
  mutable mutex mu_;
 
  // The current state of each servable version that has appeared on the bus.
  // (Entries are never removed, even when they enter state kEnd.)
  ServableMap states_ TF_GUARDED_BY(mu_);
 
  // The current state of each servable version that has not transitioned to
  // state ServableState::ManagerState::kEnd.
  ServableMap live_states_ TF_GUARDED_BY(mu_);
 
  // Deque of pairs of timestamp and ServableState, corresponding to the most
  // recent servable state events handled by the monitor. The size of this deque
  // is upper bounded by max_count_log_events in Options.
  BoundedLog log_ TF_GUARDED_BY(mu_);
 
  std::vector<ServableStateNotificationRequest>
      servable_state_notification_requests_ TF_GUARDED_BY(mu_);
 
  // Separate mutex to protect the notify_fns_ so that they can be updated
  // independently. This also allows these notify_fns_ to call other methods
  // in ServableStateMonitor which don't depend on this mutex without being
  // deadlocked.
  mutable mutex notify_mu_;
  std::vector<NotifyFn> notify_fns_ TF_GUARDED_BY(notify_mu_);
 
  TF_DISALLOW_COPY_AND_ASSIGN(ServableStateMonitor);
};
 
inline bool operator==(const ServableStateMonitor::ServableStateAndTime& a,
                       const ServableStateMonitor::ServableStateAndTime& b) {
  return a.event_time_micros == b.event_time_micros && a.state == b.state;
}
 
inline bool operator!=(const ServableStateMonitor::ServableStateAndTime& a,
                       const ServableStateMonitor::ServableStateAndTime& b) {
  return !(a == b);
}
 
inline std::ostream& operator<<(
    std::ostream& os,
    const ServableStateMonitor::ServableStateAndTime& state_and_time) {
  return os << state_and_time.DebugString();
}
 
}  // namespace serving
}  // namespace tensorflow
 
#endif  // TENSORFLOW_SERVING_CORE_SERVABLE_STATE_MONITOR_H_