tmp save

NKcqx · NKcqx · commit c8dd31fe7cf7 · 2023-06-26T12:06:46.000+08:00
diff --git a/fed/api.py b/fed/api.py
@@ -210,7 +210,7 @@ def init(
 
     logger.info(f'Started rayfed with {cluster_config}')
     set_exit_on_failure_sending(exit_on_failure_cross_silo_sending)
-    recv_actor_config = fed_config.ProxyActorConfig(
+    recv_actor_config = fed_config.CrossSiloProxyConfig(
         resource_label=cross_silo_recv_resource_label)
     # Start recv proxy
     start_recv_proxy(
@@ -222,7 +222,7 @@ def init(
         actor_config=recv_actor_config
     )
 
-    send_actor_config = fed_config.ProxyActorConfig(
+    send_actor_config = fed_config.CrossSiloProxyConfig(
         resource_label=cross_silo_send_resource_label)
     start_send_proxy(
         cluster=cluster,
diff --git a/fed/barriers.py b/fed/barriers.py
@@ -378,7 +378,7 @@ def start_recv_proxy(
     logging_level: str,
     tls_config=None,
     retry_policy=None,
-    actor_config: Optional[fed_config.ProxyActorConfig] = None
+    actor_config: Optional[fed_config.CrossSiloProxyConfig] = None
 ):
 
     # Create RecevrProxyActor
@@ -424,7 +424,7 @@ def start_send_proxy(
     tls_config: Dict = None,
     retry_policy=None,
     max_retries=None,
-    actor_config: Optional[fed_config.ProxyActorConfig] = None
+    actor_config: Optional[fed_config.CrossSiloProxyConfig] = None
 ):
     # Create SendProxyActor
     global _SEND_PROXY_ACTOR
diff --git a/fed/config.py b/fed/config.py
@@ -80,7 +80,7 @@ def get_job_config():
     return _job_config
 
 
-class ProxyActorConfig:
+class CrossSiloProxyConfig:
     """A class to store parameters used for Proxy Actor
 
     Attributes:
@@ -89,5 +89,17 @@ class ProxyActorConfig:
     """
     def __init__(
             self,
-            resource_label: Optional[Dict[str, str]] = None) -> None:
-        self.resource_label = resource_label
+            grpc_retry_policy: Dict = None,
+            send_max_retries: int = None,
+            timeout_in_seconds: int = 60,
+            messages_max_size_in_bytes: int = None,
+            serializing_allowed_list: Optional[Dict[str, str]] = None,
+            send_resource_label: Optional[Dict[str, str]] = None,
+            recv_resource_label: Optional[Dict[str, str]] = None) -> None:
+        self.grpc_retry_policy = grpc_retry_policy
+        self.send_max_retries = send_max_retries
+        self.timeout_in_seconds = timeout_in_seconds
+        self.messages_max_size_in_bytes = messages_max_size_in_bytes
+        self.serializing_allowed_list = serializing_allowed_list
+        self.send_resource_label = send_resource_label
+        self.recv_resource_label = recv_resource_label
