diff --git a/scripts/performance_test/perftest_config.yaml b/scripts/performance_test/perftest_config.yaml index 310653b..a1026ca 100644 --- a/scripts/performance_test/perftest_config.yaml +++ b/scripts/performance_test/perftest_config.yaml @@ -38,6 +38,7 @@ backend: # WARNING: When set to `true`, TQ will attempt to terminate any existing mooncake_master process. auto_init: true # Address of the metadata coordination server. + # Set to "P2PHANDSHAKE" to use peer-to-peer handshake mode (recommended for multi-NIC environments). metadata_server: localhost:50050 # Address of the Mooncake master server. master_server_address: localhost:50051 @@ -53,6 +54,47 @@ backend: # Network device name. # Set to "" to let Mooncake auto-select available devices. device_name: "" + # Whether to hard-pin objects in memory. Hard-pinned objects are never evicted. + # Set to null (default) for auto-management: hard_pin is disabled when offload is enabled, + # and enabled otherwise. Set to true/false to explicitly override. + hard_pin: null + + # SSD Offload configuration. + # When enabled, TQ will start a standalone mooncake_client process that offloads + # data from DRAM to NVMe SSD when memory usage exceeds the high watermark. + # This is essential for environments where total CPU DRAM < GPU HBM. + # + # Architecture note: + # The offload client is a **single-node centralized SSD storage pool** that runs + # on the first node where tq.init() is called. It registers itself with mooncake_master + # and serves the entire cluster — when eviction is triggered, the master moves data + # from ANY node's memory to this offload client via TCP/RDMA. + # Ensure that `local_hostname` (above) is set to a cluster-reachable address so that + # all nodes can transfer data to the offload client. + # + # TODO: Support multi-node distributed offload (one client per node) for better + # scalability in large clusters. + offload: + # Master switch for SSD offload. + enabled: false + # Path to the SSD directory for offloaded data. + file_storage_path: /tmp/mooncake_offload + # Local buffer size in bytes for the offload client (default: 2GB). + local_buffer_size_bytes: 2147483648 + # Whether to use io_uring for async I/O (requires kernel >= 5.1 and liburing). + use_uring: false + # Heartbeat interval in seconds for the offload client. + # Note: This is paired with mooncake_master's -client_ttl (default: 30s). + # The client sends heartbeats at this interval; the master considers the client + # dead if no heartbeat is received within client_ttl. Ensure: + # heartbeat_interval_seconds << client_ttl (currently 2s << 30s = tolerates ~15 missed beats) + heartbeat_interval_seconds: 2 + # Port for the standalone mooncake_client process. + client_port: 42052 + # Eviction high watermark ratio (0.0-1.0). Eviction starts when memory usage exceeds this ratio. + eviction_high_watermark_ratio: 0.9 + # Fraction of memory to free during each eviction cycle. + eviction_ratio: 0.1 # For Yuanrong: Yuanrong: diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index 8a6df3d..331cd80 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -39,6 +39,7 @@ backend: # WARNING: When set to `true`, TQ will attempt to terminate any existing mooncake_master process. auto_init: true # Address of the metadata coordination server. + # Set to "P2PHANDSHAKE" to use peer-to-peer handshake mode (recommended for multi-NIC environments). metadata_server: localhost:50050 # Address of the Mooncake master server. master_server_address: localhost:50051 @@ -54,6 +55,47 @@ backend: # Network device name. # Set to "" to let Mooncake auto-select available devices. device_name: "" + # Whether to hard-pin objects in memory. Hard-pinned objects are never evicted. + # Set to null (default) for auto-management: hard_pin is disabled when offload is enabled, + # and enabled otherwise. Set to true/false to explicitly override. + hard_pin: null + + # SSD Offload configuration. + # When enabled, TQ will start a standalone mooncake_client process that offloads + # data from DRAM to NVMe SSD when memory usage exceeds the high watermark. + # This is essential for environments where total CPU DRAM < GPU HBM. + # + # Architecture note: + # The offload client is a **single-node centralized SSD storage pool** that runs + # on the first node where tq.init() is called. It registers itself with mooncake_master + # and serves the entire cluster — when eviction is triggered, the master moves data + # from ANY node's memory to this offload client via TCP/RDMA. + # Ensure that `local_hostname` (above) is set to a cluster-reachable address so that + # all nodes can transfer data to the offload client. + # + # TODO: Support multi-node distributed offload (one client per node) for better + # scalability in large clusters. + offload: + # Master switch for SSD offload. + enabled: false + # Path to the SSD directory for offloaded data. + file_storage_path: /tmp/mooncake_offload + # Local buffer size in bytes for the offload client (default: 2GB). + local_buffer_size_bytes: 2147483648 + # Whether to use io_uring for async I/O (requires kernel >= 5.1 and liburing). + use_uring: false + # Heartbeat interval in seconds for the offload client. + # Note: This is paired with mooncake_master's -client_ttl (default: 30s). + # The client sends heartbeats at this interval; the master considers the client + # dead if no heartbeat is received within client_ttl. Ensure: + # heartbeat_interval_seconds << client_ttl (currently 2s << 30s = tolerates ~15 missed beats) + heartbeat_interval_seconds: 2 + # Port for the standalone mooncake_client process. + client_port: 42052 + # Eviction high watermark ratio (0.0-1.0). Eviction starts when memory usage exceeds this ratio. + eviction_high_watermark_ratio: 0.9 + # Fraction of memory to free during each eviction cycle. + eviction_ratio: 0.1 # For Yuanrong: Yuanrong: diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 82ceaca..c3dd801 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -246,6 +246,17 @@ def close(): f"Consider manually killing the mooncake_master." ) + # Terminate offload client process if it was started + if isinstance(value, dict): + offload_proc = value.get("offload_client_process") + if offload_proc is not None and offload_proc.poll() is None: + offload_proc.terminate() + try: + offload_proc.wait(timeout=5) + except subprocess.TimeoutExpired: + offload_proc.kill() + logger.info(f"Terminated mooncake_client offload process (PID: {offload_proc.pid}).") + if _TQ_CLIENT: try: ret = _TQ_CLIENT.storage_manager.storage_client._store.remove_all() diff --git a/transfer_queue/storage/bootstrap/mooncake_bootstrap.py b/transfer_queue/storage/bootstrap/mooncake_bootstrap.py index 536599c..c5d7ba1 100644 --- a/transfer_queue/storage/bootstrap/mooncake_bootstrap.py +++ b/transfer_queue/storage/bootstrap/mooncake_bootstrap.py @@ -27,15 +27,27 @@ @StorageBootstrapProvider.register_provider("MooncakeStore") -def initialize_mooncake_storage(conf: DictConfig) -> subprocess.Popen | None: +def initialize_mooncake_storage(conf: DictConfig) -> subprocess.Popen | dict | None: """ Initialize Mooncake store backend. + + Supports two metadata modes: + - HTTP metadata server (default): metadata_server = "host:port" + - P2P handshake (recommended for multi-NIC environments): metadata_server = "P2PHANDSHAKE" + + When ``offload.enabled`` is set to ``true`` in the config, this function also starts + a standalone ``mooncake_client`` process that offloads data from DRAM to NVMe SSD. + Args: conf (DictConfig): Configuration dictionary for the Mooncake store backend. Returns: - subprocess.Popen | None: Process object for the Mooncake store backend process. + subprocess.Popen | dict | None: + - None if auto_init is disabled. + - subprocess.Popen: master process (when offload is disabled). + - dict: {"master_process": Popen, "offload_client_process": Popen} (when offload is enabled). Raises: ValueError: If the Mooncake store is not initialized successfully. + RuntimeError: If mooncake_master or mooncake_client fails to start. """ if not conf.backend.MooncakeStore.auto_init: return None @@ -52,21 +64,28 @@ def initialize_mooncake_storage(conf: DictConfig) -> subprocess.Popen | None: else: raise RuntimeError(f"Failed to kill existing mooncake_master processes (exit code: {result}).") - # process metadata_server - metadata_server_raw_address = conf.backend.MooncakeStore.metadata_server - if "://" not in metadata_server_raw_address: - metadata_server_raw_address = "//" + metadata_server_raw_address + # process metadata_server (normalize and validate) + metadata_server_raw_address = str(conf.backend.MooncakeStore.metadata_server).strip() + use_p2p_handshake = metadata_server_raw_address.upper() == "P2PHANDSHAKE" + + if use_p2p_handshake: + metadata_server_host = None + metadata_server_port = None + logger.info("mooncake_master: Using P2PHANDSHAKE mode (no HTTP metadata server)") + else: + if "://" not in metadata_server_raw_address: + metadata_server_raw_address = "//" + metadata_server_raw_address - metadata_server_parsed = urlparse(metadata_server_raw_address) + metadata_server_parsed = urlparse(metadata_server_raw_address) - if not metadata_server_parsed.hostname or metadata_server_parsed.port is None: - raise ValueError( - f"Invalid metadata_server '{conf.backend.MooncakeStore.metadata_server}'. " - f"Host and port are required (e.g., host:port)." - ) + if not metadata_server_parsed.hostname or metadata_server_parsed.port is None: + raise ValueError( + f"Invalid metadata_server '{conf.backend.MooncakeStore.metadata_server}'. " + f"Host and port are required (e.g., host:port)." + ) - metadata_server_host = metadata_server_parsed.hostname - metadata_server_port = str(metadata_server_parsed.port) + metadata_server_host = metadata_server_parsed.hostname + metadata_server_port = str(metadata_server_parsed.port) # process master_server master_server_raw_address = conf.backend.MooncakeStore.master_server_address @@ -83,20 +102,62 @@ def initialize_mooncake_storage(conf: DictConfig) -> subprocess.Popen | None: master_server_port = str(master_server_parsed.port) + # Read offload configuration from config + offload_conf = conf.backend.MooncakeStore.get("offload", {}) + enable_offload = offload_conf.get("enabled", False) + cmd = [ "mooncake_master", "-client_ttl=30", "-default_kv_lease_ttl=999999", "-default_kv_soft_pin_ttl=999999", - "--eviction_high_watermark_ratio=1.0", - "--eviction_ratio=0.0", - "--enable_http_metadata_server=true", "--allow_evict_soft_pinned_objects=false", - f"--http_metadata_server_host={metadata_server_host}", - f"--http_metadata_server_port={metadata_server_port}", f"--rpc_port={master_server_port}", ] + if not use_p2p_handshake: + # Enable HTTP metadata server for non-P2P mode + cmd.extend( + [ + "--enable_http_metadata_server=true", + f"--http_metadata_server_host={metadata_server_host}", + f"--http_metadata_server_port={metadata_server_port}", + ] + ) + + if enable_offload: + eviction_high_watermark = offload_conf.get("eviction_high_watermark_ratio", 0.9) + eviction_ratio = offload_conf.get("eviction_ratio", 0.1) + + # Validate offload parameters + if not (0.0 < eviction_high_watermark <= 1.0): + raise ValueError( + f"offload.eviction_high_watermark_ratio must be in (0.0, 1.0], got {eviction_high_watermark}" + ) + if not (0.0 < eviction_ratio <= 1.0): + raise ValueError(f"offload.eviction_ratio must be in (0.0, 1.0], got {eviction_ratio}") + + # Enable SSD offload: lower watermark to trigger eviction, offload on evict + cmd.extend( + [ + "--enable_offload=true", + f"--eviction_high_watermark_ratio={eviction_high_watermark}", + f"--eviction_ratio={eviction_ratio}", + "--offload_on_evict=true", + "--offload_force_evict=false", + "--offloading_queue_limit=10000", + ] + ) + logger.info("mooncake_master: SSD offload enabled (offload_on_evict=true)") + else: + # Default: no eviction, no offload + cmd.extend( + [ + "--eviction_high_watermark_ratio=1.0", + "--eviction_ratio=0.0", + ] + ) + log_file_path = "/tmp/mooncake_master.log" with open(log_file_path, "w") as log_file: process = subprocess.Popen( @@ -124,4 +185,102 @@ def initialize_mooncake_storage(conf: DictConfig) -> subprocess.Popen | None: f"mooncake_master exited with error. Check {log_file_path} for detailed logs. Output:\n{error_msg}" ) + # Start standalone mooncake_client for SSD offload if enabled. + # + # Architecture: The offload client is a single-node centralized SSD storage pool. + # It runs only on the first node that calls tq.init(), but serves the entire cluster: + # when mooncake_master triggers eviction, it moves data from ANY node's DRAM segment + # to this offload client's SSD via TCP/RDMA. The client registers itself with the + # master using `local_hostname`, which must be reachable from all cluster nodes. + # + # Heartbeat relationship: + # - mooncake_master uses -client_ttl=30 (seconds) to detect dead clients. + # - mooncake_client sends heartbeats every `heartbeat_interval_seconds` (default: 2s). + # - Ratio: 30s / 2s = tolerates up to ~15 consecutive missed heartbeats. + if enable_offload: + ssd_path = offload_conf.get("file_storage_path", "/tmp/mooncake_offload") + client_port = str(offload_conf.get("client_port", 42052)) + local_buffer_size = str(offload_conf.get("local_buffer_size_bytes", 2147483648)) + use_uring = "1" if offload_conf.get("use_uring", False) else "0" + heartbeat_interval = str(offload_conf.get("heartbeat_interval_seconds", 2)) + global_segment_size = str(conf.backend.MooncakeStore.get("global_segment_size", 4294967296)) + + # Get local hostname + local_hostname = conf.backend.MooncakeStore.get("local_hostname", "") + if not local_hostname: + try: + from transfer_queue.utils.zmq_utils import get_node_ip_address + + local_hostname = get_node_ip_address() + except Exception: + import socket + + local_hostname = socket.gethostbyname(socket.gethostname()) + + os.makedirs(ssd_path, exist_ok=True) + + client_env = os.environ.copy() + client_env["MOONCAKE_OFFLOAD_FILE_STORAGE_PATH"] = ssd_path + client_env["MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES"] = local_buffer_size + client_env["MOONCAKE_OFFLOAD_USE_URING"] = use_uring + client_env["MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS"] = heartbeat_interval + + if use_p2p_handshake: + metadata_server_url = "P2PHANDSHAKE" + else: + metadata_server_url = f"http://{metadata_server_host}:{metadata_server_port}/metadata" + master_address = f"{master_server_parsed.hostname}:{master_server_port}" + + client_cmd = [ + "mooncake_client", + f"-host={local_hostname}", + f"-global_segment_size={global_segment_size}", + f"-master_server_address={master_address}", + f"-metadata_server={metadata_server_url}", + f"-protocol={conf.backend.MooncakeStore.get('protocol', 'tcp')}", + "-enable_offload=true", + f"-port={client_port}", + ] + + client_log_path = "/tmp/mooncake_client.log" + with open(client_log_path, "w") as client_log: + client_process = subprocess.Popen( + client_cmd, + stdout=client_log, + stderr=subprocess.STDOUT, + env=client_env, + start_new_session=True, + ) + time.sleep(5) + + if client_process.poll() is None: + logger.info( + f"mooncake_client started for SSD offload, PID: {client_process.pid}. " + f"SSD path: {ssd_path}. Logs: {client_log_path}" + ) + else: + # Offload client is a required component when offload is enabled. + # Hard fail to prevent silent degradation. + error_msg = "" + try: + with open(client_log_path) as f: + error_msg = f.read() + except Exception as e: + error_msg = f"Failed to read log file: {e}" + + # Terminate the master process since offload cannot work without the client + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + raise RuntimeError( + f"mooncake_client exited unexpectedly (exit code: {client_process.returncode}). " + f"SSD offload is enabled but the offload client failed to start. " + f"Check {client_log_path} for details. Output:\n{error_msg}" + ) + + # Return structured resources for lifecycle management + if enable_offload: + return {"master_process": process, "offload_client_process": client_process} return process diff --git a/transfer_queue/storage/clients/mooncake_client.py b/transfer_queue/storage/clients/mooncake_client.py index 297b7cf..e687830 100644 --- a/transfer_queue/storage/clients/mooncake_client.py +++ b/transfer_queue/storage/clients/mooncake_client.py @@ -78,13 +78,27 @@ def __init__(self, config: dict[str, Any]): if self.master_server_address is None or not isinstance(self.master_server_address, str): raise ValueError("Missing or invalid 'master_server_address' in config") - if not self.metadata_server.startswith("http://") and not self.metadata_server.startswith("etcd://"): - self.metadata_server = f"http://{self.metadata_server}" - if not self.metadata_server.startswith("etcd://") and not self.metadata_server.endswith("/metadata"): - self.metadata_server = self.metadata_server + "/metadata" + # Support P2PHANDSHAKE mode: if metadata_server is "P2PHANDSHAKE" (case-insensitive), + # normalize to the exact string "P2PHANDSHAKE" and pass it directly without adding + # http:// prefix. This avoids IP detection issues in multi-NIC environments. + if str(self.metadata_server).strip().upper() == "P2PHANDSHAKE": + self.metadata_server = "P2PHANDSHAKE" + else: + if not self.metadata_server.startswith("http://") and not self.metadata_server.startswith("etcd://"): + self.metadata_server = f"http://{self.metadata_server}" + if not self.metadata_server.startswith("etcd://") and not self.metadata_server.endswith("/metadata"): + self.metadata_server = self.metadata_server + "/metadata" self.replica_config = ReplicateConfig() - self.replica_config.with_hard_pin = True + # When offload is enabled, hard_pin must be disabled so that objects can be evicted + # and offloaded to SSD. Hard-pinned objects are never evicted by Mooncake. + offload_conf = config.get("offload", {}) + offload_enabled = offload_conf.get("enabled", False) if isinstance(offload_conf, dict) else False + hard_pin = config.get("hard_pin", None) + if hard_pin is None: + # Auto-manage: disable hard_pin when offload is enabled + hard_pin = not offload_enabled + self.replica_config.with_hard_pin = bool(hard_pin) self._store = MooncakeDistributedStore() ret = self._store.setup(