move gRPC initialization to where it's used

Closes #56
tf-encrypted · Oct 26, 2020 · 43a597e · 43a597e
1 parent c04b476
commit 43a597e
Show file tree

Hide file tree

Showing 7 changed files with 3 additions and 16 deletions.
diff --git a/examples/keras/worker.py b/examples/keras/worker.py
@@ -2,8 +2,6 @@
 import asyncio
 import logging
 
-from grpc.experimental import aio
-
 from moose.cluster.cluster_spec import load_cluster_spec
 from moose.logger import get_logger
 from moose.worker import Worker
@@ -20,8 +18,6 @@
     get_logger().setLevel(level=logging.DEBUG)
 
 if __name__ == "__main__":
-    aio.init_grpc_aio()
-
     get_logger().info(f"Starting on {args.host}:{args.port}")
     cluster_spec = load_cluster_spec(args.cluster_spec)
     worker = Worker(args.name, args.host, args.port, cluster_spec)

diff --git a/examples/mp-spdz/main.py b/examples/mp-spdz/main.py
@@ -1,8 +1,6 @@
 import argparse
 import logging
 
-from grpc.experimental import aio
-
 from moose.compiler.edsl import HostPlacement
 from moose.compiler.edsl import add
 from moose.compiler.edsl import computation
@@ -77,7 +75,6 @@ def my_comp():
 concrete_comp = my_comp.trace_func()
 
 if __name__ == "__main__":
-    aio.init_grpc_aio()
     runtime = RemoteRuntime(args.cluster_spec)
     runtime.evaluate_computation(
         computation=concrete_comp,

diff --git a/examples/mp-spdz/worker.py b/examples/mp-spdz/worker.py
@@ -2,8 +2,6 @@
 import asyncio
 import logging
 
-from grpc.experimental import aio
-
 from moose.cluster.cluster_spec import load_cluster_spec
 from moose.logger import get_logger
 from moose.worker import Worker
@@ -20,8 +18,6 @@
     get_logger().setLevel(level=logging.DEBUG)
 
 if __name__ == "__main__":
-    aio.init_grpc_aio()
-
     get_logger().info(f"Starting on {args.host}:{args.port}")
     cluster_spec = load_cluster_spec(args.cluster_spec)
     worker = Worker(args.name, args.host, args.port, cluster_spec)

diff --git a/examples/python-functions/main.py b/examples/python-functions/main.py
@@ -2,8 +2,6 @@
 import logging
 import os
 
-from grpc.experimental import aio
-
 from moose.compiler.edsl import HostPlacement
 from moose.compiler.edsl import add
 from moose.compiler.edsl import computation
@@ -63,7 +61,6 @@ def my_comp():
 concrete_comp = my_comp.trace_func()
 
 if __name__ == "__main__":
-    aio.init_grpc_aio()
     if args.runtime == "test":
         runtime = TestRuntime(workers=concrete_comp.devices())
     elif args.runtime == "remote":

diff --git a/moose/channels/grpc.py b/moose/channels/grpc.py
@@ -10,6 +10,7 @@ class Channel:
     def __init__(self, endpoint, buffer, ca_cert, ident_cert, ident_key):
         self._buffer = buffer
 
+        aio.init_grpc_aio()
         if ca_cert:
             credentials = grpc.ssl_channel_credentials(
                 root_certificates=ca_cert,

diff --git a/moose/executor/proxy.py b/moose/executor/proxy.py
@@ -7,6 +7,7 @@
 
 class RemoteExecutor:
     def __init__(self, endpoint, ca_cert, ident_cert, ident_key):
+        aio.init_grpc_aio()
         if ca_cert:
             credentials = grpc.ssl_channel_credentials(
                 root_certificates=ca_cert,

diff --git a/moose/worker.py b/moose/worker.py
@@ -79,8 +79,6 @@ def __init__(
         ident_key_filename,
         allow_insecure_networking=False,
     ):
-        aio.init_grpc_aio()
-
         ca_cert = load_certificate(ca_cert_filename)
         ident_cert = load_certificate(ident_cert_filename)
         ident_key = load_certificate(ident_key_filename)
@@ -93,6 +91,7 @@ def __init__(
         executor = AsyncExecutor(name=name, channel_manager=channel_manager)
 
         # set up server
+        aio.init_grpc_aio()
         # self._server = aio.server(interceptors=(MyInterceptor(),))
         self._server = aio.server()