teaclave-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From ms...@apache.org
Subject [incubator-teaclave] 01/02: [execution] Have a simple flow of function invocation
Date Wed, 18 Mar 2020 21:14:05 GMT
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 <bob@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


Mime
View raw message