Merge pull request #78 from pipeless-ai/kv_store

feat(kv_store): Add KV store

pipeless/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pipeless-ai"
-version = "1.0.2"
+version = "1.1.0"
 edition = "2021"
 authors = ["Miguel A. Cabrera Minagorri"]
 description = "An open-source computer vision framework to build and deploy applications in minutes"
@@ -40,6 +40,8 @@ clap = { version = "4.4.7", features = ["derive"] }
 reqwest = { version = "0.11.22", features = ["blocking", "json"] }
 openssl = { version = "0.10", features = ["vendored"] }
 json_to_table = "0.6.0"
+sled = "0.34.7"
+lazy_static = "1.4.0"
 
 [dependencies.uuid]
 version = "1.4.1"

pipeless/src/data.rs

@@ -16,13 +16,15 @@ pub struct RgbFrame {
     input_ts: std::time::Instant, // to measure processing performance
     inference_input: ndarray::ArrayBase<ndarray::OwnedRepr<f32>, ndarray::Dim<ndarray::IxDynImpl>>,
     inference_output: ndarray::ArrayBase<ndarray::OwnedRepr<f32>, ndarray::Dim<ndarray::IxDynImpl>>,
+    pipeline_id: uuid::Uuid,
 }
 impl RgbFrame {
     pub fn new(
         original: ndarray::Array3<u8>,
         width: usize, height: usize,
         pts: gst::ClockTime, dts: gst::ClockTime, duration: gst::ClockTime,
         fps: u8, input_ts: std::time::Instant,
+        pipeline_id: uuid::Uuid,
     ) -> Self {
         let modified = original.to_owned();
         RgbFrame {
@@ -33,6 +35,7 @@ impl RgbFrame {
             input_ts,
             inference_input: ndarray::ArrayBase::zeros(ndarray::IxDyn(&[1])),
             inference_output: ndarray::ArrayBase::zeros(ndarray::IxDyn(&[1])),
+            pipeline_id,
         }
     }
 
@@ -45,6 +48,7 @@ impl RgbFrame {
         fps: u8, input_ts: u64,
         inference_input: ndarray::ArrayBase<ndarray::OwnedRepr<f32>, ndarray::Dim<ndarray::IxDynImpl>>,
         inference_output: ndarray::ArrayBase<ndarray::OwnedRepr<f32>, ndarray::Dim<ndarray::IxDynImpl>>,
+        pipeline_id: &str,
     ) -> Self {
         RgbFrame {
             uuid: uuid::Uuid::from_str(uuid).unwrap(),
@@ -55,7 +59,8 @@ impl RgbFrame {
             duration: gst::ClockTime::from_mseconds(duration),
             fps,
             input_ts: std::time::Instant::now() - std::time::Duration::from_millis(input_ts),
-            inference_input, inference_output
+            inference_input, inference_output,
+            pipeline_id: uuid::Uuid::from_str(pipeline_id).unwrap(),
         }
     }
 
@@ -118,6 +123,12 @@ impl RgbFrame {
     pub fn set_inference_output(&mut self, output_data: ndarray::ArrayBase<ndarray::OwnedRepr<f32>, ndarray::Dim<ndarray::IxDynImpl>>) {
         self.inference_output = output_data;
     }
+    pub fn get_pipeline_id(&self) -> uuid::Uuid {
+        self.pipeline_id
+    }
+    pub fn set_pipeline_id(&mut self, pipeline_id: &str) {
+        self.pipeline_id = uuid::Uuid::from_str(pipeline_id).unwrap();
+    }
 }
 
 pub enum Frame {
@@ -129,11 +140,12 @@ impl Frame {
         width: usize, height: usize,
         pts: gst::ClockTime, dts: gst::ClockTime, duration: gst::ClockTime,
         fps: u8, input_ts: std::time::Instant,
+        pipeline_id: uuid::Uuid
     ) -> Self {
         let rgb = RgbFrame::new(
             original, width, height,
             pts, dts, duration,
-            fps, input_ts
+            fps, input_ts, pipeline_id
         );
         Self::RgbFrame(rgb)
     }

pipeless/src/input/pipeline.rs

@@ -37,6 +37,7 @@ impl StreamDef {
 }
 
 fn on_new_sample(
+    pipeless_pipeline_id: uuid::Uuid,
     appsink: &gst_app::AppSink,
     pipeless_bus_sender: &tokio::sync::mpsc::UnboundedSender<pipeless::events::Event>,
 ) -> Result<gst::FlowSuccess, gst::FlowError> {
@@ -97,7 +98,8 @@ fn on_new_sample(
     let frame = pipeless::data::Frame::new_rgb(
         ndframe, width, height,
         pts, dts, duration,
-        fps as u8, frame_input_instant
+        fps as u8, frame_input_instant,
+        pipeless_pipeline_id
     );
     // The event takes ownership of the frame
     pipeless::events::publish_new_frame_change_event_sync(
@@ -300,6 +302,7 @@ fn on_bus_message(
 }
 
 fn create_gst_pipeline(
+    pipeless_pipeline_id: uuid::Uuid,
     input_uri: &str,
     pipeless_bus_sender: &tokio::sync::mpsc::UnboundedSender<pipeless::events::Event>,
 ) -> gst::Pipeline {
@@ -322,8 +325,9 @@ fn create_gst_pipeline(
                 let pipeless_bus_sender = pipeless_bus_sender.clone();
                 move |appsink: &gst_app::AppSink| {
                 on_new_sample(
+                    pipeless_pipeline_id,
                     appsink,
-                    &pipeless_bus_sender
+                    &pipeless_bus_sender,
                 )
             }
         }).build();
@@ -350,7 +354,7 @@ impl Pipeline {
         pipeless_bus_sender: &tokio::sync::mpsc::UnboundedSender<pipeless::events::Event>,
     ) -> Self {
         let input_uri = stream.get_video().get_uri();
-        let gst_pipeline = create_gst_pipeline(input_uri, pipeless_bus_sender);
+        let gst_pipeline = create_gst_pipeline(id, input_uri, pipeless_bus_sender);
         let pipeline = Pipeline {
             id,
             stream,

pipeless/src/kvs/mod.rs

@@ -0,0 +1 @@
+pub mod store;

pipeless/src/kvs/store.rs

@@ -0,0 +1,61 @@
+use log::error;
+use sled;
+use lazy_static::lazy_static;
+
+// We assume the type implementing StoreInterface can be send thread safely
+pub trait StoreInterface: Sync {
+    fn get(&self, key: &str) -> String;
+    fn set(&self, key: &str, value: &str);
+}
+
+struct LocalStore {
+    backend: sled::Db,
+}
+impl LocalStore {
+    fn new() -> Self {
+        let db_path = "/tmp/.pipeless_kv_store";
+        let db = sled::open(db_path).expect("Failed to open KV store");
+        Self { backend: db }
+    }
+}
+impl StoreInterface for LocalStore {
+    /// Insert a KV pair, logs an error if it fails, but do not stop the program.
+    /// Trying to get the key will return an empty string.
+    fn set(&self, key: &str, value: &str) {
+        if let Err(err) = self.backend.insert(key, value) {
+            error!("Error inserting key {} with value {} in local KV store. Error: {}", key, value, err);
+        }
+    }
+
+    /// Returns the value or an empty string
+    fn get(&self, key: &str) -> String {
+        match self.backend.get(key) {
+            Ok(v) => {
+                match v.as_deref() {
+                    Some(v) => std::str::from_utf8(v).unwrap().into(),
+                    None => String::from("")
+                }
+            },
+            Err(err) => {
+                error!("Error getting value for key {} from local KV store. Error: {}", key, err);
+                String::from("")
+            }
+        }
+    }
+}
+
+// TODO: setup Redis or any other distributed solution.
+// Important: Note that any type implementing StoreInterface must be thread safe
+struct DistributedStore {}
+impl DistributedStore {
+    fn new() -> Self { unimplemented!() }
+}
+impl StoreInterface for DistributedStore {
+    fn get(&self, key: &str) -> String { unimplemented!() }
+    fn set(&self, key: &str, value: &str) { unimplemented!() }
+}
+
+lazy_static! {
+    // TODO: Add support for distributed store do not hardcode the local one
+    pub static ref KV_STORE: Box<dyn StoreInterface> = Box::new(LocalStore::new());
+}

pipeless/src/lib.rs

@@ -18,4 +18,5 @@ pub mod events;
 pub mod data;
 pub mod dispatcher;
 pub mod stages;
-pub mod cli;
+pub mod cli;
+pub mod kvs;