source_adapter.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_SOURCE_ADAPTER_H_
#define TENSORFLOW_SERVING_CORE_SOURCE_ADAPTER_H_
 
#include <algorithm>
#include <vector>
 
#include "tensorflow/core/lib/core/notification.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow_serving/core/loader.h"
#include "tensorflow_serving/core/servable_data.h"
#include "tensorflow_serving/core/source.h"
#include "tensorflow_serving/core/storage_path.h"
#include "tensorflow_serving/core/target.h"
#include "tensorflow_serving/util/class_registration.h"
 
namespace tensorflow {
namespace serving {
 
template <typename InputType, typename OutputType>
class SourceAdapter : public TargetBase<InputType>, public Source<OutputType> {
 public:
  ~SourceAdapter() override = 0;
 
  void SetAspiredVersions(const StringPiece servable_name,
                          std::vector<ServableData<InputType>> versions) final;
 
  void SetAspiredVersionsCallback(
      typename Source<OutputType>::AspiredVersionsCallback callback) final;
 
  virtual std::vector<ServableData<OutputType>> Adapt(
      const StringPiece servable_name,
      std::vector<ServableData<InputType>> versions) = 0;
 
  ServableData<OutputType> AdaptOneVersion(ServableData<InputType> input);
 
 protected:
  // This is an abstract class.
  SourceAdapter() = default;
 
 private:
  // The callback for emitting OutputType-based aspired-version lists.
  typename Source<OutputType>::AspiredVersionsCallback outgoing_callback_;
 
  // Has 'outgoing_callback_' been set yet, so that the SourceAdapter is ready
  // to propagate aspired versions?
  Notification outgoing_callback_set_;
};
 
// START_SKIP_DOXYGEN
 
// Define a SourceAdapter registry for the common case of adapting from a
// storage path to a loader.
using StoragePathSourceAdapter =
    SourceAdapter<StoragePath, std::unique_ptr<Loader>>;
DEFINE_CLASS_REGISTRY(StoragePathSourceAdapterRegistry,
                      StoragePathSourceAdapter);
#define REGISTER_STORAGE_PATH_SOURCE_ADAPTER(ClassCreator, ConfigProto)      \
  REGISTER_CLASS(StoragePathSourceAdapterRegistry, StoragePathSourceAdapter, \
                 ClassCreator, ConfigProto);
 
// A source adapter that converts InputType instances to OutputType instances
// one at a time (i.e. there is no interaction among members of a given aspired-
// version list). Most source adapters can subclass UnarySourceAdapter, and do
// not need the full generality of SourceAdapter.
//
// Requires OutputType to be default-constructable and updatable in-place.
//
// Implementing subclasses supply an implementation of the Convert() virtual
// method, which converts a servable from InputType to OutputType.
//
// IMPORTANT: Every leaf derived class must call Detach() at the top of its
// destructor. (See documentation on TargetBase::Detach() in target.h.) Doing so
// ensures that no Convert() calls are in flight during destruction of member
// variables.
template <typename InputType, typename OutputType>
class UnarySourceAdapter : public SourceAdapter<InputType, OutputType> {
 public:
  ~UnarySourceAdapter() override = 0;
 
 protected:
  // This is an abstract class.
  UnarySourceAdapter() = default;
 
 private:
  // This method is implemented in terms of Convert(), which the implementing
  // subclass must supply.
  std::vector<ServableData<OutputType>> Adapt(
      const StringPiece servable_name,
      std::vector<ServableData<InputType>> versions) final;
 
  // Converts a single InputType instance into a corresponding OutputType
  // instance.
  virtual Status Convert(const InputType& data, OutputType* converted_data) = 0;
};
 
// A source adapter that converts every incoming ServableData<InputType> item
// into an error-containing ServableData<OutputType>. If the incoming data item
// was already an error, the existing error is passed through; otherwise a new
// error Status given via this class's constructor is added.
//
// This class is useful in conjunction with a router, to handle servable data
// items that do not conform to any explicitly-programmed route. Specifically,
// consider a fruit router configured route apples to output port 0, oranges to
// output port 1, and anything else to a final port 2. If we only have proper
// SourceAdapters to handle apples and oranges, we might connect an
// ErrorInjectingSourceAdapter to port 2, to catch any unexpected fruits.
template <typename InputType, typename OutputType>
class ErrorInjectingSourceAdapter final
    : public SourceAdapter<InputType, OutputType> {
 public:
  explicit ErrorInjectingSourceAdapter(const Status& error);
  ~ErrorInjectingSourceAdapter() override;
 
 private:
  std::vector<ServableData<OutputType>> Adapt(
      const StringPiece servable_name,
      std::vector<ServableData<InputType>> versions) override;
 
  // The error status to inject.
  const Status error_;
 
  TF_DISALLOW_COPY_AND_ASSIGN(ErrorInjectingSourceAdapter);
};
 
// Implementation details follow. API users need not read.
 
template <typename InputType, typename OutputType>
SourceAdapter<InputType, OutputType>::~SourceAdapter() {}
 
template <typename InputType, typename OutputType>
void SourceAdapter<InputType, OutputType>::SetAspiredVersions(
    const StringPiece servable_name,
    std::vector<ServableData<InputType>> versions) {
  outgoing_callback_set_.WaitForNotification();
  outgoing_callback_(servable_name, Adapt(servable_name, std::move(versions)));
}
 
template <typename InputType, typename OutputType>
void SourceAdapter<InputType, OutputType>::SetAspiredVersionsCallback(
    typename Source<OutputType>::AspiredVersionsCallback callback) {
  outgoing_callback_ = callback;
  outgoing_callback_set_.Notify();
}
 
template <typename InputType, typename OutputType>
ServableData<OutputType> SourceAdapter<InputType, OutputType>::AdaptOneVersion(
    ServableData<InputType> input) {
  const StringPiece servable_name(input.id().name);
  std::vector<ServableData<InputType>> input_versions = {input};
  std::vector<ServableData<OutputType>> output_versions =
      Adapt(servable_name, input_versions);
  DCHECK_EQ(1, output_versions.size());
  return std::move(output_versions[0]);
}
 
template <typename InputType, typename OutputType>
UnarySourceAdapter<InputType, OutputType>::~UnarySourceAdapter() {}
 
template <typename InputType, typename OutputType>
std::vector<ServableData<OutputType>>
UnarySourceAdapter<InputType, OutputType>::Adapt(
    const StringPiece servable_name,
    std::vector<ServableData<InputType>> versions) {
  std::vector<ServableData<OutputType>> adapted_versions;
  for (const ServableData<InputType>& version : versions) {
    if (version.status().ok()) {
      OutputType adapted_data;
      Status adapt_status = Convert(version.DataOrDie(), &adapted_data);
      if (adapt_status.ok()) {
        adapted_versions.emplace_back(
            ServableData<OutputType>{version.id(), std::move(adapted_data)});
      } else {
        adapted_versions.emplace_back(
            ServableData<OutputType>{version.id(), adapt_status});
      }
    } else {
      adapted_versions.emplace_back(
          ServableData<OutputType>{version.id(), version.status()});
    }
  }
  return adapted_versions;
}
 
template <typename InputType, typename OutputType>
ErrorInjectingSourceAdapter<InputType, OutputType>::ErrorInjectingSourceAdapter(
    const Status& error)
    : error_(error) {
  DCHECK(!error.ok());
}
 
template <typename InputType, typename OutputType>
ErrorInjectingSourceAdapter<InputType,
                            OutputType>::~ErrorInjectingSourceAdapter() {
  TargetBase<InputType>::Detach();
}
 
template <typename InputType, typename OutputType>
std::vector<ServableData<OutputType>>
ErrorInjectingSourceAdapter<InputType, OutputType>::Adapt(
    const StringPiece servable_name,
    std::vector<ServableData<InputType>> versions) {
  std::vector<ServableData<OutputType>> adapted_versions;
  for (const ServableData<InputType>& version : versions) {
    if (version.status().ok()) {
      LOG(INFO) << "Injecting error for servable " << version.id() << ": "
                << error_.message();
      adapted_versions.emplace_back(
          ServableData<OutputType>{version.id(), error_});
    } else {
      adapted_versions.emplace_back(
          ServableData<OutputType>{version.id(), version.status()});
    }
  }
  return adapted_versions;
}
 
// END_SKIP_DOXYGEN
 
}  // namespace serving
}  // namespace tensorflow
 
#endif  // TENSORFLOW_SERVING_CORE_SOURCE_ADAPTER_H_