(2/n torchx-allocator)(monarch/tools) implement a way to resolve hostname to ipv6 or ipv4 address (#295)

Summary:

implement `monarch.tools.network.get_ip_addr(hostname)` that resolves a hostname to an ip address by:

1. Look up a TCP (`SOCK_STREAM`) compatible ipv6 address
2. If not found, fall-back to TCP compatible ipv4 address
3. Error if no ipv6 or ipv4

This is required since hyperactor's `ChannelAddr` for the `tcp` transport (e.g. `tcp!slurm-compute-node-0:26600`) takes a socketaddr (ip + port) not a hostname. So the caller has to resolve the hostname (typically queried from the scheduler) to an ip.

Reviewed By: ahmadsharif1

Differential Revision: D76846286

Author: kiukchung

python/monarch/tools/network.py

@@ -0,0 +1,61 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+import logging
+import socket
+from typing import Optional
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+def get_ip_addr(hostname: str) -> str:
+    """Resolves and returns the ip address of the given hostname.
+
+    This function will return an ipv6 address if one that can bind
+    `SOCK_STREAM` (TCP) socket is found. Otherwise it will fall-back
+    to resolving an ipv4 `SOCK_STREAM` address.
+
+    Raises a `RuntimeError` if neither ipv6 or ipv4 ip can be resolved from hostname.
+    """
+
+    def get_sockaddr(family: socket.AddressFamily) -> Optional[str]:
+        try:
+            addrs = socket.getaddrinfo(
+                hostname, port=None, family=family, type=socket.SOCK_STREAM
+            )  # tcp
+            if addrs:
+                # socket.getaddrinfo return a list of addr 5-tuple addr infos
+                _, _, _, _, sockaddr = addrs[0]  # use the first address
+
+                # sockaddr is a tuple (ipv4) or a 4-tuple (ipv6) where the first element is the ip addr
+                ipaddr = str(sockaddr[0])
+
+                logger.info(
+                    "Resolved %s address: `%s` for host: `%s`",
+                    family.name,
+                    ipaddr,
+                    hostname,
+                )
+                return str(ipaddr)
+            else:
+                return None
+        except socket.gaierror as e:
+            logger.info(
+                "No %s address that can bind TCP sockets for host: %s. %s",
+                family.name,
+                hostname,
+                e,
+            )
+            return None
+
+    ipaddr = get_sockaddr(socket.AF_INET6) or get_sockaddr(socket.AF_INET)
+    if not ipaddr:
+        raise RuntimeError(
+            f"Unable to resolve `{hostname}` to ipv6 or ipv4 address that can bind TCP socket."
+            " Check the network configuration on the host."
+        )
+    return ipaddr

python/tests/tools/test_network.py

@@ -0,0 +1,60 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import socket
+import unittest
+from unittest import mock
+
+from monarch.tools import network
+
+
+class TestNetwork(unittest.TestCase):
+    def test_network_ipv4_fallback(self) -> None:
+        with mock.patch(
+            "socket.getaddrinfo",
+            side_effect=[
+                socket.gaierror,
+                [
+                    (
+                        socket.AF_INET,
+                        socket.SOCK_STREAM,
+                        socket.IPPROTO_TCP,
+                        "",
+                        ("123.45.67.89", 80),
+                    )
+                ],
+            ],
+        ):
+            self.assertEqual(
+                "123.45.67.89", network.get_ip_addr("foo.bar.facebook.com")
+            )
+
+    def test_network_ipv6(self) -> None:
+        with mock.patch(
+            "socket.getaddrinfo",
+            return_value=(
+                [
+                    (
+                        socket.AF_INET6,
+                        socket.SOCK_STREAM,
+                        socket.IPPROTO_TCP,
+                        "",
+                        ("1234:ab00:567c:89d:abcd:0:328:0", 0, 0, 0),
+                    )
+                ]
+            ),
+        ):
+            self.assertEqual(
+                "1234:ab00:567c:89d:abcd:0:328:0",
+                network.get_ip_addr("foo.bar.facebook.com"),
+            )
+
+    def test_network(self) -> None:
+        # since we patched `socket.getaddrinfo` above
+        # don't patch and just make sure things don't error out
+        self.assertIsNotNone(network.get_ip_addr(socket.getfqdn()))