From 33a2320f4e805a09d4cfd7cabdfc160353fdfd4b Mon Sep 17 00:00:00 2001
From: xla authors <google-ml-automation@google.com>
Date: Fri, 24 Jan 2025 16:54:24 -0800
Subject: [PATCH] Extra debug logging

PiperOrigin-RevId: 719477321
---
 xla/pjrt/BUILD          | 1 +
 xla/pjrt/pjrt_future.cc | 7 +++++++
 2 files changed, 8 insertions(+)

diff --git a/xla/pjrt/BUILD b/xla/pjrt/BUILD
index 326cdf61b038d..94e30153b09ac 100644
--- a/xla/pjrt/BUILD
+++ b/xla/pjrt/BUILD
@@ -667,6 +667,7 @@ cc_library(
         "//xla/tsl/concurrency:async_value",
         "//xla/tsl/concurrency:ref_count",
         "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
diff --git a/xla/pjrt/pjrt_future.cc b/xla/pjrt/pjrt_future.cc
index 0d208ef8ec274..51bc0cc46d95b 100644
--- a/xla/pjrt/pjrt_future.cc
+++ b/xla/pjrt/pjrt_future.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include <utility>
 
 #include "absl/base/thread_annotations.h"
+#include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
@@ -54,6 +55,12 @@ PjRtFuture<> JoinFutures(absl::Span<const PjRtFuture<>> futures) {
     future.OnReady([state](absl::Status status) {
       if (!status.ok()) {
         absl::MutexLock lock(&state->mu);
+        if (VLOG_IS_ON(2)) {
+          if (status.code() != state->status.code()) {
+            VLOG(2) << "Ignoring status " << status
+                    << " because first error was " << state->status;
+          }
+        }
         state->status.Update(status);
       }