loader_harness.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_LOADER_HARNESS_H_
#define TENSORFLOW_SERVING_CORE_LOADER_HARNESS_H_
 
#include <memory>
 
#include "absl/status/status.h"
#include "absl/types/optional.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow_serving/core/loader.h"
#include "tensorflow_serving/core/servable_id.h"
 
namespace tensorflow {
namespace serving {
 
// START_SKIP_DOXYGEN
// Forward-declaration for the use in LoaderHarness.
template <typename T>
struct ServableStateSnapshot;
// END_SKIP_DOXYGEN
 
class LoaderHarness final {
 public:
  enum class State {
    kNew,
 
    kLoadRequested,
 
    kLoadApproved,
 
    kLoading,
 
    kReady,
 
    kUnloadRequested,
 
    kQuiescing,
 
    kQuiesced,
 
    kUnloading,
 
    kDisabled,
 
    kError
  };
 
  struct Options {
    Options() {}
 
    uint32 max_num_load_retries = 0;
 
    uint64_t load_retry_interval_micros = 0;
 
    std::function<void(const ServableId& id, const Status& error)>
        error_callback;
  };
 
  LoaderHarness(const ServableId& id, std::unique_ptr<Loader> loader,
                const Options& options = Options());
 
  template <typename T>
  LoaderHarness(const ServableId& id, std::unique_ptr<Loader> loader,
                std::unique_ptr<T> additional_state,
                const Options& options = Options())
      : id_(id),
        loader_(std::move(loader)),
        additional_state_(std::move(additional_state)),
        options_(options),
        should_retry_([&](absl::Status status) { return true; }) {}
 
  ~LoaderHarness();
 
  ServableId id() const { return id_; }
 
  State state() const TF_LOCKS_EXCLUDED(mu_);
 
  Loader* loader() const { return loader_.get(); }
 
  template <typename T = std::nullptr_t>
  ServableStateSnapshot<T> loader_state_snapshot() const TF_LOCKS_EXCLUDED(mu_);
 
  Status LoadRequested() TF_LOCKS_EXCLUDED(mu_);
 
  Status LoadApproved() TF_LOCKS_EXCLUDED(mu_);
 
  Status Load() TF_LOCKS_EXCLUDED(mu_);
 
  Status UnloadRequested() TF_LOCKS_EXCLUDED(mu_);
 
  void set_should_retry(std::function<bool(absl::Status)> should_retry)
      TF_LOCKS_EXCLUDED(mu_);
 
  bool should_retry(absl::Status status) TF_LOCKS_EXCLUDED(mu_);
 
  Status Unload() TF_LOCKS_EXCLUDED(mu_);
 
  Status StartQuiescing() TF_LOCKS_EXCLUDED(mu_);
 
  Status DoneQuiescing() TF_LOCKS_EXCLUDED(mu_);
 
  void Error(const Status& status) TF_LOCKS_EXCLUDED(mu_);
 
  Status status() const TF_LOCKS_EXCLUDED(mu_);
 
  template <typename T>
  T* additional_state() {
    return additional_state_.get<T>();
  }
 
  static string StateDebugString(State state);
 
 private:
  // Transitions the state to kError and invokes 'options_.error_callback'.
  // Private method to be used when we want to set an error from another method
  // in this class, where mu_ is already held.
  void ErrorInternal(const Status& status) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
  // Expects 'state_' to equal 'from', and if so transitions it to 'to'. If not,
  // DCHECK-fails, calls ErrorInternal() with a suitable error and returns the
  // same error.
  Status TransitionState(State from, State to) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
  Status UnloadInternal(State from_state) TF_LOCKS_EXCLUDED(mu_);
  Status UnloadDueToCancelledLoad() TF_LOCKS_EXCLUDED(mu_);
 
  const ServableId id_;
  const std::unique_ptr<Loader> loader_;
  // Additional state that the manager uses.
  const UniqueAnyPtr additional_state_;
  const Options options_;
  mutable mutex mu_;
  State state_ TF_GUARDED_BY(mu_) = State::kNew;
  // If state_ is kError, this will be non-OK.
  Status status_ TF_GUARDED_BY(mu_);
  // The retry policy for the servable. If the function returns false, we cancel
  // the next retry. This does not affect the current load action already
  // running.
  // There is no retry if the last action was successful.
  std::function<bool(absl::Status)> should_retry_ TF_GUARDED_BY(mu_);
 
  TF_DISALLOW_COPY_AND_ASSIGN(LoaderHarness);
};
 
// START_SKIP_DOXYGEN
 
// A snapshot of a servable's state and aspiredness, from the LoaderHarness's
// perspective.
template <typename T = std::nullptr_t>
struct ServableStateSnapshot final {
  ServableId id;
  LoaderHarness::State state;
  absl::optional<T> additional_state;
};
 
template <typename T>
inline bool operator==(const ServableStateSnapshot<T>& a,
                       const ServableStateSnapshot<T>& b) {
  return a.id == b.id && a.state == b.state &&
         a.additional_state == b.additional_state;
}
 
template <typename T>
inline bool operator!=(const ServableStateSnapshot<T>& a,
                       const ServableStateSnapshot<T>& b) {
  return !(a == b);
}
 
inline std::ostream& operator<<(std::ostream& os, LoaderHarness::State state) {
  os << LoaderHarness::StateDebugString(state);
  return os;
}
 
// Implementation details. API readers may skip.
 
template <typename T>
ServableStateSnapshot<T> LoaderHarness::loader_state_snapshot() const {
  mutex_lock l(mu_);
  if (additional_state_.get<T>() == nullptr) {
    return {id_, state_, {}};
  } else {
    return {id_, state_, {*additional_state_.get<T>()}};
  }
}
 
// Specialization for std::nullptr_t.
//
// We mark this inline to follow ODR.
template <>
inline ServableStateSnapshot<std::nullptr_t>
LoaderHarness::loader_state_snapshot() const {
  mutex_lock l(mu_);
  return {id_, state_, {}};
}
 
// END_SKIP_DOXYGEN
 
}  // namespace serving
}  // namespace tensorflow
 
#endif  // TENSORFLOW_SERVING_CORE_LOADER_HARNESS_H_