fix: fix concurrency problems

Cargo.toml

@@ -8,9 +8,9 @@ resolver = "2"
 [workspace.dependencies]
 anyhow = "1"
 rand = "0.8.5"
-dashmap = { git = "https://github.com/xacrimon/dashmap.git", branch = "master", features = ["inline"] }
-derive_builder = "0.12.0"
-deadpool = { version = "0.9.5", default-features = false, features = ["managed"] }
+dashmap = { version = "5.5.0", features = ["inline"] }
 chrono = "0.4.26"
-burn = { git = "https://github.com/burn-rs/burn.git", branch = "main" }
+burn = "0.9.0"
 serde = { version = "1.0.163", features = ["derive"] }
+typed-builder = "0.16.2"
+derive-getters = "0.2.1"

cml-core/Cargo.toml

@@ -8,8 +8,8 @@ edition = "2021"
 [dependencies]
 anyhow = { workspace = true }
 dashmap = { workspace = true }
-derive-getters = "0.2.1"
-derive_builder = { workspace = true }
-deadpool = { workspace = true }
+derive-getters = { workspace = true }
+typed-builder = { workspace = true }
+deadpool = { version = "0.10.0", default-features = false, features = ["managed", "rt_tokio_1"] }
 serde = { workspace = true }
 chrono = { workspace = true }

cml-core/src/core.rs

@@ -1,3 +1,4 @@
 pub mod inference;
 pub mod register;
 pub mod task;
+pub(crate) mod utils;

cml-core/src/core/inference.rs

@@ -1,36 +1,37 @@
-use crate::metadata::MetaData;
+use crate::{metadata::Metadata, SharedBatchState};
 use anyhow::Result;
 use deadpool::managed::{Manager, Pool};
 use derive_getters::Getters;
-use std::path::PathBuf;
+use std::future::Future;
+use typed_builder::TypedBuilder;
 
-#[derive(Builder, Getters)]
+#[derive(TypedBuilder, Getters)]
 pub struct NewSample<F> {
-    data_path: PathBuf,
-    #[builder(default = "None")]
-    output: Option<F>,
-    #[builder(default = "None")]
+    #[builder(default, setter(strip_option))]
+    pub output: Option<F>,
+    #[builder(default, setter(strip_option))]
     optional_fields: Option<Vec<F>>,
-    #[builder(default = "None")]
+    #[builder(default, setter(strip_option))]
     optional_tags: Option<Vec<F>>,
 }
 
 pub trait Inference<M, F, T, C: Manager> {
-    async fn init_inference(
+    fn init_inference(
         &self,
         target_type: M,
         optional_fields: Option<Vec<M>>,
         optional_tags: Option<Vec<M>>,
-    ) -> Result<()>;
+    ) -> impl Future<Output = Result<()>> + Send;
 
-    async fn inference<FN>(
+    fn inference<FN>(
         &self,
-        metadata: MetaData<F>,
-        available_status: &[&str],
-        data: &mut Vec<NewSample<F>>,
+        metadata: &Metadata<F>,
+        data: Vec<NewSample<F>>,
+        current_ts: Option<T>,
+        batch_state: Option<&SharedBatchState>,
         pool: &Pool<C>,
         inference_fn: FN,
-    ) -> Result<()>
+    ) -> impl Future<Output = Result<()>>
     where
-        FN: FnOnce(&mut Vec<NewSample<F>>, &str, T) -> Vec<NewSample<F>>;
+        FN: FnOnce(&mut [NewSample<F>], &str, T);
 }

cml-core/src/core/register.rs

@@ -1,48 +1,31 @@
-use crate::metadata::MetaData;
+use crate::{metadata::Metadata, SharedBatchState};
 use anyhow::Result;
-use dashmap::DashMap;
 use deadpool::managed::{Manager, Pool};
 use derive_getters::Getters;
-use std::{
-    path::PathBuf,
-    sync::{Arc, Mutex},
-};
+use std::future::Future;
+use typed_builder::TypedBuilder;
 
-#[derive(Builder, Getters)]
+#[derive(TypedBuilder, Getters, Clone)]
 pub struct TrainData<F> {
-    data_path: PathBuf,
     gt: F,
-    #[builder(default = "None")]
+    #[builder(default, setter(strip_option))]
     optional_fields: Option<Vec<F>>,
 }
 
-pub struct BatchState {
-    pub map: DashMap<String, bool>,
-}
-
-pub type SharedBatchState = Arc<Mutex<BatchState>>;
-
-impl BatchState {
-    pub fn create(num_shards: usize) -> SharedBatchState {
-        Arc::new(Mutex::new(BatchState {
-            map: DashMap::with_shard_amount(num_shards),
-        }))
-    }
-}
-
-pub trait Register<M, F, C: Manager> {
-    async fn init_register(
+pub trait Register<M, F, T, C: Manager> {
+    fn init_register(
         &self,
         gt_type: M,
         optional_fields: Option<Vec<M>>,
         optional_tags: Option<Vec<M>>,
-    ) -> Result<()>;
+    ) -> impl Future<Output = Result<()>> + Send;
 
-    async fn register(
+    fn register(
         &self,
-        metadata: MetaData<F>,
+        metadata: &Metadata<F>,
         train_data: Vec<TrainData<F>>,
-        batch_state: Arc<Mutex<BatchState>>,
+        current_ts: Option<T>,
+        batch_state: Option<&SharedBatchState>,
         pool: &Pool<C>,
-    ) -> Result<()>;
+    ) -> impl Future<Output = Result<()>> + Send;
 }

cml-core/src/core/task.rs

@@ -1,28 +1,11 @@
 use anyhow::Result;
-use chrono::Duration;
-use derive_getters::Getters;
-
-#[derive(Builder, Getters, Clone)]
-pub struct TaskConfig<'a> {
-    min_start_count: usize,
-    min_update_count: usize,
-    working_status: Vec<&'a str>,
-    limit_time: Duration,
-}
-
+use std::future::Future;
 pub trait Task<M> {
-    async fn init_task(
+    fn init_task(
         &self,
         optional_fields: Option<Vec<M>>,
         optional_tags: Option<Vec<M>>,
-    ) -> Result<()>;
+    ) -> impl Future<Output = Result<()>> + Send;
 
-    fn run<FN>(
-        &self,
-        task_config: TaskConfig,
-        build_from_scratch_fn: FN,
-        fining_build_fn: FN,
-    ) -> Result<()>
-    where
-        FN: Fn(&str) -> Result<()> + Send + Sync;
+    fn prepare(&self) -> Result<(Vec<String>, Vec<String>)>;
 }

cml-core/src/core/utils.rs

@@ -0,0 +1,3 @@
+pub fn get_placeholders<T>(fields: &[T]) -> String {
+    vec!["?"; fields.len()].join(", ")
+}

cml-core/src/handler.rs

@@ -1,6 +1,11 @@
 use anyhow::Result;
+use std::future::Future;
 
 pub trait Handler {
     type Database;
-    async fn init(self, client: &Self::Database, db: Option<&str>) -> Result<()>;
+    fn init(
+        self,
+        client: &Self::Database,
+        db: Option<&str>,
+    ) -> impl Future<Output = Result<()>> + Send;
 }

cml-core/src/lib.rs

@@ -1,9 +1,14 @@
-#![allow(incomplete_features)]
-#![feature(async_fn_in_trait)]
-
-#[macro_use]
-extern crate derive_builder;
+use std::{
+    collections::HashSet,
+    sync::{Arc, Condvar, Mutex},
+};
 
 pub mod core;
-pub mod handler;
-pub mod metadata;
+mod handler;
+mod metadata;
+
+pub type SharedBatchState = Arc<(Mutex<HashSet<String>>, Condvar)>;
+
+pub use core::{register::Register, utils::get_placeholders};
+pub use handler::Handler;
+pub use metadata::Metadata;

cml-core/src/metadata.rs

@@ -1,22 +1,9 @@
 use derive_getters::Getters;
+use typed_builder::TypedBuilder;
 
-#[derive(Builder, Getters)]
-pub struct MetaData<F> {
-    model_update_time: i64,
-    batch: String,
-    inherent_field_num: usize,
-    inherent_tag_num: usize,
-    optional_field_num: usize,
-    #[builder(default = "None")]
+#[derive(TypedBuilder, Getters, Clone)]
+pub struct Metadata<F> {
+    pub batch: String,
+    #[builder(default, setter(strip_option))]
     optional_tags: Option<Vec<F>>,
 }
-
-impl<F> MetaData<F> {
-    pub fn get_placeholders(&self) -> (String, String) {
-        (
-            vec!["?"; self.optional_tags.as_ref().map_or(0, |v| v.len()) + self.inherent_tag_num]
-                .join(", "),
-            vec!["?"; self.optional_field_num + self.inherent_field_num].join(", "),
-        )
-    }
-}

cml-tdengine/Cargo.toml

@@ -14,14 +14,15 @@ taos-query = { version = "*" }
 anyhow = { workspace = true }
 rand = { workspace = true }
 dashmap = { workspace = true }
-derive_builder = { workspace = true }
-deadpool = { workspace = true }
+typed-builder = { workspace = true }
 serde = { workspace = true }
 chrono = { workspace = true }
 rayon = "1.7.0"
-tokio = { version = "1.29.1", features = ["rt", "macros"] }
+tokio = { version = "1", features = ["rt", "macros"] }
+derive-getters = { workspace = true }
 
 [dev-dependencies]
 burn = { workspace = true }
-burn-autodiff = { git = "https://github.com/burn-rs/burn.git", branch = "main" }
-burn-ndarray = { git = "https://github.com/burn-rs/burn.git", branch = "main" }
+burn-autodiff = "0.9.0"
+burn-ndarray = "0.9.0"
+tempfile = "3.8.0"