[feat] simple storage support ssd offload

scripts/performance_test/README_PERFTEST.md

@@ -43,6 +43,8 @@ python perftest.py \
 | `--worker_node_ip` | Worker node IP address (required for Yuanrong) | None | No |
 | `--output_csv` | Path to output CSV file | None | No |
 | `--use_complex_case` | Use complex test case with nested tensors and NonTensorStack fields | False | No |
+| `--ssd_offload` | Enable SimpleStorage SSD offload; requires `--ssd_path` | False | No |
+| `--ssd_path` | Existing local SSD directory used by `--ssd_offload` | None | No |
 
 ## Backend Configuration
 
@@ -58,6 +60,10 @@ backend:
     num_data_storage_units: 16
 ```
 
+`--ssd_offload` injects the SSD configuration and writes `--ssd_path` into the
+in-memory configuration for that run. The path must already exist and be a
+directory on every storage node.
+
 ### Yuanrong Configuration
 
 ```yaml
@@ -171,6 +177,49 @@ HEAD_NODE_IP=192.168.0.1 WORKER_NODE_IP=192.168.0.2 DEVICE=npu ./run_perf_test.s
 
 After running the tests, `draw_figure.py` reads all CSV files from `results/` and generates a grouped bar chart comparing total throughput (Gbps) across backends and data sizes.
 
+## Running the SSD Offload Performance Test
+
+`run_ssd_offload_perf_test.sh` is the dedicated entry point for storage-tier
+features such as SimpleStorage SSD offload. The script keeps backend and size
+matrices as arrays for future extension; currently it runs only:
+
+- **Backend**: SimpleStorage
+- **Workloads**:
+  - Sample1MiB (batch=512, fields=8, seq=262144, total=4 GiB)
+  - Sample2MiB (batch=512, fields=8, seq=524288, total=8 GiB)
+  - Sample4MiB (batch=512, fields=8, seq=1048576, total=16 GiB)
+
+The workload names describe the size of one float32 field sample. All three
+sizes meet the SSD offload threshold, while the batch size and field count stay
+fixed for a direct comparison.
+For each backend and workload, the script runs SSD offload first and then runs
+the same workload without offload as the host-memory baseline.
+
+```bash
+HEAD_NODE_IP=192.168.0.1 \
+WORKER_NODE_IP=192.168.0.2 \
+SSD_OFFLOAD_PATH=/path/to/local/ssd \
+./run_ssd_offload_perf_test.sh
+```
+
+Configuration variables:
+
+| Variable | Description | Default |
+|----------|-------------|---------|
+| `HEAD_NODE_IP` | Head node IP address | `127.0.0.1` |
+| `WORKER_NODE_IP` | Worker node IP address | `127.0.0.1` |
+| `DEVICE` | Device type (`cpu`, `npu`, `gpu`) | `cpu` |
+| `NUM_TEST_ITERATIONS` | Number of iterations | `4` |
+| `USE_COMPLEX_CASE` | Use complex test data | `false` |
+| `SSD_OFFLOAD_PATH` | Existing local SSD directory | Required |
+
+Results are written to:
+
+- SSD offload: `results/simplestorage_ssd_{sample-size}.csv`
+- Host memory: `results/simplestorage_{sample-size}.csv`
+
+Here, `{sample-size}` is `sample1mib`, `sample2mib`, or `sample4mib`.
+
 ## Examples
 
 ### SimpleStorage backend (simple case)
@@ -185,6 +234,24 @@ python perftest.py --backend_config=perftest_config.yaml --backend=SimpleStorage
   --head_node_ip=192.168.0.1 --use_complex_case
 ```
 
+### SimpleStorage SSD offload
+
+Run the host-memory and SSD cases separately with identical workload options:
+
+```bash
+# Host-memory baseline
+python perftest.py --backend_config=perftest_config.yaml --backend=SimpleStorage \
+  --head_node_ip=192.168.0.1 --global_batch_size=64 --field_num=8 \
+  --seq_len=262144 --num_test_iterations=3 \
+  --output_csv=results/simple_storage_memory.csv
+
+# SSD offload
+python perftest.py --backend_config=perftest_config.yaml --backend=SimpleStorage \
+  --head_node_ip=192.168.0.1 --global_batch_size=64 --field_num=8 \
+  --seq_len=262144 --num_test_iterations=3 --ssd_offload --ssd_path=/path/to/local/ssd \
+  --output_csv=results/simple_storage_ssd.csv
+```
+
 ### Yuanrong backend (inter-node)
 ```bash
 python perftest.py --backend_config=perftest_config.yaml --backend=Yuanrong \
@@ -225,6 +292,7 @@ Throughput is shown in both Gb/s (gigabits per second) and GB/s (gigabytes per s
 |--------|-------------|
 | `backend` | Backend name |
 | `device` | Device type |
+| `ssd_offload` | Whether SimpleStorage SSD offload was enabled |
 | `total_data_size_gb` | Data size in GB |
 | `put_time` | PUT duration (seconds) |
 | `get_time` | GET duration (seconds) |

scripts/performance_test/perftest.py

@@ -19,6 +19,7 @@
 import logging
 import os
 import time
+from pathlib import Path
 from typing import Any
 
 import ray
@@ -293,6 +294,8 @@ def __init__(
         worker_node_ip: str | None = None,
         output_csv: str | None = None,
         use_complex_case: bool = False,
+        ssd_offload: bool = False,
+        ssd_path: str | None = None,
     ):
         """Initialize the throughput tester.
 
@@ -308,6 +311,8 @@ def __init__(
             worker_node_ip: Worker node IP address (required for Yuanrong)
             output_csv: Path to output CSV file (optional)
             use_complex_case: Whether to use complex test case (nested + nontensor fields)
+            ssd_offload: Enable SimpleStorage SSD offload
+            ssd_path: Existing local SSD directory used when SSD offload is enabled
         """
         self.backend_config_path = backend_config_path
         self.backend_override = backend
@@ -320,12 +325,18 @@ def __init__(
         self.worker_node_ip = worker_node_ip
         self.output_csv = output_csv
         self.use_complex_case = use_complex_case
+        self.ssd_offload_requested = ssd_offload
+        self.ssd_path_override = ssd_path
 
         # Prepare full config for tq.init()
         self.full_config = self._prepare_config()
 
         # Get backend from config
         self.backend = self.full_config["backend"]["storage_backend"]
+        simple_storage_config = self.full_config["backend"].get("SimpleStorage", {})
+        ssd_config = simple_storage_config.get("ssd_offload", {})
+        self.ssd_offload_enabled = self.backend == "SimpleStorage" and bool(ssd_config.get("enabled", False))
+        self.ssd_path = str(ssd_config["path"]) if self.ssd_offload_enabled else None
 
         # GDR is configured via backend.MooncakeStore.use_gdr (no separate CLI flag).
         self.use_gdr = bool(self.full_config["backend"].get("MooncakeStore", {}).get("use_gdr", False))
@@ -376,6 +387,26 @@ def _prepare_config(self) -> dict[str, Any]:
         if config.backend.storage_backend == "SimpleStorage":
             config.backend.SimpleStorage.total_storage_size = total_storage_size
 
+        if self.ssd_offload_requested:
+            if config.backend.storage_backend != "SimpleStorage":
+                raise ValueError("--ssd_offload is only supported by the SimpleStorage backend")
+            if not self.ssd_path_override:
+                raise ValueError("--ssd_offload requires --ssd_path")
+            config.backend.SimpleStorage.ssd_offload = {
+                "enabled": True,
+                "path": self.ssd_path_override,
+            }
+
+        if config.backend.storage_backend == "SimpleStorage":
+            ssd_config = config.backend.SimpleStorage.get("ssd_offload", None)
+            if ssd_config is not None and ssd_config.get("enabled", False):
+                ssd_path = ssd_config.get("path", None)
+                if not ssd_path or not Path(str(ssd_path)).is_dir():
+                    raise ValueError(
+                        f"SSD offload directory does not exist or is not a directory: {ssd_path!r}. "
+                        "Update --ssd_path to an existing local SSD directory."
+                    )
+
         return OmegaConf.to_container(config, resolve=True)
 
     def _initialize_clients(self) -> None:
@@ -490,6 +521,9 @@ def run_throughput_test(self, skip_dataset_create=False) -> dict[str, Any]:
         logger.info(f"Backend: {self.backend}")
         logger.info(f"Device: {self.device}")
         logger.info(f"GDR: {self.use_gdr}")
+        logger.info(f"SSD Offload: {self.ssd_offload_enabled}")
+        if self.ssd_offload_enabled:
+            logger.info(f"SSD Path: {self.ssd_path}")
         logger.info(f"Total Data Size: {self.total_data_size_gb:.6f} GB")
         logger.info(f"PUT Time: {put_time:.8f}s")
         logger.info(f"GET Time: {get_time:.8f}s")
@@ -499,17 +533,19 @@ def run_throughput_test(self, skip_dataset_create=False) -> dict[str, Any]:
         logger.info("=" * 60)
 
         # Return results (only Gb/s for CSV, not GB/s)
-        return {
+        result = {
             "backend": self.backend,
             "device": self.device,
             "use_gdr": self.use_gdr,
+            "ssd_offload": self.ssd_offload_enabled,
             "total_data_size_gb": self.total_data_size_gb,
             "put_time": put_time,
             "get_time": get_time,
             "put_gbit_per_sec": put_gbit_per_sec,
             "get_gbit_per_sec": get_gbit_per_sec,
             "total_gbit_per_sec": total_gbit_per_sec,
         }
+        return result
 
     def close(self) -> None:
         """Close the transfer_queue clients."""
@@ -607,6 +643,18 @@ def main() -> None:
         default=False,
         help="Use complex test case with nested tensors and nontensor fields (default: False, simple case)",
     )
+    parser.add_argument(
+        "--ssd_offload",
+        action="store_true",
+        default=False,
+        help=("Enable SimpleStorage SSD offload. Requires --ssd_path."),
+    )
+    parser.add_argument(
+        "--ssd_path",
+        type=str,
+        default=None,
+        help="Existing local SSD directory used by --ssd_offload.",
+    )
 
     args = parser.parse_args()
 
@@ -623,6 +671,8 @@ def main() -> None:
         worker_node_ip=args.worker_node_ip,
         output_csv=args.output_csv,
         use_complex_case=args.use_complex_case,
+        ssd_offload=args.ssd_offload,
+        ssd_path=args.ssd_path,
     )
 
     # Run test multiple times for consistent results using a for loop

scripts/performance_test/run_ssd_offload_perf_test.sh

@@ -0,0 +1,83 @@
+#!/bin/bash
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+RESULTS_DIR="${SCRIPT_DIR}/results"
+PERFTEST_PY="${SCRIPT_DIR}/perftest.py"
+CONFIG_YAML="${SCRIPT_DIR}/perftest_config.yaml"
+
+mkdir -p "${RESULTS_DIR}"
+
+# ========== User Configuration ==========
+HEAD_NODE_IP="${HEAD_NODE_IP:-127.0.0.1}"
+WORKER_NODE_IP="${WORKER_NODE_IP:-127.0.0.1}"
+DEVICE="${DEVICE:-cpu}"
+NUM_TEST_ITERATIONS="${NUM_TEST_ITERATIONS:-4}"
+USE_COMPLEX_CASE="${USE_COMPLEX_CASE:-false}"
+SSD_OFFLOAD_PATH="${SSD_OFFLOAD_PATH:-}"
+# ========================================
+
+if [[ -z "${SSD_OFFLOAD_PATH}" ]]; then
+    echo "SSD_OFFLOAD_PATH must point to an existing local SSD directory." >&2
+    exit 1
+fi
+
+if [[ ! -d "${SSD_OFFLOAD_PATH}" ]]; then
+    echo "SSD_OFFLOAD_PATH does not exist or is not a directory: ${SSD_OFFLOAD_PATH}" >&2
+    exit 1
+fi
+
+# Extension points for additional SSD-capable backends and workload sizes.
+# Currently, SSD offload is a SimpleStorage feature. Workload names describe
+# the per-field sample size that controls memory-versus-SSD routing.
+BACKENDS=("SimpleStorage")
+declare -a SETTINGS=(
+    # batch_size, field_num, seq_len, name
+    "512,8,262144,Sample1MiB"
+    "512,8,524288,Sample2MiB"
+    "512,8,1048576,Sample4MiB"
+)
+
+COMPLEX_ARGS=()
+if [[ "${USE_COMPLEX_CASE}" == "true" ]]; then
+    COMPLEX_ARGS=(--use_complex_case)
+fi
+
+for backend in "${BACKENDS[@]}"; do
+    echo "=========================================="
+    echo "Testing SSD offload comparison: ${backend}"
+    echo "SSD path: ${SSD_OFFLOAD_PATH}"
+    echo "=========================================="
+
+    for setting in "${SETTINGS[@]}"; do
+        IFS=',' read -r batch_size field_num seq_len name <<< "${setting}"
+        ssd_output_csv="${RESULTS_DIR}/${backend,,}_ssd_${name,,}.csv"
+        memory_output_csv="${RESULTS_DIR}/${backend,,}_${name,,}.csv"
+
+        echo "  SSD offload: ${name} (batch=${batch_size}, fields=${field_num}, seq=${seq_len})"
+        python "${PERFTEST_PY}" --backend_config="${CONFIG_YAML}" --backend="${backend}" \
+            --device="${DEVICE}" \
+            --global_batch_size="${batch_size}" --field_num="${field_num}" --seq_len="${seq_len}" \
+            --num_test_iterations="${NUM_TEST_ITERATIONS}" \
+            --head_node_ip="${HEAD_NODE_IP}" --worker_node_ip="${WORKER_NODE_IP}" \
+            --output_csv="${ssd_output_csv}" --ssd_offload --ssd_path="${SSD_OFFLOAD_PATH}" \
+            "${COMPLEX_ARGS[@]}"
+
+        sleep 10
+
+        echo "  Host memory: ${name} (batch=${batch_size}, fields=${field_num}, seq=${seq_len})"
+        python "${PERFTEST_PY}" --backend_config="${CONFIG_YAML}" --backend="${backend}" \
+            --device="${DEVICE}" \
+            --global_batch_size="${batch_size}" --field_num="${field_num}" --seq_len="${seq_len}" \
+            --num_test_iterations="${NUM_TEST_ITERATIONS}" \
+            --head_node_ip="${HEAD_NODE_IP}" --worker_node_ip="${WORKER_NODE_IP}" \
+            --output_csv="${memory_output_csv}" \
+            "${COMPLEX_ARGS[@]}"
+
+        sleep 10
+    done
+done
+
+echo ""
+echo "All SSD offload comparison tests completed!"