Add an option to propagate current Context in periodic functions from…

… AspiredVersionsManager.

PiperOrigin-RevId: 566682492

tensorflow_serving/core/aspired_versions_manager.cc

@@ -25,6 +25,7 @@ limitations under the License.
 
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/context.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow_serving/core/loader_harness.h"
@@ -169,7 +170,8 @@ Status AspiredVersionsManager::Create(
 
   manager->reset(new AspiredVersionsManager(
       options.manage_state_interval_micros, options.env,
-      std::move(options.aspired_version_policy), std::move(basic_manager)));
+      std::move(options.aspired_version_policy), std::move(basic_manager),
+      options.with_current_context));
   (manager->get())->enable_reload_servables_with_error_ =
       options.enable_reload_servables_with_error;
   return OkStatus();
@@ -178,7 +180,7 @@ Status AspiredVersionsManager::Create(
 AspiredVersionsManager::AspiredVersionsManager(
     int64_t manage_state_interval_micros, Env* env,
     std::unique_ptr<AspiredVersionPolicy> aspired_version_policy,
-    std::unique_ptr<BasicManager> basic_manager)
+    std::unique_ptr<BasicManager> basic_manager, bool with_current_context)
     : aspired_version_policy_(std::move(aspired_version_policy)),
       target_impl_(new internal::AspiredVersionsManagerTargetImpl(this)),
       basic_manager_(std::move(basic_manager)) {
@@ -190,13 +192,25 @@ AspiredVersionsManager::AspiredVersionsManager(
     PeriodicFunction::Options pf_options;
     pf_options.env = env;
     pf_options.thread_name_prefix = "AspiredVersionsManager_ManageState_Thread";
-    manage_state_thread_.reset(new PeriodicFunction(
-        [this]() {
-          this->FlushServables();
-          this->HandlePendingAspiredVersionsRequests();
-          this->InvokePolicyAndExecuteAction();
-        },
-        manage_state_interval_micros));
+    if (with_current_context) {
+      tensorflow::Context context(tensorflow::ContextKind::kThread);
+      manage_state_thread_.reset(new PeriodicFunction(
+          [this, context = std::move(context)]() {
+            tensorflow::WithContext wc(context);
+            this->FlushServables();
+            this->HandlePendingAspiredVersionsRequests();
+            this->InvokePolicyAndExecuteAction();
+          },
+          manage_state_interval_micros));
+    } else {
+      manage_state_thread_.reset(new PeriodicFunction(
+          [this]() {
+            this->FlushServables();
+            this->HandlePendingAspiredVersionsRequests();
+            this->InvokePolicyAndExecuteAction();
+          },
+          manage_state_interval_micros));
+    }
   }
 }
 

tensorflow_serving/core/aspired_versions_manager.h

@@ -144,6 +144,10 @@ class AspiredVersionsManager : public Manager,
     // For servables which end with LoaderHarness::State::kError, enable
     // future attempts at reload to progress.
     bool enable_reload_servables_with_error = false;
+
+    // If true, the AspiredVersionsManager will propagate its current context to
+    // the newly created periodic functions.
+    bool with_current_context = false;
   };
   static Status Create(Options options,
                        std::unique_ptr<AspiredVersionsManager>* manager);
@@ -220,7 +224,7 @@ class AspiredVersionsManager : public Manager,
   AspiredVersionsManager(
       int64_t manage_state_interval_micros, Env* env,
       std::unique_ptr<AspiredVersionPolicy> aspired_version_policy,
-      std::unique_ptr<BasicManager> basic_manager);
+      std::unique_ptr<BasicManager> basic_manager, bool with_current_context);
 
   Status GetUntypedServableHandle(
       const ServableRequest& request,

tensorflow_serving/model_servers/BUILD

@@ -87,6 +87,7 @@ cc_library(
         "//tensorflow_serving/core:source",
         "//tensorflow_serving/core:source_adapter",
         "//tensorflow_serving/core:storage_path",
+        "//tensorflow_serving/core:stream_logger",
         "//tensorflow_serving/resources:resource_values",
         "//tensorflow_serving/servables/tensorflow:predict_util",
         "//tensorflow_serving/servables/tensorflow:saved_model_bundle_source_adapter",

tensorflow_serving/model_servers/server_core.cc

@@ -751,6 +751,7 @@ Status ServerCore::CreateAspiredVersionsManager(
   manager_options.flush_filesystem_caches = options_.flush_filesystem_caches;
   manager_options.enable_reload_servables_with_error =
       options_.enable_reload_servables_with_error;
+  manager_options.with_current_context = options_.with_current_context;
   const tensorflow::Status status =
       AspiredVersionsManager::Create(std::move(manager_options), manager);
   if (!status.ok()) {

tensorflow_serving/model_servers/server_core.h

@@ -202,6 +202,10 @@ class ServerCore : public Manager {
     std::string storage_path_prefix;
 
     bool enable_cors_support = false;
+
+    // If true, propagate current context to children threads (periodic
+    // functions) in AspiredVersionsManager.
+    bool with_current_context = false;
   };
 
   virtual ~ServerCore() = default;