Skip to content

Commit 46f6f3f

Browse files
deven96deven96
authored
Fix persistence with custom serialize deserialize (#229)
* Removing verbose logging when persistence is turned on * Fix broken persistence from new grpc types
1 parent 64b3bcf commit 46f6f3f

File tree

11 files changed

+111
-73
lines changed

11 files changed

+111
-73
lines changed

‎ahnlich/Cargo.lock

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎ahnlich/Makefile

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,9 +57,22 @@ run-db: ## Run ahnlich_db bin
5757
run-db-tracing: ## Run ahnlich_db bin with tracing enabled
5858
cargo run --bin ahnlich-db run --enable-tracing
5959

60+
run-db-persistence: ## Run ahnlich_db bin with persistence
61+
cargo run --bin ahnlich-db run --enable-tracing \
62+
--enable-persistence \
63+
--persistence-interval 60 \
64+
--persist-location ahnlich-db.dat
65+
6066
run-ai: ## Run ahnlich_ai bin with default supported models
6167
cargo run --bin ahnlich-ai run --supported-models all-minilm-l6-v2,resnet-50,bge-base-en-v1.5,clip-vit-b32-text,clip-vit-b32-image
6268

69+
run-ai-persistence: ## Run ahnlich_ai bin with persistence
70+
cargo run --bin ahnlich-ai run --enable-tracing --enable-persistence \
71+
--persist-location ahnlich-ai.dat \
72+
--persistence-interval 60 \
73+
--supported-models all-minilm-l6-v2,resnet-50,bge-base-en-v1.5,clip-vit-b32-text,clip-vit-b32-image
74+
75+
6376
run-ai-coreml: ## Run ahnlich_ai bin with coreML and tracing on mac-os
6477
RUST_LOG='ort=debug' \
6578
cargo run --features coreml \

‎ahnlich/db/src/engine/predicate.rs

Lines changed: 2 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -268,38 +268,8 @@ impl PredicateIndices {
268268
/// ids. This is essential in helping us filter down the entire dataset using a predicate before
269269
/// performing similarity algorithmic search
270270
#[derive(Debug, Serialize, Deserialize)]
271-
struct PredicateIndex(#[serde(with = "custom_metadata_map")] InnerPredicateIndex);
272-
273-
mod custom_metadata_map {
274-
use super::*;
275-
use serde::{self, Deserialize, Deserializer, Serialize, Serializer};
276-
277-
pub fn serialize<S>(map: &InnerPredicateIndex, serializer: S) -> Result<S::Ok, S::Error>
278-
where
279-
S: Serializer,
280-
{
281-
let hash_map: HashMap<MetadataValue, InnerPredicateIndexVal> = map
282-
.iter(&map.guard())
283-
.map(|(k, v)| (k.clone(), v.clone()))
284-
.collect();
285-
hash_map.serialize(serializer)
286-
}
287-
288-
pub fn deserialize<'de, D>(deserializer: D) -> Result<InnerPredicateIndex, D::Error>
289-
where
290-
D: Deserializer<'de>,
291-
{
292-
let hashmap: HashMap<MetadataValue, InnerPredicateIndexVal> =
293-
HashMap::deserialize(deserializer)?;
294-
295-
// let vec: Vec<(MetadataValue, InnerPredicateIndexVal)> = Vec::deserialize(deserializer)?;
296-
let map = ConcurrentHashMap::new();
297-
for (k, v) in hashmap {
298-
map.insert(k, v, &map.guard());
299-
}
300-
Ok(map)
301-
}
302-
}
271+
#[serde(transparent)]
272+
struct PredicateIndex(InnerPredicateIndex);
303273

304274
impl PredicateIndex {
305275
#[tracing::instrument(skip(self))]

‎ahnlich/db/src/engine/store.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ use utils::persistence::AhnlichPersistenceUtils;
3030
/// We should be only able to generate a store key id from a 1D vector except during tests
3131
3232
#[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd, Hash, Serialize, Deserialize)]
33+
#[serde(transparent)]
3334
pub(crate) struct StoreKeyId(String);
3435

3536
#[cfg(test)]

‎ahnlich/db/src/tests/server_tests.rs

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -713,6 +713,10 @@ async fn test_del_key() {
713713
#[tokio::test]
714714
async fn test_server_with_persistence() {
715715
// First server instance
716+
//
717+
// Clean up - delete persistence file
718+
let _ = std::fs::remove_file(&*PERSISTENCE_FILE);
719+
716720
let server = Server::new(&CONFIG_WITH_PERSISTENCE)
717721
.await
718722
.expect("Failed to create server");
@@ -775,7 +779,7 @@ async fn test_server_with_persistence() {
775779
}),
776780
value: Some(StoreValue {
777781
value: HashMap::from_iter([(
778-
"medal".into(),
782+
"role".into(),
779783
MetadataValue {
780784
value: Some(MetadataValueEnum::Image(vec![1, 2, 3])),
781785
},
@@ -861,7 +865,7 @@ async fn test_server_with_persistence() {
861865
stores: vec![db_response_types::StoreInfo {
862866
name: "Main".to_string(),
863867
len: 2,
864-
size_in_bytes: 1232,
868+
size_in_bytes: 1320,
865869
}],
866870
},
867871
)),
@@ -885,7 +889,7 @@ async fn test_server_with_persistence() {
885889
stores: vec![db_response_types::StoreInfo {
886890
name: "Main".to_string(),
887891
len: 1,
888-
size_in_bytes: 1184,
892+
size_in_bytes: 1272,
889893
}],
890894
},
891895
)),
@@ -912,8 +916,6 @@ async fn test_server_with_persistence() {
912916
tokio::spawn(async move { server.start().await });
913917

914918
// Verify persistence file exists and is not empty
915-
// let file_metadata = std::fs::metadata(&*PERSISTENCE_FILE).unwrap();
916-
917919
let file_metadata = std::fs::metadata(
918920
&CONFIG_WITH_PERSISTENCE
919921
.common
@@ -922,7 +924,6 @@ async fn test_server_with_persistence() {
922924
.unwrap(),
923925
)
924926
.unwrap();
925-
926927
assert!(file_metadata.len() > 0, "The persistence file is empty");
927928

928929
let address = format!("http://{}", address);
@@ -996,7 +997,7 @@ async fn test_server_with_persistence() {
996997
}),
997998
value: Some(StoreValue {
998999
value: HashMap::from_iter([(
999-
"medal".into(),
1000+
"role".into(),
10001001
MetadataValue {
10011002
value: Some(MetadataValueEnum::Image(vec![1, 2, 3])),
10021003
},

‎ahnlich/types/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ tonic.workspace = true
1515
prost.workspace = true
1616
prost-types = "0.13"
1717
serde.workspace = true
18+
ascii85 = "0.2.1"
1819

1920

2021
[build-dependencies]

‎ahnlich/types/build.rs

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -71,20 +71,16 @@ fn main() -> Result<()> {
7171
)
7272
.type_attribute(
7373
"metadata.MetadataValue",
74-
"#[derive(serde::Serialize, serde::Deserialize, PartialOrd, Ord, Hash, Eq)]",
74+
"#[derive(PartialOrd, Ord, Hash, Eq)]",
7575
)
76-
.type_attribute(
77-
"keyval.StoreName",
78-
"#[derive(serde::Serialize, serde::Deserialize, Eq, Hash, Ord, PartialOrd)]",
79-
)
80-
.type_attribute("keyval.StoreName", "#[serde(transparent)]")
76+
.type_attribute("keyval.StoreName", "#[derive(Eq, Hash, Ord, PartialOrd)]")
8177
.type_attribute(
8278
"db.server.StoreInfo",
8379
"#[derive(Hash, Eq, Ord, PartialOrd)]",
8480
)
8581
.type_attribute(
8682
"metadata.MetadataValue.value",
87-
"#[derive(serde::Serialize, serde::Deserialize, PartialOrd, Ord, Hash, Eq)]",
83+
"#[derive(PartialOrd, Ord, Hash, Eq)]",
8884
)
8985
.type_attribute(
9086
"ai.models.AIModel",

‎ahnlich/types/src/keyval.rs

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
// This file is @generated by prost-build.
2-
#[derive(serde::Serialize, serde::Deserialize, Eq, Hash, Ord, PartialOrd)]
3-
#[serde(transparent)]
4-
#[derive(Clone, PartialEq, ::prost::Message)]
2+
#[derive(Eq, Hash, Ord, PartialOrd, Clone, PartialEq, ::prost::Message)]
53
pub struct StoreName {
64
#[prost(string, tag = "1")]
75
pub value: ::prost::alloc::string::String,

‎ahnlich/types/src/metadata.rs

Lines changed: 2 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,32 +1,12 @@
11
// This file is @generated by prost-build.
2-
#[derive(
3-
serde::Serialize,
4-
serde::Deserialize,
5-
PartialOrd,
6-
Ord,
7-
Hash,
8-
Eq,
9-
Clone,
10-
PartialEq,
11-
::prost::Message,
12-
)]
2+
#[derive(PartialOrd, Ord, Hash, Eq, Clone, PartialEq, ::prost::Message)]
133
pub struct MetadataValue {
144
#[prost(oneof = "metadata_value::Value", tags = "2, 3")]
155
pub value: ::core::option::Option<metadata_value::Value>,
166
}
177
/// Nested message and enum types in `MetadataValue`.
188
pub mod metadata_value {
19-
#[derive(
20-
serde::Serialize,
21-
serde::Deserialize,
22-
PartialOrd,
23-
Ord,
24-
Hash,
25-
Eq,
26-
Clone,
27-
PartialEq,
28-
::prost::Oneof,
29-
)]
9+
#[derive(PartialOrd, Ord, Hash, Eq, Clone, PartialEq, ::prost::Oneof)]
3010
pub enum Value {
3111
#[prost(string, tag = "2")]
3212
RawString(::prost::alloc::string::String),

‎ahnlich/types/src/utils/mod.rs

Lines changed: 73 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
11
use std::hash::{Hash, Hasher};
22
use std::num::NonZeroUsize;
3+
use std::str::FromStr;
4+
5+
use serde::{Deserialize, Deserializer, Serialize, Serializer};
36

47
use crate::ai::models::AiStoreInputType;
58
use crate::client::ConnectedClient;
69
use crate::keyval::store_input::Value;
7-
use crate::keyval::StoreInput;
10+
use crate::keyval::{StoreInput, StoreName};
11+
use crate::metadata::metadata_value::Value as MetadataValueInner;
812
use crate::metadata::MetadataValue;
913
use crate::predicates::{AndCondition, Equals, In, NotEquals, NotIn, OrCondition};
1014
use crate::shared::info::StoreUpsert;
@@ -40,6 +44,74 @@ impl TryFrom<MetadataValue> for StoreInput {
4044
}
4145
}
4246

47+
impl std::fmt::Display for MetadataValueInner {
48+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
49+
match self {
50+
MetadataValueInner::Image(bytes) => write!(f, "img:{}", ascii85::encode(bytes)),
51+
MetadataValueInner::RawString(s) => write!(f, "str:{}", s),
52+
}
53+
}
54+
}
55+
56+
impl Serialize for MetadataValue {
57+
fn serialize<S>(&self, s: S) -> Result<S::Ok, S::Error>
58+
where
59+
S: Serializer,
60+
{
61+
match &self.value {
62+
Some(v) => s.serialize_str(&v.to_string()),
63+
None => Err(serde::ser::Error::custom(
64+
"Metadata value struct is empty, cannot serialize",
65+
)),
66+
}
67+
}
68+
}
69+
70+
impl Serialize for StoreName {
71+
fn serialize<S>(&self, s: S) -> Result<S::Ok, S::Error>
72+
where
73+
S: Serializer,
74+
{
75+
s.serialize_str(&self.value)
76+
}
77+
}
78+
79+
impl<'de> Deserialize<'de> for StoreName {
80+
fn deserialize<D>(d: D) -> Result<Self, D::Error>
81+
where
82+
D: Deserializer<'de>,
83+
{
84+
let value = String::deserialize(d)?;
85+
Ok(StoreName { value })
86+
}
87+
}
88+
89+
// If you need deserialization back into MyKey:
90+
impl FromStr for MetadataValueInner {
91+
type Err = String;
92+
fn from_str(s: &str) -> Result<Self, Self::Err> {
93+
if let Some(rest) = s.strip_prefix("img:") {
94+
let bytes = ascii85::decode(rest).map_err(|e| e.to_string())?;
95+
Ok(MetadataValueInner::Image(bytes))
96+
} else if let Some(rest) = s.strip_prefix("str:") {
97+
Ok(MetadataValueInner::RawString(rest.to_string()))
98+
} else {
99+
Err("unknown key format".into())
100+
}
101+
}
102+
}
103+
impl<'de> Deserialize<'de> for MetadataValue {
104+
fn deserialize<D>(d: D) -> Result<Self, D::Error>
105+
where
106+
D: Deserializer<'de>,
107+
{
108+
let s = String::deserialize(d)?;
109+
Ok(MetadataValue {
110+
value: Some(MetadataValueInner::from_str(&s).map_err(serde::de::Error::custom)?),
111+
})
112+
}
113+
}
114+
43115
impl TryFrom<&StoreInput> for AiStoreInputType {
44116
type Error = ();
45117

0 commit comments

Comments
 (0)