

#include "tensorflow_serving/core/eager_unload_policy.h"

namespace tensorflow {
namespace serving {

optional<AspiredVersionPolicy::ServableAction> EagerUnloadPolicy::GetNextAction(
    const std::vector<AspiredServableStateSnapshot>& all_versions) const {
  // First iterate over all_versions and find any in kReady that are no longer
  // aspired. Unload the first if any.
  for (const auto& version : all_versions) {
    if (version.state == LoaderHarness::State::kReady && !version.is_aspired) {
      return AspiredVersionPolicy::ServableAction(
          {Action::kUnload, version.id});
    }
  }

  // Second and only if no action was found earlier, iterate over all
  // versions and find any in kNew that are aspired. Load the first if any.
  for (const auto& version : all_versions) {
    if (version.state == LoaderHarness::State::kNew && version.is_aspired) {
      return AspiredVersionPolicy::ServableAction({Action::kLoad, version.id});
    }
  }

  return nullopt;
}

}  // namespace serving
}  // namespace tensorflow
