IPv6-safe store address parsing (3 line change)

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,

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