loader.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_H_
#define TENSORFLOW_SERVING_CORE_LOADER_H_
 
#include <memory>
 
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow_serving/core/source.h"
#include "tensorflow_serving/resources/resources.pb.h"
#include "tensorflow_serving/util/any_ptr.h"
 
namespace tensorflow {
namespace serving {
 
class Loader {
 public:
  virtual ~Loader() = default;
 
  //      include: (a) replace conservative estimate with actual measurement
  //      once loaded in memory; (b) load process consumes extra transient
  //      memory that is not used in steady-state after the load completes.
  virtual Status EstimateResources(ResourceAllocation* estimate) const = 0;
 
  virtual Status Load() {
    return errors::Unimplemented("Load isn't implemented.");
  }
 
  struct Metadata {
    ServableId servable_id;
  };
  virtual Status LoadWithMetadata(const Metadata& metadata) { return Load(); }
 
  virtual void Unload() = 0;
 
  virtual AnyPtr servable() = 0;
};
 
inline bool operator==(const Loader::Metadata& a, const Loader::Metadata& b) {
  return a.servable_id == b.servable_id;
}
 
inline bool operator!=(const Loader::Metadata& a, const Loader::Metadata& b) {
  return !(a == b);
}
 
class ResourceUnsafeLoader : public Loader {
 public:
  Status EstimateResources(ResourceAllocation* estimate) const final {
    estimate->Clear();
    return Status();
  }
};
 
// A source that emits Loader unique pointers.
using LoaderSource = Source<std::unique_ptr<Loader>>;
 
}  // namespace serving
}  // namespace tensorflow
 
#endif  // TENSORFLOW_SERVING_CORE_LOADER_H_