[Multi-Job][2/N] Make proxy's RPC framework replaceable

fed/_private/constants.py

@@ -25,6 +25,8 @@
 
 KEY_OF_TLS_CONFIG = "TLS_CONFIG"
 
+KEY_OF_CROSS_SILO_COMM_CONFIG = "CROSS_SILO_COMM_CONFIG"
+
 KEY_OF_CROSS_SILO_SERIALIZING_ALLOWED_LIST = "CROSS_SILO_SERIALIZING_ALLOWED_LIST" # noqa
 
 KEY_OF_CROSS_SILO_MESSAGES_MAX_SIZE_IN_BYTES = "CROSS_SILO_MESSAGES_MAX_SIZE_IN_BYTES" # noqa

fed/_private/serialization_utils.py

@@ -63,7 +63,8 @@ def find_class(self, module, name):
 def _apply_loads_function_with_whitelist():
     global _pickle_whitelist
 
-    _pickle_whitelist = fed_config.get_cluster_config().serializing_allowed_list
+    _pickle_whitelist = fed_config.get_job_config() \
+        .cross_silo_comm_config.serializing_allowed_list
     if _pickle_whitelist is None:
         return
 

fed/api.py

@@ -15,7 +15,7 @@
 import functools
 import inspect
 import logging
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Union, Optional
 
 import cloudpickle
 import ray
@@ -33,7 +33,11 @@
     send,
     start_recv_proxy,
     start_send_proxy,
+    SendProxy,
+    RecvProxy
 )
+from fed.config import CrossSiloCommConfig
+
 from fed.fed_object import FedObject
 from fed.utils import is_ray_object_refs, setup_logger
 
@@ -45,16 +49,10 @@ def init(
     party: str = None,
     tls_config: Dict = None,
     logging_level: str = 'info',
-    cross_silo_grpc_retry_policy: Dict = None,
-    cross_silo_send_max_retries: int = None,
-    cross_silo_serializing_allowed_list: Dict = None,
-    cross_silo_send_resource_label: Dict = None,
-    cross_silo_recv_resource_label: Dict = None,
-    exit_on_failure_cross_silo_sending: bool = False,
-    cross_silo_messages_max_size_in_bytes: int = None,
-    cross_silo_timeout_in_seconds: int = 60,
     enable_waiting_for_other_parties_ready: bool = False,
-    grpc_metadata: Dict = None,
+    send_proxy_cls: SendProxy = None,
+    recv_proxy_cls: RecvProxy = None,
+    global_cross_silo_comm_config: Optional[CrossSiloCommConfig] = None,
     **kwargs,
 ):
     """
@@ -116,48 +114,20 @@ def init(
                     "cert": "bob's server cert",
                     "key": "bob's server cert key",
                 }
-
         logging_level: optional; the logging level, could be `debug`, `info`,
             `warning`, `error`, `critical`, not case sensititive.
-        cross_silo_grpc_retry_policy: a dict descibes the retry policy for
-            cross silo rpc call. If None, the following default retry policy
-            will be used. More details please refer to
-            `retry-policy <https://github.com/grpc/proposal/blob/master/A6-client-retries.md#retry-policy>`_. # noqa
-
+        enable_waiting_for_other_parties_ready: ping other parties until they
+            are all ready if True.
+        global_cross_silo_comm_config: Global cross-silo communication related
+            config that are applied to all connections. Supported configs
+            can refer to CrossSiloCommConfig in config.py.
+        dest_party_comm_config: Communication config for the destination party
+            specifed by the key. E.g.
             .. code:: python
                 {
-                    "maxAttempts": 4,
-                    "initialBackoff": "0.1s",
-                    "maxBackoff": "1s",
-                    "backoffMultiplier": 2,
-                    "retryableStatusCodes": [
-                        "UNAVAILABLE"
-                    ]
+                    'alice': alice_CrossSiloCommConfig,
+                    'bob': bob_CrossSiloCommConfig
                 }
-        cross_silo_send_max_retries: the max retries for sending data cross silo.
-        cross_silo_serializing_allowed_list: The package or class list allowed for
-            serializing(deserializating) cross silos. It's used for avoiding pickle
-            deserializing execution attack when crossing solis.
-        cross_silo_send_resource_label: Customized resource label, the SendProxyActor
-            will be scheduled based on the declared resource label. For example,
-            when setting to `{"my_label": 1}`, then the SendProxyActor will be started
-            only on Nodes with `{"resource": {"my_label": $NUM}}` where $NUM >= 1.
-        cross_silo_recv_resource_label: Customized resource label, the RecverProxyActor
-            will be scheduled based on the declared resource label. For example,
-            when setting to `{"my_label": 1}`, then the RecverProxyActor will be started
-            only on Nodes with `{"resource": {"my_label": $NUM}}` where $NUM >= 1.
-        exit_on_failure_cross_silo_sending: whether exit when failure on
-            cross-silo sending. If True, a SIGTERM will be signaled to self
-            if failed to sending cross-silo data.
-        cross_silo_messages_max_size_in_bytes: The maximum length in bytes of
-            cross-silo messages.
-            If None, the default value of 500 MB is specified.
-        cross_silo_timeout_in_seconds: The timeout in seconds of a cross-silo RPC call.
-            It's 60 by default.
-        enable_waiting_for_other_parties_ready: ping other parties until they
-            are all ready if True.
-        grpc_metadata: optional; The metadata sent with the grpc request. This won't override
-            basic tcp headers, such as `user-agent`, but aggregate them together.
 
     Examples:
         >>> import fed
@@ -182,22 +152,21 @@ def init(
         assert (
             'cert' in tls_config and 'key' in tls_config
         ), 'Cert or key are not in tls_config.'
+
+    global_cross_silo_comm_config = \
+        global_cross_silo_comm_config or CrossSiloCommConfig()
     # A Ray private accessing, should be replaced in public API.
     compatible_utils._init_internal_kv()
 
     cluster_config = {
         constants.KEY_OF_CLUSTER_ADDRESSES: cluster,
         constants.KEY_OF_CURRENT_PARTY_NAME: party,
         constants.KEY_OF_TLS_CONFIG: tls_config,
-        constants.KEY_OF_CROSS_SILO_SERIALIZING_ALLOWED_LIST:
-            cross_silo_serializing_allowed_list,
-        constants.KEY_OF_CROSS_SILO_MESSAGES_MAX_SIZE_IN_BYTES:
-            cross_silo_messages_max_size_in_bytes,
-        constants.KEY_OF_CROSS_SILO_TIMEOUT_IN_SECONDS: cross_silo_timeout_in_seconds,
     }
 
     job_config = {
-        constants.KEY_OF_GRPC_METADATA : grpc_metadata,
+        constants.KEY_OF_CROSS_SILO_COMM_CONFIG:
+            global_cross_silo_comm_config,
     }
     compatible_utils.kv.put(constants.KEY_OF_CLUSTER_CONFIG,
                             cloudpickle.dumps(cluster_config))
@@ -215,30 +184,31 @@ def init(
 
     logger.info(f'Started rayfed with {cluster_config}')
     get_global_context().get_cleanup_manager().start(
-        exit_when_failure_sending=exit_on_failure_cross_silo_sending)
+        exit_when_failure_sending=global_cross_silo_comm_config.exit_on_sending_failure)
 
-    recv_actor_config = fed_config.ProxyActorConfig(
-        resource_label=cross_silo_recv_resource_label)
+    if recv_proxy_cls is None:
+        from fed.proxy.grpc_proxy import GrpcRecvProxy
+        recv_proxy_cls = GrpcRecvProxy
     # Start recv proxy
     start_recv_proxy(
         cluster=cluster,
         party=party,
         logging_level=logging_level,
         tls_config=tls_config,
-        retry_policy=cross_silo_grpc_retry_policy,
-        actor_config=recv_actor_config
+        proxy_cls=recv_proxy_cls,
+        proxy_config=global_cross_silo_comm_config
     )
 
-    send_actor_config = fed_config.ProxyActorConfig(
-        resource_label=cross_silo_send_resource_label)
+    if send_proxy_cls is None:
+        from fed.proxy.grpc_proxy import GrpcSendProxy
+        send_proxy_cls = GrpcSendProxy
     start_send_proxy(
         cluster=cluster,
         party=party,
         logging_level=logging_level,
         tls_config=tls_config,
-        retry_policy=cross_silo_grpc_retry_policy,
-        max_retries=cross_silo_send_max_retries,
-        actor_config=send_actor_config
+        proxy_cls=send_proxy_cls,
+        proxy_config=global_cross_silo_comm_config
     )
 
     if enable_waiting_for_other_parties_ready:

fed/config.py

@@ -7,7 +7,8 @@
 import fed._private.compatible_utils as compatible_utils
 import fed._private.constants as fed_constants
 import cloudpickle
-from typing import Dict, Optional
+from typing import Dict, List, Optional
+import json
 
 
 class ClusterConfig:
@@ -48,8 +49,8 @@ def __init__(self, raw_bytes: bytes) -> None:
             self._data = cloudpickle.loads(raw_bytes)
 
     @property
-    def grpc_metadata(self):
-        return self._data.get(fed_constants.KEY_OF_GRPC_METADATA, {})
+    def cross_silo_comm_config(self):
+        return self._data.get(fed_constants.KEY_OF_CROSS_SILO_COMM_CONFIG, {})
 
 
 # A module level cache for the cluster configurations.
@@ -80,14 +81,105 @@ def get_job_config():
     return _job_config
 
 
-class ProxyActorConfig:
+class CrossSiloCommConfig:
     """A class to store parameters used for Proxy Actor
 
     Attributes:
-        resource_label: The customized resources for the actor. This will be
-            filled into the "resource" field of Ray ActorClass.options.
+        proxier_fo_max_retries: The max restart times for the send proxy.
+        serializing_allowed_list: The package or class list allowed for
+            serializing(deserializating) cross silos. It's used for avoiding pickle
+            deserializing execution attack when crossing solis.
+        send_resource_label: Customized resource label, the SendProxyActor
+            will be scheduled based on the declared resource label. For example,
+            when setting to `{"my_label": 1}`, then the SendProxyActor will be started
+            only on Nodes with `{"resource": {"my_label": $NUM}}` where $NUM >= 1.
+        recv_resource_label: Customized resource label, the RecverProxyActor
+            will be scheduled based on the declared resource label. For example,
+            when setting to `{"my_label": 1}`, then the RecverProxyActor will be started
+            only on Nodes with `{"resource": {"my_label": $NUM}}` where $NUM >= 1.
+        exit_on_sending_failure: whether exit when failure on
+            cross-silo sending. If True, a SIGTERM will be signaled to self
+            if failed to sending cross-silo data.
+        messages_max_size_in_bytes: The maximum length in bytes of
+            cross-silo messages.
+            If None, the default value of 500 MB is specified.
+        timeout_in_seconds: The timeout in seconds of a cross-silo RPC call.
+            It's 60 by default.
+        http_header: The HTTP header, e.g. metadata in grpc, sent with the RPC request.
+            This won't override basic tcp headers, such as `user-agent`, but concat
+            them together.
     """
     def __init__(
             self,
-            resource_label: Optional[Dict[str, str]] = None) -> None:
-        self.resource_label = resource_label
+            proxier_fo_max_retries: int = None,
+            timeout_in_seconds: int = 60,
+            messages_max_size_in_bytes: int = None,
+            exit_on_sending_failure: Optional[bool] = False,
+            serializing_allowed_list: Optional[Dict[str, str]] = None,
+            send_resource_label: Optional[Dict[str, str]] = None,
+            recv_resource_label: Optional[Dict[str, str]] = None,
+            http_header: Optional[Dict[str, str]] = None) -> None:
+        self.proxier_fo_max_retries = proxier_fo_max_retries
+        self.timeout_in_seconds = timeout_in_seconds
+        self.messages_max_size_in_bytes = messages_max_size_in_bytes
+        self.exit_on_sending_failure = exit_on_sending_failure
+        self.serializing_allowed_list = serializing_allowed_list
+        self.send_resource_label = send_resource_label
+        self.recv_resource_label = recv_resource_label
+        self.http_header = http_header
+
+    def __json__(self):
+        return json.dumps(self.__dict__)
+
+    @classmethod
+    def from_json(cls, json_str):
+        data = json.loads(json_str)
+        return cls(**data)
+
+
+class CrossSiloGrpcCommConfig(CrossSiloCommConfig):
+    """A class to store parameters used for GRPC communication
+
+    Attributes:
+        grpc_retry_policy: a dict descibes the retry policy for
+            cross silo rpc call. If None, the following default retry policy
+            will be used. More details please refer to
+            `retry-policy <https://github.com/grpc/proposal/blob/master/A6-client-retries.md#retry-policy>`_. # noqa
+
+            .. code:: python
+                {
+                    "maxAttempts": 4,
+                    "initialBackoff": "0.1s",
+                    "maxBackoff": "1s",
+                    "backoffMultiplier": 2,
+                    "retryableStatusCodes": [
+                        "UNAVAILABLE"
+                    ]
+                }
+        grpc_channel_options: A list of tuples to store GRPC channel options,
+            e.g. [
+                    ('grpc.enable_retries', 1),
+                    ('grpc.max_send_message_length', 50 * 1024 * 1024)
+                ]
+    """
+    def __init__(self,
+                 grpc_channel_options: List = None,
+                 grpc_retry_policy: Dict[str, str] = None,
+                 *args,
+                 **kwargs):
+        super().__init__(*args, **kwargs)
+        self.grpc_retry_policy = grpc_retry_policy
+        self.grpc_channel_options = grpc_channel_options
+
+
+class CrossSiloBRPCConfig(CrossSiloCommConfig):
+    """A class to store parameters used for GRPC communication
+
+    Attributes:
+    """
+    def __init__(self,
+                 brpc_options,
+                 *args,
+                 **kwargs):
+        super().__init__(*args, **kwargs)
+        self.brpc_options = brpc_options