serving_session.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_SERVABLES_TENSORFLOW_SERVING_SESSION_H_
#define TENSORFLOW_SERVING_SERVABLES_TENSORFLOW_SERVING_SESSION_H_
 
#include <memory>
#include <string>
#include <utility>
#include <vector>
 
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/threadpool_options.h"
#include "tensorflow/core/public/session.h"
 
namespace tensorflow {
namespace serving {
 
class ServingSession : public Session {
 public:
  ServingSession() = default;
  ~ServingSession() override = default;
 
  // Methods that return errors.
  Status Create(const GraphDef& graph) final;
  Status Extend(const GraphDef& graph) final;
  Status Close() final;
 
  // (Subclasses just implement Run().)
};
 
class ServingSessionWrapper : public ServingSession {
 public:
  explicit ServingSessionWrapper(std::unique_ptr<Session> wrapped)
      : wrapped_(std::move(wrapped)) {
    VLOG(2) << "Created the ServingSessionWrapper around the Session.";
  }
 
  ~ServingSessionWrapper() override = default;
 
  Status Run(const std::vector<std::pair<string, Tensor>>& inputs,
             const std::vector<string>& output_tensor_names,
             const std::vector<string>& target_node_names,
             std::vector<Tensor>* outputs) override {
    return wrapped_->Run(inputs, output_tensor_names, target_node_names,
                         outputs);
  }
 
  Status Run(const RunOptions& run_options,
             const std::vector<std::pair<string, Tensor>>& inputs,
             const std::vector<string>& output_tensor_names,
             const std::vector<string>& target_node_names,
             std::vector<Tensor>* outputs, RunMetadata* run_metadata) override {
    return wrapped_->Run(run_options, inputs, output_tensor_names,
                         target_node_names, outputs, run_metadata);
  }
 
  Status Run(const RunOptions& run_options,
             const std::vector<std::pair<string, Tensor>>& inputs,
             const std::vector<string>& output_tensor_names,
             const std::vector<string>& target_node_names,
             std::vector<Tensor>* outputs, RunMetadata* run_metadata,
             const thread::ThreadPoolOptions& thread_pool_options) override {
    return wrapped_->Run(run_options, inputs, output_tensor_names,
                         target_node_names, outputs, run_metadata,
                         thread_pool_options);
  }
 
  Status ListDevices(std::vector<DeviceAttributes>* response) override {
    return wrapped_->ListDevices(response);
  }
 
 private:
  std::unique_ptr<Session> wrapped_;
 
  TF_DISALLOW_COPY_AND_ASSIGN(ServingSessionWrapper);
};
 
// Subclass of SessionWrapper which reroutes Run() calls with
// thread_pool_options to Run() without those options.  This is to provide
// support for RemoteSession::Run which does not implement the overloaded Run()
// method with thread pool options.
class SessionWrapperIgnoreThreadPoolOptions : public ServingSessionWrapper {
 public:
  explicit SessionWrapperIgnoreThreadPoolOptions(
      std::unique_ptr<Session> wrapped)
      : ServingSessionWrapper(std::move(wrapped)) {
    VLOG(2) << "Created the SessionWrapperIgnoreThreadPoolOptions around the "
               "Session.";
  }
 
  Status Run(const RunOptions& run_options,
             const std::vector<std::pair<string, Tensor>>& inputs,
             const std::vector<string>& output_tensor_names,
             const std::vector<string>& target_node_names,
             std::vector<Tensor>* outputs, RunMetadata* run_metadata,
             const thread::ThreadPoolOptions& thread_pool_options) override {
    return ServingSessionWrapper::Run(run_options, inputs, output_tensor_names,
                                      target_node_names, outputs, run_metadata);
  }
 
 private:
  TF_DISALLOW_COPY_AND_ASSIGN(SessionWrapperIgnoreThreadPoolOptions);
};
 
}  // namespace serving
}  // namespace tensorflow
 
#endif  // TENSORFLOW_SERVING_SERVABLES_TENSORFLOW_SERVING_SESSION_H_