Improve Jax distributed connect errors with possible user remedies.

PiperOrigin-RevId: 666524115
openxla · Aug 22, 2024 · 511141e · 511141e
1 parent 7dcc8e9
commit 511141e
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 4 deletions.
diff --git a/xla/pjrt/distributed/client.cc b/xla/pjrt/distributed/client.cc
@@ -133,8 +133,18 @@ absl::Status DistributedRuntimeCoordinationServiceClient::Connect() {
   }
   if (s.ok()) {
     LOG(INFO) << "Connected to distributed JAX controller";
+  } else if (absl::IsDeadlineExceeded(s)) {
+    LOG(ERROR)
+        << "Failed to connect to distributed JAX controller: waited too "
+           "long for some tasks to show up. This may be due to 1) some "
+           "tasks crashed earlier before connecting, 2) some tasks were never "
+           "scheduled, or 3) scheduling delays. Consider setting a longer "
+           "initialization timeout if such delays are expected, the timeout is "
+           "currently set to: "
+        << absl::Milliseconds(config_.cluster_register_timeout_in_ms())
+        << ".\n\nOriginal runtime error: " << s;
   } else {
-    LOG(INFO) << "Failed to connect to distributed JAX controller: " << s;
+    LOG(ERROR) << "Failed to connect to distributed JAX controller: " << s;
   }
   return s;
 }

diff --git a/xla/tsl/distributed_runtime/coordination/coordination_service.cc b/xla/tsl/distributed_runtime/coordination/coordination_service.cc
@@ -547,7 +547,7 @@ void CoordinationServiceStandaloneImpl::CheckBarrierTimeout() {
                         "Total Number of tasks already at the barrier: ",
                         barrier->tasks_at_barrier.size() - pending_task_count,
                         "/", barrier->tasks_at_barrier.size(),
-                        ". Timed out task names:\n%s", pending_tasks);
+                        ". Timed out task names:\n", pending_tasks);
       }
       const absl::Status error =
           MakeCoordinationError(absl::DeadlineExceededError(error_message));

diff --git a/xla/tsl/distributed_runtime/coordination/coordination_service_agent.cc b/xla/tsl/distributed_runtime/coordination/coordination_service_agent.cc
@@ -944,8 +944,9 @@ void CoordinationServiceAgentImpl::WaitAtBarrierAsync(
   leader_client_->BarrierAsync(
       request.get(), response.get(),
       [request, response, done = std::move(done)](const absl::Status& s) {
-        done(s);
-        VLOG(3) << "WaitAtBarrierResponse: " << s;
+        auto status = TrimCoordinationErrorMessage(s);
+        done(status);
+        VLOG(3) << "WaitAtBarrierResponse: " << status;
       });
 }