jax-ml/jax · error · xla::XlaRuntimeError

make_gloo_tcp_collectives only implemented for linux and mac

Error message

make_gloo_tcp_collectives only implemented for linux and macos

What it means

make_gloo_tcp_collectives compiles the gloo TCP transport only on Linux (and macOS); on the #else branch (e.g. Windows) the binding exists but throws at call time.

Source

Thrown at jaxlib/jax.cc:427

        std::shared_ptr<xla::KeyValueStoreInterface> kv_store = nullptr;
        if (distributed_client != nullptr) {
          kv_store = GetDistributedKeyValueStore(distributed_client,
                                                 /*key_prefix=*/"cpu:");
        }
        auto gloo_kv_store =
            std::make_unique<xla::cpu::GlooKeyValueStore>(kv_store);
        auto uv_attrs = gloo::transport::uv::attr();
        if (hostname) {
          uv_attrs.hostname = *hostname;
        }
        if (interface) {
          uv_attrs.iface = *interface;
        }
        auto uv_device = gloo::transport::uv::CreateDevice(uv_attrs);
        return std::make_shared<xla::cpu::GlooCollectives>(
            std::move(gloo_kv_store), std::move(uv_device));
#else   // defined(__linux__)
        throw xla::XlaRuntimeError(
            "make_gloo_tcp_collectives only implemented for linux and macos");
#endif  // defined(__linux__)
      },
      nb::arg("distributed_client"), nb::arg("hostname").none() = std::nullopt,
      nb::arg("interface").none() = std::nullopt);

#if !defined(_WIN32) && !defined(PLATFORM_GOOGLE)
  nb::class_<xla::cpu::MpiCollectives> mpi_collectives(m, "MpiCollectives",
                                                       cpu_collectives);
  mpi_collectives.def("Init", &xla::cpu::MpiCollectives::Init);
  mpi_collectives.def("Finalize", &xla::cpu::MpiCollectives::Finalize);
  m.def("make_mpi_collectives",
        []() -> std::shared_ptr<xla::cpu::MpiCollectives> {
          return std::make_shared<xla::cpu::MpiCollectives>();
        });
#else   // !_WIN32 && !PLATFORM_GOOGLE
  m.def("make_mpi_collectives",
        []() -> std::shared_ptr<xla::cpu::CpuCollectives> {

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Run under WSL2 or a Linux container
  2. Use a different collective backend available on your platform (e.g. MPI build of jaxlib)
  3. Check platform before calling and skip distributed CPU collectives on Windows

Example fix

# before
import jaxlib
kv = jaxlib.make_gloo_tcp_collectives(client, hostname)
# after
import sys, jaxlib
kv = jaxlib.make_gloo_tcp_collectives(client, hostname) if sys.platform == 'linux' else None
Defensive patterns

Strategy: type-guard

Validate before calling

import sys
assert sys.platform in ('linux', 'darwin'), 'gloo TCP collectives unsupported here'

Type guard

def gloo_supported() -> bool:
    import sys; return sys.platform.startswith('linux')

Try / catch

try:
    coll = jaxlib.make_gloo_tcp_collectives(...)
except Exception:
    coll = None  # degrade to single-process

Prevention

When it happens

Trigger: Calling jaxlib.make_gloo_tcp_collectives(...) on Windows (or any platform where the gloo uv transport is not compiled in).

Common situations: Running jax multi-host CPU collectives on Windows, where gloo's transport layer is not built.

Related errors


AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/689e6fd08af2ff28. Report an issue: GitHub.