Revert unintended modifications.

fed/_private/fed_call_holder.py

@@ -24,17 +24,6 @@
 # Set config in the very beginning to avoid being overwritten by other packages.
 logging.basicConfig(level=logging.INFO)
 
-from fed._private.global_context import get_global_context
-from fed.fed_object import FedObject
-from fed.proxy.barriers import send
-from fed.utils import resolve_dependencies
-
-try:
-    from jax.tree_util import tree_flatten
-except ImportError:
-    from fed.tree_util import tree_flatten
-
-import fed.config as fed_config
 
 logger = logging.getLogger(__name__)
 

fed/_private/global_context.py

@@ -76,11 +76,15 @@ def acquire_shutdown_flag(self) -> bool:
 
 
 def init_global_context(
-    current_party: str, job_name: str, sending_failure_handler: Callable[[], None] = None
+    current_party: str,
+    job_name: str,
+    sending_failure_handler: Callable[[], None] = None,
 ) -> None:
     global _global_context
     if _global_context is None:
-        _global_context = GlobalContext(job_name, current_party, sending_failure_handler)
+        _global_context = GlobalContext(
+            job_name, current_party, sending_failure_handler
+        )
 
 
 def get_global_context():

fed/api.py

@@ -168,7 +168,9 @@ def init(
 
     fed_utils.validate_addresses(addresses)
     init_global_context(
-        current_party=party, job_name=job_name, failure_handler=sending_failure_handler
+        current_party=party,
+        job_name=job_name,
+        sending_failure_handler=sending_failure_handler,
     )
     tls_config = {} if tls_config is None else tls_config
     if tls_config:

fed/proxy/barriers.py

@@ -20,6 +20,7 @@
 import ray
 
 import fed.config as fed_config
+from fed.exceptions import FedRemoteError
 from fed._private import constants
 from fed._private.global_context import get_global_context
 from fed.proxy.base_proxy import ReceiverProxy, SenderProxy, SenderReceiverProxy
@@ -223,7 +224,7 @@ async def get_data(self, src_party, upstream_seq_id, curr_seq_id):
         data = await self._proxy_instance.get_data(
             src_party, upstream_seq_id, curr_seq_id
         )
-        if isinstance(data, Exception):
+        if isinstance(data, FedRemoteError):
             logger.debug(
                 f"Receiving exception: {type(data)}, {data} from {src_party}, "
                 f"upstream_seq_id: {upstream_seq_id}, "