aspired_versions_manager.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_ASPIRED_VERSIONS_MANAGER_H_
#define TENSORFLOW_SERVING_CORE_ASPIRED_VERSIONS_MANAGER_H_
 
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
 
#include "absl/types/optional.h"
#include "tensorflow/core/kernels/batching_util/periodic_function.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/hash/hash.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow_serving/core/aspired_version_policy.h"
#include "tensorflow_serving/core/basic_manager.h"
#include "tensorflow_serving/core/loader.h"
#include "tensorflow_serving/core/manager.h"
#include "tensorflow_serving/core/servable_data.h"
#include "tensorflow_serving/core/servable_handle.h"
#include "tensorflow_serving/core/servable_id.h"
#include "tensorflow_serving/core/servable_state.h"
#include "tensorflow_serving/core/target.h"
#include "tensorflow_serving/util/event_bus.h"
#include "tensorflow_serving/util/observer.h"
 
namespace tensorflow {
namespace serving {
 
class AspiredVersionsManager;
 
namespace internal {
 
class AspiredVersionsManagerTargetImpl;
 
uint32 GetManagerNumLoadThreads(AspiredVersionsManager* manager);
 
// Returns the Notifier function of the manager's Observer, which forwards
// SetNumLoadThreads().  This indirection is to prevent callers from using
// SetNumLoadThreads() on a deleted manager.
std::function<void(uint32)> SetManagerNumLoadThreadsNotifier(
    AspiredVersionsManager* manager);
 
}  // namespace internal
 
namespace test_util {
class AspiredVersionsManagerTestAccess;
}  // namespace test_util
 
class AspiredVersionsManager : public Manager,
                               public Target<std::unique_ptr<Loader>> {
 public:
  using PreLoadHook = BasicManager::PreLoadHook;
 
  using CustomSortActionsFn =
      std::function<bool(const AspiredVersionPolicy::ServableAction&,
                         const AspiredVersionPolicy::ServableAction&)>;
 
  struct Options {
    std::unique_ptr<ResourceTracker> resource_tracker;
 
    int64_t manage_state_interval_micros = 100 * 1000;
 
    EventBus<ServableState>* servable_event_bus = nullptr;
 
    std::unique_ptr<AspiredVersionPolicy> aspired_version_policy;
 
    CustomSortActionsFn custom_sort_actions;
 
    uint32 num_load_threads = 0;
 
    uint32 num_unload_threads = 0;
 
    uint32 max_num_load_retries = 5;
 
    int64_t load_retry_interval_micros = 1LL * 60 * 1000 * 1000;
 
    // Defines how we want to retry when model loading fails.
    std::function<bool(absl::Status)> should_retry_model_load;
 
    // If true, and there are not multiple load threads, filesystem caches will
    // be flushed after each servable is loaded. (Cache flush is skipped when
    // multiple load threads are active, in order to avoid setting back a
    // concurrent load on another thread.)
    bool flush_filesystem_caches = false;
 
    Env* env = Env::Default();
 
    PreLoadHook pre_load_hook;
 
    // For servables which end with LoaderHarness::State::kError, enable
    // future attempts at reload to progress.
    bool enable_reload_servables_with_error = false;
 
    // If true, the AspiredVersionsManager will propagate its current context to
    // the newly created periodic functions.
    bool with_current_context = false;
  };
  static Status Create(Options options,
                       std::unique_ptr<AspiredVersionsManager>* manager);
  ~AspiredVersionsManager() override;
 
  std::vector<ServableId> ListAvailableServableIds() const override;
 
  //
  // AspiredVersionsManager's semantics with respect to this callback are as
  // follows:
  //
  // 1. OMITTING A VERSION INSTRUCTS THE MANAGER TO UNLOAD IT
  //
  // An invocation of the callback for servable stream S specifies all the
  // versions of S (if any) the manager should aim to have loaded. Each callback
  // invocation for S supercedes any prior invocations for S. Versions of S
  // supplied in previous invocations that are omitted from the latest
  // invocation will be unloaded. An invocation for S supplying an empty version
  // list causes the manager to unload all versions of S.
  //
  // First example call sequence:
  //  callback(A, {A1})      // Aspire to load version 1 of servable A.
  //  callback(B, {B1, B2})  // Aspire to load versions 1 and 2 of servable B.
  //  callback(A, {A2})      // Aspire to unload A1 and load A2.
  //  callback(B, {})        // Aspire to unload all versions of servable B.
  //
  // Second example call sequence:
  //  callback(A, {A1})      // Aspire to load version 1 of servable A.
  //  callback(A, {A1, A2})  // Aspire to load versions 1 and 2 of servable A.
  //  callback(A, {A2})      // Aspire to unload A1.
  //
  //
  // 2. Load()/Unload() CALLS GO TO A SINGLE LOADER OBJECT
  //
  // In general, multiple callback calls may supply a loader object for a given
  // servable id. Once the manager calls Load() on one of those loaders, its
  // next call for that id will be to the same loader's Unload() method. (In
  // other words, bracketed Load() and Unload() calls will be to the same loader
  // object.)
  //
  //
  // 3. NO SPONTANEOUS UNLOADING
  //
  // The manager aims to evolve the loadedness states of the servable objects it
  // manages to match the aspired list, but at a given point in time the two may
  // not coincide. That is because (a) loading/unloading are not instantaneous
  // operations, (b) loading can fail, and (c) the manager reserves the right to
  // refuse to load a servable version in the aspired list e.g. due to resource
  // limitations.
  //
  // However, the manager does obey the following constraint: Once it has loaded
  // a given servable version V, as long as V is present in the latest aspired
  // list it cannot unload V. One purpose of this guarantee is to facilitate
  // incremental loading, in which version V's Load() implementation arranges to
  // copy state from (or share state with) and already-loaded version V-1 (or
  // any prior version(s) that are loaded, for that matter). As long as V-1 is
  // currently loaded, and remains part of the aspired list, V can rely on V-1
  // remaining loaded.
  //
  Source<std::unique_ptr<Loader>>::AspiredVersionsCallback
  GetAspiredVersionsCallback() override;
 
 private:
  friend class internal::AspiredVersionsManagerTargetImpl;
  friend class test_util::AspiredVersionsManagerTestAccess;
  friend class ServerCore;
  friend uint32 internal::GetManagerNumLoadThreads(
      AspiredVersionsManager* manager);
  friend std::function<void(uint32)> internal::SetManagerNumLoadThreadsNotifier(
      AspiredVersionsManager* manager);
 
  AspiredVersionsManager(
      int64_t manage_state_interval_micros, Env* env,
      std::unique_ptr<AspiredVersionPolicy> aspired_version_policy,
      CustomSortActionsFn custom_sort_actions,
      std::unique_ptr<BasicManager> basic_manager, bool with_current_context);
 
  Status GetUntypedServableHandle(
      const ServableRequest& request,
      std::unique_ptr<UntypedServableHandle>* untyped_handle) override;
 
  std::map<ServableId, std::unique_ptr<UntypedServableHandle>>
  GetAvailableUntypedServableHandles() const override;
 
  // Enqueues an incoming aspired-versions request to be processed later,
  // asynchronously.
  void EnqueueAspiredVersionsRequest(
      const StringPiece servable_name,
      std::vector<ServableData<std::unique_ptr<Loader>>> versions)
      TF_LOCKS_EXCLUDED(pending_aspired_versions_requests_mu_);
 
  // Processes an aspired-versions request. It assumes the request doesn't
  // re-aspire any servables currently marked as not aspired in
  // 'basic_manager_'.
  void ProcessAspiredVersionsRequest(
      const StringPiece servable_name,
      std::vector<ServableData<std::unique_ptr<Loader>>> versions)
      TF_EXCLUSIVE_LOCKS_REQUIRED(basic_manager_read_modify_write_mu_);
 
  // Determines whether an aspired-versions request contains any versions that
  // are currently being managed in 'basic_manager_' with is_aspired==false.
  bool ContainsAnyReaspiredVersions(
      const StringPiece servable_name,
      const std::vector<ServableData<std::unique_ptr<Loader>>>& versions) const
      TF_SHARED_LOCKS_REQUIRED(basic_manager_read_modify_write_mu_);
 
  // Performs the action on the harness.
  void PerformAction(const AspiredVersionPolicy::ServableAction action)
      TF_EXCLUSIVE_LOCKS_REQUIRED(basic_manager_read_modify_write_mu_);
 
  // Goes through the harness map and calls the configured servable_policy with
  // the state snapshots to get a list of suggested actions. The actions are
  // then ordered and finally the topmost one is performed.
  absl::optional<AspiredVersionPolicy::ServableAction> GetNextAction()
      TF_EXCLUSIVE_LOCKS_REQUIRED(basic_manager_read_modify_write_mu_);
 
  // Checks for servables that are not aspired and at some final state and tells
  // 'basic_manager_' to forget about them. This method is intended to be
  // invoked periodically, interleaved with InvokePolicyAndExecuteAction() and
  // HandlePendingAspiredVersionsRequests().
  void FlushServables() TF_LOCKS_EXCLUDED(basic_manager_read_modify_write_mu_);
 
  // Handles enqueued aspired-versions requests. This method is intended to be
  // invoked periodically, interleaved with InvokePolicyAndExecuteAction().
  void HandlePendingAspiredVersionsRequests()
      TF_LOCKS_EXCLUDED(basic_manager_read_modify_write_mu_,
                        pending_aspired_versions_requests_mu_);
 
  // Invokes the aspired-version policy and executes any returned policy action.
  // This method is intended to be invoked periodically.
  void InvokePolicyAndExecuteAction()
      TF_LOCKS_EXCLUDED(basic_manager_read_modify_write_mu_);
 
  // Sets the number of load threads.
  //
  // This may block all new load requests, or temporarily allow more threads to
  // start, before it returns. See BasicManager::SetNumLoadThreads for details
  void SetNumLoadThreads(uint32 num_load_threads);
  uint32 num_load_threads() const;
 
  std::unique_ptr<AspiredVersionPolicy> aspired_version_policy_;
  CustomSortActionsFn custom_sort_actions_;
 
  // Aspired-versions requests pending to be processed, keyed by servable name.
  //
  // We stage incoming aspired-versions requests here and process them
  // asynchronously from the SetAspiredVersions() call, to avoid blocking in
  // SetAspiredVersions() to handle re-aspiring versions.
  //
  // For a given servable name we to need store at most pending request, since
  // each new request we receive supercedes the prior one.
  using AspiredVersionsMap =
      std::map<string, std::vector<ServableData<std::unique_ptr<Loader>>>>;
  AspiredVersionsMap pending_aspired_versions_requests_
      TF_GUARDED_BY(pending_aspired_versions_requests_mu_);
  mutable mutex pending_aspired_versions_requests_mu_;
 
  // To lock basic_manager_ to perform atomic read/modify/write operations on
  // the set of managed servables and their state (in particular, aspiredness).
  mutable mutex basic_manager_read_modify_write_mu_;
 
  // Periodically runs HandlePendingAspiredVersionsRequests() and
  // InvokePolicyAndExecuteAction() in a background thread.
  std::unique_ptr<PeriodicFunction> manage_state_thread_;
 
  // The object that implements the Target API on behalf of this manager.
  std::unique_ptr<TargetBase<std::unique_ptr<Loader>>> target_impl_;
 
  // This is where the servables "live" while they are being managed.
  std::unique_ptr<BasicManager> basic_manager_;
 
  // An observer object that forwards to SetNumLoadThreads(), if not detached.
  // This is declared last here so that it is deleted before basic_manager_.
  std::unique_ptr<Observer<const uint32>> set_num_load_threads_observer_;
 
  // For servables which end with LoaderHarness::State::kError, enable
  // future attempts at reload to progress.
  bool enable_reload_servables_with_error_ = false;
 
  TF_DISALLOW_COPY_AND_ASSIGN(AspiredVersionsManager);
};
 
}  // namespace serving
}  // namespace tensorflow
 
#endif  // TENSORFLOW_SERVING_CORE_ASPIRED_VERSIONS_MANAGER_H_