Automated Code Change

PiperOrigin-RevId: 624410019
tensorflow · Apr 13, 2024 · 11a994a · 11a994a
1 parent 415ba2b
commit 11a994a
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 14 deletions.
diff --git a/tensorflow/core/tfrt/run_handler_thread_pool/run_handler.cc b/tensorflow/core/tfrt/run_handler_thread_pool/run_handler.cc
@@ -159,12 +159,12 @@ ThreadWorkSource::~ThreadWorkSource() {
 Task ThreadWorkSource::EnqueueTask(Task t, bool is_blocking,
                                    bool enable_wake_up) {
   uint64_t id = t.f->trace_id;
-  tensorflow::profiler::TraceMe activity(
+  tsl::profiler::TraceMe activity(
       [id, is_blocking] {
-        return tensorflow::profiler::TraceMeEncode(
+        return tsl::profiler::TraceMeEncode(
             "Enqueue", {{"id", id}, {"is_blocking", is_blocking}});
       },
-      tensorflow::profiler::TraceMeLevel::kInfo);
+      tsl::profiler::TraceMeLevel::kInfo);
   tensorflow::mutex* mu = nullptr;
   Queue* task_queue = nullptr;
   thread_local int64_t closure_counter = 0;
@@ -586,12 +586,12 @@ void RunHandlerThreadPool::WorkerLoop(int thread_id,
       tws->DecrementInflightTaskCount(task_from_blocking_queue);
       tws->DecrementPendingTaskCount();
     } else {
-      tensorflow::profiler::TraceMe activity(
+      tsl::profiler::TraceMe activity(
           [thread_id] {
-            return tensorflow::profiler::TraceMeEncode(
-                "Sleeping", {{"thread_id", thread_id}});
+            return tsl::profiler::TraceMeEncode("Sleeping",
+                                                {{"thread_id", thread_id}});
           },
-          tensorflow::profiler::TraceMeLevel::kInfo);
+          tsl::profiler::TraceMeLevel::kInfo);
       if (VLOG_IS_ON(4)) {
         for (int i = 0; i < thread_work_sources->size(); ++i) {
           VLOG(4) << "source id " << i << " "
@@ -789,12 +789,12 @@ class RunHandlerPool::Impl {
     {
       tensorflow::mutex_lock l(mu_);
       if (!has_free_handler()) {
-        tensorflow::profiler::TraceMe activity(
+        tsl::profiler::TraceMe activity(
             [step_id] {
-              return tensorflow::profiler::TraceMeEncode(
-                  "WaitingForHandler", {{"step_id", step_id}});
+              return tsl::profiler::TraceMeEncode("WaitingForHandler",
+                                                  {{"step_id", step_id}});
             },
-            tensorflow::profiler::TraceMeLevel::kInfo);
+            tsl::profiler::TraceMeLevel::kInfo);
         if (timeout_in_ms == 0) {
           mu_.Await(tensorflow::Condition(this, &Impl::has_free_handler));
         } else if (!mu_.AwaitWithDeadline(

diff --git a/tensorflow/core/tfrt/run_handler_thread_pool/run_handler_concurrent_work_queue.cc b/tensorflow/core/tfrt/run_handler_thread_pool/run_handler_concurrent_work_queue.cc
@@ -59,7 +59,7 @@ RunHandlerThreadWorkQueue::RunHandlerThreadWorkQueue(const Options& options)
   handler_pool_ = std::make_unique<RunHandlerPool>(pool_options);
 }
 
-tensorflow::StatusOr<std::unique_ptr<tensorflow::tfrt_stub::WorkQueueInterface>>
+absl::StatusOr<std::unique_ptr<tensorflow::tfrt_stub::WorkQueueInterface>>
 RunHandlerThreadWorkQueue::InitializeRequest(int64_t request_id) const {
   RunHandlerOptions options;
   std::unique_ptr<RunHandler> handler =

diff --git a/tensorflow/core/tfrt/run_handler_thread_pool/run_handler_concurrent_work_queue.h b/tensorflow/core/tfrt/run_handler_thread_pool/run_handler_concurrent_work_queue.h
@@ -91,8 +91,7 @@ class RunHandlerThreadWorkQueue
         " complementary threads)");
   }
 
-  tensorflow::StatusOr<
-      std::unique_ptr<tensorflow::tfrt_stub::WorkQueueInterface>>
+  absl::StatusOr<std::unique_ptr<tensorflow::tfrt_stub::WorkQueueInterface>>
   InitializeRequest(int64_t request_id) const override;
 
   int GetParallelismLevel() const override {