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
The following commit(s) were added to refs/heads/develop by this push:
new c4e71e2 [execution] A preliminary working implementation to execute
staged task
c4e71e2 is described below
commit c4e71e2505da4e082dd1880370afaae8464d5b4b
Author: Mingshen Sun <[email protected]>
AuthorDate: Fri Mar 20 15:49:10 2020 -0700
[execution] A preliminary working implementation to execute staged task
---
services/execution/enclave/src/lib.rs | 3 +-
services/execution/enclave/src/service.rs | 346 ++++++++++++---------
services/management/enclave/src/service.rs | 2 +-
services/proto/Cargo.toml | 1 +
services/proto/src/teaclave_scheduler_service.rs | 21 +-
services/scheduler/enclave/src/service.rs | 4 +-
tests/fixtures/functions/gbdt_training/train.enc | Bin 0 -> 4096 bytes
.../enclave/src/teaclave_scheduler_service.rs | 6 +-
types/Cargo.toml | 1 +
types/src/crypto.rs | 12 +-
types/src/staged_task.rs | 34 +-
types/src/worker.rs | 18 +-
worker/src/function/echo.rs | 6 +-
worker/src/function/gbdt_training.rs | 3 +
worker/src/runtime/default.rs | 2 +
worker/src/runtime/raw_io.rs | 1 +
16 files changed, 284 insertions(+), 176 deletions(-)
diff --git a/services/execution/enclave/src/lib.rs
b/services/execution/enclave/src/lib.rs
index 49c50f9..bcef69c 100644
--- a/services/execution/enclave/src/lib.rs
+++ b/services/execution/enclave/src/lib.rs
@@ -103,9 +103,8 @@ pub mod tests {
pub fn run_tests() -> bool {
run_tests!(
ocall::tests::test_handle_file_request,
- service::tests::test_invoke_echo_function,
+ service::tests::test_invoke_echo,
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 48430f4..c8b7027 100644
--- a/services/execution/enclave/src/service.rs
+++ b/services/execution/enclave/src/service.rs
@@ -18,14 +18,18 @@
#[cfg(feature = "mesalock_sgx")]
use std::prelude::v1::*;
+use crate::ocall::handle_file_request;
+
+use std::collections::HashMap;
use std::sync::{Arc, SgxMutex as Mutex};
use teaclave_proto::teaclave_scheduler_service::*;
use teaclave_rpc::endpoint::Endpoint;
-use teaclave_types::{StagedTask, TeaclaveFunctionArguments,
WorkerInvocationResult};
+use teaclave_types::{StagedTask, TaskStatus, WorkerInvocation,
WorkerInvocationResult};
use teaclave_worker::Worker;
use anyhow::Result;
+use uuid::Uuid;
#[derive(Clone)]
pub(crate) struct TeaclaveExecutionService {
@@ -76,191 +80,239 @@ impl TeaclaveExecutionService {
}
};
log::debug!("response: {:?}", response);
- let _result = self.invoke_task(response.staged_task);
- // self.update_task(result);
+ let staged_task = response.staged_task;
+ let result = self.invoke_task(&staged_task).unwrap();
+ match self.update_task_status(&staged_task.task_id,
TaskStatus::Finished) {
+ Ok(_) => (),
+ Err(e) => {
+ log::error!("Error: {:?}", e);
+ continue;
+ }
+ }
+ match self.update_task_result(&staged_task.task_id, result) {
+ Ok(_) => (),
+ Err(e) => {
+ log::error!("Error: {:?}", e);
+ continue;
+ }
+ }
}
}
- 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!()
+ fn invoke_task(&mut self, task: &StagedTask) ->
Result<WorkerInvocationResult> {
+ self.update_task_status(&task.task_id, TaskStatus::Running)?;
+ let invocation = prepare_task(&task);
+ let worker = Worker::default();
+ let summary = worker.invoke_function(invocation)?;
+ finalize_task(&task)?;
+ let mut result = WorkerInvocationResult::default();
+ result.return_value = summary.as_bytes().to_vec();
+
+ Ok(result)
+ }
+
+ fn update_task_result(&mut self, task_id: &Uuid, result:
WorkerInvocationResult) -> Result<()> {
+ let request = UpdateTaskResultRequest::new(
+ task_id.to_owned(),
+ &result.return_value,
+ result.output_file_hash,
+ );
+ let _response = self
+ .scheduler_client
+ .clone()
+ .lock()
+ .map_err(|_| anyhow::anyhow!("Cannot lock scheduler client"))?
+ .update_task_result(request)?;
+
+ Ok(())
}
- #[allow(unused)]
- fn update_task(&mut self, _result: WorkerInvocationResult) {
- unimplemented!()
+ fn update_task_status(&mut self, task_id: &Uuid, task_status: TaskStatus)
-> Result<()> {
+ let request = UpdateTaskStatusRequest::new(task_id.to_owned(),
task_status);
+ let _response = self
+ .scheduler_client
+ .clone()
+ .lock()
+ .map_err(|_| anyhow::anyhow!("Cannot lock scheduler client"))?
+ .update_task_status(request)?;
+
+ Ok(())
}
}
-#[cfg(test_mode)]
-mod test_mode {
- use super::*;
+fn finalize_task(task: &StagedTask) -> Result<()> {
+ use std::path::Path;
+ use teaclave_types::*;
+
+ let agent_dir = format!("/tmp/teaclave_agent/{}", task.task_id);
+ let agent_dir_path = Path::new(&agent_dir);
+
+ let mut file_request_info = vec![];
+ for (key, value) in task.output_map.iter() {
+ let mut src = agent_dir_path.to_path_buf();
+ src.push(&format!("{}.out", key));
+ let handle_file_info = HandleFileInfo::new(&src, &value.url);
+ file_request_info.push(handle_file_info);
+ }
+ let request = FileAgentRequest::new(HandleFileCommand::Upload,
file_request_info);
+ handle_file_request(request)?;
+
+ Ok(())
+}
+
+fn prepare_task(task: &StagedTask) -> WorkerInvocation {
+ use std::path::Path;
+ use std::path::PathBuf;
+ use std::untrusted::fs;
+ use std::untrusted::path::PathEx;
+ use teaclave_types::*;
+
+ let runtime_name = "default".to_string();
+ let executor_type = task.executor_type();
+ let function_name = task.function_name.clone();
+ let function_payload =
String::from_utf8_lossy(&task.function_payload).to_string();
+ let function_args = TeaclaveFunctionArguments::new(&task.arg_list);
+
+ let agent_dir = format!("/tmp/teaclave_agent/{}", task.task_id);
+ let agent_dir_path = Path::new(&agent_dir);
+ if !agent_dir_path.exists() {
+ fs::create_dir_all(agent_dir_path).unwrap();
+ }
+
+ let mut input_file_map: HashMap<String, (PathBuf, TeaclaveFileCryptoInfo)>
= HashMap::new();
+ let mut file_request_info = vec![];
+ for (key, value) in task.input_map.iter() {
+ let mut dest = agent_dir_path.to_path_buf();
+ dest.push(&format!("{}.in", key));
+ let info = HandleFileInfo::new(&dest, &value.url);
+ file_request_info.push(info);
+ input_file_map.insert(key.to_string(), (dest, value.crypto_info));
+ }
+ let request = FileAgentRequest::new(HandleFileCommand::Download,
file_request_info);
+ handle_file_request(request).unwrap();
+
+ let mut converted_input_file_map: HashMap<String,
TeaclaveWorkerInputFileInfo> = HashMap::new();
+ for (key, value) in input_file_map.iter() {
+ let (from, crypto_info) = value;
+ let mut dest = from.clone();
+ let mut file_name = dest.file_name().unwrap().to_os_string();
+ file_name.push(".converted");
+ dest.set_file_name(file_name);
+ let input_file_info = convert_encrypted_input_file(from, *crypto_info,
&dest).unwrap();
+ converted_input_file_map.insert(key.to_string(), input_file_info);
+ }
+ let input_files =
TeaclaveWorkerFileRegistry::new(converted_input_file_map);
+
+ let mut output_file_map: HashMap<String, TeaclaveWorkerOutputFileInfo> =
HashMap::new();
+ for (key, value) in task.output_map.iter() {
+ let mut dest = agent_dir_path.to_path_buf();
+ dest.push(&format!("{}.out", key));
+ let crypto = match value.crypto_info {
+ TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(crypto) => crypto,
+ _ => unimplemented!(),
+ };
+ let output_info =
+
TeaclaveWorkerOutputFileInfo::new(dest.to_string_lossy().to_string(), crypto);
+ output_file_map.insert(key.to_string(), output_info);
+ }
+ let output_files = TeaclaveWorkerFileRegistry::new(output_file_map);
+
+ WorkerInvocation {
+ runtime_name,
+ executor_type,
+ function_name,
+ function_payload,
+ function_args,
+ input_files,
+ output_files,
+ }
}
#[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_gbdt_training() {
- let function_args = TeaclaveFunctionArguments::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 plain_input = "fixtures/functions/gbdt_training/train.txt";
- let enc_output = "fixtures/functions/gbdt_training/model.enc.out";
-
- let input_info =
-
TeaclaveWorkerInputFileInfo::create_with_plaintext_file(plain_input).unwrap();
- let input_files = TeaclaveWorkerFileRegistry::new(hashmap!(
- "training_data".to_string() => input_info));
+ pub fn test_invoke_echo() {
+ let task_id = Uuid::new_v4();
+ let mut arg_map = HashMap::new();
+ arg_map.insert("message".to_string(), "Hello, Teaclave!".to_string());
+ let input_map = HashMap::new();
+ let output_map = HashMap::new();
+ let staged_task = StagedTask::new()
+ .task_id(task_id)
+ .function_name("echo")
+ .args(arg_map)
+ .input(input_map)
+ .output(output_map);
- let output_info =
- TeaclaveWorkerOutputFileInfo::new(enc_output,
TeaclaveFileRootKey128::default());
- let output_files = TeaclaveWorkerFileRegistry::new(hashmap!(
- "trained_model".to_string() => output_info));
- let invocation = WorkerInvocation {
- runtime_name: "default".to_string(),
- executor_type: "native".try_into().unwrap(),
- function_name: "gbdt_training".to_string(),
- function_payload: String::new(),
- function_args,
- input_files,
- output_files,
- };
+ let invocation = prepare_task(&staged_task);
let worker = Worker::default();
let result = worker.invoke_function(invocation);
- assert!(result.is_ok());
- log::debug!("summary: {:?}", result.unwrap());
- }
-
- pub fn test_invoke_echo_function() {
- let invocation = WorkerInvocation {
- runtime_name: "default".to_string(),
- executor_type: "native".try_into().unwrap(),
- function_name: "echo".to_string(),
- function_payload: String::new(),
- function_args: TeaclaveFunctionArguments::new(&hashmap!(
- "payload" => "Hello Teaclave!"
- )),
- input_files: TeaclaveWorkerFileRegistry::default(),
- output_files: TeaclaveWorkerFileRegistry::default(),
- };
-
- let worker = Worker::default();
- let result = worker.invoke_function(invocation).unwrap();
- assert_eq!(result, "Hello Teaclave!");
+ if result.is_ok() {
+ finalize_task(&staged_task).unwrap();
+ }
+ assert_eq!(result.unwrap(), "Hello, Teaclave!");
}
- pub fn test_invoke_gbdt_prediction() {
+ pub fn test_invoke_gbdt_training() {
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,
+ 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 fixture_dir = format!(
+ "file:///{}/fixtures/functions/gbdt_training",
+ env!("TEACLAVE_TEST_INSTALL_DIR")
+ );
+ let input_url = Url::parse(&format!("{}/train.enc",
fixture_dir)).unwrap();
+ let output_url = Url::parse(&format!(
+ "{}/model-{}.enc.out",
+ fixture_dir,
+ task_id.to_string()
+ ))
+ .unwrap();
+ let crypto = TeaclaveFileRootKey128::new(&[0; 16]).unwrap();
+ let crypto_info =
TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(crypto);
+
+ let input_data = InputData {
+ url: input_url,
hash: "".to_string(),
- crypto_info: crypto_info.clone(),
+ crypto_info,
};
- 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,
+ url: output_url,
crypto_info,
};
- let mut output_map = HashMap::new();
- output_map.insert("of_result".to_string(), output_data);
+ let input_map = hashmap!("training_data".to_string() => input_data);
+ let output_map = hashmap!("trained_model".to_string() => output_data);
let staged_task = StagedTask::new()
.task_id(task_id)
- .function(&function)
- .args(arg_list)
+ .function_name("gbdt_training")
+ .args(arg_map)
.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 invocation = prepare_task(&staged_task);
let worker = Worker::default();
let result = worker.invoke_function(invocation);
- log::debug!("result: {:?}", result);
+ if result.is_ok() {
+ finalize_task(&staged_task).unwrap();
+ }
+ log::debug!("summary: {:?}", result);
assert!(result.is_ok());
- log::debug!("summary: {:?}", result.unwrap());
}
}
diff --git a/services/management/enclave/src/service.rs
b/services/management/enclave/src/service.rs
index ab55ffe..ae4e717 100644
--- a/services/management/enclave/src/service.rs
+++ b/services/management/enclave/src/service.rs
@@ -770,7 +770,7 @@ pub mod tests {
let input_data = InputData {
url: url.clone(),
hash,
- crypto_info: crypto_info.clone(),
+ crypto_info,
};
let output_data = OutputData { url, crypto_info };
let mut input_map = HashMap::new();
diff --git a/services/proto/Cargo.toml b/services/proto/Cargo.toml
index b1d898d..9ec14d6 100644
--- a/services/proto/Cargo.toml
+++ b/services/proto/Cargo.toml
@@ -23,6 +23,7 @@ rand = { version = "0.7.0" }
serde = { version = "1.0.39", features = ["derive"] }
serde_json = { version = "1.0.39" }
url = { version = "2.1.1" }
+uuid = { version = "0.8.1", features = ["v4"] }
sgx_cov = { version = "1.1.0", optional = true }
sgx_tstd = { version = "1.1.0", features = ["net", "backtrace"], optional
= true }
diff --git a/services/proto/src/teaclave_scheduler_service.rs
b/services/proto/src/teaclave_scheduler_service.rs
index 2677367..f298444 100644
--- a/services/proto/src/teaclave_scheduler_service.rs
+++ b/services/proto/src/teaclave_scheduler_service.rs
@@ -14,6 +14,7 @@ pub use proto::TeaclaveSchedulerRequest;
pub use proto::TeaclaveSchedulerResponse;
use teaclave_rpc::into_request;
use teaclave_types::{StagedTask, TaskStatus};
+use uuid::Uuid;
#[into_request(TeaclaveSchedulerRequest::Subscribe)]
pub struct SubscribeRequest {}
@@ -40,19 +41,19 @@ impl PullTaskResponse {
#[into_request(TeaclaveSchedulerRequest::UpdateTaskResult)]
pub struct UpdateTaskResultRequest {
- pub task_id: String,
+ pub task_id: Uuid,
pub return_value: Vec<u8>,
pub output_file_hash: HashMap<String, String>,
}
impl UpdateTaskResultRequest {
pub fn new(
- task_id: impl Into<String>,
+ task_id: Uuid,
return_value: &[u8],
output_file_hash: HashMap<String, String>,
) -> Self {
Self {
- task_id: task_id.into(),
+ task_id,
return_value: return_value.to_vec(),
output_file_hash,
}
@@ -64,14 +65,14 @@ pub struct UpdateTaskResultResponse {}
#[into_request(TeaclaveSchedulerRequest::UpdateTaskStatus)]
pub struct UpdateTaskStatusRequest {
- pub task_id: String,
+ pub task_id: Uuid,
pub task_status: TaskStatus,
}
impl UpdateTaskStatusRequest {
- pub fn new(task_id: impl Into<String>, task_status: TaskStatus) -> Self {
+ pub fn new(task_id: Uuid, task_status: TaskStatus) -> Self {
Self {
- task_id: task_id.into(),
+ task_id,
task_status,
}
}
@@ -153,7 +154,7 @@ impl std::convert::TryFrom<proto::UpdateTaskResultRequest>
for UpdateTaskResultR
type Error = Error;
fn try_from(proto: proto::UpdateTaskResultRequest) -> Result<Self> {
let ret = Self {
- task_id: proto.task_id,
+ task_id: Uuid::parse_str(&proto.task_id)?,
return_value: proto.return_value,
output_file_hash: proto.output_file_hash,
};
@@ -164,7 +165,7 @@ impl std::convert::TryFrom<proto::UpdateTaskResultRequest>
for UpdateTaskResultR
impl std::convert::From<UpdateTaskResultRequest> for
proto::UpdateTaskResultRequest {
fn from(req: UpdateTaskResultRequest) -> Self {
proto::UpdateTaskResultRequest {
- task_id: req.task_id,
+ task_id: req.task_id.to_string(),
return_value: req.return_value,
output_file_hash: req.output_file_hash,
}
@@ -190,7 +191,7 @@ impl std::convert::TryFrom<proto::UpdateTaskStatusRequest>
for UpdateTaskStatusR
fn try_from(proto: proto::UpdateTaskStatusRequest) -> Result<Self> {
let task_status = i32_to_task_status(proto.task_status)?;
let ret = Self {
- task_id: proto.task_id,
+ task_id: Uuid::parse_str(&proto.task_id)?,
task_status,
};
Ok(ret)
@@ -201,7 +202,7 @@ impl std::convert::From<UpdateTaskStatusRequest> for
proto::UpdateTaskStatusRequ
fn from(req: UpdateTaskStatusRequest) -> Self {
let task_status = i32_from_task_status(req.task_status);
proto::UpdateTaskStatusRequest {
- task_id: req.task_id,
+ task_id: req.task_id.to_string(),
task_status,
}
}
diff --git a/services/scheduler/enclave/src/service.rs
b/services/scheduler/enclave/src/service.rs
index 5a843d7..5a4a5fe 100644
--- a/services/scheduler/enclave/src/service.rs
+++ b/services/scheduler/enclave/src/service.rs
@@ -161,7 +161,7 @@ impl TeaclaveScheduler for TeaclaveSchedulerService {
request: Request<UpdateTaskStatusRequest>,
) -> TeaclaveServiceResponseResult<UpdateTaskStatusResponse> {
let request = request.message;
- let mut task = self.get_task(&request.task_id)?;
+ let mut task = self.get_task(&request.task_id.to_string())?;
task.status = request.task_status;
self.put_task(&task)?;
Ok(UpdateTaskStatusResponse {})
@@ -172,7 +172,7 @@ impl TeaclaveScheduler for TeaclaveSchedulerService {
request: Request<UpdateTaskResultRequest>,
) -> TeaclaveServiceResponseResult<UpdateTaskResultResponse> {
let request = request.message;
- let mut task = self.get_task(&request.task_id)?;
+ let mut task = self.get_task(&request.task_id.to_string())?;
task.return_value = Some(request.return_value);
task.output_file_hash = request.output_file_hash;
self.put_task(&task)?;
diff --git a/tests/fixtures/functions/gbdt_training/train.enc
b/tests/fixtures/functions/gbdt_training/train.enc
new file mode 100644
index 0000000..2bb15f7
Binary files /dev/null and b/tests/fixtures/functions/gbdt_training/train.enc
differ
diff --git a/tests/functional/enclave/src/teaclave_scheduler_service.rs
b/tests/functional/enclave/src/teaclave_scheduler_service.rs
index 5ef46d1..7a6caa1 100644
--- a/tests/functional/enclave/src/teaclave_scheduler_service.rs
+++ b/tests/functional/enclave/src/teaclave_scheduler_service.rs
@@ -57,14 +57,14 @@ fn test_update_task_status_result() {
let request = PullTaskRequest {};
let response = client.pull_task(request).unwrap();
log::debug!("response: {:?}", response);
- let task_id = response.staged_task.task_id.to_string();
+ let task_id = response.staged_task.task_id;
- let request = UpdateTaskStatusRequest::new(&task_id, TaskStatus::Finished);
+ let request = UpdateTaskStatusRequest::new(task_id, TaskStatus::Finished);
let response = client.update_task_status(request);
assert!(response.is_ok());
let request =
- UpdateTaskResultRequest::new(&task_id,
"return".to_string().as_bytes(), HashMap::new());
+ UpdateTaskResultRequest::new(task_id, "return".to_string().as_bytes(),
HashMap::new());
let response = client.update_task_result(request);
assert!(response.is_ok());
diff --git a/types/Cargo.toml b/types/Cargo.toml
index f590ca8..8c74088 100644
--- a/types/Cargo.toml
+++ b/types/Cargo.toml
@@ -19,6 +19,7 @@ enclave_unit_test = ["teaclave_test_utils/mesalock_sgx"]
[dependencies]
protected_fs_rs = { path = "../common/protected_fs_rs", default-features =
false}
+log = { version = "0.4.6" }
anyhow = { version = "1.0.26" }
sgx_types = { version = "1.1.0" }
rand = { version = "0.7.0" }
diff --git a/types/src/crypto.rs b/types/src/crypto.rs
index 954abf4..bba434d 100644
--- a/types/src/crypto.rs
+++ b/types/src/crypto.rs
@@ -10,7 +10,7 @@ use std::format;
const AES_GCM_256_KEY_LENGTH: usize = 32;
const AES_GCM_256_IV_LENGTH: usize = 12;
-#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
+#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct AesGcm256CryptoInfo {
pub key: [u8; AES_GCM_256_KEY_LENGTH],
pub iv: [u8; AES_GCM_256_IV_LENGTH],
@@ -61,7 +61,7 @@ impl Default for AesGcm256CryptoInfo {
const AES_GCM_128_KEY_LENGTH: usize = 16;
const AES_GCM_128_IV_LENGTH: usize = 12;
-#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
+#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct AesGcm128CryptoInfo {
pub key: [u8; AES_GCM_128_KEY_LENGTH],
pub iv: [u8; AES_GCM_128_IV_LENGTH],
@@ -113,12 +113,16 @@ impl Default for AesGcm128CryptoInfo {
const TEACLAVE_FILE_ROOT_KEY_128_LENGTH: usize = 16;
-#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
+#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct TeaclaveFileRootKey128 {
pub key: [u8; TEACLAVE_FILE_ROOT_KEY_128_LENGTH],
}
impl TeaclaveFileRootKey128 {
+ pub fn random() -> Self {
+ Self::default()
+ }
+
pub fn new(in_key: &[u8]) -> anyhow::Result<Self> {
anyhow::ensure!(
in_key.len() == TEACLAVE_FILE_ROOT_KEY_128_LENGTH,
@@ -140,7 +144,7 @@ impl Default for TeaclaveFileRootKey128 {
}
}
-#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
+#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq)]
pub enum TeaclaveFileCryptoInfo {
AesGcm128(AesGcm128CryptoInfo),
AesGcm256(AesGcm256CryptoInfo),
diff --git a/types/src/staged_task.rs b/types/src/staged_task.rs
index 54fbdd1..e38912d 100644
--- a/types/src/staged_task.rs
+++ b/types/src/staged_task.rs
@@ -1,4 +1,7 @@
-use crate::{Function, Storable, TeaclaveFileCryptoInfo, TeaclaveInputFile,
TeaclaveOutputFile};
+use crate::{
+ Function, Storable, TeaclaveExecutorSelector, TeaclaveFileCryptoInfo,
TeaclaveInputFile,
+ TeaclaveOutputFile,
+};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::prelude::v1::*;
@@ -79,6 +82,27 @@ impl StagedTask {
}
}
+ pub fn function_id(self, function_id: impl Into<String>) -> Self {
+ Self {
+ function_id: function_id.into(),
+ ..self
+ }
+ }
+
+ pub fn function_name(self, function_name: impl Into<String>) -> Self {
+ Self {
+ function_name: function_name.into(),
+ ..self
+ }
+ }
+
+ pub fn function_payload(self, function_payload: impl Into<Vec<u8>>) ->
Self {
+ Self {
+ function_payload: function_payload.into(),
+ ..self
+ }
+ }
+
pub fn args(self, args: HashMap<String, String>) -> Self {
Self {
arg_list: args,
@@ -103,4 +127,12 @@ impl StagedTask {
pub fn get_queue_key() -> &'static str {
QUEUE_KEY
}
+
+ pub fn executor_type(&self) -> TeaclaveExecutorSelector {
+ if self.function_payload.is_empty() {
+ TeaclaveExecutorSelector::Native
+ } else {
+ TeaclaveExecutorSelector::Python
+ }
+ }
}
diff --git a/types/src/worker.rs b/types/src/worker.rs
index de46505..502e265 100644
--- a/types/src/worker.rs
+++ b/types/src/worker.rs
@@ -87,6 +87,8 @@ impl TeaclaveWorkerInputFileInfo {
}
pub fn get_readable_io(&self) -> anyhow::Result<Box<dyn io::Read>> {
+ log::debug!("path: {:?}", self.path);
+ log::debug!("key: {:?}", self.crypto_info.key);
let f = ProtectedFile::open_ex(&self.path, &self.crypto_info.key)?;
Ok(Box::new(f))
}
@@ -104,7 +106,8 @@ impl TeaclaveWorkerInputFileInfo {
path: impl AsRef<std::path::Path>,
bytes: &[u8],
) -> anyhow::Result<TeaclaveWorkerInputFileInfo> {
- let crypto = TeaclaveFileRootKey128::default();
+ // let crypto = TeaclaveFileRootKey128::default();
+ let crypto = TeaclaveFileRootKey128::new(&[0; 16]).unwrap();
let mut f = ProtectedFile::create_ex(&path, &crypto.key)?;
f.write_all(bytes)?;
Ok(Self::new(path.as_ref(), crypto))
@@ -123,6 +126,8 @@ impl TeaclaveWorkerOutputFileInfo {
}
pub fn get_writable_io(&self) -> anyhow::Result<Box<dyn io::Write>> {
+ log::debug!("path: {:?}", self.path);
+ log::debug!("key: {:?}", self.crypto_info.key);
let f = ProtectedFile::create_ex(&self.path, &self.crypto_info.key)?;
Ok(Box::new(f))
}
@@ -148,6 +153,11 @@ pub fn convert_encrypted_input_file(
crypto_info: TeaclaveFileCryptoInfo,
dst: impl AsRef<std::path::Path>,
) -> anyhow::Result<TeaclaveWorkerInputFileInfo> {
+ log::debug!("from: {:?}, to: {:?}", path.as_ref(), dst.as_ref());
+ #[cfg(not(feature = "mesalock_sgx"))]
+ use std::fs;
+ #[cfg(feature = "mesalock_sgx")]
+ use std::untrusted::fs;
let plain_text = match crypto_info {
TeaclaveFileCryptoInfo::AesGcm128(crypto) => {
let mut bytes = read_all_bytes(path)?;
@@ -160,8 +170,9 @@ pub fn convert_encrypted_input_file(
bytes
}
TeaclaveFileCryptoInfo::TeaclaveFileRootKey128(crypto) => {
- let path = path.as_ref().to_owned();
- return Ok(TeaclaveWorkerInputFileInfo::new(path, crypto));
+ fs::copy(path, dst.as_ref())?;
+ let dst = dst.as_ref().to_owned();
+ return Ok(TeaclaveWorkerInputFileInfo::new(dst, crypto));
}
TeaclaveFileCryptoInfo::Raw => read_all_bytes(path)?,
};
@@ -268,6 +279,7 @@ pub struct WorkerInvocation {
pub output_files: TeaclaveWorkerFileRegistry<TeaclaveWorkerOutputFileInfo>,
}
+#[derive(Default)]
pub struct WorkerInvocationResult {
pub return_value: Vec<u8>,
pub output_file_hash: HashMap<String, String>,
diff --git a/worker/src/function/echo.rs b/worker/src/function/echo.rs
index 359980f..7857c72 100644
--- a/worker/src/function/echo.rs
+++ b/worker/src/function/echo.rs
@@ -32,8 +32,8 @@ impl TeaclaveFunction for Echo {
_runtime: Box<dyn TeaclaveRuntime + Send + Sync>,
args: TeaclaveFunctionArguments,
) -> anyhow::Result<String> {
- let payload: String = args.try_get("payload")?;
- Ok(payload)
+ let message: String = args.try_get("message")?;
+ Ok(message)
}
}
@@ -55,7 +55,7 @@ pub mod tests {
fn test_echo() {
let func_args = TeaclaveFunctionArguments::new(&hashmap!(
- "payload" => "Hello Teaclave!"
+ "message" => "Hello Teaclave!"
));
let input_files = TeaclaveWorkerFileRegistry::default();
diff --git a/worker/src/function/gbdt_training.rs
b/worker/src/function/gbdt_training.rs
index e7c48cd..340fde8 100644
--- a/worker/src/function/gbdt_training.rs
+++ b/worker/src/function/gbdt_training.rs
@@ -44,6 +44,7 @@ impl TeaclaveFunction for GbdtTraining {
runtime: Box<dyn TeaclaveRuntime + Send + Sync>,
args: TeaclaveFunctionArguments,
) -> anyhow::Result<String> {
+ log::debug!("start traning...");
let feature_size: usize = args.try_get("feature_size")?;
let max_depth: u32 = args.try_get("max_depth")?;
let iterations: usize = args.try_get("iterations")?;
@@ -54,6 +55,7 @@ impl TeaclaveFunction for GbdtTraining {
let loss: String = args.try_get("loss")?;
let training_optimization_level: u8 =
args.try_get("training_optimization_level")?;
+ log::debug!("open input...");
// read input
let training_file = runtime.open_input(IN_DATA)?;
let mut train_dv = parse_training_data(training_file, feature_size)?;
@@ -77,6 +79,7 @@ impl TeaclaveFunction for GbdtTraining {
gbdt_train_mod.fit(&mut train_dv);
let model_json = serde_json::to_string(&gbdt_train_mod)?;
+ log::debug!("create file...");
// save the model to output
let mut model_file = runtime.create_output(OUT_MODEL)?;
model_file.write_all(model_json.as_bytes())?;
diff --git a/worker/src/runtime/default.rs b/worker/src/runtime/default.rs
index 21adeae..00e1ce3 100644
--- a/worker/src/runtime/default.rs
+++ b/worker/src/runtime/default.rs
@@ -51,6 +51,7 @@ impl TeaclaveRuntime for DefaultRuntime {
.get(identifier)
.ok_or_else(|| anyhow::anyhow!("Invalid input file identifier."))?;
+ log::debug!("open_input: {:?}", file_info.path);
let readable = file_info.get_readable_io()?;
Ok(readable)
}
@@ -62,6 +63,7 @@ impl TeaclaveRuntime for DefaultRuntime {
.get(identifier)
.ok_or_else(|| anyhow::anyhow!("Invalide output file
identifier"))?;
+ log::debug!("create_output: {:?}", file_info.path);
let writable = file_info.get_writable_io()?;
Ok(writable)
}
diff --git a/worker/src/runtime/raw_io.rs b/worker/src/runtime/raw_io.rs
index 7b64a97..50d5a29 100644
--- a/worker/src/runtime/raw_io.rs
+++ b/worker/src/runtime/raw_io.rs
@@ -52,6 +52,7 @@ impl TeaclaveRuntime for RawIoRuntime {
.entries
.get(identifier)
.ok_or_else(|| anyhow::anyhow!("Invalid input file identifier."))?;
+ log::debug!("open_input: {:?}", file_info.path);
let f = File::open(&file_info.path)?;
Ok(Box::new(f))
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
