[feat] Support P2PHANDSHAKE metadata mode and SSD offload for MooncakeStore

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:

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:

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()

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