[train] set Neuron rendezvous env vars in the XLA backend

[violivei]

Description

_TorchAwsNeuronXLABackend cannot bring up a Neuron collective that spans instances. On
two trn1.32xlarge hosts with 64 workers at one NeuronCore each, backend startup fails
inside _setup_xla_torch_process_group with

Nrt::BuildGlobalComm failed on NeuronCores 0-1(2):  nrt_status=1, message="Non specific failure".

The Neuron runtime builds a cross-instance communicator by having every rank meet at the
endpoint named in NEURON_RT_ROOT_COMM_ID, and nothing in Ray sets it. on_start already
asks worker 0 for get_address_and_port and uses the answer for MASTER_ADDR and
MASTER_PORT, so the address is already in hand; it is just never exported under the name
the runtime reads. A single instance has nothing to rendezvous with, which is why this is
invisible until the job crosses a host boundary.

This patch takes a second free port from worker 0 and sets
NEURON_RT_ROOT_COMM_ID=<worker 0 address>:<port> in the existing set_env_vars closure,
so every worker gets the same string. A second get_address_and_port call rather than
MASTER_PORT plus an offset, because the runtime needs an endpoint of its own and
get_address_and_port returns a port it has just verified free.

Setting that alone was not sufficient in my environment, which is the part I did not
expect. With the root comm id correct on every rank, the run then died in the bootstrap:

CCOM WARN Net : No interface found in the same subnet as remote address <peer-ip><48820>
CCOM WARN NET/Socket : No usable listening interface found

CCOM selects its bootstrap socket by enumerating local interfaces and looking for one whose
subnet contains the peer address. Under the AWS VPC CNI each pod holds a /32, so no local
interface is ever in the same subnet as a peer and the search returns nothing. Naming the
interface skips the search. I want to be clear that this half is deployment specific: the
/32 comes from the CNI, and on bare EC2 instances where the host address sits in a real
subnet I would expect the match to succeed and this failure never to appear. I have not
tested that and cannot confirm it. The NEURON_RT_ROOT_COMM_ID half is not conditional on
any of that.

So the interface is an opt-in socket_ifname field on TorchXLAConfig defaulting to
None, not a hardcoded eth0. eth0 is correct for a VPC CNI pod and wrong in general,
and #42808 already removed DEFAULT_NCCL_SOCKET_IFNAME from Ray Train because users prefer
these left alone; defaulting to None keeps that promise while making the knob
discoverable next to neuron_parallel_compile. When the field is set the backend uses
os.environ.setdefault, so a value the user already exported wins, matching the
TORCH_NCCL_ASYNC_ERROR_HANDLING guard in python/ray/train/torch/config.py.

Open question for the reviewer, since I can see the argument both ways. The two interface
variables can already be set through runtime_env env_vars, so the field buys
discoverability rather than capability, and you may prefer to drop it and document the
runtime_env route instead. NEURON_RT_ROOT_COMM_ID genuinely cannot be set that way,
because its correct value is not known until the worker group exists. Happy to cut the
field if you would rather keep the config surface small.

NEURON_RT_ROOT_COMM_ID itself is set unconditionally, matching how MASTER_ADDR and
MASTER_PORT are handled two lines above. If you would rather it also yield to a
pre-existing value, say so and I will change it.

Related issues

Closes #65070

Additional information

Tests:

• Verified on 2 x trn1.32xlarge (32 NeuronCores each, 64 Ray Train workers at one core
  per worker) on Kubernetes with the AWS VPC CNI and EFA enabled, running ray 2.56.1,
  torch-neuronx 2.8.0.2.10.16998, torch-xla 2.8.1, neuronx-cc 2.21.33363.0, Python 3.11.11.
  Before the change, backend startup fails with Nrt::BuildGlobalComm failed. After it, the
  process group forms and a 64 rank job trains across both hosts: 24 steady state steps, no
  NaN, and throughput around 90 percent of twice the single instance rate measured on the
  same two hosts in the same session.
• The equivalent of this change was carried as a TorchXLAConfig subclass overriding
  on_start before being written as a patch, and that is what the run above used. The
  in-tree version differs only in taking the interface name from the config instead of
  hardcoding it.
• black and ruff clean on the changed file, at the repo's 88 column line length.
• No unit test added. There is no existing test module for this backend, and the behaviour
  that matters here is not observable without Trainium hardware: the env vars are trivially
  assertable with a mocked worker group, but that only tests that the strings are spelled
  correctly. I did not find a Trainium multi-instance runner in Ray's CI, so this path
  cannot be covered there. If you want a mock-based unit test asserting the env vars land on
  every worker, say the word and I will add one.
• Single instance behaviour is unchanged in substance. NEURON_RT_ROOT_COMM_ID is set but
  unused when there is nothing to rendezvous with, and socket_ifname defaults to None so
  no interface variable is touched unless asked for.

[gemini-code-assist]

Code Review

This pull request introduces support for configuring the network interface and allocating a dedicated rendezvous port for the Neuron runtime in TorchXLAConfig. Specifically, it adds a socket_ifname parameter and sets the NEURON_RT_ROOT_COMM_ID, NCCL_SOCKET_IFNAME, and CCOM_SOCKET_IFNAME environment variables. The feedback highlights two main improvements: first, sequentially calling get_address_and_port may result in port conflicts, so both ports should be allocated simultaneously in a single remote call; second, variables like root_comm_port and backend_config should be passed explicitly to set_env_vars rather than captured in a closure to avoid serialization issues.

[gemini-code-assist]

Calling get_address_and_port twice sequentially on the same worker can easily return the exact same port because the first port is closed before the second call binds. This would cause master_port and root_comm_port to conflict, leading to bind errors during training.

Additionally, executing two separate synchronous remote tasks (execute_single) adds unnecessary round-trip latency. We can solve both issues by executing a single remote function that binds two sockets simultaneously to guarantee distinct ports.

        # Get master address and two free ports from the first worker.
        # The Neuron runtime rendezvous cannot share MASTER_PORT with
        # torch.distributed, so we must allocate two distinct ports.
        # We allocate them simultaneously to guarantee they are distinct.
        def get_ports():
            import socket
            from ray.train._internal.utils import is_ipv6
            addr = ray.util.get_node_ip_address()
            family = socket.AF_INET6 if is_ipv6(addr) else socket.AF_INET
            with socket.socket(family, socket.SOCK_STREAM) as s1, socket.socket(family, socket.SOCK_STREAM) as s2:
                s1.bind(("", 0))
                s2.bind(("", 0))
                return addr, s1.getsockname()[1], s2.getsockname()[1]

        master_addr, master_port, root_comm_port = worker_group.execute_single(0, get_ports)

[gemini-code-assist]

Instead of implicitly capturing root_comm_port and backend_config in the set_env_vars closure, it is cleaner and more robust to pass them explicitly as arguments. This avoids potential serialization issues with larger objects in remote worker tasks.

        def set_env_vars(addr, port, root_comm_port, socket_ifname):
            os.environ["MASTER_ADDR"] = addr
            os.environ["MASTER_PORT"] = str(port)
            # To trigger the xrt server
            os.environ["TORCHELASTIC_RUN_ID"] = self.unique_run_id
            # Rendezvous endpoint the Neuron runtime uses to build a communicator
            # spanning instances. Unset, each rank builds its own local
            # communicator and a multi-instance collective never forms.
            os.environ["NEURON_RT_ROOT_COMM_ID"] = build_address(addr, root_comm_port)
            if socket_ifname:
                # CCOM chooses its bootstrap socket by looking for a local
                # interface whose subnet contains the peer address. That search
                # cannot succeed where hosts carry a /32, so allow the interface
                # to be named outright.
                os.environ.setdefault("NCCL_SOCKET_IFNAME", socket_ifname)
                os.environ.setdefault("CCOM_SOCKET_IFNAME", socket_ifname)

        # Set the env vars on all workers.
        worker_group.execute(
            set_env_vars,
            addr=master_addr,
            port=master_port,
            root_comm_port=root_comm_port,
            socket_ifname=backend_config.socket_ifname,
        )

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 47371b7. Configure here.}

[violivei]

Thanks for the review, both points are right and I have pushed a fix for each.

Two ports could be the same. This is a real bug and I missed it. find_free_port
binds, reads getsockname(), and closes inside a closing() block, so the port is back in
the pool before the second call probes. Nothing made the two results differ, and if they had
matched then NEURON_RT_ROOT_COMM_ID would have pointed at torch.distributed's port, which
is the exact thing the second call exists to avoid. Linux's ephemeral allocator rotates so I
could not get it to collide in a loop locally, but "hard to hit" is not "cannot happen" and
this is not something to leave to the allocator.

Fixed the way you suggested, with one call that holds both sockets bound while it reads both
port numbers, so the ports cannot be equal:

def _get_address_and_ports() -> Tuple[str, int, int]:
    """Returns this node's address and two distinct free ports.

    Both sockets stay bound until both port numbers have been read, so the
    ports cannot come back equal. Two ``get_address_and_port`` calls can, since
    each one releases its port before the next one probes.
    """
    addr = ray.util.get_node_ip_address()
    family = socket.AF_INET6 if is_ipv6(addr) else socket.AF_INET
    with closing(socket.socket(family, socket.SOCK_STREAM)) as master_sock:
        with closing(socket.socket(family, socket.SOCK_STREAM)) as comm_sock:
            master_sock.bind(("", 0))
            comm_sock.bind(("", 0))
            return addr, master_sock.getsockname()[1], comm_sock.getsockname()[1]

Three small deviations from the sketch, all so this matches what is already in the tree:

• is_ipv6 lives in ray._common.network_utils, not ray.train._internal.utils. The
  latter only re-exports it as a side effect of its own import, so importing it from there
  would work by accident. I import it from ray._common.network_utils, next to
  build_address which this file already needed.
• closing() and the address-family selection are copied verbatim from find_free_port and
  get_address_and_port, so the new helper is the same shape as the code it replaces. No
  SO_REUSEADDR anywhere in that path, so I did not add one.
• Module level rather than nested in on_start, matching how get_address_and_port was
  passed to execute_single before.

I looked for an existing "give me N free ports" helper to use instead and there is not one
in the tree; find_free_port is the only such utility and it is single port by
construction. Happy to move _get_address_and_ports into ray/train/_internal/utils.py
next to get_address_and_port if you would rather it were shared, but I left it local since
nothing else needs it yet. master_addr and master_port are unchanged: same
ray.util.get_node_ip_address(), same bind-to-port-0.

Implicit capture in the closure. Agreed, changed. set_env_vars now takes
root_comm_port and socket_ifname as arguments and receives them through the existing
worker_group.execute(..., **kwargs) path, so the whole TorchXLAConfig is no longer
dragged into the closure:

def set_env_vars(addr, port, root_comm_port, socket_ifname):
    ...

worker_group.execute(
    set_env_vars,
    addr=master_addr,
    port=master_port,
    root_comm_port=root_comm_port,
    socket_ifname=backend_config.socket_ifname,
)

self is still captured for unique_run_id, same as before this PR.

black and ruff are clean on the file at the repo's settings. The two B904 hits in
ruff output are pre-existing and in the repo ignore list.

One thing I want to keep flagged, since it has not changed: the exact diff here still has
not been run on Trainium. What ran on two trn1.32xlarge hosts was the equivalent logic as
a TorchXLAConfig subclass overriding on_start. The rewrite above does not change which
strings land in the environment, but I am not going to claim hardware coverage for code that
has not been on hardware.

[gemini-code-assist]

Code Review

This pull request updates the Torch XLA configuration to support multi-instance Neuron training by introducing a socket_ifname parameter and allocating two distinct free ports to prevent port collisions between the Neuron runtime rendezvous and torch.distributed. It also configures relevant environment variables like NEURON_RT_ROOT_COMM_ID, NCCL_SOCKET_IFNAME, and CCOM_SOCKET_IFNAME. The feedback suggests using "::" instead of "" as the bind address when the socket family is IPv6 (socket.AF_INET6) to ensure cross-platform compatibility and robust IPv6 support.

[gemini-code-assist]

When family is socket.AF_INET6, binding to "" can fail with OSError on certain platforms (such as macOS or Windows) because "" is not always recognized as a valid wildcard address for IPv6. To ensure cross-platform compatibility and robust IPv6 support, use "::" as the bind address when family is socket.AF_INET6.

Suggested change

	family = socket.AF_INET6 if is_ipv6(addr) else socket.AF_INET
	with closing(socket.socket(family, socket.SOCK_STREAM)) as master_sock:
	with closing(socket.socket(family, socket.SOCK_STREAM)) as comm_sock:
	master_sock.bind(("", 0))
	comm_sock.bind(("", 0))
	return addr, master_sock.getsockname()[1], comm_sock.getsockname()[1]
	family = socket.AF_INET6 if is_ipv6(addr) else socket.AF_INET
	bind_addr = "::" if family == socket.AF_INET6 else ""
	with closing(socket.socket(family, socket.SOCK_STREAM)) as master_sock:
	with closing(socket.socket(family, socket.SOCK_STREAM)) as comm_sock:
	master_sock.bind((bind_addr, 0))
	comm_sock.bind((bind_addr, 0))
	return addr, master_sock.getsockname()[1], comm_sock.getsockname()[1]

[violivei]

Thanks, but I would rather leave this one as it is, and I think the suggestion belongs
somewhere else if it is a real problem.

find_free_port in ray/_common/network_utils.py already does exactly this, for both
families:

def find_free_port(family: socket.AddressFamily = socket.AF_INET) -> int:
    with closing(socket.socket(family, socket.SOCK_STREAM)) as s:
        s.bind(("", 0))
        return s.getsockname()[1]

get_address_and_port calls it with AF_INET6 whenever is_ipv6(addr), so the IPv6 path
through bind(("", 0)) is already the one Ray takes today. My helper only exists because
that function cannot hand back two ports at once, and I copied its body so the two read the
same. Special casing the bind address here would make the new helper differ from the one it
was modelled on, while leaving the original untouched, which seems like the worse outcome.

On the portability claim itself: CPython passes "" through getaddrinfo with AI_PASSIVE,
which resolves to in6addr_any for AF_INET6, so it is not platform specific in the way the
comment suggests. I checked it locally as well and bind(("", 0)) succeeds on both families.

If you do consider "" unsafe for AF_INET6 on some platform, then find_free_port has the
same exposure and every caller of get_address_and_port inherits it, so it would be worth
fixing there in its own change rather than only in this new function. Happy to send that
separately if you want it. And if you would still prefer the explicit "::" here regardless,
say so and I will add it, it is a one line change and not worth blocking on.