[modify] get ip

Author: cuixiaojin

checkpoint_engine/device_utils.py

@@ -2,9 +2,26 @@
 import re
 import socket
 import subprocess
-
 import torch
 
+from functools import lru_cache
+from loguru import logger
+
+
+@lru_cache(maxsize=1)
+def get_ip() -> str:
+    try:
+        # try to get ip from network interface
+        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
+            s.connect(("8.8.8.8", 80))
+            return s.getsockname()[0]
+    except Exception as e:  # noqa: BLE001
+        # fallback to get ip from hostname
+        logger.warning(
+            f"fail to get ip from network interface, fallback to get ip from hostname: {e}"
+        )
+        return socket.gethostbyname(socket.gethostname())
+
 
 def npu_generate_uuid() -> str:
     str_pid = str(os.getpid())
@@ -21,8 +38,7 @@ def npu_generate_uuid() -> str:
                 search_after_pid = str_result[str_result.find(str_pid) + len(str_pid) :]
                 match_chip_id = re.search(r"Chip ID[^\d]*(\d+)", search_after_pid)
                 chip_id = int(match_chip_id.group(1))
-                server_ip = socket.gethostbyname(socket.gethostname())
-                return f"{server_ip}-{npu_id * chip_count + chip_id}"
+                return f"{get_ip()}-{npu_id * chip_count + chip_id}"
         ValueError("The current process is not running on the npu device")
     except subprocess.CalledProcessError:
         ValueError("The current process is not running on the npu device")

checkpoint_engine/ps.py

@@ -4,13 +4,11 @@
 import os
 import pickle
 import random
-import socket
 import threading
 import time
 from collections import defaultdict
 from collections.abc import Callable
 from datetime import timedelta
-from functools import lru_cache
 from typing import TYPE_CHECKING, Annotated, Any, BinaryIO, NamedTuple
 
 import httpx
@@ -23,7 +21,7 @@
 from safetensors.torch import safe_open
 from torch.multiprocessing.reductions import reduce_tensor
 
-from checkpoint_engine.device_utils import DeviceManager, npu_generate_uuid
+from checkpoint_engine.device_utils import DeviceManager, get_ip, npu_generate_uuid
 
 
 if TYPE_CHECKING:
@@ -261,21 +259,6 @@ def _get_physical_gpu_id(device_manager: DeviceManager, device_index: int | None
         raise ValueError(f"fail to get physical gpu id {device_index}") from e
 
 
-@lru_cache(maxsize=1)
-def _get_ip() -> str:
-    try:
-        # try to get ip from network interface
-        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
-            s.connect(("8.8.8.8", 80))
-            return s.getsockname()[0]
-    except Exception as e:  # noqa: BLE001
-        # fallback to get ip from hostname
-        logger.warning(
-            f"fail to get ip from network interface, fallback to get ip from hostname: {e}"
-        )
-        return socket.gethostbyname(socket.gethostname())
-
-
 def _ibv_get_device_list() -> list[str]:
     lib = ctypes.CDLL("libibverbs.so.1")
     lib.ibv_get_device_list.argtypes = [ctypes.POINTER(ctypes.c_int)]  # int *num_devices
@@ -600,7 +583,7 @@ def __init__(self, device_manager: DeviceManager):
         gpu_count = device_manager.device_module.device_count()
         local_rank = self.rank % gpu_count
         device = _get_my_rdma_device(local_rank, gpu_count, _get_rdma_devices())
-        self.ip = _get_ip()
+        self.ip = get_ip()
 
         # we will start at most 8 ps processes, so we use 8 retries to avoid port conflicts in extreme cases
         retry_count = 8
@@ -792,7 +775,7 @@ def gather_metas(self, checkpoint_name: str):
                 for x in self._memory_pool.get(checkpoint_name, [])
             ],
             p2p_store_addr=None if self._p2p_store is None else self._p2p_store.addr,
-            host_ip=_get_ip(),
+            host_ip=get_ip(),
             device_uuid=self._device_uuid,
         )