Add ip_address as argument to listener

tests/test_listener.py

@@ -0,0 +1,38 @@
+import pytest
+
+import ucp
+
+
+@pytest.mark.asyncio
+async def test_bind_ip():
+    listener = ucp.create_listener(lambda: None, ip_address="127.0.0.1")
+
+    assert isinstance(listener.port, int)
+    assert listener.port >= 1024
+
+    assert isinstance(listener.ip, str)
+    assert listener.ip == "127.0.0.1"
+
+
+@pytest.mark.asyncio
+async def test_bind_port():
+    listener = None
+    port = None
+
+    for i in range(10000, 20000):
+        try:
+            listener = ucp.create_listener(
+                lambda: None,
+                i,
+            )
+        except ucp.UCXError:
+            # Port already in use, try another
+            continue
+        else:
+            port = i
+            break
+
+    assert isinstance(listener.port, int)
+    assert listener.port == port
+
+    assert isinstance(listener.ip, str)

ucp/_libs/ucx_listener.pyx

@@ -48,6 +48,9 @@ cdef class UCXListener(UCXObject):
         Extra arguments to the call-back function
     cb_kwargs: dict, optional
         Extra keyword arguments to the call-back function
+    ip_address: str, optional
+        IP address to bind the listener to. Binds to `0.0.0.0` if not
+        specified.
 
     Returns
     -------
@@ -69,7 +72,8 @@ cdef class UCXListener(UCXObject):
         uint16_t port,
         cb_func,
         tuple cb_args=None,
-        dict cb_kwargs=None
+        dict cb_kwargs=None,
+        str ip_address=None,
     ):
         if cb_args is None:
             cb_args = ()
@@ -90,7 +94,13 @@ cdef class UCXListener(UCXObject):
         )
         params.conn_handler.cb = _listener_cb
         params.conn_handler.arg = <void*> self.cb_data
-        if c_util_set_sockaddr(&params.sockaddr, NULL, port):
+
+        cdef alloc_sockaddr_ret = (
+            c_util_set_sockaddr(&params.sockaddr, NULL, port)
+            if ip_address is None else
+            c_util_set_sockaddr(&params.sockaddr, ip_address.encode(), port)
+        )
+        if alloc_sockaddr_ret:
             raise MemoryError("Failed allocation of sockaddr")
 
         cdef ucs_status_t status = ucp_listener_create(

ucp/core.py

@@ -234,6 +234,7 @@ def create_listener(
         callback_func,
         port=0,
         endpoint_error_handling=True,
+        ip_address=None,
     ):
         """Create and start a listener to accept incoming connections
 
@@ -257,6 +258,9 @@ def create_listener(
             but prevents a process from terminating unexpectedly that may
             happen when disabled. If `False` endpoint endpoint error handling
             is disabled.
+        ip_address: str, optional
+            IP address to bind the listener to. Binds to `0.0.0.0` if not
+            specified.
 
         Returns
         -------
@@ -274,6 +278,7 @@ def create_listener(
                 port=port,
                 cb_func=_listener_handler,
                 cb_args=(callback_func, self, endpoint_error_handling),
+                ip_address=ip_address,
             )
         )
         return ret
@@ -992,11 +997,14 @@ def register_am_allocator(allocator, allocator_type):
     return _get_ctx().register_am_allocator(allocator, allocator_type)
 
 
-def create_listener(callback_func, port=None, endpoint_error_handling=True):
+def create_listener(
+    callback_func, port=None, endpoint_error_handling=True, ip_address=None
+):
     return _get_ctx().create_listener(
         callback_func,
         port,
         endpoint_error_handling=endpoint_error_handling,
+        ip_address=ip_address,
     )