[jax] Migrate jaxlib to xla::Future

PiperOrigin-RevId: 824307165

Author: ezhulenev

jaxlib/BUILD

@@ -1021,6 +1021,7 @@ cc_library(
         "@tsl//tsl/platform:ml_dtypes",
         "@tsl//tsl/profiler/lib:traceme",
         "@xla//third_party/python_runtime:headers",  # buildcleaner: keep
+        "@xla//xla:future",
         "@xla//xla:literal",
         "@xla//xla:shape_util",
         "@xla//xla:status_macros",
@@ -1034,7 +1035,6 @@ cc_library(
         "@xla//xla/pjrt:pjrt_client",
         "@xla//xla/pjrt:pjrt_compiler",
         "@xla//xla/pjrt:pjrt_executable",
-        "@xla//xla/pjrt:pjrt_future",
         "@xla//xla/pjrt:pjrt_layout",
         "@xla//xla/pjrt:status_casters",
         "@xla//xla/python:nb_absl_span",
@@ -1172,9 +1172,9 @@ cc_library(
         "@llvm-project//llvm:Support",
         "@nanobind",
         "@tsl//tsl/platform:casts",
+        "@xla//xla:future",
         "@xla//xla:util",
         "@xla//xla/pjrt:pjrt_client",
-        "@xla//xla/pjrt:pjrt_future",
         "@xla//xla/pjrt:status_casters",
         "@xla//xla/pjrt/distributed:client",
         "@xla//xla/pjrt/distributed:key_value_store_interface",
@@ -1353,8 +1353,8 @@ cc_library(
         "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:span",
         "@nanobind",
+        "@xla//xla:future",
         "@xla//xla:util",
-        "@xla//xla/pjrt:pjrt_future",
         "@xla//xla/python:version",
         "@xla//xla/python/ifrt",
         "@xla//xla/tsl/concurrency:async_value",

jaxlib/py_array.cc

@@ -69,13 +69,13 @@ limitations under the License.
 #include "jaxlib/to_ifrt_sharding.h"
 #include "jaxlib/traceback.h"
 #include "jaxlib/util.h"
+#include "xla/future.h"
 #include "xla/layout.h"
 #include "xla/layout_util.h"
 #include "xla/pjrt/exceptions.h"
 #include "xla/pjrt/lru_cache.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_compiler.h"
-#include "xla/pjrt/pjrt_future.h"
 #include "xla/pjrt/pjrt_layout.h"
 #include "xla/pjrt/status_casters.h"
 #include "xla/primitive_util.h"
@@ -475,7 +475,7 @@ PyArray_Storage::PyArray_Storage(nb::object aval, bool weak_type,
                                  nb::object sharding, bool committed,
                                  nb_class_ptr<PyClient> py_client,
                                  ifrt::ArrayRef ifrt_array,
-                                 xla::PjRtFuture<> result_status)
+                                 xla::Future<> result_status)
     : aval(std::move(aval)),
       weak_type(weak_type),
       dtype(std::move(dtype)),
@@ -512,7 +512,7 @@ void PyInit_helper(PyArray self, nb::object aval, nb::object sharding,
   Construct(reinterpret_cast<PyArrayObject*>(self.ptr()), aval,
             nb::cast<bool>(aval.attr("weak_type")), std::move(dtype),
             std::move(shape), std::move(sharding), committed, py_client,
-            std::move(ifrt_array), xla::PjRtFuture<>());
+            std::move(ifrt_array), xla::Future<>());
 }
 
 void PyArray::PyInit(PyArray self, nb::object aval, nb::object sharding,
@@ -533,7 +533,7 @@ void PyArray::PyInit(PyArray self, nb::object aval, nb::object sharding,
 PyArray PyArray::MakeFromSingleDeviceArray(nb_class_ptr<PyClient> py_client,
                                            ifrt::ArrayRef ifrt_array,
                                            bool weak_type, bool committed,
-                                           xla::PjRtFuture<> result_status) {
+                                           xla::Future<> result_status) {
   if (!llvm::isa<ifrt::SingleDeviceSharding>(ifrt_array->sharding())) {
     throw xla::XlaRuntimeError(xla::InvalidArgument(
         "Constructing single device jax.Array from non-single "
@@ -606,27 +606,27 @@ PyArray PyArrayResultHandler::Call(absl::Span<const PyArray> py_arrays) const {
   return Call(py_device_list.value()->py_client(),
               CreateIfRtArrayFromSingleDeviceShardedPyArrays(
                   dtype_, shape_, py_arrays, sharding_),
-              xla::PjRtFuture<>());
+              xla::Future<>());
 }
 
 PyArray PyArrayResultHandler::Call(nb_class_ptr<PyClient> py_client,
                                    ifrt::ArrayRef ifrt_array,
-                                   xla::PjRtFuture<> result_status) const {
+                                   xla::Future<> result_status) const {
   return PyArray(aval_, weak_type_, dtype_, shape_, sharding_,
                  std::move(py_client), std::move(ifrt_array), committed_,
                  skip_checks_, std::move(result_status));
 }
 
 PyArray PyArrayResultHandler::Call(PyArray py_array) const {
   return Call(py_array.py_client(), tsl::FormRef(py_array.ifrt_array()),
-              xla::PjRtFuture<>());
+              xla::Future<>());
 }
 
 PyArray::PyArray(nb::object aval, bool weak_type, xla::nb_dtype dtype,
                  std::vector<int64_t> shape, nb::object sharding,
                  nb_class_ptr<PyClient> py_client, ifrt::ArrayRef ifrt_array,
                  bool committed, bool skip_checks,
-                 xla::PjRtFuture<> result_status) {
+                 xla::Future<> result_status) {
   auto* self =
       PyArray_tp_new(reinterpret_cast<PyTypeObject*>(type_), nullptr, nullptr);
   m_ptr = self;

jaxlib/py_array.h

@@ -38,9 +38,9 @@ limitations under the License.
 #include "jaxlib/py_client.h"
 #include "jaxlib/py_user_context.h"
 #include "jaxlib/traceback.h"
+#include "xla/future.h"
 #include "xla/pjrt/exceptions.h"
 #include "xla/pjrt/pjrt_client.h"
-#include "xla/pjrt/pjrt_future.h"
 #include "xla/pjrt/pjrt_layout.h"
 #include "xla/python/ifrt/array.h"
 #include "xla/python/ifrt/device_list.h"
@@ -99,8 +99,7 @@ struct PyArray_Storage {
   PyArray_Storage(nanobind::object aval, bool weak_type, xla::nb_dtype dtype,
                   std::vector<int64_t> shape, nanobind::object sharding,
                   bool committed, nb_class_ptr<PyClient> py_client,
-                  xla::ifrt::ArrayRef ifrt_array,
-                  xla::PjRtFuture<> result_status);
+                  xla::ifrt::ArrayRef ifrt_array, xla::Future<> result_status);
 
   ~PyArray_Storage();
   nanobind::handle AsHandle();
@@ -125,7 +124,7 @@ struct PyArray_Storage {
   // Only set if this Array was generated by a computation that has effects.
   // This is the result status of the XLA computation that generated this
   // array.
-  xla::PjRtFuture<> result_status;
+  xla::Future<> result_status;
 
   // Doubly-linked list of all PyArrays known to the client. Protected by the
   // GIL. Since multiple PyArrays may share the same PjRtBuffer, there may be
@@ -158,12 +157,12 @@ class PyArray : public nanobind::object {
           std::vector<int64_t> shape, nanobind::object sharding,
           nb_class_ptr<PyClient> py_client, xla::ifrt::ArrayRef ifrt_array,
           bool committed, bool skip_checks,
-          xla::PjRtFuture<> result_status = xla::PjRtFuture<>());
+          xla::Future<> result_status = xla::Future<>());
 
   static PyArray MakeFromSingleDeviceArray(
       nb_class_ptr<PyClient> py_client, xla::ifrt::ArrayRef ifrt_array,
       bool weak_type, bool committed,
-      xla::PjRtFuture<> result_status = xla::PjRtFuture<>());
+      xla::Future<> result_status = xla::Future<>());
 
   static PyArray MakeFromIfrtArrayAndSharding(nb_class_ptr<PyClient> py_client,
                                               xla::ifrt::ArrayRef ifrt_array,
@@ -226,7 +225,7 @@ class PyArray : public nanobind::object {
   }
 
   // Returns xla::InvalidArgument if the buffer has been deleted.
-  // See `PjRtFuture` for the semantics of `IsReady` and `IsKnownReady`.
+  // See `Future` for the semantics of `IsReady` and `IsKnownReady`.
   absl::StatusOr<bool> IsReady() {
     xla::ifrt::Array* ifrt_array_ptr = ifrt_array();
     if (ifrt_array_ptr->IsDeleted()) {
@@ -237,7 +236,7 @@ class PyArray : public nanobind::object {
     return ifrt_array_ptr->GetReadyFuture().IsReady();
   }
 
-  const xla::PjRtFuture<>& result_status() const {
+  const xla::Future<>& result_status() const {
     return GetStorage().result_status;
   }
 
@@ -365,7 +364,7 @@ class PyArrayResultHandler {
   PyArray Call(PyArray py_array) const;
 
   PyArray Call(nb_class_ptr<PyClient> py_client, xla::ifrt::ArrayRef ifrt_array,
-               xla::PjRtFuture<> result_status = xla::PjRtFuture<>()) const;
+               xla::Future<> result_status = xla::Future<>()) const;
 
  private:
   nanobind::object aval_;

jaxlib/py_executable.cc

@@ -43,8 +43,8 @@ limitations under the License.
 #include "jaxlib/py_client.h"
 #include "jaxlib/py_device.h"
 #include "jaxlib/py_user_context.h"
+#include "xla/future.h"
 #include "xla/hlo/ir/hlo_module.h"
-#include "xla/pjrt/pjrt_future.h"
 #include "xla/pjrt/pjrt_layout.h"
 #include "xla/python/ifrt/array.h"
 #include "xla/python/ifrt/device.h"
@@ -220,7 +220,7 @@ static ifrt::ArrayRef GetIfRtArray(const ExecuteShardedArg& arg) {
 
 void PopulateExecuteShardedResults(const nb_class_ptr<PyClient>& client,
                                    std::vector<ifrt::ArrayRef> ifrt_arrays,
-                                   const xla::PjRtFuture<>& result_status,
+                                   const xla::Future<>& result_status,
                                    int num_computations,
                                    std::vector<std::vector<PyArray>>& outputs) {
   DCHECK_GT(num_computations, 0);
@@ -246,11 +246,11 @@ absl::StatusOr<PyExecuteResults> ExecuteShardedOnLocalDevicesInternal(
     const ifrt::ExecuteOptions& options, const nb_class_ptr<PyClient>& client,
     ifrt::LoadedExecutable* ifrt_loaded_executable,
     absl::Span<const ExecuteShardedArg> args,
-    std::optional<std::vector<xla::PjRtFuture<>>>& returned_futures) {
+    std::optional<std::vector<xla::Future<>>>& returned_futures) {
   std::vector<ifrt::ArrayRef> output_arrays;
   std::unique_ptr<tsl::Future<>> returned_future;
   int num_computations = ifrt_loaded_executable->addressable_devices().size();
-  xla::PjRtFuture<> result_status;
+  xla::Future<> result_status;
   {
     nb::gil_scoped_release gil_release;
     for (const auto& arg : args) {
@@ -301,7 +301,7 @@ absl::StatusOr<PyExecuteResults> ExecuteShardedOnLocalDevicesInternal(
 PyExecuteResults::PyExecuteResults(const nb_class_ptr<PyClient>& client,
                                    std::vector<ifrt::ArrayRef> ifrt_arrays,
                                    int num_computations, PyShardedToken token,
-                                   xla::PjRtFuture<> result_status)
+                                   xla::Future<> result_status)
     : client_(client),
       ifrt_arrays_(std::move(ifrt_arrays)),
       num_computations_(num_computations),
@@ -333,7 +333,7 @@ PyExecuteResults::DisassembleIntoSingleDeviceArrays() {
   std::vector<std::vector<PyArray>> outputs;
   PopulateExecuteShardedResults(
       client_, Consume(),
-      result_status_.IsValid() ? result_status_ : xla::PjRtFuture<>(),
+      result_status_.IsValid() ? result_status_ : xla::Future<>(),
       num_computations_, outputs);
   return outputs;
 }
@@ -357,7 +357,7 @@ PyExecuteResults::DisassemblePrefixIntoSingleDeviceArrays(size_t n) {
   std::vector<std::vector<PyArray>> outputs;
   PopulateExecuteShardedResults(
       client_, std::move(ifrt_arrays),
-      result_status_.IsValid() ? result_status_ : xla::PjRtFuture<>(),
+      result_status_.IsValid() ? result_status_ : xla::Future<>(),
       num_computations_, outputs);
   return outputs;
 }
@@ -382,7 +382,7 @@ std::vector<nb::object> PyExecuteResults::ConsumeWithHandlers(
     if (std::holds_alternative<const PyArrayResultHandler*>(handler)) {
       outputs.push_back(std::get<const PyArrayResultHandler*>(handler)->Call(
           client_, std::move(ifrt_arrays[buffer_id]),
-          result_status_.IsValid() ? result_status_ : xla::PjRtFuture<>()));
+          result_status_.IsValid() ? result_status_ : xla::Future<>()));
     } else {
       tsl::profiler::TraceMe traceme("ConsumeWithHandlers fallback.");
       auto disassembled_arrays =
@@ -396,7 +396,7 @@ std::vector<nb::object> PyExecuteResults::ConsumeWithHandlers(
       for (auto& disassembled_array : *disassembled_arrays) {
         nb::object array = PyArray::MakeFromSingleDeviceArray(
             client_, std::move(disassembled_array), false, true,
-            result_status_.IsValid() ? result_status_ : xla::PjRtFuture<>());
+            result_status_.IsValid() ? result_status_ : xla::Future<>());
         PyList_SET_ITEM(bufs.ptr(), i, array.release().ptr());
         ++i;
       }
@@ -417,7 +417,7 @@ absl::StatusOr<PyExecuteResults> PyLoadedExecutable::ExecuteSharded(
   }
   PyUserContextScope user_context_scope;
   PopulateCallLocation(options, xla::ifrt::UserContextScope::current().get());
-  std::optional<std::vector<xla::PjRtFuture<>>> returned_futures;
+  std::optional<std::vector<xla::Future<>>> returned_futures;
   if (with_tokens) {
     returned_futures.emplace();
   }

jaxlib/py_executable.h

@@ -36,11 +36,11 @@ limitations under the License.
 #include "jaxlib/py_array.h"
 #include "jaxlib/py_client.h"
 #include "jaxlib/traceback.h"
+#include "xla/future.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/pjrt/exceptions.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_executable.h"
-#include "xla/pjrt/pjrt_future.h"
 #include "xla/pjrt/pjrt_layout.h"
 #include "xla/python/ifrt/array.h"
 #include "xla/python/ifrt/attribute_map.h"
@@ -53,24 +53,24 @@ namespace jax {
 class PyToken {
  public:
   PyToken() = default;
-  explicit PyToken(xla::PjRtFuture<> future) : future_(std::move(future)) {}
+  explicit PyToken(xla::Future<> future) : future_(std::move(future)) {}
 
   static PyToken ReadyPyToken() {
-    return PyToken(xla::PjRtFuture<>(absl::OkStatus()));
+    return PyToken(xla::Future<>(absl::OkStatus()));
   }
 
   absl::Status Await();
 
  private:
-  xla::PjRtFuture<> future_;
+  xla::Future<> future_;
 };
 
 // PyShardedToken contains a PyToken for each device's execution.
 class PyShardedToken {
  public:
   // Default construction creates a always-ready token.
   PyShardedToken() = default;
-  explicit PyShardedToken(std::vector<xla::PjRtFuture<>> futures)
+  explicit PyShardedToken(std::vector<xla::Future<>> futures)
       : futures_(std::move(futures)) {}
 
   PyToken GetPyToken(int device_id) const {
@@ -81,15 +81,15 @@ class PyShardedToken {
   absl::Status Await();
 
  private:
-  std::vector<xla::PjRtFuture<>> futures_;
+  std::vector<xla::Future<>> futures_;
 };
 
 class PyExecuteResults {
  public:
   PyExecuteResults(const nb_class_ptr<PyClient>& client,
                    std::vector<xla::ifrt::ArrayRef> ifrt_arrays,
                    int num_computations, PyShardedToken token,
-                   xla::PjRtFuture<> result_status = xla::PjRtFuture<>());
+                   xla::Future<> result_status = xla::Future<>());
 
   std::vector<std::vector<PyArray>> DisassembleIntoSingleDeviceArrays();
 
@@ -119,7 +119,7 @@ class PyExecuteResults {
   int num_computations_;
   PyShardedToken token_;
   // Only set if the computation has tokens.
-  xla::PjRtFuture<> result_status_;
+  xla::Future<> result_status_;
 };
 
 using ExecuteShardedArg = std::variant<PyArray, std::vector<PyArray>>;

jaxlib/py_socket_transfer.cc

@@ -44,10 +44,10 @@ limitations under the License.
 #include "jaxlib/py_executable.h"
 #include "jaxlib/py_user_context.h"
 #include "jaxlib/to_ifrt_sharding.h"
+#include "xla/future.h"
 #include "xla/pjrt/distributed/client.h"
 #include "xla/pjrt/distributed/key_value_store_interface.h"
 #include "xla/pjrt/pjrt_client.h"
-#include "xla/pjrt/pjrt_future.h"
 #include "xla/pjrt/status_casters.h"
 #include "xla/python/ifrt/array.h"
 #include "xla/python/ifrt/array_spec.h"
@@ -424,12 +424,12 @@ void RegisterTransferServerTypes(nanobind::module_& m) {
                                 nb::repr(slice).c_str(), device_size)
                     .c_str());
           }
-          std::vector<xla::PjRtFuture<>> futures_per_array;
+          std::vector<xla::Future<>> futures_per_array;
           for (auto& buffer : arrs[i]->pjrt_buffers()) {
             auto raw_buffer = xla::ValueOrThrow(
                 xla::PjRtRawBuffer::CreateRawAliasOfBuffer(buffer.get()));
             tsl::RCReference<ChunkDestination> dest;
-            xla::PjRtFuture<> future;
+            xla::Future<> future;
             std::tie(dest, future) = xla::ValueOrThrow(
                 CreateSlicedRawBufferDest(raw_buffer, start, total_size));
             futures_per_array.push_back(std::move(future));

jaxlib/util.cc

@@ -23,7 +23,7 @@ limitations under the License.
 #include "absl/time/time.h"
 #include "absl/types/span.h"
 #include "nanobind/nanobind.h"
-#include "xla/pjrt/pjrt_future.h"
+#include "xla/future.h"
 #include "xla/python/ifrt/array.h"
 #include "xla/python/ifrt/client.h"
 #include "xla/python/ifrt/value.h"
@@ -37,7 +37,7 @@ namespace ifrt = xla::ifrt;
 
 namespace jax {
 
-void BlockUntilReadyWithCancel(xla::PjRtFuture<>& future) {
+void BlockUntilReadyWithCancel(xla::Future<>& future) {
   future.BlockUntilReady([](tsl::AsyncValue* value) {
     auto state = std::make_shared<absl::Notification>();
     value->AndThen([state]() { state->Notify(); });

jaxlib/util.h

@@ -18,13 +18,13 @@ limitations under the License.
 
 #include "absl/status/status.h"
 #include "absl/types/span.h"
-#include "xla/pjrt/pjrt_future.h"
+#include "xla/future.h"
 #include "xla/python/ifrt/array.h"
 
 namespace jax {
 
 // Waits until future is ready but will cancel if ctrl-c is pressed.
-void BlockUntilReadyWithCancel(xla::PjRtFuture<>& future);
+void BlockUntilReadyWithCancel(xla::Future<>& future);
 
 // Requests if given buffers are ready, awaits for results and returns OK if
 // all of the buffers are ready or the last non-ok status.