diff --git a/torchft/process_group.py b/torchft/process_group.py index fc88eac4..dc9aec08 100644 --- a/torchft/process_group.py +++ b/torchft/process_group.py @@ -114,8 +114,12 @@ def create_store_client(store_addr: str, timeout: timedelta) -> Store: Ex: localhost:1234/my/prefix """ - host, _, rest = store_addr.partition(":") - port, _, prefix = rest.partition("/") + # IPv6-safe: the address is host:port/prefix (optionally [host]:port/...). + # Split off the prefix, then take the port as the field after the LAST colon + # so a bare IPv6 address (e.g. a flat-IPv6 fabric) is not mangled. + addr_part, _, prefix = store_addr.partition("/") + host, _, port = addr_part.rpartition(":") + host = host.strip("[]") store = TCPStore( host_name=host, diff --git a/torchft/process_group_test.py b/torchft/process_group_test.py index a82fa299..b0fe8deb 100644 --- a/torchft/process_group_test.py +++ b/torchft/process_group_test.py @@ -37,6 +37,7 @@ from torchft.manager import Manager from torchft.process_group import ( _ErrorSwallowingWork, + create_store_client, ErrorSwallowingProcessGroupWrapper, ManagedProcessGroup, ProcessGroup, @@ -790,6 +791,45 @@ def test_managed_process_group(self) -> None: self.assertEqual(manager.allreduce.call_count, 2) +class CreateStoreClientTest(TestCase): + @parameterized.expand( + [ + # (store_addr, expected_host, expected_port, expected_prefix) + ("localhost:1234/my/prefix", "localhost", 1234, "my/prefix"), + ("127.0.0.1:29500/torchft/0/1", "127.0.0.1", 29500, "torchft/0/1"), + # Bare IPv6: the host itself contains colons, so only the field + # after the final colon is the port. + ( + "fdaa:0:1:2:3:4:5:6:29521/torchft/0/1", + "fdaa:0:1:2:3:4:5:6", + 29521, + "torchft/0/1", + ), + ("::1:29521/torchft/0/1", "::1", 29521, "torchft/0/1"), + # Bracketed IPv6: brackets are stripped from the host. + ( + "[fdaa:0:1:2:3:4:5:6]:29521/torchft/0/1", + "fdaa:0:1:2:3:4:5:6", + 29521, + "torchft/0/1", + ), + ] + ) + def test_parses_addr( + self, store_addr: str, host: str, port: int, prefix: str + ) -> None: + with ( + patch("torchft.process_group.TCPStore") as tcp_store, + patch("torchft.process_group.PrefixStore") as prefix_store, + ): + create_store_client(store_addr, timeout=timedelta(seconds=1)) + + _, kwargs = tcp_store.call_args + self.assertEqual(kwargs["host_name"], host) + self.assertEqual(kwargs["port"], port) + self.assertEqual(prefix_store.call_args.args[0], prefix) + + class MultiPgBaseTest(TestCase): """ A base test that creates N processes (via ThreadPoolExecutor) sharing