You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@teaclave.apache.org by ms...@apache.org on 2020/03/18 21:14:05 UTC
[incubator-teaclave] 01/02: [execution] Have a simple flow of
function invocation
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 2e24cd16c9241435d78c77344abde6baa8fb77e1
Author: Mingshen Sun <bo...@mssun.me>
AuthorDate: Wed Mar 18 13:21:38 2020 -0700
[execution] Have a simple flow of function invocation
---
services/execution/enclave/Cargo.toml | 2 +
services/execution/enclave/src/lib.rs | 3 +-
services/execution/enclave/src/service.rs | 113 ++++++++++++++++++++++++++++--
types/src/crypto.rs | 4 ++
types/src/staged_task.rs | 2 +
types/src/worker.rs | 3 +-
worker/src/runtime/mod.rs | 4 +-
worker/src/worker.rs | 7 ++
8 files changed, 129 insertions(+), 9 deletions(-)
diff --git a/services/execution/enclave/Cargo.toml b/services/execution/enclave/Cargo.toml
index 4fbc9c3..b282bf3 100644
--- a/services/execution/enclave/Cargo.toml
+++ b/services/execution/enclave/Cargo.toml
@@ -33,6 +33,8 @@ serde_json = { version = "1.0.39" }
serde = { version = "1.0.92", features = ["derive"] }
thiserror = { version = "1.0.9" }
gbdt = { version = "0.1.0", features = ["input", "enable_training"] }
+uuid = { version = "0.8.1", features = ["v4"] }
+url = { version = "2.1.1", features = ["serde"]}
teaclave_attestation = { path = "../../../attestation" }
teaclave_config = { path = "../../../config" }
diff --git a/services/execution/enclave/src/lib.rs b/services/execution/enclave/src/lib.rs
index d0ba237..a21c37e 100644
--- a/services/execution/enclave/src/lib.rs
+++ b/services/execution/enclave/src/lib.rs
@@ -100,8 +100,9 @@ pub mod tests {
pub fn run_tests() -> bool {
run_tests!(
- service::tests::test_invoke_function,
service::tests::test_invoke_echo_function,
+ service::tests::test_invoke_gbdt_training,
+ service::tests::test_invoke_gbdt_prediction
)
}
}
diff --git a/services/execution/enclave/src/service.rs b/services/execution/enclave/src/service.rs
index 503c6e7..48430f4 100644
--- a/services/execution/enclave/src/service.rs
+++ b/services/execution/enclave/src/service.rs
@@ -22,7 +22,7 @@ use std::sync::{Arc, SgxMutex as Mutex};
use teaclave_proto::teaclave_scheduler_service::*;
use teaclave_rpc::endpoint::Endpoint;
-use teaclave_types::{StagedTask, WorkerInvocationResult};
+use teaclave_types::{StagedTask, TeaclaveFunctionArguments, WorkerInvocationResult};
use teaclave_worker::Worker;
use anyhow::Result;
@@ -76,16 +76,18 @@ impl TeaclaveExecutionService {
}
};
log::debug!("response: {:?}", response);
- let result = self.invoke_task(response.staged_task);
- self.update_task(result);
+ let _result = self.invoke_task(response.staged_task);
+ // self.update_task(result);
}
}
- fn invoke_task(&mut self, _task: StagedTask) -> WorkerInvocationResult {
+ fn invoke_task(&mut self, task: StagedTask) -> WorkerInvocationResult {
+ let _function_args = TeaclaveFunctionArguments::new(&task.arg_list);
// TODO: convert task to function, i.e., needs help from agent
unimplemented!()
}
+ #[allow(unused)]
fn update_task(&mut self, _result: WorkerInvocationResult) {
unimplemented!()
}
@@ -99,10 +101,15 @@ mod test_mode {
#[cfg(feature = "enclave_unit_test")]
pub mod tests {
use super::*;
+ use std::collections::HashMap;
use std::convert::TryInto;
+ use std::format;
+ use std::vec;
use teaclave_types::*;
+ use url::Url;
+ use uuid::Uuid;
- pub fn test_invoke_function() {
+ pub fn test_invoke_gbdt_training() {
let function_args = TeaclaveFunctionArguments::new(&hashmap!(
"feature_size" => "4",
"max_depth" => "4",
@@ -160,4 +167,100 @@ pub mod tests {
let result = worker.invoke_function(invocation).unwrap();
assert_eq!(result, "Hello Teaclave!");
}
+
+ pub fn test_invoke_gbdt_prediction() {
+ let task_id = Uuid::new_v4();
+ let function = Function {
+ function_id: Uuid::new_v4(),
+ name: "gbdt_prediction".to_string(),
+ description: "".to_string(),
+ payload: b"".to_vec(),
+ is_public: false,
+ arg_list: vec![],
+ input_list: vec![],
+ output_list: vec![],
+ owner: "mock_user".to_string(),
+ is_native: true,
+ };
+ let arg_list = HashMap::new();
+ let test_install_dir = env!("TEACLAVE_TEST_INSTALL_DIR");
+ let fixture_dir = format!("{}/fixtures/functions/gbdt_prediction", test_install_dir);
+ let model_url = Url::parse(&format!("file:///{}/model.txt", fixture_dir)).unwrap();
+ let test_data_url = Url::parse(&format!("file:///{}/test_data.txt", fixture_dir)).unwrap();
+ let crypto_info = TeaclaveFileCryptoInfo::Raw;
+
+ let input_data_model = InputData {
+ url: model_url,
+ hash: "".to_string(),
+ crypto_info: crypto_info.clone(),
+ };
+ let input_data_test_data = InputData {
+ url: test_data_url,
+ hash: "".to_string(),
+ crypto_info: crypto_info.clone(),
+ };
+ let mut input_map = HashMap::new();
+ input_map.insert("if_model".to_string(), input_data_model);
+ input_map.insert("if_data".to_string(), input_data_test_data);
+
+ let result_url = Url::parse(&format!("file:///{}/result.txt.out", fixture_dir)).unwrap();
+ let output_data = OutputData {
+ url: result_url,
+ crypto_info,
+ };
+ let mut output_map = HashMap::new();
+ output_map.insert("of_result".to_string(), output_data);
+ let staged_task = StagedTask::new()
+ .task_id(task_id)
+ .function(&function)
+ .args(arg_list)
+ .input(input_map)
+ .output(output_map);
+
+ // StagedTask => WorkerInvocation
+
+ let function_args = TeaclaveFunctionArguments::new(&staged_task.arg_list);
+
+ let plain_if_model = "fixtures/functions/gbdt_prediction/model.txt";
+ let plain_if_data = "fixtures/functions/gbdt_prediction/test_data.txt";
+ let plain_output = "fixtures/functions/gbdt_prediction/result.txt.out";
+
+ // for (key, value) in staged_task.input_map.iter() {
+ // }
+
+ // for (key, value) in staged_task.output_map.iter() {
+ // }
+
+ let input_files = TeaclaveWorkerFileRegistry::new(hashmap!(
+ "if_model".to_string() =>
+ TeaclaveWorkerInputFileInfo::new(plain_if_model, TeaclaveFileRootKey128::default()),
+ "if_data".to_string() =>
+ TeaclaveWorkerInputFileInfo::new(plain_if_data, TeaclaveFileRootKey128::default())
+ ));
+
+ let output_info =
+ TeaclaveWorkerOutputFileInfo::new(plain_output, TeaclaveFileRootKey128::default());
+ let output_files = TeaclaveWorkerFileRegistry::new(hashmap!(
+ "of_result".to_string() => output_info
+ ));
+
+ let function_name = staged_task.function_name;
+ let function_payload = String::from_utf8_lossy(&staged_task.function_payload).to_string();
+
+ let invocation = WorkerInvocation {
+ runtime_name: "raw-io".to_string(),
+ executor_type: "native".try_into().unwrap(),
+ function_name,
+ function_payload,
+ function_args,
+ input_files,
+ output_files,
+ };
+
+ let worker = Worker::default();
+ let result = worker.invoke_function(invocation);
+ log::debug!("result: {:?}", result);
+ assert!(result.is_ok());
+ log::debug!("summary: {:?}", result.unwrap());
+ }
}
diff --git a/types/src/crypto.rs b/types/src/crypto.rs
index 466361a..954abf4 100644
--- a/types/src/crypto.rs
+++ b/types/src/crypto.rs
@@ -145,6 +145,7 @@ pub enum TeaclaveFileCryptoInfo {
AesGcm128(AesGcm128CryptoInfo),
AesGcm256(AesGcm256CryptoInfo),
TeaclaveFileRootKey128(TeaclaveFileRootKey128),
+ Raw,
}
impl TeaclaveFileCryptoInfo {
@@ -166,6 +167,7 @@ impl TeaclaveFileCryptoInfo {
let info = TeaclaveFileRootKey128::new(key)?;
TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(info)
}
+ "raw" => TeaclaveFileCryptoInfo::Raw,
_ => anyhow::bail!("Invalid crypto schema: {}", schema),
};
Ok(info)
@@ -178,6 +180,7 @@ impl TeaclaveFileCryptoInfo {
TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(_) => {
"teaclave_file_root_key_128".to_string()
}
+ TeaclaveFileCryptoInfo::Raw => "raw".to_string(),
}
}
@@ -188,6 +191,7 @@ impl TeaclaveFileCryptoInfo {
TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(crypto) => {
(crypto.key.to_vec(), Vec::new())
}
+ TeaclaveFileCryptoInfo::Raw => (vec![], vec![]),
}
}
}
diff --git a/types/src/staged_task.rs b/types/src/staged_task.rs
index 20430fa..54fbdd1 100644
--- a/types/src/staged_task.rs
+++ b/types/src/staged_task.rs
@@ -25,6 +25,7 @@ pub struct OutputData {
pub struct StagedTask {
pub task_id: Uuid,
pub function_id: String,
+ pub function_name: String,
pub function_payload: Vec<u8>,
pub arg_list: HashMap<String, String>,
pub input_map: HashMap<String, InputData>,
@@ -72,6 +73,7 @@ impl StagedTask {
pub fn function(self, function: &Function) -> Self {
Self {
function_id: function.external_id(),
+ function_name: function.name.clone(),
function_payload: function.payload.clone(),
..self
}
diff --git a/types/src/worker.rs b/types/src/worker.rs
index 3a107b6..de46505 100644
--- a/types/src/worker.rs
+++ b/types/src/worker.rs
@@ -163,6 +163,7 @@ pub fn convert_encrypted_input_file(
let path = path.as_ref().to_owned();
return Ok(TeaclaveWorkerInputFileInfo::new(path, crypto));
}
+ TeaclaveFileCryptoInfo::Raw => read_all_bytes(path)?,
};
TeaclaveWorkerInputFileInfo::create_with_bytes(dst.as_ref(), &plain_text)
}
@@ -233,7 +234,7 @@ impl TeaclaveFunctionArguments {
pub fn try_get<T: std::str::FromStr>(&self, key: &str) -> anyhow::Result<T> {
self.args
.get(key)
- .ok_or_else(|| anyhow::anyhow!("Cannot find function argument"))
+ .ok_or_else(|| anyhow::anyhow!("Cannot find function argument: {}", key))
.and_then(|s| {
s.parse::<T>()
.map_err(|_| anyhow::anyhow!("parse argument error"))
diff --git a/worker/src/runtime/mod.rs b/worker/src/runtime/mod.rs
index 5fb32c4..948d15b 100644
--- a/worker/src/runtime/mod.rs
+++ b/worker/src/runtime/mod.rs
@@ -29,9 +29,9 @@ pub trait TeaclaveRuntime {
mod default;
pub use default::DefaultRuntime;
-#[cfg(feature = "enclave_unit_test")]
+#[cfg(any(feature = "enclave_unit_test", test_mode))]
mod raw_io;
-#[cfg(feature = "enclave_unit_test")]
+#[cfg(any(feature = "enclave_unit_test", test_mode))]
pub use raw_io::RawIoRuntime;
#[cfg(feature = "enclave_unit_test")]
diff --git a/worker/src/worker.rs b/worker/src/worker.rs
index ff8c4ef..fbf811e 100644
--- a/worker/src/worker.rs
+++ b/worker/src/worker.rs
@@ -122,6 +122,13 @@ fn setup_runtimes() -> HashMap<String, RuntimeBuilder> {
Box::new(runtime::DefaultRuntime::new(input_files, output_files))
}),
);
+ #[cfg(test_mode)]
+ runtimes.insert(
+ "raw-io".to_string(),
+ Box::new(|input_files, output_files| {
+ Box::new(runtime::RawIoRuntime::new(input_files, output_files))
+ }),
+ );
runtimes
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@teaclave.apache.org
For additional commands, e-mail: commits-help@teaclave.apache.org