This is an automated email from the ASF dual-hosted git repository. mssun pushed a commit to branch develop in repository https://gitbox.apache.org/repos/asf/incubator-teaclave.git
commit 9a4de971c74ea909807993163bc975825ddc2b13 Author: Mingshen Sun <[email protected]> AuthorDate: Sun Mar 22 15:54:07 2020 -0700 [types] Use FunctionArguments in other proto definitions --- services/execution/enclave/src/service.rs | 28 ++++++------- services/management/enclave/src/service.rs | 12 +++--- services/management/enclave/src/task.rs | 6 +-- .../src/proto/teaclave_frontend_service.proto | 9 +--- services/proto/src/teaclave_frontend_service.rs | 43 ++++++------------- .../enclave/src/teaclave_frontend_service.rs | 30 ++++++-------- .../enclave/src/teaclave_management_service.rs | 13 +++--- tests/integration/enclave/src/teaclave_worker.rs | 6 +-- types/src/macros.rs | 2 +- types/src/staged_function.rs | 48 ++++++++++++++++------ worker/src/function/echo.rs | 2 +- worker/src/function/gbdt_training.rs | 2 +- worker/src/function/mesapy.rs | 12 +++--- worker/src/worker.rs | 11 ++--- 14 files changed, 110 insertions(+), 114 deletions(-) diff --git a/services/execution/enclave/src/service.rs b/services/execution/enclave/src/service.rs index 50bf2e8..11ea963 100644 --- a/services/execution/enclave/src/service.rs +++ b/services/execution/enclave/src/service.rs @@ -260,17 +260,17 @@ pub mod tests { pub fn test_invoke_gbdt_training() { let task_id = Uuid::new_v4(); - let arg_map = hashmap!( - "feature_size".to_string() => "4".to_string(), - "max_depth".to_string() => "4".to_string(), - "iterations".to_string() => "100".to_string(), - "shrinkage".to_string() => "0.1".to_string(), - "feature_sample_ratio".to_string() => "1.0".to_string(), - "data_sample_ratio".to_string() => "1.0".to_string(), - "min_leaf_size".to_string() => "1".to_string(), - "loss".to_string() => "LAD".to_string(), - "training_optimization_level".to_string() => "2".to_string(), - ); + let function_arguments = FunctionArguments::new(hashmap!( + "feature_size" => "4", + "max_depth" => "4", + "iterations" => "100", + "shrinkage" => "0.1", + "feature_sample_ratio" => "1.0", + "data_sample_ratio" => "1.0", + "min_leaf_size" => "1", + "loss" => "LAD", + "training_optimization_level" => "2", + )); let fixture_dir = format!( "file:///{}/fixtures/functions/gbdt_training", env!("TEACLAVE_TEST_INSTALL_DIR") @@ -288,13 +288,13 @@ pub mod tests { let training_input_data = FunctionInputFile::new(input_url, "", crypto_info); let model_output_data = FunctionOutputFile::new(output_url, crypto_info); - let input_data = hashmap!("training_data".to_string() => training_input_data); - let output_data = hashmap!("trained_model".to_string() => model_output_data); + let input_data = hashmap!("training_data" => training_input_data); + let output_data = hashmap!("trained_model" => model_output_data); let staged_task = StagedTask::new() .task_id(task_id) .function_name("gbdt_training") - .function_arguments(arg_map.into()) + .function_arguments(function_arguments) .input_data(input_data) .output_data(output_data); diff --git a/services/management/enclave/src/service.rs b/services/management/enclave/src/service.rs index d48ce2f..3c527f1 100644 --- a/services/management/enclave/src/service.rs +++ b/services/management/enclave/src/service.rs @@ -315,7 +315,7 @@ impl TeaclaveManagement for TeaclaveManagementService { let task = crate::task::create_task( function, user_id, - request.arg_list, + request.function_arguments, request.input_data_owner_list, request.output_data_owner_list, ) @@ -352,7 +352,7 @@ impl TeaclaveManagement for TeaclaveManagementService { creator: task.creator, function_id: task.function_id, function_owner: task.function_owner, - arg_list: task.function_arguments.into(), + function_arguments: task.function_arguments, input_data_owner_list: task.input_data_owner_list, output_data_owner_list: task.output_data_owner_list, participants: task.participants, @@ -663,7 +663,8 @@ pub mod tests { use super::*; use std::collections::{HashMap, HashSet}; use teaclave_types::{ - FunctionInput, FunctionOutput, TeaclaveFileCryptoInfo, TeaclaveFileRootKey128, + hashmap, FunctionArguments, FunctionInput, FunctionOutput, TeaclaveFileCryptoInfo, + TeaclaveFileRootKey128, }; use url::Url; @@ -730,13 +731,12 @@ pub mod tests { owner: "mock_user".to_string(), is_native: false, }; - let mut arg_list = HashMap::new(); - arg_list.insert("arg".to_string(), "data".to_string()); + let function_arguments = FunctionArguments::new(hashmap!("arg" => "data")); let task = crate::task::create_task( function, "mock_user".to_string(), - arg_list, + function_arguments, HashMap::new(), HashMap::new(), ) diff --git a/services/management/enclave/src/task.rs b/services/management/enclave/src/task.rs index c5b804f..5d0b6be 100644 --- a/services/management/enclave/src/task.rs +++ b/services/management/enclave/src/task.rs @@ -20,14 +20,15 @@ use std::collections::HashSet; use std::prelude::v1::*; use teaclave_types::Function; use teaclave_types::{ - DataOwnerList, Storable, Task, TaskStatus, TeaclaveInputFile, TeaclaveOutputFile, + DataOwnerList, FunctionArguments, Storable, Task, TaskStatus, TeaclaveInputFile, + TeaclaveOutputFile, }; use uuid::Uuid; pub(crate) fn create_task( function: Function, creator: String, - arg_list: HashMap<String, String>, + function_arguments: FunctionArguments, input_data_owner_list: HashMap<String, DataOwnerList>, output_data_owner_list: HashMap<String, DataOwnerList>, ) -> Result<Task> { @@ -47,7 +48,6 @@ pub(crate) fn create_task( participants.insert(user_id.clone()); } } - let function_arguments = arg_list.into(); let task = Task { task_id, creator, diff --git a/services/proto/src/proto/teaclave_frontend_service.proto b/services/proto/src/proto/teaclave_frontend_service.proto index b2434f4..0ecbce2 100644 --- a/services/proto/src/proto/teaclave_frontend_service.proto +++ b/services/proto/src/proto/teaclave_frontend_service.proto @@ -106,14 +106,9 @@ message DataMap { string data_id = 2; } -message Argument { - string arg_name = 1; - string arg_value = 2; -} - message CreateTaskRequest { string function_id = 1; - repeated Argument arg_list = 2; + map<string, string> function_arguments = 2; repeated DataOwnerList input_data_owner_list = 3; repeated DataOwnerList output_data_owner_list = 4; } @@ -131,7 +126,7 @@ message GetTaskResponse { string creator = 2; string function_id = 3; string function_owner = 4; - repeated Argument arg_list = 5; + map<string, string> function_arguments = 5; repeated DataOwnerList input_data_owner_list = 6; repeated DataOwnerList output_data_owner_list = 7; repeated string participants = 8; diff --git a/services/proto/src/teaclave_frontend_service.rs b/services/proto/src/teaclave_frontend_service.rs index 123371d..8878eb6 100644 --- a/services/proto/src/teaclave_frontend_service.rs +++ b/services/proto/src/teaclave_frontend_service.rs @@ -10,7 +10,8 @@ use std::collections::{HashMap, HashSet}; use std::prelude::v1::*; use teaclave_rpc::into_request; use teaclave_types::{ - DataOwnerList, FunctionInput, FunctionOutput, TaskStatus, TeaclaveFileCryptoInfo, + DataOwnerList, FunctionArguments, FunctionInput, FunctionOutput, TaskStatus, + TeaclaveFileCryptoInfo, }; use url::Url; @@ -264,7 +265,7 @@ pub struct GetFunctionResponse { #[derive(Debug)] pub struct CreateTaskRequest { pub function_id: String, - pub arg_list: HashMap<String, String>, + pub function_arguments: FunctionArguments, pub input_data_owner_list: HashMap<String, DataOwnerList>, pub output_data_owner_list: HashMap<String, DataOwnerList>, } @@ -305,7 +306,7 @@ pub struct GetTaskResponse { pub creator: String, pub function_id: String, pub function_owner: String, - pub arg_list: HashMap<String, String>, + pub function_arguments: FunctionArguments, pub input_data_owner_list: HashMap<String, DataOwnerList>, pub output_data_owner_list: HashMap<String, DataOwnerList>, pub participants: HashSet<String>, @@ -845,36 +846,16 @@ pub fn data_owner_list_to_proto<S: std::hash::BuildHasher>( ret } -fn arg_list_from_proto(vector: Vec<proto::Argument>) -> Result<HashMap<String, String>> { - let mut ret = HashMap::with_capacity(vector.len()); - for item in vector.into_iter() { - ret.insert(item.arg_name, item.arg_value); - } - Ok(ret) -} - -fn arg_list_to_proto(map: HashMap<String, String>) -> Vec<proto::Argument> { - let mut ret = Vec::with_capacity(map.len()); - for (arg_name, arg_value) in map.into_iter() { - let argument = proto::Argument { - arg_name, - arg_value, - }; - ret.push(argument); - } - ret -} - impl std::convert::TryFrom<proto::CreateTaskRequest> for CreateTaskRequest { type Error = Error; fn try_from(proto: proto::CreateTaskRequest) -> Result<Self> { - let arg_list = arg_list_from_proto(proto.arg_list)?; + let function_arguments = proto.function_arguments.into(); let input_data_owner_list = data_owner_list_from_proto(proto.input_data_owner_list)?; let output_data_owner_list = data_owner_list_from_proto(proto.output_data_owner_list)?; let ret = Self { function_id: proto.function_id, - arg_list, + function_arguments, input_data_owner_list, output_data_owner_list, }; @@ -884,13 +865,13 @@ impl std::convert::TryFrom<proto::CreateTaskRequest> for CreateTaskRequest { impl From<CreateTaskRequest> for proto::CreateTaskRequest { fn from(request: CreateTaskRequest) -> Self { - let arg_list = arg_list_to_proto(request.arg_list); + let function_arguments = request.function_arguments.into(); let input_data_owner_list = data_owner_list_to_proto(request.input_data_owner_list); let output_data_owner_list = data_owner_list_to_proto(request.output_data_owner_list); Self { function_id: request.function_id, - arg_list, + function_arguments, input_data_owner_list, output_data_owner_list, } @@ -958,7 +939,7 @@ impl std::convert::TryFrom<proto::GetTaskResponse> for GetTaskResponse { type Error = Error; fn try_from(proto: proto::GetTaskResponse) -> Result<Self> { - let arg_list = arg_list_from_proto(proto.arg_list)?; + let function_arguments = proto.function_arguments.into(); let input_data_owner_list = data_owner_list_from_proto(proto.input_data_owner_list)?; let output_data_owner_list = data_owner_list_from_proto(proto.output_data_owner_list)?; let input_map = data_map_from_proto(proto.input_map)?; @@ -970,7 +951,7 @@ impl std::convert::TryFrom<proto::GetTaskResponse> for GetTaskResponse { creator: proto.creator, function_id: proto.function_id, function_owner: proto.function_owner, - arg_list, + function_arguments, input_data_owner_list, output_data_owner_list, participants: proto.participants.into_iter().collect(), @@ -986,7 +967,7 @@ impl std::convert::TryFrom<proto::GetTaskResponse> for GetTaskResponse { impl From<GetTaskResponse> for proto::GetTaskResponse { fn from(response: GetTaskResponse) -> Self { - let arg_list = arg_list_to_proto(response.arg_list); + let function_arguments = response.function_arguments.into(); let input_data_owner_list = data_owner_list_to_proto(response.input_data_owner_list); let output_data_owner_list = data_owner_list_to_proto(response.output_data_owner_list); let input_map = data_map_to_proto(response.input_map); @@ -997,7 +978,7 @@ impl From<GetTaskResponse> for proto::GetTaskResponse { creator: response.creator, function_id: response.function_id, function_owner: response.function_owner, - arg_list, + function_arguments, input_data_owner_list, output_data_owner_list, participants: response.participants.into_iter().collect(), diff --git a/tests/functional/enclave/src/teaclave_frontend_service.rs b/tests/functional/enclave/src/teaclave_frontend_service.rs index b670133..5982981 100644 --- a/tests/functional/enclave/src/teaclave_frontend_service.rs +++ b/tests/functional/enclave/src/teaclave_frontend_service.rs @@ -276,11 +276,10 @@ fn test_create_task() { let function_id = "function-00000000-0000-0000-0000-000000000002"; let mut output_data_owner_list = HashMap::new(); output_data_owner_list.insert("output".to_string(), data_owner_id_list); + let function_arguments = FunctionArguments::new(hashmap!("arg1" => "data1")); let request = CreateTaskRequest { function_id: function_id.to_string(), - arg_list: vec![("arg1".to_string(), "data1".to_string())] - .into_iter() - .collect(), + function_arguments, input_data_owner_list: HashMap::new(), output_data_owner_list: output_data_owner_list.clone(), }; @@ -288,11 +287,10 @@ fn test_create_task() { assert!(response.is_ok()); assert!(!response.unwrap().task_id.is_empty()); + let function_arguments = FunctionArguments::new(hashmap!("arg1" => "data1")); let request = CreateTaskRequest { function_id: function_id.to_string(), - arg_list: vec![("arg1".to_string(), "data1".to_string())] - .into_iter() - .collect(), + function_arguments, input_data_owner_list: HashMap::new(), output_data_owner_list, }; @@ -313,11 +311,10 @@ fn test_get_task() { }; let mut output_data_owner_list = HashMap::new(); output_data_owner_list.insert("output".to_string(), data_owner_id_list); + let function_arguments = FunctionArguments::new(hashmap!("arg1" => "data1")); let request = CreateTaskRequest { function_id: "function-00000000-0000-0000-0000-000000000002".to_string(), - arg_list: vec![("arg1".to_string(), "data1".to_string())] - .into_iter() - .collect(), + function_arguments, input_data_owner_list: HashMap::new(), output_data_owner_list, }; @@ -345,11 +342,10 @@ fn test_assign_data() { }; let mut output_data_owner_list = HashMap::new(); output_data_owner_list.insert("output".to_string(), data_owner_id_list); + let function_arguments = FunctionArguments::new(hashmap!("arg1" => "data1")); let request = CreateTaskRequest { function_id: "function-00000000-0000-0000-0000-000000000002".to_string(), - arg_list: vec![("arg1".to_string(), "data1".to_string())] - .into_iter() - .collect(), + function_arguments, input_data_owner_list: HashMap::new(), output_data_owner_list, }; @@ -399,11 +395,10 @@ fn test_approve_task() { }; let mut output_data_owner_list = HashMap::new(); output_data_owner_list.insert("output".to_string(), data_owner_id_list); + let function_arguments = FunctionArguments::new(hashmap!("arg1" => "data1")); let request = CreateTaskRequest { function_id: "function-00000000-0000-0000-0000-000000000002".to_string(), - arg_list: vec![("arg1".to_string(), "data1".to_string())] - .into_iter() - .collect(), + function_arguments, input_data_owner_list: HashMap::new(), output_data_owner_list, }; @@ -450,11 +445,10 @@ fn test_invoke_task() { }; let mut output_data_owner_list = HashMap::new(); output_data_owner_list.insert("output".to_string(), data_owner_id_list); + let function_arguments = FunctionArguments::new(hashmap!("arg1" => "data1")); let request = CreateTaskRequest { function_id: "function-00000000-0000-0000-0000-000000000002".to_string(), - arg_list: vec![("arg1".to_string(), "data1".to_string())] - .into_iter() - .collect(), + function_arguments, input_data_owner_list: HashMap::new(), output_data_owner_list, }; diff --git a/tests/functional/enclave/src/teaclave_management_service.rs b/tests/functional/enclave/src/teaclave_management_service.rs index d802df1..c97fc75 100644 --- a/tests/functional/enclave/src/teaclave_management_service.rs +++ b/tests/functional/enclave/src/teaclave_management_service.rs @@ -220,9 +220,10 @@ fn test_get_function() { } fn get_correct_create_task() -> CreateTaskRequest { - let mut arg_list = HashMap::new(); - arg_list.insert("arg1".to_string(), "data1".to_string()); - arg_list.insert("arg2".to_string(), "data2".to_string()); + let function_arguments = FunctionArguments::new(hashmap!( + "arg1" => "data1", + "arg2" => "data2", + )); let data_owner_id_list = DataOwnerList { user_id_list: vec!["mock_user1".to_string()].into_iter().collect(), }; @@ -240,7 +241,7 @@ fn get_correct_create_task() -> CreateTaskRequest { CreateTaskRequest { function_id: "function-00000000-0000-0000-0000-000000000001".to_string(), - arg_list, + function_arguments, input_data_owner_list, output_data_owner_list, } @@ -249,7 +250,7 @@ fn get_correct_create_task() -> CreateTaskRequest { fn test_create_task() { let request = CreateTaskRequest { function_id: "invalid_function".to_string(), - arg_list: HashMap::new(), + function_arguments: HashMap::new().into(), input_data_owner_list: HashMap::new(), output_data_owner_list: HashMap::new(), }; @@ -262,7 +263,7 @@ fn test_create_task() { assert!(response.is_ok()); let mut request = get_correct_create_task(); - request.arg_list.remove("arg1"); + request.function_arguments.inner_mut().remove("arg1"); let response = client.create_task(request); assert!(response.is_err()); diff --git a/tests/integration/enclave/src/teaclave_worker.rs b/tests/integration/enclave/src/teaclave_worker.rs index 7d2feb4..9956587 100644 --- a/tests/integration/enclave/src/teaclave_worker.rs +++ b/tests/integration/enclave/src/teaclave_worker.rs @@ -7,7 +7,7 @@ use teaclave_types::{ use teaclave_worker::Worker; fn test_start_worker() { - let arguments = FunctionArguments::from_map(&hashmap!( + let arguments = FunctionArguments::new(hashmap!( "feature_size" => "4", "max_depth" => "4", "iterations" => "100", @@ -26,12 +26,12 @@ fn test_start_worker() { let input_info = StagedInputFile::create_with_plaintext_file(plain_input).unwrap(); let input_files = StagedFiles::new(hashmap!( - "training_data".to_string() => input_info)); + "training_data" => input_info)); let output_info = StagedOutputFile::new(enc_output, TeaclaveFileRootKey128::random()); let output_files = StagedFiles::new(hashmap!( - "trained_model".to_string() => output_info.clone())); + "trained_model" => output_info.clone())); let staged_function = StagedFunction::new() .name("gbdt_training") diff --git a/types/src/macros.rs b/types/src/macros.rs index 34e0ab0..db59a38 100644 --- a/types/src/macros.rs +++ b/types/src/macros.rs @@ -3,7 +3,7 @@ macro_rules! hashmap { ($( $key: expr => $value: expr,)+) => { hashmap!($($key => $value),+) }; ($( $key: expr => $value: expr ),*) => {{ let mut map = ::std::collections::HashMap::new(); - $( map.insert($key, $value); )* + $( map.insert($key.into(), $value.into()); )* map }} } diff --git a/types/src/staged_function.rs b/types/src/staged_function.rs index cff442c..8048e14 100644 --- a/types/src/staged_function.rs +++ b/types/src/staged_function.rs @@ -12,6 +12,30 @@ pub struct ArgumentValue { inner: String, } +impl From<String> for ArgumentValue { + fn from(value: String) -> Self { + ArgumentValue::new(value) + } +} + +impl From<&str> for ArgumentValue { + fn from(value: &str) -> Self { + ArgumentValue::new(value.into()) + } +} + +impl From<&String> for ArgumentValue { + fn from(value: &String) -> Self { + ArgumentValue::new(value.into()) + } +} + +impl From<ArgumentValue> for String { + fn from(value: ArgumentValue) -> Self { + value.as_str().to_owned() + } +} + impl ArgumentValue { pub fn new(value: String) -> Self { Self { inner: value } @@ -72,28 +96,28 @@ impl<S: core::default::Default + std::hash::BuildHasher> From<FunctionArguments> impl From<HashMap<String, String>> for FunctionArguments { fn from(map: HashMap<String, String>) -> Self { - FunctionArguments::from_map(&map) - } -} - -impl FunctionArguments { - pub fn from_map<K, V>(input: &HashMap<K, V>) -> Self - where - K: std::string::ToString, - V: std::string::ToString, - { - let inner = input.iter().fold(HashMap::new(), |mut acc, (k, v)| { - acc.insert(k.to_string(), ArgumentValue::new(v.to_string())); + let inner = map.iter().fold(HashMap::new(), |mut acc, (k, v)| { + acc.insert(k.into(), v.into()); acc }); Self { inner } } +} + +impl FunctionArguments { + pub fn new(map: HashMap<String, ArgumentValue>) -> Self { + Self { inner: map } + } pub fn inner(&self) -> &HashMap<String, ArgumentValue> { &self.inner } + pub fn inner_mut(&mut self) -> &mut HashMap<String, ArgumentValue> { + &mut self.inner + } + pub fn get(&self, key: &str) -> anyhow::Result<&ArgumentValue> { self.inner .get(key) diff --git a/worker/src/function/echo.rs b/worker/src/function/echo.rs index a98ee16..5d46bb2 100644 --- a/worker/src/function/echo.rs +++ b/worker/src/function/echo.rs @@ -54,7 +54,7 @@ pub mod tests { } fn test_echo() { - let func_args = FunctionArguments::from_map(&hashmap!( + let func_args = FunctionArguments::new(hashmap!( "message" => "Hello Teaclave!" )); diff --git a/worker/src/function/gbdt_training.rs b/worker/src/function/gbdt_training.rs index df45c6f..5287435 100644 --- a/worker/src/function/gbdt_training.rs +++ b/worker/src/function/gbdt_training.rs @@ -150,7 +150,7 @@ pub mod tests { } fn test_gbdt_training() { - let func_arguments = FunctionArguments::from_map(&hashmap!( + let func_arguments = FunctionArguments::new(hashmap!( "feature_size" => "4", "max_depth" => "4", "iterations" => "100", diff --git a/worker/src/function/mesapy.rs b/worker/src/function/mesapy.rs index b1d4828..2ba6a1a 100644 --- a/worker/src/function/mesapy.rs +++ b/worker/src/function/mesapy.rs @@ -119,7 +119,7 @@ pub mod tests { } fn test_mesapy() { - let py_args = FunctionArguments::from_map(&hashmap!("--name" => "Teaclave")); + let py_args = FunctionArguments::new(hashmap!("--name" => "Teaclave")); let py_payload = r#" def entrypoint(argv): in_file_id = "in_f1" @@ -157,8 +157,8 @@ def entrypoint(argv): teaclave_open("invalid_key", "wb") except RuntimeError as e: assert e.message == "fileio_init: teaclave_ffi_error" - - # open invalid option + + # open invalid option try: teaclave_open(out_file_id, "w") except RuntimeError as e: @@ -181,9 +181,9 @@ def entrypoint(argv): }; let runtime = Box::new(RawIoRuntime::new(input_files, output_files)); - let func_args = FunctionArguments::from_map(&hashmap!( - "py_payload".to_string() => py_payload.to_string(), - "py_args".to_string() => serde_json::to_string(&py_args).unwrap() + let func_args = FunctionArguments::new(hashmap!( + "py_payload" => py_payload.to_string(), + "py_args" => serde_json::to_string(&py_args).unwrap() )); let function = Mesapy; diff --git a/worker/src/worker.rs b/worker/src/worker.rs index bbcdf68..6166778 100644 --- a/worker/src/worker.rs +++ b/worker/src/worker.rs @@ -25,7 +25,7 @@ use anyhow; use serde_json; use teaclave_types::{ - ExecutorType, FunctionArguments, StagedFiles, StagedFunction, StagedInputFile, + hashmap, ExecutorType, FunctionArguments, StagedFiles, StagedFunction, StagedInputFile, StagedOutputFile, WorkerCapability, }; @@ -162,11 +162,12 @@ fn prepare_arguments( !function_payload.is_empty(), "Python function payload must not be empty!" ); - let mut wrap_args = HashMap::new(); let req_args = serde_json::to_string(&function_arguments)?; - wrap_args.insert("py_payload".to_string(), function_payload); - wrap_args.insert("py_args".to_string(), req_args); - FunctionArguments::from_map(&wrap_args) + let wrap_args = hashmap!( + "py_payload" => function_payload, + "py_args" => req_args, + ); + FunctionArguments::new(wrap_args) } }; --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
