Add multi layer test and fix bug.

Signed-off-by: vegetableysm <[email protected]>
v6d-io · Feb 22, 2024 · f446e46 · f446e46
1 parent 1e84581
commit f446e46
Show file tree

Hide file tree

Showing 7 changed files with 46 additions and 27 deletions.
diff --git a/modules/kv-state-cache/ds/kv_state_cache.cc b/modules/kv-state-cache/ds/kv_state_cache.cc
@@ -69,11 +69,12 @@ KVStateCache::~KVStateCache() {
 }
 
 KVStateCacheBuilder::KVStateCacheBuilder(Client& client, int dimension,
-                                         int cacheCapacity) {
+                                         int cacheCapacity, int layer) {
   this->dimension = dimension;
   this->version = 0;
+  this->layer = layer;
   KVStateCacheBlockBuilder* builder =
-      new KVStateCacheBlockBuilder(client, this->dimension);
+      new KVStateCacheBlockBuilder(client, this->dimension, layer);
 
   this->rootTree = std::make_shared<RadixTree>(cacheCapacity);
 
@@ -125,7 +126,7 @@ KVStateCacheBlockBuilder* KVStateCacheBuilder::Split(
   // Split the tree if the list of kvState is full.
   VINEYARD_ASSERT(nodeDataList.size() > 0);
   KVStateCacheBlockBuilder* childKVStateCacheBlockBuilder =
-      new KVStateCacheBlockBuilder(client, this->dimension);
+      new KVStateCacheBlockBuilder(client, this->dimension, this->layer);
   for (size_t i = 0; i < nodeDataList.size(); i++) {
     OffsetData* data = (OffsetData*) nodeDataList[i]->nodeData->data;
     if (data == nullptr)

diff --git a/modules/kv-state-cache/ds/kv_state_cache.h b/modules/kv-state-cache/ds/kv_state_cache.h
@@ -83,7 +83,8 @@ class KVStateCacheBuilder : public vineyard::ObjectBuilder {
   uint64_t version;
 
  public:
-  KVStateCacheBuilder(Client& client, int dimension, int cacheCapacity);
+  KVStateCacheBuilder(Client& client, int dimension, int cacheCapacity,
+                      int layer);
 
   KVStateCacheBuilder(Client& client, std::shared_ptr<KVStateCache> cache);
 

diff --git a/modules/kv-state-cache/ds/kv_state_cache_block.h b/modules/kv-state-cache/ds/kv_state_cache_block.h
@@ -114,12 +114,12 @@ class KVStateCacheBlockBuilder : public ObjectBuilder {
   // support more than 64 kv-state cache slots
   uint64_t bitmap;
   int dimension;
-  int layer = 1;
+  int layer;
 
   int FindEmptySlot();
 
  public:
-  KVStateCacheBlockBuilder(Client& client, int dimension, int layer = 1);
+  KVStateCacheBlockBuilder(Client& client, int dimension, int layer);
 
   KVStateCacheBlockBuilder(
       Client& client, std::shared_ptr<KVStateCacheBlock> kv_state_cache_block);

diff --git a/modules/kv-state-cache/utils/kv_state_cache_utils.cc b/modules/kv-state-cache/utils/kv_state_cache_utils.cc
@@ -56,7 +56,7 @@ void signalHandler(int signum) {
   exit(signum);
 }
 
-void InitKVStateCache(int dimension, int cacheCapacity) {
+void InitKVStateCache(int dimension, int cacheCapacity, int layer) {
   if (kvStateCacheBuilder == nullptr) {
     std::string socket = std::string(getenv("VINEYARD_IPC_SOCKET"));
     LOG(INFO) << "socket:" << socket;
@@ -94,7 +94,7 @@ void InitKVStateCache(int dimension, int cacheCapacity) {
       // if failed, create a new cache object
       LOG(INFO) << "failed to get the cache object, create a new one";
       kvStateCacheBuilder = std::make_shared<KVStateCacheBuilder>(
-          client, dimension, cacheCapacity);
+          client, dimension, cacheCapacity, layer);
     }
 
     // // release the lock

diff --git a/modules/kv-state-cache/utils/kv_state_cache_utils.h b/modules/kv-state-cache/utils/kv_state_cache_utils.h
@@ -18,7 +18,8 @@ limitations under the License.
 #ifndef MODULES_KV_STATE_CACHE_UTILS_H_
 #define MODULES_KV_STATE_CACHE_UTILS_H_
 
-void InitKVStateCache(int dimension = 10, int cacheCapacity = 10);
+void InitKVStateCache(int dimension = 10, int cacheCapacity = 10,
+                      int layer = 1);
 
 void Update(const std::vector<int>& tokenList, int nextToken,
             const KV_STATE_WITH_LAYER& kvState);

diff --git a/test/kv_state_cache_test.cc b/test/kv_state_cache_test.cc
@@ -26,8 +26,9 @@ using namespace vineyard;
 
 #define DEMENSION 10
 #define CAPACITY 20
+#define LAYER 3
 
-void init() { InitKVStateCache(DEMENSION, CAPACITY); }
+void init() { InitKVStateCache(DEMENSION, CAPACITY, LAYER); }
 
 void print_current_tokens(const std::vector<int>& prefix, int next_token) {
   std::string tokens_str = "";
@@ -49,23 +50,30 @@ void print_kv_state(
       key_state_str += std::to_string(iter->second.first[i]) + " ";
       value_state_str += std::to_string(iter->second.second[i]) + " ";
     }
+    LOG(INFO) << "layer " << iter->first << ":";
     LOG(INFO) << "key_state: " << key_state_str;
     LOG(INFO) << "value_state: " << value_state_str;
+    LOG(INFO) << "---------------------";
   }
 }
 
 // we do not consider the layer.
 std::map<int, std::pair<std::vector<double>, std::vector<double>>>
 generate_kv_state(int token) {
-  std::vector<double> key_state;
-  std::vector<double> value_state;
-  for (int i = 0; i < DEMENSION; ++i) {
-    key_state.push_back(((double) token) / DEMENSION * (i + 1));
-    value_state.push_back(((double) token) / DEMENSION * (i + 1) * 2);
-  }
-
   std::map<int, std::pair<std::vector<double>, std::vector<double>>> kv_state;
-  kv_state.insert(std::make_pair(0, std::make_pair(key_state, value_state)));
+  for (int currentLayer = 0; currentLayer < LAYER; currentLayer++) {
+    std::vector<double> key_state;
+    std::vector<double> value_state;
+    for (int i = 0; i < DEMENSION; ++i) {
+      key_state.push_back(((double) token) / DEMENSION * (i + 1) +
+                          currentLayer * 10);
+      value_state.push_back(((double) token) / DEMENSION * (i + 1) * 2 +
+                            currentLayer * 10);
+    }
+
+    kv_state.insert(
+        std::make_pair(currentLayer, std::make_pair(key_state, value_state)));
+  }
   return kv_state;
 }
 

diff --git a/test/kv_state_cache_test_2.cc b/test/kv_state_cache_test_2.cc
@@ -26,8 +26,9 @@ using namespace vineyard;
 
 #define DEMENSION 10
 #define CAPACITY 20
+#define LAYER 3
 
-void init() { InitKVStateCache(DEMENSION, CAPACITY); }
+void init() { InitKVStateCache(DEMENSION, CAPACITY, LAYER); }
 
 void print_current_tokens(const std::vector<int>& prefix, int next_token) {
   std::string tokens_str = "";
@@ -49,23 +50,30 @@ void print_kv_state(
       key_state_str += std::to_string(iter->second.first[i]) + " ";
       value_state_str += std::to_string(iter->second.second[i]) + " ";
     }
+    LOG(INFO) << "layer " << iter->first << ":";
     LOG(INFO) << "key_state: " << key_state_str;
     LOG(INFO) << "value_state: " << value_state_str;
+    LOG(INFO) << "---------------------";
   }
 }
 
 // we do not consider the layer.
 std::map<int, std::pair<std::vector<double>, std::vector<double>>>
 generate_kv_state(int token) {
-  std::vector<double> key_state;
-  std::vector<double> value_state;
-  for (int i = 0; i < DEMENSION; ++i) {
-    key_state.push_back(((double) token) / DEMENSION * (i + 1));
-    value_state.push_back(((double) token) / DEMENSION * (i + 1) * 2);
-  }
-
   std::map<int, std::pair<std::vector<double>, std::vector<double>>> kv_state;
-  kv_state.insert(std::make_pair(0, std::make_pair(key_state, value_state)));
+  for (int currentLayer = 0; currentLayer < LAYER; currentLayer++) {
+    std::vector<double> key_state;
+    std::vector<double> value_state;
+    for (int i = 0; i < DEMENSION; ++i) {
+      key_state.push_back(((double) token) / DEMENSION * (i + 1) +
+                          currentLayer * 10);
+      value_state.push_back(((double) token) / DEMENSION * (i + 1) * 2 +
+                            currentLayer * 10);
+    }
+
+    kv_state.insert(
+        std::make_pair(currentLayer, std::make_pair(key_state, value_state)));
+  }
   return kv_state;
 }