init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

tests/e2e/nightly/multi_node/scripts/utils.py

@@ -0,0 +1,183 @@
+import logging
+import os
+import socket
+import time
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def temp_env(env_dict: dict[str, Any]):
+    old_env = {}
+    for key, value in env_dict.items():
+        old_env[key] = os.environ.get(key)
+        os.environ[key] = str(value)
+    try:
+        yield
+    finally:
+        for key, value in old_env.items():
+            if value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = value
+
+
+def setup_logger() -> None:
+    logging.basicConfig(
+        level=logging.INFO,
+        format="[%(asctime)s] [%(levelname)s] %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+
+def load_yaml_mapping(
+    yaml_path: str | None,
+    *,
+    default_name: str,
+    default_base_path: str,
+    description: str,
+) -> dict[str, Any]:
+    if not yaml_path:
+        yaml_path = os.getenv("CONFIG_YAML_PATH", default_name)
+
+    path = Path(yaml_path)
+    if not path.is_absolute() and not path.exists():
+        base_path = os.getenv("CONFIG_BASE_PATH") or default_base_path
+        path = Path(base_path) / yaml_path
+
+    logger.info("Loading %s yaml: %s", description, path)
+    with path.open(encoding="utf-8") as f:
+        data = yaml.safe_load(f)
+    if not isinstance(data, dict):
+        raise TypeError(f"{description} must be a mapping: {path}")
+    return data
+
+
+def dns_resolver(retries: int = 240, base_delay: float = 0.5):
+    def resolve(dns: str) -> str:
+        delay = base_delay
+        for attempt in range(retries):
+            try:
+                return socket.gethostbyname(dns)
+            except socket.gaierror:
+                if attempt == retries - 1:
+                    raise
+                time.sleep(delay)
+                delay = min(delay * 1.5, 5)
+        raise RuntimeError(f"Unable to resolve DNS: {dns}")
+
+    return resolve
+
+
+def get_cluster_dns_list(world_size: int) -> list[str]:
+    if world_size < 1:
+        raise ValueError(f"world_size must be >= 1, got {world_size}")
+
+    leader_dns = os.getenv("LWS_LEADER_ADDRESS")
+    if not leader_dns:
+        raise RuntimeError("environment variable LWS_LEADER_ADDRESS is not set")
+
+    parts = leader_dns.split(".")
+    if len(parts) < 3:
+        raise ValueError(f"invalid leader DNS format: {leader_dns}")
+
+    leader_name, group_name, namespace = parts[0], parts[1], parts[2]
+    worker_dns_list = [f"{leader_name}-{idx}.{group_name}.{namespace}" for idx in range(1, world_size)]
+    return [leader_dns, *worker_dns_list]
+
+
+def get_cluster_ips(world_size: int = 2) -> list[str]:
+    resolver = dns_resolver()
+    return [resolver(dns) for dns in get_cluster_dns_list(world_size)]
+
+
+def resolve_cluster_ips(
+    raw_config: dict[str, Any],
+    num_nodes: int,
+    explicit_cluster_ips: list[str] | None = None,
+    *,
+    cluster_hosts_log_message: str | None = None,
+    dns_log_message: str = "Resolving cluster IPs via DNS...",
+) -> list[str]:
+    if explicit_cluster_ips is not None:
+        if len(explicit_cluster_ips) != num_nodes:
+            raise AssertionError("cluster_ips size mismatch")
+        return explicit_cluster_ips
+
+    cluster_hosts = raw_config.get("cluster_hosts")
+    if cluster_hosts:
+        if cluster_hosts_log_message:
+            logger.info(cluster_hosts_log_message)
+        if len(cluster_hosts) != num_nodes:
+            raise AssertionError("cluster_hosts size mismatch")
+        return list(cluster_hosts)
+
+    logger.info(dns_log_message)
+    return get_cluster_ips(num_nodes)
+
+
+def get_available_port(start_port: int = 6000, end_port: int = 7000) -> int:
+    for port in range(start_port, end_port):
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            try:
+                s.bind(("", port))
+                return port
+            except OSError:
+                continue
+    raise RuntimeError("No available port found")
+
+
+def get_cur_ip(retries: int = 20, base_delay: float = 0.5) -> str:
+    delay = base_delay
+    for attempt in range(retries):
+        try:
+            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
+                s.connect(("8.8.8.8", 80))
+                return s.getsockname()[0]
+        except Exception:
+            try:
+                return socket.gethostbyname(socket.gethostname())
+            except Exception:
+                if attempt == retries - 1:
+                    raise RuntimeError("Failed to determine local IP address")
+                time.sleep(delay)
+                delay = min(delay * 1.5, 5)
+    raise RuntimeError("Failed to determine local IP address")
+
+
+def get_net_interface(ip: str | None = None) -> str:
+    import psutil
+
+    if ip is None:
+        ip = get_cur_ip()
+
+    for iface, addrs in psutil.net_if_addrs().items():
+        for addr in addrs:
+            if addr.family == socket.AF_INET and addr.address == ip:
+                return iface
+    raise RuntimeError(f"No network interface found for IP {ip}")
+
+
+def get_all_ipv4() -> list[str]:
+    ipv4s = {"127.0.0.1"}
+    hostname = socket.gethostname()
+    for info in socket.getaddrinfo(hostname, None, family=socket.AF_INET):
+        ipv4s.add(info[4][0])
+    return list(ipv4s)
+
+
+def resolve_current_node_index(cluster_ips: list[str]) -> int:
+    worker_index = os.environ.get("LWS_WORKER_INDEX")
+    if worker_index:
+        return int(worker_index)
+
+    local_ips = set(get_all_ipv4())
+    for index, ip in enumerate(cluster_ips):
+        if ip in local_ips:
+            return index
+    raise RuntimeError("Unable to determine current node index")