saved_model_bundle_factory.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_SAVED_MODEL_BUNDLE_FACTORY_H_
#define TENSORFLOW_SERVING_SERVABLES_TENSORFLOW_SAVED_MODEL_BUNDLE_FACTORY_H_
 
#include "absl/types/optional.h"
#include "tensorflow/cc/saved_model/loader.h"
#include "tensorflow/core/kernels/batching_util/shared_batch_scheduler.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow_serving/batching/batching_session.h"
#include "tensorflow_serving/core/loader.h"
#include "tensorflow_serving/resources/resources.pb.h"
#include "tensorflow_serving/servables/tensorflow/session_bundle_config.pb.h"
 
namespace tensorflow {
namespace serving {
 
class SavedModelBundleFactory {
 public:
  static Status Create(const SessionBundleConfig& config,
                       std::unique_ptr<SavedModelBundleFactory>* factory);
 
  Status CreateSavedModelBundle(const string& path,
                                std::unique_ptr<SavedModelBundle>* bundle);
 
  Status CreateSavedModelBundleWithMetadata(
      const Loader::Metadata& metadata, const string& path,
      std::unique_ptr<SavedModelBundle>* bundle);
 
  Status EstimateResourceRequirement(const string& path,
                                     ResourceAllocation* estimate) const;
 
  const SessionBundleConfig& config() const { return config_; }
  SessionBundleConfig& mutable_config() { return config_; }
 
 private:
  using Batcher = SharedBatchScheduler<BatchingSessionTask>;
 
  SavedModelBundleFactory(const SessionBundleConfig& config,
                          std::shared_ptr<Batcher> batch_scheduler);
 
  Status InternalCreateSavedModelBundle(
      const absl::optional<Loader::Metadata>& metadata, const string& path,
      std::unique_ptr<SavedModelBundle>* bundle);
 
  SessionBundleConfig config_;
 
  // A shared batch scheduler. One queue is used for each session this factory
  // emits. If batching is not configured, this remains null.
  std::shared_ptr<Batcher> batch_scheduler_;
 
  TF_DISALLOW_COPY_AND_ASSIGN(SavedModelBundleFactory);
};
 
}  // namespace serving
}  // namespace tensorflow
 
#endif  // TENSORFLOW_SERVING_SERVABLES_TENSORFLOW_SAVED_MODEL_BUNDLE_FACTORY_H_