Add GrpcRuntime to PyMoose (#1083)

modules/Cargo.toml

@@ -21,7 +21,6 @@ notify = "4.0"
 prost = "0.9"
 serde = { version = "~1.0", features = ["derive"] }
 serde_json = "1.0"
-tokio = { version = "1.0", features = ["macros", "rt-multi-thread"] }
 toml = "0.5"
 tonic = { version = "~0.6", features = ["tls"] }
 tracing = { version="0.1", features = ["log"] }

modules/src/choreography/filesystem.rs

@@ -2,6 +2,7 @@ use super::parse_session_config_file_with_computation;
 use crate::choreography::{NetworkingStrategy, StorageStrategy};
 use crate::execution::ExecutionContext;
 use moose::prelude::*;
+use moose::tokio;
 use notify::{DebouncedEvent, Watcher};
 use std::path::Path;
 

modules/src/choreography/grpc.rs

@@ -15,6 +15,7 @@ use dashmap::mapref::entry::Entry;
 use dashmap::DashMap;
 use moose::computation::{SessionId, Value};
 use moose::execution::Identity;
+use moose::tokio;
 use std::collections::HashMap;
 use std::sync::Arc;
 

modules/src/networking/tcpstream.rs

@@ -1,4 +1,5 @@
 use async_trait::async_trait;
+use moose::tokio;
 use moose::{
     computation::{RendezvousKey, SessionId, Value},
     error::Error,

modules/src/storage/csv.rs

@@ -195,6 +195,7 @@ where
 #[cfg(test)]
 mod tests {
     use super::*;
+    use moose::tokio;
     use ndarray::array;
     use std::io::Write;
     use tempfile::NamedTempFile;

modules/src/storage/local_file_storage.rs

@@ -81,6 +81,7 @@ pub fn parse_columns(query: &str) -> Result<Vec<String>> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use moose::tokio;
     use moose::types::HostFloat64Tensor;
     use ndarray::array;
     use std::convert::TryFrom;

modules/src/storage/numpy.rs

@@ -301,6 +301,7 @@ fn extract_dtype(npy_filename: &str) -> Result<Ty> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use moose::tokio;
     use ndarray::array;
     use std::io::Write;
     use tempfile::NamedTempFile;

moose/src/lib.rs

@@ -2795,3 +2795,4 @@ pub mod types;
 
 #[doc(inline)]
 pub use error::{Error, Result};
+pub use tokio;

pymoose/Cargo.toml

@@ -20,6 +20,7 @@ approx = "~0.5"
 numpy = "~0.16"
 ndarray = "~0.15"
 moose = { path = "../moose" }
+moose-modules = { path = "../modules" }
 pyo3 = "~0.16"
 rmp-serde = "~1.0"
 serde = { version = "~1.0", features = ["derive"] }

pymoose/examples/grpc/README.md

@@ -0,0 +1,11 @@
+## Running
+
+Make sure to have three gRPC workers at the right endpoints before running this example.
+
+This can be done using Comet as follows (run each command in its own terminal):
+
+```sh
+$ cargo run --bin comet -- --identity localhost:50000 --port 50000
+$ cargo run --bin comet -- --identity localhost:50001 --port 50001
+$ cargo run --bin comet -- --identity localhost:50002 --port 50002
+```

pymoose/examples/grpc/grpc.py

@@ -0,0 +1,45 @@
+import argparse
+import logging
+
+import numpy as np
+
+import pymoose as pm
+from pymoose.logger import get_logger
+
+
+@pm.computation
+def my_computation():
+    alice = pm.host_placement("alice")
+    bob = pm.host_placement("bob")
+    carole = pm.host_placement("carole")
+
+    with alice:
+        x = pm.constant(np.array([1.0, 2.0], dtype=np.float64))
+
+    with bob:
+        y = pm.constant(np.array([3.0, 4.0], dtype=np.float64))
+
+    with carole:
+        z = pm.add(x, y)
+
+    return z
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Example")
+    parser.add_argument("--verbose", action="store_true")
+    args = parser.parse_args()
+
+    if args.verbose:
+        get_logger().setLevel(level=logging.DEBUG)
+
+    runtime = pm.GrpcMooseRuntime(
+        {
+            "alice": "localhost:50000",
+            "bob": "localhost:50001",
+            "carole": "localhost:50002",
+        }
+    )
+
+    results = runtime.evaluate_computation(my_computation)
+    print(results)

pymoose/examples/linear-regression/linreg_test.py

@@ -9,7 +9,6 @@
 import pymoose as pm
 from pymoose.computation import utils
 from pymoose.logger import get_logger
-from pymoose.testing import LocalMooseRuntime
 
 FIXED = pm.fixed(8, 27)
 # Rust compiler currently supports only limited set of alternative precisions:
@@ -132,7 +131,7 @@ def _linear_regression_eval(self, metric_name):
             "y-owner": {"y_data": y_data},
             "model-owner": {},
         }
-        runtime = LocalMooseRuntime(storage_mapping=executors_storage)
+        runtime = pm.LocalMooseRuntime(storage_mapping=executors_storage)
         traced = pm.trace(linear_comp)
         _ = runtime.evaluate_computation(
             computation=traced,

pymoose/examples/logistic-regression/logistic_regression_test.py

@@ -8,7 +8,6 @@
 import pymoose as pm
 from pymoose.computation import utils
 from pymoose.logger import get_logger
-from pymoose.testing import LocalMooseRuntime
 
 
 class ReplicatedExample(parameterized.TestCase):
@@ -61,7 +60,7 @@ def test_logistic_regression_example_execute(self):
             "bob": {},
             "carole": {},
         }
-        runtime = LocalMooseRuntime(storage_mapping=storage)
+        runtime = pm.LocalMooseRuntime(storage_mapping=storage)
         _ = runtime.evaluate_computation(
             computation=traced_model_comp,
             role_assignment={"alice": "alice", "bob": "bob", "carole": "carole"},

pymoose/examples/replicated/aes_test.py

@@ -9,7 +9,6 @@
 import pymoose as pm
 from pymoose.computation import utils
 from pymoose.logger import get_logger
-from pymoose.testing import LocalMooseRuntime
 
 
 class ReplicatedExample(parameterized.TestCase):
@@ -62,7 +61,7 @@ def test_aes_example_execute(self, host_decrypt):
             "bob": {},
             "carole": {},
         }
-        runtime = LocalMooseRuntime(storage_mapping=storage)
+        runtime = pm.LocalMooseRuntime(storage_mapping=storage)
         _ = runtime.evaluate_computation(
             computation=traced_aes_comp,
             role_assignment={"alice": "alice", "bob": "bob", "carole": "carole"},

pymoose/examples/replicated/division_test.py

@@ -6,7 +6,6 @@
 
 import pymoose as pm
 from pymoose.logger import get_logger
-from pymoose.testing import LocalMooseRuntime
 
 FIXED = pm.fixed(14, 23)
 
@@ -46,7 +45,7 @@ def my_comp():
             "carole": {},
             "dave": {},
         }
-        runtime = LocalMooseRuntime(storage_mapping=executors_storage)
+        runtime = pm.LocalMooseRuntime(storage_mapping=executors_storage)
 
         logical_comp = pm.trace(my_comp)
         runtime.evaluate_computation(

pymoose/examples/replicated/identity_test.py

@@ -8,7 +8,6 @@
 import pymoose as pm
 from pymoose.computation import utils
 from pymoose.logger import get_logger
-from pymoose.testing import LocalMooseRuntime
 
 
 class TensorIdentityExample(parameterized.TestCase):
@@ -108,7 +107,7 @@ def test_identity_example_execute(self, f, t, e):
             "bob-1": {},
             "carole-1": {},
         }
-        runtime = LocalMooseRuntime(storage_mapping=storage)
+        runtime = pm.LocalMooseRuntime(storage_mapping=storage)
         result_dict = runtime.evaluate_computation(
             computation=traced_identity_comp,
             role_assignment={

pymoose/pymoose/__init__.py

@@ -64,6 +64,8 @@
 from pymoose.edsl.tracer import trace_and_compile
 from pymoose.rust import elk_compiler
 from pymoose.rust import moose_runtime as _moose_runtime
+from pymoose.testing import GrpcMooseRuntime
+from pymoose.testing import LocalMooseRuntime
 
 MooseComputation = _moose_runtime.MooseComputation
 
@@ -93,6 +95,7 @@
     FloatType,
     host_placement,
     greater,
+    GrpcMooseRuntime,
     identity,
     index_axis,
     int32,
@@ -101,6 +104,7 @@
     IntType,
     less,
     load,
+    LocalMooseRuntime,
     log,
     log2,
     logical_or,

pymoose/pymoose/predictors/tree_ensemble_test.py

@@ -14,7 +14,6 @@
 from pymoose.logger import get_logger
 from pymoose.predictors import predictor_utils
 from pymoose.predictors import tree_ensemble
-from pymoose.testing import LocalMooseRuntime
 
 _XGB_REGRESSOR_MODELS = [("xgboost_regressor", [14.121551, 14.121551, 113.279236])]
 _SK_REGRESSOR_MODELS = [
@@ -165,7 +164,7 @@ def test_tree_ensemble_logic(self, test_case, predictor_cls):
 
         traced_model_comp = pm.trace(predictor_logic)
         storage = {plc.name: {} for plc in predictor.host_placements}
-        runtime = LocalMooseRuntime(storage_mapping=storage)
+        runtime = pm.LocalMooseRuntime(storage_mapping=storage)
         role_assignment = {plc.name: plc.name for plc in predictor.host_placements}
         result_dict = runtime.evaluate_computation(
             computation=traced_model_comp,

pymoose/pymoose/testing.py

@@ -1,4 +1,6 @@
 from pymoose.computation import utils
+from pymoose.edsl.base import AbstractComputation
+from pymoose.edsl.tracer import trace
 from pymoose.rust import moose_runtime
 
 
@@ -41,3 +43,22 @@ def read_value_from_storage(self, identity, key):
 
     def write_value_to_storage(self, identity, key, value):
         return super().write_value_to_storage(identity, key, value)
+
+
+class GrpcMooseRuntime(moose_runtime.GrpcRuntime):
+    def __new__(cls, role_assignment):
+        return moose_runtime.GrpcRuntime.__new__(GrpcMooseRuntime, role_assignment)
+
+    def evaluate_computation(
+        self,
+        computation,
+        arguments=None,
+    ):
+        if isinstance(computation, AbstractComputation):
+            computation = trace(computation)
+
+        if arguments is None:
+            arguments = {}
+
+        comp_bin = utils.serialize_computation(computation)
+        return super().evaluate_computation(comp_bin, arguments)

pymoose/rust_integration_tests/add_n_test.py

@@ -9,7 +9,6 @@
 from pymoose import elk_compiler
 from pymoose.computation import utils
 from pymoose.logger import get_logger
-from pymoose.testing import LocalMooseRuntime
 
 player0 = pm.host_placement("player0")
 player1 = pm.host_placement("player1")
@@ -53,7 +52,7 @@ def my_comp():
             "player1": {},
             "player2": {},
         }
-        runtime = LocalMooseRuntime(storage_mapping=executors_storage)
+        runtime = pm.LocalMooseRuntime(storage_mapping=executors_storage)
         concrete_comp = pm.trace(my_comp)
 
         comp_bin = utils.serialize_computation(concrete_comp)

pymoose/rust_integration_tests/argmax_test.py

@@ -8,7 +8,6 @@
 import pymoose as pm
 from pymoose.computation import types as ty
 from pymoose.logger import get_logger
-from pymoose.testing import LocalMooseRuntime
 
 
 class ArgmaxExample(parameterized.TestCase):
@@ -77,7 +76,7 @@ def test_example_execute(self, x, axis, axis_idx_max):
             "bob": {"x_arg": x_arg},
         }
 
-        runtime = LocalMooseRuntime(storage_mapping=storage)
+        runtime = pm.LocalMooseRuntime(storage_mapping=storage)
         _ = runtime.evaluate_computation(
             computation=traced_less_comp,
             role_assignment={"alice": "alice", "bob": "bob", "carole": "carole"},

pymoose/rust_integration_tests/boolean_ops_test.py

@@ -8,7 +8,6 @@
 import pymoose as pm
 from pymoose.computation import types as ty
 from pymoose.logger import get_logger
-from pymoose.testing import LocalMooseRuntime
 
 
 class BooleanLogicExample(parameterized.TestCase):
@@ -79,7 +78,7 @@ def test_bool_example_execute(self, x, y):
             "bob": {"x_arg": x, "y_arg": y},
         }
 
-        runtime = LocalMooseRuntime(storage_mapping=storage)
+        runtime = pm.LocalMooseRuntime(storage_mapping=storage)
         _ = runtime.evaluate_computation(
             computation=traced_less_comp,
             role_assignment={"alice": "alice", "bob": "bob", "carole": "carole"},

pymoose/rust_integration_tests/concat_test.py

@@ -8,7 +8,6 @@
 import pymoose as pm
 from pymoose.computation import utils
 from pymoose.logger import get_logger
-from pymoose.testing import LocalMooseRuntime
 
 player0 = pm.host_placement("player0")
 player1 = pm.host_placement("player1")
@@ -52,7 +51,7 @@ def my_comp():
             "player1": {},
             "player2": {},
         }
-        runtime = LocalMooseRuntime(storage_mapping=executors_storage)
+        runtime = pm.LocalMooseRuntime(storage_mapping=executors_storage)
         concrete_comp = pm.trace(my_comp)
 
         comp_bin = utils.serialize_computation(concrete_comp)

pymoose/rust_integration_tests/dtype_conversions_test.py

@@ -7,7 +7,6 @@
 
 import pymoose as pm
 from pymoose.logger import get_logger
-from pymoose.testing import LocalMooseRuntime
 
 
 class DTypeConversionTest(parameterized.TestCase):
@@ -94,7 +93,7 @@ def test_host_dtype_conversions(self, x_array, from_dtype, to_dtype):
             "bob": {},
             "carole": {},
         }
-        runtime = LocalMooseRuntime(storage_mapping=storage)
+        runtime = pm.LocalMooseRuntime(storage_mapping=storage)
         _ = runtime.evaluate_computation(
             computation=traced_comp,
             role_assignment={"alice": "alice", "bob": "bob", "carole": "carole"},

pymoose/rust_integration_tests/exp_test.py

@@ -7,7 +7,6 @@
 
 import pymoose as pm
 from pymoose.logger import get_logger
-from pymoose.testing import LocalMooseRuntime
 
 
 class ReplicatedExample(parameterized.TestCase):
@@ -66,7 +65,7 @@ def test_exp_example_execute(self, x):
             "bob": {},
             "carole": {},
         }
-        runtime = LocalMooseRuntime(storage_mapping=storage)
+        runtime = pm.LocalMooseRuntime(storage_mapping=storage)
         _ = runtime.evaluate_computation(
             computation=traced_exp_comp,
             role_assignment={"alice": "alice", "bob": "bob", "carole": "carole"},
@@ -90,7 +89,7 @@ def test_float_exp_execute(self, x, moose_dtype):
             "bob": {},
             "carole": {},
         }
-        runtime = LocalMooseRuntime(storage_mapping=storage)
+        runtime = pm.LocalMooseRuntime(storage_mapping=storage)
         _ = runtime.evaluate_computation(
             computation=traced_exp_comp,
             role_assignment={"alice": "alice", "bob": "bob", "carole": "carole"},