Error Handling: refactor `XlaCoordinator` to use status types. #9386

ysiraichi · 2025-06-18T19:34:01Z

This PR does 3 different things:

Adds functions and macros to be used with Status and StatusOr<T> constructs, so as to make it easier to propagate non-ok status.
Introduces XLA_SHOW_CPP_ERROR_CONTEXT environment variable for toggling error context.
Refactors XlaCoordinator to use the new status propagation constructs as an example

Mainly, inspired by tensorflow implementation, it introduces the following status propagation functions:

XLA_ASSIGN_OR_RETURN(LHS, REXPR, ...) , which either assigns the value held by REXPR to LHS if it holds a non-ok status code, or return the non-ok status.
XLA_RETURN_IF_ERROR(EXPR, ...), which early-returns if EXPR is a non-ok status, propagating the error to the caller.
XLA_ERROR_WITH_LOCATION(STATUS), which builds a new absl::Status from STATUS, by maybe appending the current file location

C++ Error Handling

Idea: by default, print only the user targeted message, with no extra C++ details. Then, by setting the XLA_SHOW_CPP_ERROR_CONTEXT environment variable, also print extra context information, such as C++ source location, and other contextual messages.

Callers may overwrite the error message by specifying the last optional argument of XLA_ASSIGN_OR_RETURN and XLA_RETURN_IF_ERROR. Rationale: they might have more context to create a more user-friendly message.
If XLA_SHOW_CPP_ERROR_CONTEXT is set, the callee overwritten error messages are shown in the following lines

Example errors:

When XLA_SHOW_CPP_ERROR_CONTEXT=0 (default):

RuntimeError: <message>

When XLA_SHOW_CPP_ERROR_CONTEXT=1:

RuntimeError: <message> (at <file>:<line>)
From Error: <previous-message> (at <file>:<line>)
...

ysiraichi · 2025-06-18T19:34:27Z

Should be merged only after #9384.

zhanyong-wan · 2025-06-23T16:42:42Z

This PR is a draft? Is it ready for review? If yes, please take it out of draft. Thanks!

ysiraichi · 2025-06-23T16:45:38Z

This is a draft. I'm still working on it.

Add new Status and StatusOr utility functions that provide better error handling and contextual information: - `XLA_ERROR_WITH_LOCATION`: Creates Status with file location info - `XLA_RETURN_IF_ERROR` and `XLA_ASSIGN_OR_RETURN`: Enhanced macros for status propagation with optional custom error messages - `MaybeWithLocation` and `MaybeWithNewMessage`: Functions for enriching error messages - `XLA_SHOW_CPP_ERROR_CONTEXT` environment variable to control error context display

Replace direct constructor with `Create()` factory method that returns `StatusOr<unique_ptr<XlaCoordinator>>`: - Uses new `XLA_ASSIGN_OR_RETURN` and `XLA_RETURN_IF_ERROR` macros - Provides better error handling with contextual messages - Maintains same functionality with improved error propagation

ysiraichi · 2025-06-25T16:49:55Z

torch_xla/csrc/runtime/xla_coordinator.cpp

+    XLA_ASSIGN_OR_RETURN(
+        dist_runtime_service_,
+        xla::GetDistributedRuntimeService(dist_service_addr, service_options),
+        "Failed to initialize distributed runtime service.");


An error status from xla::GetDistributedRuntimeService() call will get its error message overwritten by this macro. In the end, the user will see:

RuntimeError: Failed to initialize distributed runtime service.

zhanyong-wan

Nice work!

torch_xla/csrc/status.h

zhanyong-wan · 2025-06-25T21:55:07Z

torch_xla/csrc/runtime/env_vars.h

@@ -37,6 +37,7 @@ extern const char* const kEnvPjrtDynamicPlugins;
 extern const char* const kEnvDistSvcHeartbeatIntervalInSec;
 extern const char* const kEnvDistSvcMaxMissingHeartbeats;
 extern const char* const kEnvDistSvcShutdownTimeoutInMin;
+extern const char* const kEnvShowCppErrorContext;


Document how this env var affects the torch_xla behavior?

zhanyong-wan · 2025-06-25T21:57:02Z

torch_xla/csrc/runtime/xla_coordinator.h

 private:
+  // Convenience function called by `Create()` that initializes the current
+  // XlaCoordinator.
+  absl::Status Initialize(int global_rank, int world_size,


Document the parameters?

zhanyong-wan · 2025-06-25T21:57:39Z

torch_xla/csrc/runtime/xla_coordinator.h

@@ -41,7 +48,17 @@ class XlaCoordinator {
  // false otherwise.
  bool ReachedSyncPoint(int step);

+  // Creates a new instance of XlaCoordinator, and initializes it.
+  static absl::StatusOr<absl_nonnull std::unique_ptr<XlaCoordinator>> Create(


Document the parameters?

zhanyong-wan · 2025-06-25T22:04:15Z

torch_xla/csrc/runtime/xla_coordinator.h

@@ -12,11 +13,17 @@ namespace runtime {
 // XlaCoordinator serves as the point of entry for all operations which
 // required the XLA distributed runtime, such as preemption coordination.
 class XlaCoordinator {
+ private:
+  // Private struct for making the constructor private, but still callable
+  // with std::make_unique<T>() function.


Replace with

// with std::make_unique<XlaCoordinator>(PrivateUse()) within this class?

zhanyong-wan · 2025-06-25T22:27:14Z

torch_xla/csrc/status.h

@@ -0,0 +1,81 @@
+#ifndef XLA_TORCH_XLA_CSRC_STATUS_H_


Unit tests for this library?

zhanyong-wan · 2025-06-25T22:28:04Z

torch_xla/csrc/status.cpp

+
+namespace torch_xla {
+
+bool showCppErrorContext() {


Style: CamelCase function name.

zhanyong-wan · 2025-06-25T22:28:15Z

torch_xla/csrc/status.cpp

+
+namespace torch_xla {
+
+bool showCppErrorContext() {


Style: mark this static.

zhanyong-wan · 2025-06-25T22:29:20Z

torch_xla/csrc/status.cpp

+namespace torch_xla {
+
+bool showCppErrorContext() {
+  return runtime::sys_util::GetEnvBool(runtime::env::kEnvShowCppErrorContext,


This is called at every Status error construction site. It must be efficient. Use a local static const bool variable to remember the result so that we only call GetEnvBool once?

zhanyong-wan · 2025-06-25T22:33:09Z

torch_xla/csrc/status.cpp

+  //
+  // The idea is that whenever `new_message` is given, it should have more
+  // context to give a better error message to the user.
+  std::string_view message = (new_message.empty()) ? old_message : new_message;


Nit: remove () around new_message.empty().

ghpvnist · 2025-06-26T00:57:59Z

torch_xla/csrc/runtime/xla_coordinator.cpp

+absl::StatusOr<absl_nonnull std::unique_ptr<XlaCoordinator>>
+XlaCoordinator::Create(int global_rank, int world_size, std::string master_addr,
+                       std::string port) {
+  auto coordinator = std::make_unique<XlaCoordinator>(PrivateUse());


Can this be achieved with

auto coordinator = std::make_unique<XlaCoordinator>(new XlaCoordinator());

and place the default constructor privately?

private: XlaCoordinator();

It cannot. (The code won't compile as there's no XlaCoordinator(XlaCoordinator*) ctor.)

ghpvnist · 2025-06-26T00:59:17Z

torch_xla/csrc/status.h

+// This function also appends file location information to the error message, if
+// `kEnvShowCppErrorContext` is set.
+absl::Status MaybeWithNewMessage(absl::Status&& status, const char* file,
+                                 const int32_t line,


const in const int32_t is redundant. Similarly, std::string_view is read-only so const keyword isn't needed.

ysiraichi requested review from ghpvnist and zhanyong-wan June 18, 2025 19:34

ysiraichi marked this pull request as draft June 18, 2025 20:45

ysiraichi force-pushed the ysiraichi/status-qol-functions branch 3 times, most recently from 3bd1d2c to 1b2ebca Compare June 19, 2025 16:15

ysiraichi force-pushed the ysiraichi/status-for-getcomputationclient branch from a351a69 to 7e0da02 Compare June 19, 2025 16:34

ysiraichi force-pushed the ysiraichi/status-qol-functions branch 5 times, most recently from eaac1c3 to 87d2d12 Compare June 21, 2025 16:08

ysiraichi force-pushed the ysiraichi/status-qol-functions branch 3 times, most recently from 510dcce to a1b26a4 Compare June 23, 2025 18:31

zhanyong-wan removed request for ghpvnist and zhanyong-wan June 23, 2025 19:28

ysiraichi force-pushed the ysiraichi/status-qol-functions branch 2 times, most recently from a78bc1a to 4b71836 Compare June 25, 2025 14:56

ysiraichi changed the base branch from ysiraichi/status-for-getcomputationclient to master June 25, 2025 14:56

ysiraichi changed the title ~~Add Status and StatusOr<T> QOL functions.~~ Error Handling: refactor XlaCoordinator to use status types. Jun 25, 2025

ysiraichi added 2 commits June 25, 2025 12:38

ysiraichi force-pushed the ysiraichi/status-qol-functions branch from 4b71836 to 55214ac Compare June 25, 2025 15:38

ysiraichi commented Jun 25, 2025

View reviewed changes

ysiraichi requested review from zhanyong-wan and ghpvnist June 25, 2025 16:51

ysiraichi marked this pull request as ready for review June 25, 2025 16:51

Make consumed status into rvalue.

6395fa1

zhanyong-wan requested changes Jun 25, 2025

View reviewed changes

ghpvnist reviewed Jun 26, 2025

View reviewed changes

Error Handling: refactor XlaCoordinator to use status types. #9386

Are you sure you want to change the base?

Error Handling: refactor XlaCoordinator to use status types. #9386

Uh oh!

Conversation

ysiraichi commented Jun 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

C++ Error Handling

Uh oh!

ysiraichi commented Jun 18, 2025

Uh oh!

zhanyong-wan commented Jun 23, 2025

Uh oh!

ysiraichi commented Jun 23, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zhanyong-wan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Error Handling: refactor `XlaCoordinator` to use status types. #9386

Error Handling: refactor `XlaCoordinator` to use status types. #9386

ysiraichi commented Jun 18, 2025 •

edited

Loading