diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 14ff656c..98475e06 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -27,11 +27,10 @@ jobs: python3 -m virtualenv -p python3 py3 . py3/bin/activate which python - pip install pytest torch cloudpickle cryptography pip install ray==2.0.0 - pip install flake8 # For code style checking + pip install black==23.1 - name: Lint run: | . py3/bin/activate - flake8 + black -S --check --diff . --exclude='fed/grpc|py3' diff --git a/.isort.cfg b/.isort.cfg new file mode 100644 index 00000000..eb16c8b5 --- /dev/null +++ b/.isort.cfg @@ -0,0 +1,13 @@ +[settings] +# This is to make isort compatible with Black. See +# https://black.readthedocs.io/en/stable/the_black_code_style.html#how-black-wraps-lines. +line_length=88 +profile=black +multi_line_output=3 +include_trailing_comma=True +use_parentheses=True +float_to_top=True +filter_files=True + +known_local_folder=fed +sections=FUTURE,STDLIB,THIRDPARTY,FIRSTPARTY,LOCALFOLDER \ No newline at end of file diff --git a/benchmarks/many_tiny_tasks_benchmark.py b/benchmarks/many_tiny_tasks_benchmark.py index 5fbbf387..ae450041 100644 --- a/benchmarks/many_tiny_tasks_benchmark.py +++ b/benchmarks/many_tiny_tasks_benchmark.py @@ -12,9 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -import ray -import time import sys +import time + +import ray + import fed @@ -53,8 +55,8 @@ def main(party): if i % 100 == 0: print(f"Running {i}th call") print(f"num calls: {num_calls}") - print("total time (ms) = ", (time.time() - start)*1000) - print("per task overhead (ms) =", (time.time() - start)*1000/num_calls) + print("total time (ms) = ", (time.time() - start) * 1000) + print("per task overhead (ms) =", (time.time() - start) * 1000 / num_calls) fed.shutdown() ray.shutdown() diff --git a/fed/__init__.py b/fed/__init__.py index 7636dd4a..f7870bc4 100644 --- a/fed/__init__.py +++ b/fed/__init__.py @@ -12,11 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -from fed.api import (get, init, kill, remote, - shutdown) -from fed.proxy.barriers import recv, send -from fed.fed_object import FedObject +from fed.api import get, init, kill, remote, shutdown from fed.exceptions import FedRemoteError +from fed.fed_object import FedObject +from fed.proxy.barriers import recv, send __all__ = [ "get", @@ -27,5 +26,5 @@ "recv", "send", "FedObject", - "FedRemoteError" + "FedRemoteError", ] diff --git a/fed/_private/compatible_utils.py b/fed/_private/compatible_utils.py index 83d2a098..ab5eb154 100644 --- a/fed/_private/compatible_utils.py +++ b/fed/_private/compatible_utils.py @@ -13,10 +13,11 @@ # limitations under the License. import abc -import ray -import fed._private.constants as fed_constants +import ray import ray.experimental.internal_kv as ray_internal_kv + +import fed._private.constants as fed_constants from fed._private import constants @@ -41,15 +42,14 @@ def _compare_version_strings(version1, version2): def _ray_version_less_than_2_0_0(): - """ Whther the current ray version is less 2.0.0. - """ + """Whther the current ray version is less 2.0.0.""" return _compare_version_strings( - fed_constants.RAY_VERSION_2_0_0_STR, ray.__version__) + fed_constants.RAY_VERSION_2_0_0_STR, ray.__version__ + ) def init_ray(address: str = None, **kwargs): - """A compatible API to init Ray. - """ + """A compatible API to init Ray.""" if address == 'local' and _ray_version_less_than_2_0_0(): # Ignore the `local` when ray < 2.0.0 ray.init(**kwargs) @@ -58,8 +58,7 @@ def init_ray(address: str = None, **kwargs): def _get_gcs_address_from_ray_worker(): - """A compatible API to get the gcs address from the ray worker module. - """ + """A compatible API to get the gcs address from the ray worker module.""" try: return ray._private.worker._global_node.gcs_address except AttributeError: @@ -67,19 +66,19 @@ def _get_gcs_address_from_ray_worker(): def wrap_kv_key(job_name, key: str): - """Add an prefix to the key to avoid conflict with other jobs. - """ - assert isinstance(key, str), \ - f"The key of KV data must be `str` type, got {type(key)}." + """Add an prefix to the key to avoid conflict with other jobs.""" + assert isinstance( + key, str + ), f"The key of KV data must be `str` type, got {type(key)}." - return constants.RAYFED_JOB_KV_DATA_KEY_FMT.format( - job_name, key) + return constants.RAYFED_JOB_KV_DATA_KEY_FMT.format(job_name, key) class AbstractInternalKv(abc.ABC): - """ An abstract class that represents for bridging Ray internal kv in + """An abstract class that represents for bridging Ray internal kv in both Ray client mode and non Ray client mode. """ + def __init__(self) -> None: pass @@ -105,8 +104,8 @@ def reset(self): class InternalKv(AbstractInternalKv): - """The internal kv class for non Ray client mode. - """ + """The internal kv class for non Ray client mode.""" + def __init__(self, job_name: str) -> None: super().__init__() self._job_name = job_name @@ -120,21 +119,18 @@ def initialize(self): from ray._raylet import GcsClient gcs_client = GcsClient( - address=_get_gcs_address_from_ray_worker(), - nums_reconnect_retry=10) + address=_get_gcs_address_from_ray_worker(), nums_reconnect_retry=10 + ) return ray_internal_kv._initialize_internal_kv(gcs_client) def put(self, k, v): - return ray_internal_kv._internal_kv_put( - wrap_kv_key(self._job_name, k), v) + return ray_internal_kv._internal_kv_put(wrap_kv_key(self._job_name, k), v) def get(self, k): - return ray_internal_kv._internal_kv_get( - wrap_kv_key(self._job_name, k)) + return ray_internal_kv._internal_kv_get(wrap_kv_key(self._job_name, k)) def delete(self, k): - return ray_internal_kv._internal_kv_del( - wrap_kv_key(self._job_name, k)) + return ray_internal_kv._internal_kv_del(wrap_kv_key(self._job_name, k)) def reset(self): return ray_internal_kv._internal_kv_reset() @@ -144,8 +140,8 @@ def _ping(self): class ClientModeInternalKv(AbstractInternalKv): - """The internal kv class for Ray client mode. - """ + """The internal kv class for Ray client mode.""" + def __init__(self) -> None: super().__init__() self._internal_kv_actor = ray.get_actor("_INTERNAL_KV_ACTOR") @@ -176,9 +172,13 @@ def _init_internal_kv(job_name): global kv if kv is None: from ray._private.client_mode_hook import is_client_mode_enabled + if is_client_mode_enabled: - kv_actor = ray.remote(InternalKv).options( - name="_INTERNAL_KV_ACTOR").remote(job_name) + kv_actor = ( + ray.remote(InternalKv) + .options(name="_INTERNAL_KV_ACTOR") + .remote(job_name) + ) response = kv_actor._ping.remote() ray.get(response) kv = ClientModeInternalKv() if is_client_mode_enabled else InternalKv(job_name) @@ -192,6 +192,7 @@ def _clear_internal_kv(): kv.delete(constants.KEY_OF_JOB_CONFIG) kv.reset() from ray._private.client_mode_hook import is_client_mode_enabled + if is_client_mode_enabled: _internal_kv_actor = ray.get_actor("_INTERNAL_KV_ACTOR") ray.kill(_internal_kv_actor) diff --git a/fed/_private/constants.py b/fed/_private/constants.py index f21f3d06..a77201b7 100644 --- a/fed/_private/constants.py +++ b/fed/_private/constants.py @@ -27,7 +27,7 @@ KEY_OF_CROSS_SILO_COMM_CONFIG_DICT = "CROSS_SILO_COMM_CONFIG_DICT" -RAYFED_LOG_FMT = "%(asctime)s.%(msecs)03d %(levelname)s %(filename)s:%(lineno)s [%(party)s] -- [%(jobname)s] %(message)s" # noqa +RAYFED_LOG_FMT = "%(asctime)s.%(msecs)03d %(levelname)s %(filename)s:%(lineno)s [%(party)s] -- [%(jobname)s] %(message)s" # noqa RAYFED_DATE_FMT = "%Y-%m-%d %H:%M:%S" diff --git a/fed/_private/fed_actor.py b/fed/_private/fed_actor.py index dd880579..6e84bfa2 100644 --- a/fed/_private/fed_actor.py +++ b/fed/_private/fed_actor.py @@ -16,6 +16,7 @@ import ray from ray.util.client.common import ClientActorHandle + from fed._private.fed_call_holder import FedCallHolder from fed.fed_object import FedObject @@ -90,19 +91,17 @@ def _execute_impl(self, cls_args, cls_kwargs): ) def _execute_remote_method( - self, - method_name, - options, - _ray_wrappered_method, - args, - kwargs, + self, + method_name, + options, + _ray_wrappered_method, + args, + kwargs, ): num_returns = 1 if options and 'num_returns' in options: num_returns = options['num_returns'] - logger.debug( - f"Actor method call: {method_name}, num_returns: {num_returns}" - ) + logger.debug(f"Actor method call: {method_name}, num_returns: {num_returns}") return _ray_wrappered_method.options( name='', diff --git a/fed/_private/fed_call_holder.py b/fed/_private/fed_call_holder.py index 9e349c46..80c2a350 100644 --- a/fed/_private/fed_call_holder.py +++ b/fed/_private/fed_call_holder.py @@ -14,15 +14,16 @@ import logging -# Set config in the very beginning to avoid being overwritten by other packages. -logging.basicConfig(level=logging.INFO) - +import fed.config as fed_config from fed._private.global_context import get_global_context -from fed.proxy.barriers import send from fed.fed_object import FedObject -from fed.utils import resolve_dependencies +from fed.proxy.barriers import send from fed.tree_util import tree_flatten -import fed.config as fed_config +from fed.utils import resolve_dependencies + +# Set config in the very beginning to avoid being overwritten by other packages. +logging.basicConfig(level=logging.INFO) + logger = logging.getLogger(__name__) diff --git a/fed/_private/global_context.py b/fed/_private/global_context.py index cd1337fc..8e5c5506 100644 --- a/fed/_private/global_context.py +++ b/fed/_private/global_context.py @@ -12,22 +12,24 @@ # See the License for the specific language governing permissions and # limitations under the License. -from fed.cleanup import CleanupManager -from typing import Callable import threading +from typing import Callable + +from fed.cleanup import CleanupManager class GlobalContext: - def __init__(self, job_name: str, - current_party: str, - failure_handler: Callable[[], None]) -> None: + def __init__( + self, job_name: str, current_party: str, failure_handler: Callable[[], None] + ) -> None: self._job_name = job_name self._seq_count = 0 self._failure_handler = failure_handler self._atomic_shutdown_flag_lock = threading.Lock() self._atomic_shutdown_flag = True self._cleanup_manager = CleanupManager( - current_party, self.acquire_shutdown_flag) + current_party, self.acquire_shutdown_flag + ) def next_seq_id(self) -> int: self._seq_count += 1 @@ -65,9 +67,9 @@ def acquire_shutdown_flag(self) -> bool: _global_context = None -def init_global_context(current_party: str, - job_name: str, - failure_handler: Callable[[], None] = None) -> None: +def init_global_context( + current_party: str, job_name: str, failure_handler: Callable[[], None] = None +) -> None: global _global_context if _global_context is None: _global_context = GlobalContext(job_name, current_party, failure_handler) diff --git a/fed/_private/message_queue.py b/fed/_private/message_queue.py index c4a1b6eb..9fdff7dc 100644 --- a/fed/_private/message_queue.py +++ b/fed/_private/message_queue.py @@ -12,11 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging import threading -from collections import deque import time -import logging - +from collections import deque logger = logging.getLogger(__name__) @@ -55,7 +54,8 @@ def _loop(): if self._thread is None or not self._thread.is_alive(): logger.debug( - f"Starting new thread[{self._thread_name}] for message polling.") + f"Starting new thread[{self._thread_name}] for message polling." + ) self._queue = deque() self._thread = threading.Thread(target=_loop, name=self._thread_name) self._thread.start() @@ -79,9 +79,11 @@ def stop(self): If False: forcelly kill the for-loop sub-thread. """ if threading.current_thread() == self._thread: - logger.error(f"Can't stop the message queue in the message " - f"polling thread[{self._thread_name}]. Ignore it as this" - f"could bring unknown time sequence problems.") + logger.error( + f"Can't stop the message queue in the message " + f"polling thread[{self._thread_name}]. Ignore it as this" + f"could bring unknown time sequence problems." + ) raise RuntimeError("Thread can't kill itself") # TODO(NKcqx): Force kill sub-thread by calling `._stop()` will diff --git a/fed/_private/serialization_utils.py b/fed/_private/serialization_utils.py index c1b2e73e..0c0f3c61 100644 --- a/fed/_private/serialization_utils.py +++ b/fed/_private/serialization_utils.py @@ -13,11 +13,11 @@ # limitations under the License. import io + import cloudpickle import fed.config as fed_config - _pickle_whitelist = None diff --git a/fed/cleanup.py b/fed/cleanup.py index d68b290b..91a0739a 100644 --- a/fed/cleanup.py +++ b/fed/cleanup.py @@ -16,11 +16,12 @@ import os import signal import threading -from fed._private.message_queue import MessageQueueManager -from fed.exceptions import FedRemoteError -from ray.exceptions import RayError import ray +from ray.exceptions import RayError + +from fed._private.message_queue import MessageQueueManager +from fed.exceptions import FedRemoteError logger = logging.getLogger(__name__) @@ -44,11 +45,13 @@ class CleanupManager: def __init__(self, current_party, acquire_shutdown_flag) -> None: self._sending_data_q = MessageQueueManager( lambda msg: self._process_data_sending_task_return(msg), - thread_name='DataSendingQueueThread') + thread_name='DataSendingQueueThread', + ) self._sending_error_q = MessageQueueManager( lambda msg: self._process_error_sending_task_return(msg), - thread_name="ErrorSendingQueueThread") + thread_name="ErrorSendingQueueThread", + ) self._monitor_thread = None @@ -80,12 +83,14 @@ def stop(self): self._sending_data_q.stop() self._sending_error_q.stop() - def push_to_sending(self, - obj_ref: ray.ObjectRef, - dest_party: str = None, - upstream_seq_id: int = -1, - downstream_seq_id: int = -1, - is_error: bool = False): + def push_to_sending( + self, + obj_ref: ray.ObjectRef, + dest_party: str = None, + upstream_seq_id: int = -1, + downstream_seq_id: int = -1, + is_error: bool = False, + ): """ Push the sending remote task's return value, i.e. `obj_ref` to the corresponding message queue. @@ -104,7 +109,7 @@ def push_to_sending(self, queue instead. """ msg_pack = (obj_ref, dest_party, upstream_seq_id, downstream_seq_id) - if (is_error): + if is_error: self._sending_error_q.append(msg_pack) else: self._sending_data_q.append(msg_pack) @@ -123,7 +128,7 @@ def _signal_exit(self): # will cause dead lock. In order to ensure executing `shutdown` exactly # once and avoid dead lock, the lock must be checked before sending # signals. - if (self._acquire_shutdown_flag()): + if self._acquire_shutdown_flag(): logger.debug("Signal SIGINT to exit.") os.kill(os.getpid(), signal.SIGINT) @@ -151,16 +156,24 @@ def _process_data_sending_task_return(self, message): try: res = ray.get(obj_ref) except Exception as e: - logger.warn(f'Failed to send {obj_ref} to {dest_party}, error: {e},' - f'upstream_seq_id: {upstream_seq_id}, ' - f'downstream_seq_id: {downstream_seq_id}.') - if (isinstance(e, RayError)): + logger.warn( + f'Failed to send {obj_ref} to {dest_party}, error: {e},' + f'upstream_seq_id: {upstream_seq_id}, ' + f'downstream_seq_id: {downstream_seq_id}.' + ) + if isinstance(e, RayError): logger.info(f"Sending error {e.cause} to {dest_party}.") from fed.proxy.barriers import send + # TODO(NKcqx): Cascade broadcast to all parties error_trace = e.cause if self._expose_error_trace else None - send(dest_party, FedRemoteError(self._current_party, error_trace), - upstream_seq_id, downstream_seq_id, True) + send( + dest_party, + FedRemoteError(self._current_party, error_trace), + upstream_seq_id, + downstream_seq_id, + True, + ) res = False @@ -183,10 +196,12 @@ def _process_error_sending_task_return(self, error_msg): res = False if not res: - logger.warning(f"Failed to send error {error_ref} to {dest_party}, " - f"upstream_seq_id: {upstream_seq_id} " - f"downstream_seq_id: {downstream_seq_id}. " - "In this case, other parties won't sense " - "this error and may cause unknown behaviour.") + logger.warning( + f"Failed to send error {error_ref} to {dest_party}, " + f"upstream_seq_id: {upstream_seq_id} " + f"downstream_seq_id: {downstream_seq_id}. " + "In this case, other parties won't sense " + "this error and may cause unknown behaviour." + ) # Return True so that remaining error objects can be sent return True diff --git a/fed/exceptions.py b/fed/exceptions.py index dad4abfc..bd21f9a6 100644 --- a/fed/exceptions.py +++ b/fed/exceptions.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. + class FedRemoteError(Exception): def __init__(self, src_party: str, cause: Exception) -> None: self._src_party = src_party diff --git a/fed/fed_object.py b/fed/fed_object.py index 6e62faa9..aaa7e1dd 100644 --- a/fed/fed_object.py +++ b/fed/fed_object.py @@ -17,6 +17,7 @@ class FedObjectSendingContext: """The class that's used for holding the all contexts about sending side.""" + def __init__(self) -> None: # This field holds the target(downstream) parties that this fed object # is sending or sent to. @@ -33,6 +34,7 @@ def was_sending_or_sent_to_party(self, target_party: str): class FedObjectReceivingContext: """The class that's used for holding the all contexts about receiving side.""" + pass diff --git a/fed/proxy/barriers.py b/fed/proxy/barriers.py index aa207ae4..ee614463 100644 --- a/fed/proxy/barriers.py +++ b/fed/proxy/barriers.py @@ -51,9 +51,9 @@ def set_receiver_proxy_actor_name(name: str): _RECEIVER_PROXY_ACTOR_NAME = name -def set_proxy_actor_name(job_name: str, - use_global_proxy: bool, - sender_recvr_proxy: bool = False): +def set_proxy_actor_name( + job_name: str, use_global_proxy: bool, sender_recvr_proxy: bool = False +): """ Generate the name of the proxy actor. @@ -136,7 +136,8 @@ def __init__( job_config = fed_config.get_job_config(job_name) cross_silo_comm_config = job_config.cross_silo_comm_config_dict self._proxy_instance: SenderProxy = proxy_cls( - addresses, party, job_name, tls_config, cross_silo_comm_config) + addresses, party, job_name, tls_config, cross_silo_comm_config + ) async def is_ready(self): res = await self._proxy_instance.is_ready() @@ -207,7 +208,8 @@ def __init__( job_config = fed_config.get_job_config(job_name) cross_silo_comm_config = job_config.cross_silo_comm_config_dict self._proxy_instance: ReceiverProxy = proxy_cls( - listening_address, party, job_name, tls_config, cross_silo_comm_config) + listening_address, party, job_name, tls_config, cross_silo_comm_config + ) async def start(self): await self._proxy_instance.start() @@ -222,9 +224,11 @@ async def get_data(self, src_party, upstream_seq_id, curr_seq_id): src_party, upstream_seq_id, curr_seq_id ) if isinstance(data, Exception): - logger.debug(f"Receiving exception: {type(data)}, {data} from {src_party}, " - f"upstream_seq_id: {upstream_seq_id}, " - f"curr_seq_id: {curr_seq_id}. Re-raise it.") + logger.debug( + f"Receiving exception: {type(data)}, {data} from {src_party}, " + f"upstream_seq_id: {upstream_seq_id}, " + f"curr_seq_id: {curr_seq_id}. Re-raise it." + ) raise data return data @@ -437,13 +441,14 @@ def _start_sender_receiver_proxy( global _SENDER_RECEIVER_PROXY_ACTOR _SENDER_RECEIVER_PROXY_ACTOR = SenderReceiverProxyActor.options( - **actor_options).remote( - addresses=addresses, - party=party, - job_name=job_name, - tls_config=tls_config, - logging_level=logging_level, - proxy_cls=proxy_cls, + **actor_options + ).remote( + addresses=addresses, + party=party, + job_name=job_name, + tls_config=tls_config, + logging_level=logging_level, + proxy_cls=proxy_cls, ) _SENDER_RECEIVER_PROXY_ACTOR.start.remote() server_state = ray.get( @@ -453,13 +458,7 @@ def _start_sender_receiver_proxy( logger.info("Succeeded to create receiver proxy actor.") -def send( - dest_party, - data, - upstream_seq_id, - downstream_seq_id, - is_error=False -): +def send(dest_party, data, upstream_seq_id, downstream_seq_id, is_error=False): """ Args: is_error: Whether the `data` is an error object or not. Default is False. @@ -473,7 +472,8 @@ def send( downstream_seq_id=downstream_seq_id, ) get_global_context().get_cleanup_manager().push_to_sending( - res, dest_party, upstream_seq_id, downstream_seq_id, is_error) + res, dest_party, upstream_seq_id, downstream_seq_id, is_error + ) return res diff --git a/fed/proxy/base_proxy.py b/fed/proxy/base_proxy.py index b2eba265..c62285b8 100644 --- a/fed/proxy/base_proxy.py +++ b/fed/proxy/base_proxy.py @@ -51,7 +51,7 @@ def __init__( party: str, job_name: str, tls_config: Dict, - proxy_config: CrossSiloMessageConfig = None + proxy_config: CrossSiloMessageConfig = None, ) -> None: self._listen_addr = listen_addr self._party = party diff --git a/fed/proxy/grpc/grpc_options.py b/fed/proxy/grpc/grpc_options.py index 6e4b2d14..064824b3 100644 --- a/fed/proxy/grpc/grpc_options.py +++ b/fed/proxy/grpc/grpc_options.py @@ -14,7 +14,6 @@ import json - _GRPC_SERVICE = "GrpcService" _DEFAULT_GRPC_RETRY_POLICY = { @@ -34,17 +33,16 @@ 'grpc.so_reuseport': 0, 'grpc.max_send_message_length': _DEFAULT_GRPC_MAX_SEND_MESSAGE_LENGTH, 'grpc.max_receive_message_length': _DEFAULT_GRPC_MAX_RECEIVE_MESSAGE_LENGTH, - 'grpc.service_config': - json.dumps( - { - 'methodConfig': [ - { - 'name': [{'service': _GRPC_SERVICE}], - 'retryPolicy': _DEFAULT_GRPC_RETRY_POLICY, - } - ] - } - ), + 'grpc.service_config': json.dumps( + { + 'methodConfig': [ + { + 'name': [{'service': _GRPC_SERVICE}], + 'retryPolicy': _DEFAULT_GRPC_RETRY_POLICY, + } + ] + } + ), } diff --git a/fed/proxy/grpc/grpc_proxy.py b/fed/proxy/grpc/grpc_proxy.py index 1fdfa4e2..e47c62dd 100644 --- a/fed/proxy/grpc/grpc_proxy.py +++ b/fed/proxy/grpc/grpc_proxy.py @@ -14,27 +14,29 @@ import asyncio import copy -import cloudpickle -import grpc +import json import logging import threading -import json from typing import Dict -import fed.utils as fed_utils +import cloudpickle +import grpc -from fed.config import CrossSiloMessageConfig, GrpcCrossSiloMessageConfig import fed._private.compatible_utils as compatible_utils -from fed.proxy.grpc.grpc_options import _DEFAULT_GRPC_CHANNEL_OPTIONS, _GRPC_SERVICE +import fed.utils as fed_utils +from fed.config import CrossSiloMessageConfig, GrpcCrossSiloMessageConfig from fed.proxy.barriers import ( add_two_dim_dict, get_from_two_dim_dict, - pop_from_two_dim_dict, key_exists_in_two_dim_dict, + pop_from_two_dim_dict, ) -from fed.proxy.base_proxy import SenderProxy, ReceiverProxy +from fed.proxy.base_proxy import ReceiverProxy, SenderProxy +from fed.proxy.grpc.grpc_options import _DEFAULT_GRPC_CHANNEL_OPTIONS, _GRPC_SERVICE + if compatible_utils._compare_version_strings( - fed_utils.get_package_version('protobuf'), '4.0.0'): + fed_utils.get_package_version('protobuf'), '4.0.0' +): from fed.grpc.pb4 import fed_pb2 as fed_pb2 from fed.grpc.pb4 import fed_pb2_grpc as fed_pb2_grpc else: @@ -67,43 +69,44 @@ def parse_grpc_options(proxy_config: CrossSiloMessageConfig): # However, `GrpcCrossSiloMessageConfig` provides a more flexible way # to configure grpc channel options, i.e. the `grpc_channel_options` # field, which may override the `messages_max_size_in_bytes` field. - if (isinstance(proxy_config, CrossSiloMessageConfig)): - if (proxy_config.messages_max_size_in_bytes is not None): - grpc_channel_options.update({ - 'grpc.max_send_message_length': - proxy_config.messages_max_size_in_bytes, - 'grpc.max_receive_message_length': - proxy_config.messages_max_size_in_bytes, - }) + if isinstance(proxy_config, CrossSiloMessageConfig): + if proxy_config.messages_max_size_in_bytes is not None: + grpc_channel_options.update( + { + 'grpc.max_send_message_length': proxy_config.messages_max_size_in_bytes, + 'grpc.max_receive_message_length': proxy_config.messages_max_size_in_bytes, + } + ) if isinstance(proxy_config, GrpcCrossSiloMessageConfig): if proxy_config.grpc_channel_options is not None: grpc_channel_options.update(proxy_config.grpc_channel_options) if proxy_config.grpc_retry_policy is not None: - grpc_channel_options.update({ - 'grpc.service_config': - json.dumps( - { - 'methodConfig': [ - { - 'name': [{'service': _GRPC_SERVICE}], - 'retryPolicy': proxy_config.grpc_retry_policy, - } - ] - } - ), - }) + grpc_channel_options.update( + { + 'grpc.service_config': json.dumps( + { + 'methodConfig': [ + { + 'name': [{'service': _GRPC_SERVICE}], + 'retryPolicy': proxy_config.grpc_retry_policy, + } + ] + } + ), + } + ) return grpc_channel_options class GrpcSenderProxy(SenderProxy): def __init__( - self, - cluster: Dict, - party: str, - job_name: str, - tls_config: Dict, - proxy_config: Dict = None + self, + cluster: Dict, + party: str, + job_name: str, + tls_config: Dict, + proxy_config: Dict = None, ) -> None: proxy_config = GrpcCrossSiloMessageConfig.from_dict(proxy_config) super().__init__(cluster, party, job_name, tls_config, proxy_config) @@ -113,29 +116,27 @@ def __init__( # Mapping the destination party name to the reused client stub. self._stubs = {} - async def send( - self, - dest_party, - data, - upstream_seq_id, - downstream_seq_id): + async def send(self, dest_party, data, upstream_seq_id, downstream_seq_id): dest_addr = self._addresses[dest_party] grpc_metadata, grpc_channel_options = self.get_grpc_config_by_party(dest_party) tls_enabled = fed_utils.tls_enabled(self._tls_config) if dest_party not in self._stubs: if tls_enabled: ca_cert, private_key, cert_chain = fed_utils.load_cert_config( - self._tls_config) + self._tls_config + ) credentials = grpc.ssl_channel_credentials( certificate_chain=cert_chain, private_key=private_key, root_certificates=ca_cert, ) channel = grpc.aio.secure_channel( - dest_addr, credentials, options=grpc_channel_options) + dest_addr, credentials, options=grpc_channel_options + ) else: channel = grpc.aio.insecure_channel( - dest_addr, options=grpc_channel_options) + dest_addr, options=grpc_channel_options + ) stub = fed_pb2_grpc.GrpcServiceStub(channel) self._stubs[dest_party] = stub @@ -153,8 +154,7 @@ async def send( return response.result def get_grpc_config_by_party(self, dest_party): - """Overide global config by party specific config - """ + """Overide global config by party specific config""" grpc_metadata = self._grpc_metadata grpc_options = self._grpc_options @@ -162,14 +162,9 @@ def get_grpc_config_by_party(self, dest_party): if dest_party_msg_config is not None: if dest_party_msg_config.http_header is not None: dest_party_grpc_metadata = dict(dest_party_msg_config.http_header) - grpc_metadata = { - **grpc_metadata, - **dest_party_grpc_metadata - } + grpc_metadata = {**grpc_metadata, **dest_party_grpc_metadata} dest_party_grpc_options = parse_grpc_options(dest_party_msg_config) - grpc_options = { - **grpc_options, **dest_party_grpc_options - } + grpc_options = {**grpc_options, **dest_party_grpc_options} return grpc_metadata, fed_utils.dict2tuple(grpc_options) async def get_proxy_config(self, dest_party=None): @@ -188,8 +183,10 @@ def handle_response_error(self, response): if 400 <= response.code < 500: # Request error should also be identified as a sending failure, # though the request was physically sent. - logger.warning(f"Request was successfully sent but got error response, " - f"code: {response.code}, message: {response.result}.") + logger.warning( + f"Request was successfully sent but got error response, " + f"code: {response.code}, message: {response.result}." + ) raise RuntimeError(response.result) @@ -225,12 +222,12 @@ async def send_data_grpc( class GrpcReceiverProxy(ReceiverProxy): def __init__( - self, - listen_addr: str, - party: str, - job_name: str, - tls_config: Dict, - proxy_config: Dict + self, + listen_addr: str, + party: str, + job_name: str, + tls_config: Dict, + proxy_config: Dict, ) -> None: proxy_config = GrpcCrossSiloMessageConfig.from_dict(proxy_config) super().__init__(listen_addr, party, job_name, tls_config, proxy_config) @@ -260,9 +257,11 @@ async def start(self): fed_utils.dict2tuple(self._grpc_options), ) except RuntimeError as err: - msg = f'Grpc server failed to listen to port: {port}' \ - f' Try another port by setting `listen_addr` into `cluster` config' \ - f' when calling `fed.init`. Grpc error msg: {err}' + msg = ( + f'Grpc server failed to listen to port: {port}' + f' Try another port by setting `listen_addr` into `cluster` config' + f' when calling `fed.init`. Grpc error msg: {err}' + ) self._server_ready_future.set_result((False, msg)) async def is_ready(self): @@ -289,6 +288,7 @@ async def get_data(self, src_party, upstream_seq_id, curr_seq_id): # NOTE(qwang): This is used to avoid the conflict with pickle5 in Ray. import fed._private.serialization_utils as fed_ser_utils + fed_ser_utils._apply_loads_function_with_whitelist() return cloudpickle.loads(data) @@ -309,12 +309,15 @@ def __init__(self, all_events, all_data, party, lock, job_name): async def SendData(self, request, context): job_name = request.job_name if job_name != self._job_name: - logger.warning(f"Receive data from job {job_name}, ignore it. " - f"The reason may be that the ReceiverProxy is listening " - f"on the same address with that job.") + logger.warning( + f"Receive data from job {job_name}, ignore it. " + f"The reason may be that the ReceiverProxy is listening " + f"on the same address with that job." + ) return fed_pb2.SendDataResponse( code=417, - result=f"JobName mis-match, expected {self._job_name}, got {job_name}.") + result=f"JobName mis-match, expected {self._job_name}, got {job_name}.", + ) upstream_seq_id = request.upstream_seq_id downstream_seq_id = request.downstream_seq_id logger.debug( @@ -340,8 +343,15 @@ async def SendData(self, request, context): async def _run_grpc_server( - port, event, all_data, party, lock, job_name, - server_ready_future, tls_config=None, grpc_options=None + port, + event, + all_data, + party, + lock, + job_name, + server_ready_future, + tls_config=None, + grpc_options=None, ): logger.info(f"ReceiverProxy binding port {port}, options: {grpc_options}...") server = grpc.aio.server(options=grpc_options) diff --git a/fed/tests/client_mode_tests/test_basic_client_mode.py b/fed/tests/client_mode_tests/test_basic_client_mode.py index 98078029..35f3d851 100644 --- a/fed/tests/client_mode_tests/test_basic_client_mode.py +++ b/fed/tests/client_mode_tests/test_basic_client_mode.py @@ -16,9 +16,10 @@ import pytest import ray + import fed import fed._private.compatible_utils as compatible_utils -from fed.tests.test_utils import ray_client_mode_setup # noqa +from fed.tests.test_utils import ray_client_mode_setup # noqa @fed.remote @@ -49,10 +50,13 @@ def mean(x, y): def run(party): import time + if party == 'alice': time.sleep(1.4) - address = 'ray://127.0.0.1:21012' if party == 'alice' else 'ray://127.0.0.1:21011' # noqa + address = ( + 'ray://127.0.0.1:21012' if party == 'alice' else 'ray://127.0.0.1:21011' + ) # noqa compatible_utils.init_ray(address=address) addresses = { @@ -83,7 +87,7 @@ def run(party): ray.shutdown() -def test_fed_get_in_2_parties(ray_client_mode_setup): # noqa +def test_fed_get_in_2_parties(ray_client_mode_setup): # noqa p_alice = multiprocessing.Process(target=run, args=('alice',)) p_bob = multiprocessing.Process(target=run, args=('bob',)) p_alice.start() diff --git a/fed/tests/multi-jobs/test_ignore_other_job_msg.py b/fed/tests/multi-jobs/test_ignore_other_job_msg.py index a75d8f78..d83acff3 100644 --- a/fed/tests/multi-jobs/test_ignore_other_job_msg.py +++ b/fed/tests/multi-jobs/test_ignore_other_job_msg.py @@ -13,32 +13,30 @@ # limitations under the License. import multiprocessing -import fed -import ray + import grpc import pytest -import fed.utils as fed_utils +import ray + +import fed import fed._private.compatible_utils as compatible_utils +import fed.utils as fed_utils from fed.proxy.grpc.grpc_proxy import GrpcSenderProxy, send_data_grpc + if compatible_utils._compare_version_strings( - fed_utils.get_package_version('protobuf'), '4.0.0'): + fed_utils.get_package_version('protobuf'), '4.0.0' +): from fed.grpc.pb4 import fed_pb2_grpc as fed_pb2_grpc else: from fed.grpc.pb3 import fed_pb2_grpc as fed_pb2_grpc class TestGrpcSenderProxy(GrpcSenderProxy): - async def send( - self, - dest_party, - data, - upstream_seq_id, - downstream_seq_id): + async def send(self, dest_party, data, upstream_seq_id, downstream_seq_id): dest_addr = self._addresses[dest_party] grpc_metadata, grpc_channel_options = self.get_grpc_config_by_party(dest_party) if dest_party not in self._stubs: - channel = grpc.aio.insecure_channel( - dest_addr, options=grpc_channel_options) + channel = grpc.aio.insecure_channel(dest_addr, options=grpc_channel_options) stub = fed_pb2_grpc.GrpcServiceStub(channel) self._stubs[dest_party] = stub @@ -81,14 +79,17 @@ def agg_fn(obj1, obj2): def run(party, job_name): ray.init(address='local') - fed.init(addresses=addresses, - party=party, - job_name=job_name, - sender_proxy_cls=TestGrpcSenderProxy, - config={ - 'cross_silo_comm': { - 'exit_on_sending_failure': True, - }}) + fed.init( + addresses=addresses, + party=party, + job_name=job_name, + sender_proxy_cls=TestGrpcSenderProxy, + config={ + 'cross_silo_comm': { + 'exit_on_sending_failure': True, + } + }, + ) # 'bob' only needs to start the proxy actors if party == 'alice': ds1, ds2 = [123, 789] @@ -103,6 +104,7 @@ def run(party, job_name): fed.shutdown() ray.shutdown() import time + # Wait for SIGTERM as failure on sending. time.sleep(86400) diff --git a/fed/tests/multi-jobs/test_multi_proxy_actor.py b/fed/tests/multi-jobs/test_multi_proxy_actor.py index 5021ec0d..a540086c 100644 --- a/fed/tests/multi-jobs/test_multi_proxy_actor.py +++ b/fed/tests/multi-jobs/test_multi_proxy_actor.py @@ -13,32 +13,30 @@ # limitations under the License. import multiprocessing -import fed -import ray + import grpc import pytest -import fed.utils as fed_utils +import ray + +import fed import fed._private.compatible_utils as compatible_utils +import fed.utils as fed_utils from fed.proxy.grpc.grpc_proxy import GrpcSenderProxy, send_data_grpc + if compatible_utils._compare_version_strings( - fed_utils.get_package_version('protobuf'), '4.0.0'): + fed_utils.get_package_version('protobuf'), '4.0.0' +): from fed.grpc.pb4 import fed_pb2_grpc as fed_pb2_grpc else: from fed.grpc.pb3 import fed_pb2_grpc as fed_pb2_grpc class TestGrpcSenderProxy(GrpcSenderProxy): - async def send( - self, - dest_party, - data, - upstream_seq_id, - downstream_seq_id): + async def send(self, dest_party, data, upstream_seq_id, downstream_seq_id): dest_addr = self._addresses[dest_party] grpc_metadata, grpc_channel_options = self.get_grpc_config_by_party(dest_party) if dest_party not in self._stubs: - channel = grpc.aio.insecure_channel( - dest_addr, options=grpc_channel_options) + channel = grpc.aio.insecure_channel(dest_addr, options=grpc_channel_options) stub = fed_pb2_grpc.GrpcServiceStub(channel) self._stubs[dest_party] = stub @@ -87,16 +85,19 @@ def agg_fn(obj1, obj2): def run(party, job_name): ray.init(address='local') - fed.init(addresses=addresses[job_name], - party=party, - job_name=job_name, - sender_proxy_cls=TestGrpcSenderProxy, - config={ - 'cross_silo_comm': { - 'exit_on_sending_failure': True, - # Create unique proxy for current job - 'use_global_proxy': False - }}) + fed.init( + addresses=addresses[job_name], + party=party, + job_name=job_name, + sender_proxy_cls=TestGrpcSenderProxy, + config={ + 'cross_silo_comm': { + 'exit_on_sending_failure': True, + # Create unique proxy for current job + 'use_global_proxy': False, + } + }, + ) sender_proxy_actor_name = f"SenderProxyActor_{job_name}" receiver_proxy_actor_name = f"ReceiverProxyActor_{job_name}" diff --git a/fed/tests/simple_example.py b/fed/tests/simple_example.py index 4ca1095f..88c9ce00 100644 --- a/fed/tests/simple_example.py +++ b/fed/tests/simple_example.py @@ -13,9 +13,11 @@ # limitations under the License. import multiprocessing -import fed + import ray +import fed + @fed.remote class MyActor: diff --git a/fed/tests/test_api.py b/fed/tests/test_api.py index 9a3b00af..c9eaec16 100644 --- a/fed/tests/test_api.py +++ b/fed/tests/test_api.py @@ -13,10 +13,12 @@ # limitations under the License. import multiprocessing + import pytest +import ray + import fed import fed._private.compatible_utils as compatible_utils -import ray import fed.config as fed_config diff --git a/fed/tests/test_async_startup_2_clusters.py b/fed/tests/test_async_startup_2_clusters.py index 9542f87e..af04776b 100644 --- a/fed/tests/test_async_startup_2_clusters.py +++ b/fed/tests/test_async_startup_2_clusters.py @@ -15,8 +15,8 @@ import multiprocessing import pytest - import ray + import fed import fed._private.compatible_utils as compatible_utils diff --git a/fed/tests/test_basic_pass_fed_objects.py b/fed/tests/test_basic_pass_fed_objects.py index ebca911f..223097af 100644 --- a/fed/tests/test_basic_pass_fed_objects.py +++ b/fed/tests/test_basic_pass_fed_objects.py @@ -16,6 +16,7 @@ import pytest import ray + import fed import fed._private.compatible_utils as compatible_utils diff --git a/fed/tests/test_cross_silo_error.py b/fed/tests/test_cross_silo_error.py index 17547a09..f2f623eb 100644 --- a/fed/tests/test_cross_silo_error.py +++ b/fed/tests/test_cross_silo_error.py @@ -12,15 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -import ray import multiprocessing +import sys +from unittest.mock import Mock import pytest +import ray + import fed import fed._private.compatible_utils as compatible_utils -import sys - -from unittest.mock import Mock from fed.exceptions import FedRemoteError @@ -59,10 +59,10 @@ def run(party): 'cross_silo_comm': { 'exit_on_sending_failure': True, 'timeout_ms': 20 * 1000, - 'expose_error_trace': True + 'expose_error_trace': True, }, }, - failure_handler=my_failure_handler + failure_handler=my_failure_handler, ) # Both party should catch the error @@ -107,10 +107,10 @@ def run2(party): 'cross_silo_comm': { 'exit_on_sending_failure': True, 'timeout_ms': 20 * 1000, - 'expose_error_trace': True + 'expose_error_trace': True, }, }, - failure_handler=my_failure_handler + failure_handler=my_failure_handler, ) # Both party should catch the error @@ -162,7 +162,7 @@ def run3(party): 'timeout_ms': 20 * 1000, }, }, - failure_handler=my_failure_handler + failure_handler=my_failure_handler, ) # Both party should catch the error diff --git a/fed/tests/test_enable_tls_across_parties.py b/fed/tests/test_enable_tls_across_parties.py index 5c3c1c38..59921a1b 100644 --- a/fed/tests/test_enable_tls_across_parties.py +++ b/fed/tests/test_enable_tls_across_parties.py @@ -16,8 +16,8 @@ import os import pytest - import ray + import fed import fed._private.compatible_utils as compatible_utils diff --git a/fed/tests/test_exit_on_failure_sending.py b/fed/tests/test_exit_on_failure_sending.py index ab1a8613..3da50ae7 100644 --- a/fed/tests/test_exit_on_failure_sending.py +++ b/fed/tests/test_exit_on_failure_sending.py @@ -13,17 +13,16 @@ # limitations under the License. import multiprocessing +import os +import signal +import sys import pytest import ray + import fed import fed._private.compatible_utils as compatible_utils -import signal - -import os -import sys - def signal_handler(sig, frame): if sig == signal.SIGTERM.value: @@ -73,7 +72,7 @@ def run(party): 'timeout_ms': 20 * 1000, }, }, - failure_handler=lambda : os.kill(os.getpid(), signal.SIGTERM) + failure_handler=lambda: os.kill(os.getpid(), signal.SIGTERM), ) o = f.party("alice").remote() diff --git a/fed/tests/test_fed_get.py b/fed/tests/test_fed_get.py index 5752f778..8a7e72d1 100644 --- a/fed/tests/test_fed_get.py +++ b/fed/tests/test_fed_get.py @@ -16,6 +16,7 @@ import pytest import ray + import fed import fed._private.compatible_utils as compatible_utils @@ -48,6 +49,7 @@ def mean(x, y): def run(party): import time + if party == 'alice': time.sleep(1.4) diff --git a/fed/tests/test_grpc_options_on_proxies.py b/fed/tests/test_grpc_options_on_proxies.py index cb14e922..22e31592 100644 --- a/fed/tests/test_grpc_options_on_proxies.py +++ b/fed/tests/test_grpc_options_on_proxies.py @@ -13,11 +13,12 @@ # limitations under the License. import multiprocessing + import pytest -import fed -import fed._private.compatible_utils as compatible_utils import ray +import fed +import fed._private.compatible_utils as compatible_utils from fed.proxy.barriers import receiver_proxy_actor_name, sender_proxy_actor_name @@ -131,7 +132,7 @@ def run3(party): "messages_max_size_in_bytes": 100, "grpc_channel_options": [ ('grpc.max_send_message_length', 200), - ], + ], }, }, ) diff --git a/fed/tests/test_internal_kv.py b/fed/tests/test_internal_kv.py index bb048239..65a6e4fa 100644 --- a/fed/tests/test_internal_kv.py +++ b/fed/tests/test_internal_kv.py @@ -1,10 +1,12 @@ import multiprocessing +import time + import pytest import ray +import ray.experimental.internal_kv as ray_internal_kv + import fed -import time import fed._private.compatible_utils as compatible_utils -import ray.experimental.internal_kv as ray_internal_kv def run(party): @@ -21,8 +23,10 @@ def run(party): # Test that a prefix key name is added under the hood. assert ray_internal_kv._internal_kv_get(b"test_key") is None - assert ray_internal_kv._internal_kv_get( - b"RAYFED#test_job_name#test_key") == b"test_val" + assert ( + ray_internal_kv._internal_kv_get(b"RAYFED#test_job_name#test_key") + == b"test_val" + ) time.sleep(5) fed.shutdown() diff --git a/fed/tests/test_listening_address.py b/fed/tests/test_listening_address.py index a2787226..5f9891ca 100644 --- a/fed/tests/test_listening_address.py +++ b/fed/tests/test_listening_address.py @@ -16,6 +16,7 @@ import pytest import ray + import fed import fed._private.compatible_utils as compatible_utils @@ -36,9 +37,7 @@ def _run(party): s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) s.bind(("127.0.0.1", occupied_port)) - addresses = { - 'alice': f'127.0.0.1:{occupied_port}' - } + addresses = {'alice': f'127.0.0.1:{occupied_port}'} # Starting grpc server on an used port will cause AssertionError with pytest.raises(AssertionError): diff --git a/fed/tests/test_options.py b/fed/tests/test_options.py index 13bbb6f3..6ab9107b 100644 --- a/fed/tests/test_options.py +++ b/fed/tests/test_options.py @@ -16,6 +16,7 @@ import pytest import ray + import fed import fed._private.compatible_utils as compatible_utils diff --git a/fed/tests/test_pass_fed_objects_in_containers_in_actor.py b/fed/tests/test_pass_fed_objects_in_containers_in_actor.py index 08d3a2d6..4d9e4c82 100644 --- a/fed/tests/test_pass_fed_objects_in_containers_in_actor.py +++ b/fed/tests/test_pass_fed_objects_in_containers_in_actor.py @@ -15,8 +15,8 @@ import multiprocessing import pytest - import ray + import fed import fed._private.compatible_utils as compatible_utils diff --git a/fed/tests/test_pass_fed_objects_in_containers_in_normal_tasks.py b/fed/tests/test_pass_fed_objects_in_containers_in_normal_tasks.py index 7b78cb60..866370d2 100644 --- a/fed/tests/test_pass_fed_objects_in_containers_in_normal_tasks.py +++ b/fed/tests/test_pass_fed_objects_in_containers_in_normal_tasks.py @@ -15,8 +15,8 @@ import multiprocessing import pytest - import ray + import fed import fed._private.compatible_utils as compatible_utils diff --git a/fed/tests/test_ping_others.py b/fed/tests/test_ping_others.py index 4753ddec..a734104e 100644 --- a/fed/tests/test_ping_others.py +++ b/fed/tests/test_ping_others.py @@ -12,15 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -import pytest import multiprocessing +import time + +import pytest +import ray + import fed import fed._private.compatible_utils as compatible_utils -import ray -import time from fed.proxy.barriers import ping_others - addresses = { 'alice': '127.0.0.1:11012', 'bob': '127.0.0.1:11011', @@ -31,7 +32,7 @@ def test_ping_non_started_party(): def run(party): compatible_utils.init_ray(address='local') fed.init(addresses=addresses, party=party) - if (party == 'alice'): + if party == 'alice': with pytest.raises(RuntimeError): ping_others(addresses, party, 5) @@ -47,7 +48,7 @@ def test_ping_started_party(): def run(party): compatible_utils.init_ray(address='local') fed.init(addresses=addresses, party=party) - if (party == 'alice'): + if party == 'alice': ping_success = ping_others(addresses, party, 5) assert ping_success is True else: diff --git a/fed/tests/test_repeat_init.py b/fed/tests/test_repeat_init.py index 8926c786..1e68622d 100644 --- a/fed/tests/test_repeat_init.py +++ b/fed/tests/test_repeat_init.py @@ -16,9 +16,10 @@ import multiprocessing import pytest +import ray + import fed import fed._private.compatible_utils as compatible_utils -import ray @fed.remote diff --git a/fed/tests/test_reset_context.py b/fed/tests/test_reset_context.py index 95c6e53f..6350851e 100644 --- a/fed/tests/test_reset_context.py +++ b/fed/tests/test_reset_context.py @@ -1,8 +1,10 @@ import multiprocessing -import fed + +import pytest import ray + +import fed import fed._private.compatible_utils as compatible_utils -import pytest addresses = { 'alice': '127.0.0.1:11012', @@ -21,9 +23,7 @@ def get(self): def run(party): compatible_utils.init_ray(address='local') - fed.init( - addresses=addresses, - party=party) + fed.init(addresses=addresses, party=party) actor = A.party('alice').remote(10) alice_fed_obj = actor.get.remote() @@ -45,9 +45,7 @@ def run(party): compatible_utils.kv.put("key2", "val2") compatible_utils.init_ray(address='local') - fed.init( - addresses=addresses, - party=party) + fed.init(addresses=addresses, party=party) actor = A.party('alice').remote(10) alice_fed_obj = actor.get.remote() @@ -70,8 +68,8 @@ def run(party): def test_reset_context(): - p_alice = multiprocessing.Process(target=run, args=('alice', )) - p_bob = multiprocessing.Process(target=run, args=('bob', )) + p_alice = multiprocessing.Process(target=run, args=('alice',)) + p_bob = multiprocessing.Process(target=run, args=('bob',)) p_alice.start() import time @@ -85,4 +83,5 @@ def test_reset_context(): if __name__ == "__main__": import sys + sys.exit(pytest.main(["-sv", __file__])) diff --git a/fed/tests/test_retry_policy.py b/fed/tests/test_retry_policy.py index 574ce9fb..093d7631 100644 --- a/fed/tests/test_retry_policy.py +++ b/fed/tests/test_retry_policy.py @@ -15,6 +15,7 @@ import multiprocessing from unittest import TestCase + import pytest import ray diff --git a/fed/tests/test_transport_proxy.py b/fed/tests/test_transport_proxy.py index bb6f3f2d..de37e865 100644 --- a/fed/tests/test_transport_proxy.py +++ b/fed/tests/test_transport_proxy.py @@ -31,7 +31,8 @@ from fed.proxy.grpc.grpc_proxy import GrpcReceiverProxy, GrpcSenderProxy if compatible_utils._compare_version_strings( - fed_utils.get_package_version('protobuf'), '4.0.0'): + fed_utils.get_package_version('protobuf'), '4.0.0' +): from fed.grpc.pb4 import fed_pb2 as fed_pb2 from fed.grpc.pb4 import fed_pb2_grpc as fed_pb2_grpc else: @@ -99,8 +100,9 @@ def test_n_to_1_transport(): class TestSendDataService(fed_pb2_grpc.GrpcServiceServicer): - def __init__(self, all_events, all_data, party, lock, - expected_metadata, expected_jobname): + def __init__( + self, all_events, all_data, party, lock, expected_metadata, expected_jobname + ): self.expected_metadata = expected_metadata or {} self._expected_jobname = expected_jobname or "" @@ -109,8 +111,9 @@ async def SendData(self, request, context): assert self._expected_jobname == job_name metadata = dict(context.invocation_metadata()) for k, v in self.expected_metadata.items(): - assert k in metadata, \ - f"The expected key {k} is not in the metadata keys: {metadata.keys()}." + assert ( + k in metadata + ), f"The expected key {k} is not in the metadata keys: {metadata.keys()}." assert v == metadata[k] event = asyncio.Event() event.set() @@ -129,9 +132,10 @@ async def _test_run_grpc_server( ): server = grpc.aio.server(options=grpc_options) fed_pb2_grpc.add_GrpcServiceServicer_to_server( - TestSendDataService(event, all_data, party, lock, - expected_metadata, expected_jobname), - server + TestSendDataService( + event, all_data, party, lock, expected_metadata, expected_jobname + ), + server, ) server.add_insecure_port(f'[::]:{port}') await server.start() @@ -154,13 +158,13 @@ def __init__( async def run_grpc_server(self): return await _test_run_grpc_server( - self._listen_addr[self._listen_addr.index(':') + 1:], + self._listen_addr[self._listen_addr.index(':') + 1 :], None, None, self._party, None, expected_metadata=self._expected_metadata, - expected_jobname=self._expected_jobname + expected_jobname=self._expected_jobname, ) async def is_ready(self): @@ -178,9 +182,12 @@ def _test_start_receiver_proxy( address = addresses[party] receiver_proxy_actor = TestReceiverProxyActor.options( name=receiver_proxy_actor_name(), max_concurrency=1000 - ).remote(listen_addr=address, party=party, - expected_metadata=expected_metadata, - expected_jobname=expected_jobname) + ).remote( + listen_addr=address, + party=party, + expected_metadata=expected_metadata, + expected_jobname=expected_jobname, + ) receiver_proxy_actor.run_grpc_server.remote() assert ray.get(receiver_proxy_actor.is_ready.remote()) @@ -214,7 +221,7 @@ def test_send_grpc_with_meta(): addresses, party_name, expected_metadata=metadata, - expected_jobname=test_job_name + expected_jobname=test_job_name, ) _start_sender_proxy( addresses, diff --git a/fed/tests/test_utils.py b/fed/tests/test_utils.py index f17f1a6f..d42bde8b 100644 --- a/fed/tests/test_utils.py +++ b/fed/tests/test_utils.py @@ -13,15 +13,16 @@ # limitations under the License. import time + import pytest import fed.utils as fed_utils def start_ray_cluster( - ray_port, - client_server_port, - dashboard_port, + ray_port, + client_server_port, + dashboard_port, ): command = [ 'ray', @@ -45,8 +46,9 @@ def start_ray_cluster( # container, you can increase /dev/shm size by passing '--shm-size=1.97gb' to # 'docker run' (or add it to the run_options list in a Ray cluster config). # Make sure to set this to more than 0% of available RAM. - assert 'Overwriting previous Ray address' in str(e) \ - or 'WARNING: The object store is using /tmp instead of /dev/shm' in str(e) + assert 'Overwriting previous Ray address' in str( + e + ) or 'WARNING: The object store is using /tmp instead of /dev/shm' in str(e) @pytest.fixture diff --git a/fed/tests/without_ray_tests/test_tree_utils.py b/fed/tests/without_ray_tests/test_tree_utils.py index 41729ad0..416b5d3a 100644 --- a/fed/tests/without_ray_tests/test_tree_utils.py +++ b/fed/tests/without_ray_tests/test_tree_utils.py @@ -1,4 +1,3 @@ - # Copyright 2023 The RayFed Team # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,9 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any, Dict, List, Tuple, Union + import pytest -from typing import Any, Union, List, Tuple, Dict import fed.tree_util as tree_utils @@ -28,7 +28,6 @@ def test_flatten_none(): def test_flatten_single_primivite_elements(): - def _assert_flatten_single_element(target: Any): li, tree_def = tree_utils.tree_flatten(target) assert isinstance(li, list) diff --git a/fed/tests/without_ray_tests/test_utils.py b/fed/tests/without_ray_tests/test_utils.py index b00e042e..98233250 100644 --- a/fed/tests/without_ray_tests/test_utils.py +++ b/fed/tests/without_ray_tests/test_utils.py @@ -1,4 +1,3 @@ - # Copyright 2023 The RayFed Team # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -18,18 +17,21 @@ import fed -@pytest.mark.parametrize("input_address, is_valid_address", [ - ("192.168.0.1:8080", True), - ("sa127032as:80", True), - ("https://www.example.com", True), - ("http://www.example.com", True), - ("local", True), - ("localhost", True), - (None, False), - ("invalid_string", False), - ("http", False), - ("example.com", False), -]) +@pytest.mark.parametrize( + "input_address, is_valid_address", + [ + ("192.168.0.1:8080", True), + ("sa127032as:80", True), + ("https://www.example.com", True), + ("http://www.example.com", True), + ("local", True), + ("localhost", True), + (None, False), + ("invalid_string", False), + ("http", False), + ("example.com", False), + ], +) def test_validate_address(input_address, is_valid_address): if is_valid_address: fed.utils.validate_address(input_address) @@ -43,4 +45,5 @@ def test_validate_address(input_address, is_valid_address): if __name__ == "__main__": import sys + sys.exit(pytest.main(["-sv", __file__])) diff --git a/fed/tree_util.py b/fed/tree_util.py index 14a99fc3..033960d2 100644 --- a/fed/tree_util.py +++ b/fed/tree_util.py @@ -14,10 +14,9 @@ # Most codes are copied from https://github.com/pytorch/pytorch/blob/c263bd43e8e8502d4726643bc6fd046f0130ac0e/torch/utils/_pytree.py # noqa -from typing import NamedTuple, Callable, Any, Tuple, List, Dict, Type, cast, TypeVar -from collections import namedtuple, OrderedDict +from collections import OrderedDict, namedtuple from dataclasses import dataclass - +from typing import Any, Callable, Dict, List, NamedTuple, Tuple, Type, TypeVar, cast T = TypeVar('T') S = TypeVar('S') @@ -63,9 +62,8 @@ class NodeDef(NamedTuple): def _register_pytree_node( - typ: Any, - flatten_fn: FlattenFunc, - unflatten_fn: UnflattenFunc) -> None: + typ: Any, flatten_fn: FlattenFunc, unflatten_fn: UnflattenFunc +) -> None: SUPPORTED_NODES[typ] = NodeDef(flatten_fn, unflatten_fn) @@ -161,8 +159,11 @@ def __repr__(self, indent: int = 0) -> str: children_specs_str += self.children_specs[0].__repr__(indent) children_specs_str += ',' if len(self.children_specs) > 1 else '' children_specs_str += ','.join( - ['\n' + ' ' * indent + child.__repr__(indent) - for child in self.children_specs[1:]]) + [ + '\n' + ' ' * indent + child.__repr__(indent) + for child in self.children_specs[1:] + ] + ) repr_suffix: str = f'{children_specs_str}])' return repr_prefix + repr_suffix @@ -188,8 +189,8 @@ def tree_flatten(pytree: PyTree) -> Tuple[List[Any], TreeSpec]: child_pytrees, context = flatten_fn(pytree) # Recursively flatten the children - result : List[Any] = [] - children_specs : List['TreeSpec'] = [] + result: List[Any] = [] + children_specs: List['TreeSpec'] = [] for child in child_pytrees: flat, child_spec = tree_flatten(child) result += flat @@ -205,12 +206,14 @@ def tree_unflatten(values: List[Any], spec: TreeSpec) -> PyTree: if not isinstance(spec, TreeSpec): raise ValueError( f'tree_unflatten(values, spec): Expected `spec` to be instance of ' - f'TreeSpec but got item of type {type(spec)}.') + f'TreeSpec but got item of type {type(spec)}.' + ) if len(values) != spec.num_leaves: raise ValueError( f'tree_unflatten(values, spec): `values` has length {len(values)} ' f'but the spec refers to a pytree that holds {spec.num_leaves} ' - f'items ({spec}).') + f'items ({spec}).' + ) if isinstance(spec, LeafSpec): return values[0] diff --git a/fed/utils.py b/fed/utils.py index b5450f26..59b6fe79 100644 --- a/fed/utils.py +++ b/fed/utils.py @@ -14,8 +14,8 @@ import logging import re -import sys import subprocess +import sys import ray @@ -239,17 +239,16 @@ def validate_addresses(addresses: dict): validate_address(address) -def start_command(command: str, timeout=60) : +def start_command(command: str, timeout=60): """ A util to start a shell command. """ process = subprocess.Popen( - command, - shell=True, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE) + command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) output, error = process.communicate(timeout=timeout) if len(error) != 0: raise RuntimeError( - f'Failed to start command [{command}], the error is:\n {error.decode()}') + f'Failed to start command [{command}], the error is:\n {error.decode()}' + ) return output diff --git a/setup.py b/setup.py index 468ff6ad..1305b97d 100644 --- a/setup.py +++ b/setup.py @@ -66,8 +66,10 @@ def run(self): name=package_name, version=VERSION, license='Apache 2.0', - description='A multiple parties joint, distributed execution engine based on Ray,' - 'to help build your own federated learning frameworks in minutes.', + description=( + 'A multiple parties joint, distributed execution engine based on Ray,' + 'to help build your own federated learning frameworks in minutes.' + ), long_description=long_description, long_description_content_type='text/markdown', author='RayFed Team', diff --git a/tool/generate_tls_certs.py b/tool/generate_tls_certs.py index 7b9a7897..a0372007 100644 --- a/tool/generate_tls_certs.py +++ b/tool/generate_tls_certs.py @@ -13,9 +13,9 @@ # limitations under the License. import datetime +import errno import os import socket -import errno def try_make_directory_shared(directory_path):