Consolidate find_free_port implementations to ray._common.network_utils with type hints

Signed-off-by: Yicheng-Lu-llll <[email protected]>

Author: Yicheng-Lu-llll

python/ray/_common/network_utils.py

@@ -1,4 +1,5 @@
 import socket
+from contextlib import closing
 from functools import lru_cache
 from typing import Optional, Tuple, Union
 
@@ -93,3 +94,19 @@ def is_localhost(host: str) -> bool:
         True if the host is a localhost address, False otherwise.
     """
     return host in ("localhost", "127.0.0.1", "::1")
+
+
+def find_free_port(family: socket.AddressFamily = socket.AF_INET) -> int:
+    """Find a free port on the local machine.
+
+    Args:
+        family: The socket address family (AF_INET for IPv4, AF_INET6 for IPv6).
+            Defaults to AF_INET.
+
+    Returns:
+        An available port number.
+    """
+    with closing(socket.socket(family, socket.SOCK_STREAM)) as s:
+        s.bind(("", 0))
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        return s.getsockname()[1]

python/ray/_private/test_utils.py

@@ -1826,14 +1826,6 @@ def job_hook(**kwargs):
     sys.exit(0)
 
 
-def find_free_port() -> int:
-    sock = socket.socket()
-    sock.bind(("", 0))
-    port = sock.getsockname()[1]
-    sock.close()
-    return port
-
-
 def wandb_setup_api_key_hook():
     """
     Example external hook to set up W&B API key in

python/ray/air/_internal/util.py

@@ -2,9 +2,7 @@
 import logging
 import os
 import queue
-import socket
 import threading
-from contextlib import closing
 from typing import Optional
 
 import numpy as np
@@ -14,13 +12,6 @@
 logger = logging.getLogger(__name__)
 
 
-def find_free_port(family=socket.AF_INET):
-    with closing(socket.socket(family, socket.SOCK_STREAM)) as s:
-        s.bind(("", 0))
-        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        return s.getsockname()[1]
-
-
 def is_nan(value):
     return np.isnan(value)
 

python/ray/dashboard/modules/aggregator/tests/test_aggregator_agent.py

@@ -8,11 +8,9 @@
 from google.protobuf.timestamp_pb2 import Timestamp
 
 import ray.dashboard.consts as dashboard_consts
+from ray._common.network_utils import find_free_port
 from ray._private import ray_constants
-from ray._private.test_utils import (
-    find_free_port,
-    wait_for_condition,
-)
+from ray._private.test_utils import wait_for_condition
 from ray._private.utils import init_grpc_channel
 from ray._raylet import GcsClient
 from ray.core.generated.common_pb2 import (

python/ray/dashboard/modules/reporter/tests/test_healthz.py

@@ -4,8 +4,8 @@
 import requests
 
 import ray._private.ray_constants as ray_constants
+from ray._common.network_utils import find_free_port
 from ray._common.test_utils import wait_for_condition
-from ray._private.test_utils import find_free_port
 from ray.tests.conftest import *  # noqa: F401 F403
 
 

python/ray/tests/conftest.py

@@ -22,14 +22,13 @@
 
 import ray
 import ray._private.ray_constants as ray_constants
-from ray._common.network_utils import build_address
+from ray._common.network_utils import build_address, find_free_port
 from ray._common.test_utils import wait_for_condition
 from ray._private.conftest_utils import set_override_dashboard_url  # noqa: F401
 from ray._private.runtime_env import virtualenv_utils
 from ray._private.test_utils import (
     RayletKiller,
     external_redis_test_enabled,
-    find_free_port,
     get_and_run_resource_killer,
     get_redis_cli,
     init_error_pubsub,

python/ray/tests/test_gcs_utils.py

@@ -11,11 +11,10 @@
 import ray
 import ray._private.gcs_utils as gcs_utils
 import ray._private.ray_constants as ray_constants
-from ray._common.network_utils import parse_address
+from ray._common.network_utils import find_free_port, parse_address
 from ray._common.test_utils import async_wait_for_condition
 from ray._private.test_utils import (
     external_redis_test_enabled,
-    find_free_port,
     generate_system_config_map,
 )
 from ray._raylet import GcsClient, NodeID

python/ray/tests/test_metrics_agent.py

@@ -16,7 +16,7 @@
 from google.protobuf.timestamp_pb2 import Timestamp
 
 import ray
-from ray._common.network_utils import build_address
+from ray._common.network_utils import build_address, find_free_port
 from ray._common.test_utils import SignalActor, wait_for_condition
 from ray._private.metrics_agent import (
     Gauge as MetricsAgentGauge,
@@ -27,7 +27,6 @@
     PrometheusTimeseries,
     fetch_prometheus_metric_timeseries,
     fetch_prometheus_timeseries,
-    find_free_port,
     get_log_batch,
     raw_metric_timeseries,
 )

python/ray/tests/test_ray_event_export_task_events.py

@@ -8,9 +8,10 @@
 
 import ray
 import ray.dashboard.consts as dashboard_consts
+from ray._common.network_utils import find_free_port
 from ray._common.test_utils import wait_for_condition
 from ray._private import ray_constants
-from ray._private.test_utils import find_free_port, run_string_as_driver_nonblocking
+from ray._private.test_utils import run_string_as_driver_nonblocking
 from ray._raylet import GcsClient
 
 logger = logging.getLogger(__name__)

python/ray/tests/test_runtime_env_working_dir_3.py

@@ -9,12 +9,12 @@
 
 import ray
 import ray.experimental.internal_kv as kv
+from ray._common.network_utils import find_free_port
 from ray._common.test_utils import wait_for_condition
 from ray._private.ray_constants import RAY_RUNTIME_ENV_URI_PIN_EXPIRATION_S_ENV_VAR
 from ray._private.test_utils import (
     chdir,
     check_local_files_gced,
-    find_free_port,
 )
 from ray._private.utils import get_directory_size_bytes