prevent unloading the serving versions.

tensorflow_serving/core/BUILD

@@ -659,6 +659,8 @@ cc_library(
         ":loader_harness",
         ":servable_id",
         "@com_google_absl//absl/types:optional",
+        "//tensorflow_serving/sources/storage_path:file_system_storage_path_source",
+        "//tensorflow_serving/sources/storage_path:file_system_storage_path_source_proto",
         "@org_tensorflow//tensorflow/core:lib",
     ],
 )

tensorflow_serving/core/aspired_version_policy.cc

@@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <unordered_set>
+
 #include "tensorflow_serving/core/aspired_version_policy.h"
 
 namespace tensorflow {
@@ -32,5 +34,16 @@ absl::optional<ServableId> AspiredVersionPolicy::GetHighestAspiredNewServableId(
   return highest_version_id;
 }
 
+std::unordered_set<int64>
+AspiredVersionPolicy::GetSpecificVersionsInConfig(
+    const std::string& servable_name) const {
+  mutex_lock l(mu_);
+  if (storage_path_source_ == nullptr) {
+    std::unordered_set<int64> empty_set;
+    return empty_set;
+  }
+  return storage_path_source_->GetSpecificVersionsInConfig(servable_name);
+}
+
 }  // namespace serving
 }  // namespace tensorflow

tensorflow_serving/core/aspired_version_policy.h

@@ -18,12 +18,15 @@ limitations under the License.
 
 #include <string>
 #include <vector>
+#include <unordered_set>
 
 #include "absl/types/optional.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow_serving/core/loader_harness.h"
 #include "tensorflow_serving/core/servable_id.h"
+#include "tensorflow_serving/sources/storage_path/file_system_storage_path_source.h"
+
 namespace tensorflow {
 namespace serving {
 
@@ -77,6 +80,15 @@ class AspiredVersionPolicy {
   virtual absl::optional<ServableAction> GetNextAction(
       const std::vector<AspiredServableStateSnapshot>& all_versions) const = 0;
 
+  void set_storage_path_source(
+      FileSystemStoragePathSource* storage_path_source) {
+    mutex_lock l(mu_);
+    storage_path_source_ = storage_path_source;
+  }
+
+  std::unordered_set<int64> GetSpecificVersionsInConfig(
+      const std::string& servable_name) const;
+
  protected:
   /// Returns the aspired ServableId with the highest version that matches
   /// kNew state, if any exists.
@@ -85,6 +97,8 @@ class AspiredVersionPolicy {
 
  private:
   friend class AspiredVersionPolicyTest;
+  mutable mutex mu_;
+  FileSystemStoragePathSource* storage_path_source_ GUARDED_BY(mu_) = nullptr;
 };
 
 inline bool operator==(const AspiredVersionPolicy::ServableAction& lhs,

tensorflow_serving/core/aspired_versions_manager.h

@@ -203,6 +203,11 @@ class AspiredVersionsManager : public Manager,
   //
   Source<std::unique_ptr<Loader>>::AspiredVersionsCallback
   GetAspiredVersionsCallback() override;
+
+  void SetPolicyStoragePathSource(
+      FileSystemStoragePathSource* storage_path_source) {
+    aspired_version_policy_->set_storage_path_source(storage_path_source);
+  }
 
  private:
   friend class internal::AspiredVersionsManagerTargetImpl;

tensorflow_serving/core/availability_preserving_policy.cc

@@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <unordered_set>
+
 #include "tensorflow_serving/core/availability_preserving_policy.h"
 
 #include "absl/types/optional.h"
@@ -67,7 +69,11 @@ AvailabilityPreservingPolicy::GetNextAction(
     absl::optional<ServableId> version_to_unload =
         GetLowestServableId(unaspired_serving_versions);
     if (version_to_unload) {
-      return {{Action::kUnload, version_to_unload.value()}};
+      std::unordered_set<int64> specific_versions =
+          GetSpecificVersionsInConfig(version_to_unload.value().name);
+      if (specific_versions.count(version_to_unload.value().version) == 0) {
+        return {{Action::kUnload, version_to_unload.value()}};
+      }
     }
   }
 

tensorflow_serving/model_servers/server_core.cc

@@ -347,6 +347,7 @@ Status ServerCore::AddModelsViaModelConfigList() {
 
     // Stow the source components.
     storage_path_source_and_router_ = {source.get(), router.get()};
+    manager_->SetPolicyStoragePathSource(storage_path_source_and_router_->source);
     manager_.AddDependency(std::move(source));
     if (prefix_source_adapter != nullptr) {
       manager_.AddDependency(std::move(prefix_source_adapter));

tensorflow_serving/sources/storage_path/file_system_storage_path_source.cc

@@ -386,6 +386,24 @@ void FileSystemStoragePathSource::SetAspiredVersionsCallback(
   }
 }
 
+std::unordered_set<int64>
+FileSystemStoragePathSource::GetSpecificVersionsInConfig(
+    const std::string& servable_name) const {
+  mutex_lock l(mu_);
+
+  std::unordered_set<int64> specific_versions;
+  for (const FileSystemStoragePathSourceConfig::ServableToMonitor& servable :
+       config_.servables()) {
+    if (servable.servable_name() == servable_name) {
+       specific_versions.insert(
+           servable.servable_version_policy().specific().versions().begin(),
+           servable.servable_version_policy().specific().versions().end());
+       break;
+    }
+  }
+  return specific_versions;
+}
+
 Status FileSystemStoragePathSource::PollFileSystemAndInvokeCallback() {
   mutex_lock l(mu_);
   std::map<string, std::vector<ServableData<StoragePath>>>

tensorflow_serving/sources/storage_path/file_system_storage_path_source.h

@@ -19,6 +19,7 @@ limitations under the License.
 #include <functional>
 #include <memory>
 #include <utility>
+#include <unordered_set>
 
 #include "absl/types/variant.h"
 #include "tensorflow/core/kernels/batching_util/periodic_function.h"
@@ -77,6 +78,9 @@ class FileSystemStoragePathSource : public Source<StoragePath> {
     mutex_lock l(mu_);
     return config_;
   }
+
+  std::unordered_set<int64> GetSpecificVersionsInConfig(
+      const std::string& servable_name) const;
 
  private:
   friend class internal::FileSystemStoragePathSourceTestAccess;