Merge pull request #127 from k82cn/sdk_rs_fix

Enhance rust sdk.
flame-sh · Feb 9, 2025 · b0c0247 · b0c0247
2 parents e32421e + 10b3436
commit b0c0247
Show file tree

Hide file tree

Showing 29 changed files with 276 additions and 437 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -8,8 +8,7 @@ members = [
     "session_manager",
     "executor_manager",
     "rpc",
-    "sdk/rust/service",
-    "sdk/rust/client",
+    "sdk/rust",
 ]
 
 [workspace.dependencies]

diff --git a/Makefile b/Makefile
@@ -19,7 +19,6 @@ ci-image:
 	sudo docker build -t xflops/flame-console -f docker/Dockerfile.console .
 
 update_protos:
-	cp rpc/protos/frontend.proto sdk/rust/client/protos
-	cp rpc/protos/types.proto sdk/rust/client/protos
-	cp rpc/protos/shim.proto sdk/rust/service/protos
-	cp rpc/protos/types.proto sdk/rust/service/protos
+	cp rpc/protos/frontend.proto sdk/rust/protos
+	cp rpc/protos/types.proto sdk/rust/protos
+	cp rpc/protos/shim.proto sdk/rust/protos
diff --git a/flmctl/Cargo.toml b/flmctl/Cargo.toml
@@ -6,8 +6,7 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-flame-client = { path = "../sdk/rust/client" }
-common = { path = "../common" }
+flame-rs = { path = "../sdk/rust" }
 
 tokio = { workspace = true }
 tonic = { workspace = true }

diff --git a/flmctl/src/create.rs b/flmctl/src/create.rs
@@ -13,13 +13,11 @@ limitations under the License.
 
 use std::error::Error;
 
-use common::ctx::FlameContext;
-
-use self::flame::SessionAttributes;
-use flame_client as flame;
+use flame_rs as flame;
+use flame_rs::{apis::FlameContext, client::SessionAttributes};
 
 pub async fn run(ctx: &FlameContext, app: &str, slots: &i32) -> Result<(), Box<dyn Error>> {
-    let conn = flame::connect(&ctx.endpoint).await?;
+    let conn = flame::client::connect(&ctx.endpoint).await?;
     let attr = SessionAttributes {
         application: app.to_owned(),
         slots: *slots,

diff --git a/flmctl/src/list.rs b/flmctl/src/list.rs
@@ -14,16 +14,16 @@ limitations under the License.
 use std::cmp::Ordering;
 use std::error::Error;
 
-use common::ctx::FlameContext;
-use flame_client::SessionState;
-use flame_client::{self as flame, Connection, FlameError};
+use flame_rs as flame;
+use flame_rs::apis::{FlameContext, FlameError, SessionState};
+use flame_rs::client::Connection;
 
 pub async fn run(
     ctx: &FlameContext,
     application: bool,
     session: bool,
 ) -> Result<(), Box<dyn Error>> {
-    let conn = flame::connect(&ctx.endpoint).await?;
+    let conn = flame::client::connect(&ctx.endpoint).await?;
     match (application, session) {
         (true, false) => list_application(conn).await,
         (false, true) => list_session(conn).await,

diff --git a/flmctl/src/main.rs b/flmctl/src/main.rs
@@ -14,7 +14,7 @@ limitations under the License.
 use std::error::Error;
 
 use clap::{Parser, Subcommand};
-use common::ctx::FlameContext;
+use flame_rs::apis::FlameContext;
 
 mod create;
 mod helper;

diff --git a/flmctl/src/migrate.rs b/flmctl/src/migrate.rs
@@ -13,11 +13,10 @@ limitations under the License.
 
 use std::error::Error;
 
+use flame_rs::apis::FlameContext;
 use sqlx::{migrate::MigrateDatabase, Sqlite, SqlitePool};
 use url::Url;
 
-use common::ctx::FlameContext;
-
 pub async fn run(_: &FlameContext, url: &str, sql: &str) -> Result<(), Box<dyn Error>> {
     let uri = Url::parse(url)?;
 

diff --git a/flmctl/src/register.rs b/flmctl/src/register.rs
@@ -13,10 +13,14 @@ limitations under the License.
 
 use std::{fs, path::Path};
 
-use serde_derive::{Deserialize, Serialize};
+use flame_rs as flame;
+use flame_rs::apis::Shim;
+use flame_rs::{
+    apis::{FlameContext, FlameError},
+    client::ApplicationAttributes,
+};
 
-use common::ctx::FlameContext;
-use flame_client::{self as flame, ApplicationAttributes, FlameError, Shim};
+use serde_derive::{Deserialize, Serialize};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 struct MetadataYaml {
@@ -53,7 +57,7 @@ pub async fn run(ctx: &FlameContext, path: &String) -> Result<(), FlameError> {
 
     let app_attr = ApplicationAttributes::try_from(&app)?;
 
-    let conn = flame::connect(&ctx.endpoint).await?;
+    let conn = flame::client::connect(&ctx.endpoint).await?;
 
     conn.register_application(app.metadata.name, app_attr)
         .await?;

diff --git a/flmctl/src/unregister.rs b/flmctl/src/unregister.rs
@@ -13,7 +13,7 @@ limitations under the License.
 
 use std::error::Error;
 
-use common::ctx::FlameContext;
+use flame_rs::apis::FlameContext;
 
 pub async fn run(_: &FlameContext, _: &String) -> Result<(), Box<dyn Error>> {
     todo!()

diff --git a/flmctl/src/view.rs b/flmctl/src/view.rs
@@ -13,7 +13,7 @@ limitations under the License.
 
 use std::error::Error;
 
-use common::ctx::FlameContext;
+use flame_rs::apis::FlameContext;
 
 pub async fn run(_: &FlameContext, _: &String) -> Result<(), Box<dyn Error>> {
     todo!()

diff --git a/flmexec/Cargo.toml b/flmexec/Cargo.toml
@@ -6,9 +6,7 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-flame-client = { path = "../sdk/rust/client" }
-rpc = { path = "../rpc" }
-common = { path = "../common" }
+flame-rs = { path = "../sdk/rust" }
 
 tokio = { workspace = true }
 tonic = { workspace = true }

diff --git a/flmexec/src/main.rs b/flmexec/src/main.rs
@@ -19,9 +19,9 @@ use chrono::Local;
 use clap::Parser;
 use indicatif::HumanCount;
 
-use self::flame::FlameError;
-use common::ctx::FlameContext;
-use flame_client as flame;
+use flame_rs as flame;
+use flame_rs::apis::{FlameContext, FlameError};
+use flame_rs::client::{SessionAttributes, Task, TaskInformer};
 
 #[derive(Parser)]
 #[command(name = "flmexec")]
@@ -54,10 +54,10 @@ async fn main() -> Result<(), Box<dyn Error>> {
     let slots = cli.slots.unwrap_or(DEFAULT_SLOTS);
     let task_num = cli.task_num.unwrap_or(DEFAULT_TASK_NUM);
 
-    let conn = flame::connect(&ctx.endpoint).await?;
+    let conn = flame::client::connect(&ctx.endpoint).await?;
 
     let ssn_creation_start_time = Local::now();
-    let ssn_attr = flame::SessionAttributes {
+    let ssn_attr = SessionAttributes {
         application: DEFAULT_APP.to_string(),
         slots,
         common_data: None,
@@ -97,8 +97,8 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
 struct ExecInfo {}
 
-impl flame::TaskInformer for ExecInfo {
-    fn on_update(&mut self, task: flame::Task) {
+impl TaskInformer for ExecInfo {
+    fn on_update(&mut self, task: Task) {
         if task.is_completed() {
             println!(
                 "Task {:<10}: {:?}",

diff --git a/flmping/Cargo.toml b/flmping/Cargo.toml
@@ -6,10 +6,7 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-flame-client = { path = "../sdk/rust/client" }
-flame-service = { path = "../sdk/rust/service" }
-
-common = { path = "../common" }
+flame-rs = { path = "../sdk/rust" }
 
 tokio = { workspace = true }
 tonic = { workspace = true }

diff --git a/flmping/src/client.rs b/flmping/src/client.rs
@@ -16,12 +16,13 @@ use std::sync::{Arc, Mutex};
 
 use chrono::Local;
 use clap::Parser;
+use flame_rs::apis::FlameError;
+use flame_rs::client::{SessionAttributes, Task, TaskInformer};
 use futures::future::try_join_all;
 use indicatif::HumanCount;
 
-use self::flame::FlameError;
-use common::ctx::FlameContext;
-use flame_client::{self as flame, new_ptr};
+use flame::apis::FlameContext;
+use flame_rs::{self as flame, new_ptr};
 
 #[derive(Parser)]
 #[command(name = "flmping")]
@@ -51,10 +52,10 @@ async fn main() -> Result<(), Box<dyn Error>> {
     let slots = cli.slots.unwrap_or(DEFAULT_SLOTS);
     let task_num = cli.task_num.unwrap_or(DEFAULT_TASK_NUM);
 
-    let conn = flame::connect(&ctx.endpoint).await?;
+    let conn = flame::client::connect(&ctx.endpoint).await?;
 
     let ssn_creation_start_time = Local::now();
-    let ssn_attr = flame::SessionAttributes {
+    let ssn_attr = SessionAttributes {
         application: DEFAULT_APP.to_string(),
         slots,
         common_data: None,
@@ -102,8 +103,8 @@ impl OutputInfor {
     }
 }
 
-impl flame::TaskInformer for OutputInfor {
-    fn on_update(&mut self, task: flame::Task) {
+impl TaskInformer for OutputInfor {
+    fn on_update(&mut self, task: Task) {
         if task.is_completed() {
             let output = task.output.unwrap_or_default();
             println!(

diff --git a/flmping/src/service.rs b/flmping/src/service.rs
@@ -12,13 +12,17 @@ limitations under the License.
 */
 use gethostname::gethostname;
 
-use flame_service::{self as flame, FlameError, SessionContext, TaskContext, TaskOutput};
+use flame_rs::{
+    self as flame,
+    apis::{FlameError, TaskOutput},
+    service::{SessionContext, TaskContext},
+};
 
 #[derive(Clone)]
 pub struct FlmpingService {}
 
 #[tonic::async_trait]
-impl flame::FlameService for FlmpingService {
+impl flame::service::FlameService for FlmpingService {
     async fn on_session_enter(&self, _: SessionContext) -> Result<(), FlameError> {
         Ok(())
     }
@@ -41,7 +45,7 @@ impl flame::FlameService for FlmpingService {
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     env_logger::init();
 
-    flame_service::run(FlmpingService {}).await?;
+    flame::service::run(FlmpingService {}).await?;
 
     log::debug!("FlmpingService was stopped.");