[Perf] Refactor performance test for different kv store backends

.github/workflows/perftest.yml

@@ -0,0 +1,60 @@
+# This workflow runs the SimpleStorage performance test
+name: Performance Test
+
+on:
+  push:
+    branches:
+      - main
+      - v0.*
+  pull_request:
+    branches:
+      - main
+      - v0.*
+
+jobs:
+  perftest:
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.11"]
+
+    steps:
+    - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+        pip install -e .
+    - name: Start Ray cluster
+      run: |
+        # Get the host IP
+        HOST_IP=$(hostname -I | awk '{print $1}')
+        echo "Host IP: $HOST_IP"
+        # Start Ray with node resource
+        ray start --head --resources="{\"node:$HOST_IP\":1}"
+        ray status
+    - name: Run SimpleStorage performance test
+      run: |
+        # Get the host IP
+        HOST_IP=$(hostname -I | awk '{print $1}')
+        echo "Host IP: $HOST_IP"
+        # Run the perftest with small batch size for quick test
+        cd scripts/performance_test
+        python perftest.py \
+          --backend_config=../../transfer_queue/config.yaml \
+          --device=cpu \
+          --global_batch_size=128 \
+          --field_num=4 \
+          --seq_len=1024 \
+          --head_node_ip=$HOST_IP \
+          --output_csv=results.csv
+    - name: Stop Ray cluster
+      run: |
+        ray stop
+      if: always()

.gitignore

@@ -220,3 +220,6 @@ __marimo__/
 
 #MacOS
 **/.DS_Store
+
+# Perftest
+scripts/performance_test/results/

docs/storage_backends/openyuanrong_datasystem.md

@@ -132,11 +132,11 @@ from transfer_queue import (
     TransferQueueController,
     process_zmq_server_info,
 )
-# host, port, manager_type and client_name are the config for booting the datasystem.
+# port, manager_type and client_name are the config for booting the datasystem.
+# host will be auto-detected by checking local IP addresses.
 config_str = """
   manager_type: YuanrongStorageManager
   client_name: YuanrongStorageClient
-  host: 127.0.0.1
   port: 31501
 """
 dict_conf = OmegaConf.create(config_str, flags={"allow_objects": True})
@@ -360,26 +360,22 @@ def main():
     config_str = """
         manager_type: YuanrongStorageManager
         client_name: YuanrongStorageClient
-        host: 10.170.27.24
         port: 31501
     """
     dict_conf = OmegaConf.create(config_str, flags={"allow_objects": True})
     # It is important to pay attention to the controller's lifecycle.
     controller, dict_conf.controller_info = initialize_controller()
-
-    conf_writer = dict_conf.copy()
-    conf_writer.host = HEAD_NODE_IP
-    conf_reader = dict_conf.copy()
-    conf_reader.host = WORKER_NODE_IP
+
+    # Note: host is auto-detected on each node, no need to configure explicitly
     data = TensorDict({ "prompt": torch.ones(3, 512), "big_tensor": torch.randn(3,1024,1024)}, batch_size=[3])
     # you could assign npu or gpu devices by 'resources'
     # resources={f"node:{HEAD_NODE_IP}": 0.001} could Force the actor to run on HEAD_NODE
     writer = TransferQueueClientActor.options(
             resources={f"node:{HEAD_NODE_IP}": 0.001},
-    ).remote(conf_writer, "train")
+    ).remote(dict_conf, "train")
     reader = TransferQueueClientActor.options(
             resources={f"node:{WORKER_NODE_IP}": 0.001}
-    ).remote(conf_reader, "rollout")
+    ).remote(dict_conf, "rollout")
 
     ray.get(writer.put.remote(data=data, partition_id="train_0"))