AI Proxy Persistence (deven96#70)

* Add persistence for ahnlich ai and add new command(destroy_database) for ai_proxy

* update destroydb command to purge stores

* add ai purge commands to python ai client and update readme

Author: Iamdavidonuh

ahnlich/ai/src/engine/store.rs

@@ -15,6 +15,7 @@ use serde::Serialize;
 use std::collections::HashMap as StdHashMap;
 use std::collections::HashSet as StdHashSet;
 use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
 use std::sync::Arc;
 
 /// Contains all the stores that have been created in memory
@@ -34,6 +35,24 @@ impl AIStoreHandler {
             write_flag,
         }
     }
+    pub(crate) fn get_stores(&self) -> AIStores {
+        self.stores.clone()
+    }
+
+    #[cfg(test)]
+    pub fn write_flag(&self) -> Arc<AtomicBool> {
+        self.write_flag.clone()
+    }
+
+    fn set_write_flag(&self) {
+        let _ = self
+            .write_flag
+            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst);
+    }
+
+    pub(crate) fn use_snapshot(&mut self, stores_snapshot: AIStores) {
+        self.stores = stores_snapshot;
+    }
 
     #[tracing::instrument(skip(self))]
     pub(crate) fn create_store(
@@ -42,7 +61,8 @@ impl AIStoreHandler {
         store_type: AIStoreType,
         model: AIModel,
     ) -> Result<(), AIProxyError> {
-        self.stores
+        if self
+            .stores
             .try_insert(
                 store_name.clone(),
                 Arc::new(AIStore::create(
@@ -52,8 +72,11 @@ impl AIStoreHandler {
                 )),
                 &self.stores.guard(),
             )
-            .map_err(|_| AIProxyError::StoreAlreadyExists(store_name.clone()))?;
-
+            .is_err()
+        {
+            return Err(AIProxyError::StoreAlreadyExists(store_name.clone()));
+        }
+        self.set_write_flag();
         Ok(())
     }
 
@@ -161,11 +184,20 @@ impl AIStoreHandler {
         let removed = if !removed {
             0
         } else {
-            //self.set_write_flag();
+            self.set_write_flag();
             1
         };
         Ok(removed)
     }
+
+    /// Matches DestroyDatabase - Drops all the stores in the database
+    #[tracing::instrument(skip(self))]
+    pub(crate) fn purge_stores(&self) -> usize {
+        let store_length = self.stores.pin().len();
+        let guard = self.stores.guard();
+        self.stores.clear(&guard);
+        store_length
+    }
 }
 
 #[derive(Debug, Serialize, Deserialize)]

ahnlich/ai/src/server/handler.rs

@@ -13,7 +13,7 @@ use tokio::net::{TcpListener, TcpStream};
 use tokio::select;
 use tokio_graceful::Shutdown;
 use tracing::Instrument;
-use utils::{client::ClientHandler, protocol::AhnlichProtocol};
+use utils::{client::ClientHandler, persistence::Persistence, protocol::AhnlichProtocol};
 
 use ahnlich_client_rs::db::{DbClient, DbConnManager};
 use deadpool::managed::Pool;
@@ -54,8 +54,30 @@ impl<'a> AIProxyServer<'a> {
         }
         let write_flag = Arc::new(AtomicBool::new(false));
         let db_client = Self::build_db_client(config).await;
-        let store_handler = AIStoreHandler::new(write_flag.clone());
+        let mut store_handler = AIStoreHandler::new(write_flag.clone());
         let client_handler = Arc::new(ClientHandler::new(config.maximum_clients));
+
+        // persistence
+        if let Some(persist_location) = &config.persist_location {
+            match Persistence::load_snapshot(persist_location) {
+                Err(e) => {
+                    tracing::error!("Failed to load snapshot from persist location {e}");
+                }
+                Ok(snapshot) => {
+                    store_handler.use_snapshot(snapshot);
+                }
+            }
+            // spawn the persistence task
+            let mut persistence_task = Persistence::task(
+                write_flag,
+                config.persistence_interval,
+                persist_location,
+                store_handler.get_stores(),
+            );
+            shutdown_token
+                .spawn_task_fn(|guard| async move { persistence_task.monitor(guard).await });
+        };
+
         Ok(Self {
             listener,
             shutdown_token,
@@ -148,4 +170,9 @@ impl<'a> AIProxyServer<'a> {
 
         DbClient::new_with_pool(pool)
     }
+
+    #[cfg(test)]
+    pub fn write_flag(&self) -> Arc<AtomicBool> {
+        self.store_handler.write_flag()
+    }
 }

ahnlich/ai/src/server/task.rs

@@ -285,6 +285,10 @@ impl AhnlichProtocol for AIProxyTask {
                         );
                     }
                 }
+                AIQuery::PurgeStores => {
+                    let destoryed = self.store_handler.purge_stores();
+                    Ok(AIServerResponse::Del(destoryed))
+                }
             })
         }
         result

ahnlich/ai/src/tests/aiproxy_test.rs

@@ -13,12 +13,13 @@ use ahnlich_types::{
 
 use once_cell::sync::Lazy;
 use pretty_assertions::assert_eq;
-use std::{collections::HashSet, num::NonZeroUsize};
+use std::{collections::HashSet, num::NonZeroUsize, sync::atomic::Ordering};
 
 use crate::cli::AIProxyConfig;
 use crate::server::handler::AIProxyServer;
 use ahnlich_types::bincode::BinCodeSerAndDeser;
 use std::net::SocketAddr;
+use std::path::PathBuf;
 use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader};
 use tokio::net::TcpStream;
 use tokio::time::{timeout, Duration};
@@ -31,6 +32,16 @@ static AI_CONFIG: Lazy<AIProxyConfig> = Lazy::new(|| {
     ai_proxy
 });
 
+static PERSISTENCE_FILE: Lazy<PathBuf> =
+    Lazy::new(|| PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("ahnlich_ai_proxy.dat"));
+
+static AI_CONFIG_WITH_PERSISTENCE: Lazy<AIProxyConfig> = Lazy::new(|| {
+    AIProxyConfig::default()
+        .os_select_port()
+        .set_persistence_interval(200)
+        .set_persist_location((*PERSISTENCE_FILE).clone())
+});
+
 async fn get_server_response(
     reader: &mut BufReader<TcpStream>,
     query: AIServerQuery,
@@ -446,3 +457,129 @@ async fn test_ai_proxy_fails_db_server_unavailable() {
     let err = res.err().unwrap();
     assert!(err.contains(" kind: ConnectionRefused,"))
 }
+
+#[tokio::test]
+async fn test_ai_proxy_test_with_persistence() {
+    let server = Server::new(&CONFIG)
+        .await
+        .expect("Could not initialize server");
+    let ai_server = AIProxyServer::new(&AI_CONFIG_WITH_PERSISTENCE)
+        .await
+        .expect("Could not initialize ai proxy");
+
+    let address = ai_server.local_addr().expect("Could not get local addr");
+    let _ = tokio::spawn(async move { server.start().await });
+    let write_flag = ai_server.write_flag();
+    // start up ai proxy
+    let _ = tokio::spawn(async move { ai_server.start().await });
+    // Allow some time for the servers to start
+    tokio::time::sleep(Duration::from_millis(200)).await;
+
+    let store_name = StoreName(String::from("Main"));
+    let store_name_2 = StoreName(String::from("Main2"));
+    let first_stream = TcpStream::connect(address).await.unwrap();
+
+    let message = AIServerQuery::from_queries(&[
+        AIQuery::CreateStore {
+            r#type: AIStoreType::RawString,
+            store: store_name.clone(),
+            model: AIModel::Llama3,
+            predicates: HashSet::from_iter([]),
+            non_linear_indices: HashSet::new(),
+        },
+        AIQuery::CreateStore {
+            r#type: AIStoreType::Binary,
+            store: store_name_2.clone(),
+            model: AIModel::Llama3,
+            predicates: HashSet::from_iter([]),
+            non_linear_indices: HashSet::new(),
+        },
+        AIQuery::DropStore {
+            store: store_name,
+            error_if_not_exists: true,
+        },
+    ]);
+
+    let mut expected = AIServerResult::with_capacity(3);
+
+    expected.push(Ok(AIServerResponse::Unit));
+    expected.push(Ok(AIServerResponse::Unit));
+    expected.push(Ok(AIServerResponse::Del(1)));
+
+    let mut reader = BufReader::new(first_stream);
+    query_server_assert_result(&mut reader, message, expected).await;
+
+    // write flag should show that a write has occured
+    assert!(write_flag.load(Ordering::SeqCst));
+    // Allow some time for persistence to kick in
+    tokio::time::sleep(Duration::from_millis(200)).await;
+    // start another server with existing persistence
+
+    let persisted_server = AIProxyServer::new(&AI_CONFIG_WITH_PERSISTENCE)
+        .await
+        .unwrap();
+
+    // Allow some time for the server to start
+    tokio::time::sleep(Duration::from_millis(100)).await;
+
+    let address = persisted_server
+        .local_addr()
+        .expect("Could not get local addr");
+    let write_flag = persisted_server.write_flag();
+    let _ = tokio::spawn(async move { persisted_server.start().await });
+    let second_stream = TcpStream::connect(address).await.unwrap();
+    let mut reader = BufReader::new(second_stream);
+
+    let message = AIServerQuery::from_queries(&[AIQuery::ListStores]);
+
+    let mut expected = AIServerResult::with_capacity(1);
+
+    expected.push(Ok(AIServerResponse::StoreList(HashSet::from_iter([
+        AIStoreInfo {
+            name: store_name_2.clone(),
+            r#type: AIStoreType::Binary,
+            model: AIModel::Llama3,
+            embedding_size: AIModel::Llama3.embedding_size().into(),
+        },
+    ]))));
+
+    query_server_assert_result(&mut reader, message, expected).await;
+    assert!(!write_flag.load(Ordering::SeqCst));
+    // delete persistence
+    let _ = std::fs::remove_file(&*PERSISTENCE_FILE);
+}
+
+#[tokio::test]
+async fn test_ai_proxy_destroy_database() {
+    let address = provision_test_servers().await;
+    let second_stream = TcpStream::connect(address).await.unwrap();
+    let store_name = StoreName(String::from("Deven Kicks"));
+    let message = AIServerQuery::from_queries(&[
+        AIQuery::CreateStore {
+            r#type: AIStoreType::RawString,
+            store: store_name.clone(),
+            model: AIModel::Llama3,
+            predicates: HashSet::from_iter([]),
+            non_linear_indices: HashSet::new(),
+        },
+        AIQuery::ListStores,
+        AIQuery::PurgeStores,
+        AIQuery::ListStores,
+    ]);
+    let mut expected = AIServerResult::with_capacity(4);
+
+    expected.push(Ok(AIServerResponse::Unit));
+    expected.push(Ok(AIServerResponse::StoreList(HashSet::from_iter([
+        AIStoreInfo {
+            name: store_name,
+            r#type: AIStoreType::RawString,
+            model: AIModel::Llama3,
+            embedding_size: AIModel::Llama3.embedding_size().into(),
+        },
+    ]))));
+    expected.push(Ok(AIServerResponse::Del(1)));
+    expected.push(Ok(AIServerResponse::StoreList(HashSet::from_iter([]))));
+
+    let mut reader = BufReader::new(second_stream);
+    query_server_assert_result(&mut reader, message, expected).await
+}

ahnlich/types/src/ai/query.rs

@@ -52,6 +52,7 @@ pub enum AIQuery {
     },
     InfoServer,
     ListStores,
+    PurgeStores,
     Ping,
 }
 

docs/ai.md

@@ -21,3 +21,4 @@ Here is a rudimentary list of commands for the AI proxy to accept
 - `LISTSTORES`: List all the stores on the server. It also returns information like store length/size, embedding size, AI model, e.t.c.
 - `PING`: Test server if the server is reachable
 - `DROPSTORE`: takes in a store and deletes it. Destroys everything pertaining the store
+- `PURGESTORES`: Destroys all created stores.

sdk/ahnlich-client-py/README.md

@@ -508,7 +508,7 @@ response = client.delete_key(
 
 
 ## Bulk Requests
-The clients has the ability to send multiple requests at once, and these requests will be handled sequentially. The builder class takes care of this. The response is a list of all individual request responses.
+Clients have the ability to send multiple requests at once, and these requests will be handled sequentially. The builder class takes care of this. The response is a list of all individual request responses.
 
 
 ```py

sdk/ahnlich-client-py/ahnlich_client_py/builders/ai.py

@@ -98,6 +98,9 @@ def drop_store(self, store_name: str, error_if_not_exists: bool):
             )
         )
 
+    def purge_stores(self):
+        self.queries.append(ai_query.AIQuery__PurgeStores())
+
     def info_server(self):
         self.queries.append(ai_query.AIQuery__InfoServer())
 

sdk/ahnlich-client-py/ahnlich_client_py/clients/ai.py

@@ -106,6 +106,10 @@ def drop_store(self, store_name: str, error_if_not_exists: bool):
         )
         return self.process_request(self.builder.to_server_query())
 
+    def purge_stores(self):
+        self.builder.purge_stores()
+        return self.process_request(self.builder.to_server_query())
+
     def info_server(self):
         self.builder.info_server()
         return self.process_request(self.builder.to_server_query())

sdk/ahnlich-client-py/ahnlich_client_py/internals/ai_query.py

@@ -123,11 +123,17 @@ class AIQuery__ListStores(AIQuery):
 
 
 @dataclass(frozen=True)
-class AIQuery__Ping(AIQuery):
+class AIQuery__PurgeStores(AIQuery):
     INDEX = 10  # type: int
     pass
 
 
+@dataclass(frozen=True)
+class AIQuery__Ping(AIQuery):
+    INDEX = 11  # type: int
+    pass
+
+
 AIQuery.VARIANTS = [
     AIQuery__CreateStore,
     AIQuery__GetPred,
@@ -139,6 +145,7 @@ class AIQuery__Ping(AIQuery):
     AIQuery__DropStore,
     AIQuery__InfoServer,
     AIQuery__ListStores,
+    AIQuery__PurgeStores,
     AIQuery__Ping,
 ]