Use gloo only on x86_64 MacOS

This should fix JAX nightly MacOS ARM64 builds PiperOrigin-RevId: 659722407
openxla · Aug 5, 2024 · a13fec1 · a13fec1
1 parent c123549
commit a13fec1
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/xla/python/BUILD b/xla/python/BUILD
@@ -1402,7 +1402,8 @@ tsl_pybind_extension(
         "@tsl//tsl/platform/cloud:gcs_file_system",
     ] + select({
         # gloo transport only builds on linux
-        "//xla/tsl:macos": [
+        "//xla/tsl:macos_arm64": [],
+        "//xla/tsl:macos_x86_64": [
             "//xla/pjrt/cpu:gloo_collectives",
             "//xla/pjrt/cpu:gloo_kv_store",
         ] + if_oss(["@gloo//:transport_uv"]),

diff --git a/xla/python/xla.cc b/xla/python/xla.cc
@@ -71,7 +71,7 @@ limitations under the License.
 #include "gloo/transport/tcp/device.h"
 #include "xla/pjrt/cpu/gloo_collectives.h"
 #include "xla/pjrt/cpu/gloo_kv_store.h"
-#elif __APPLE__
+#elif __APPLE__ && defined(__x86_64__)
 #include "gloo/transport/uv/device.h"
 #include "xla/pjrt/cpu/gloo_collectives.h"
 #include "xla/pjrt/cpu/gloo_kv_store.h"
@@ -278,7 +278,7 @@ NB_MODULE(xla_extension, m_nb) {
         auto tcp_device = gloo::transport::tcp::CreateDevice(tcp_attrs);
         return std::make_shared<cpu::GlooCollectives>(std::move(gloo_kv_store),
                                                       std::move(tcp_device));
-#elif defined(__APPLE__)
+#elif defined(__APPLE__) && defined(__x86_64__)
         std::shared_ptr<KeyValueStoreInterface> kv_store = nullptr;
         if (distributed_client != nullptr) {
           kv_store = GetDistributedKeyValueStore(distributed_client,
@@ -297,7 +297,8 @@ NB_MODULE(xla_extension, m_nb) {
                                                       std::move(uv_device));
 #else   // __linux__
         throw xla::XlaRuntimeError(
-            "make_gloo_tcp_collectives only implemented for linux and macos");
+            "make_gloo_tcp_collectives only implemented for linux and x86_64 "
+            "macos");
 #endif  // __linux__
       },
       nb::arg("distributed_client"), nb::arg("hostname").none() = std::nullopt,