diff --git a/dlrover/python/elastic_agent/master_client.py b/dlrover/python/elastic_agent/master_client.py index c1b2b4905..e646da113 100644 --- a/dlrover/python/elastic_agent/master_client.py +++ b/dlrover/python/elastic_agent/master_client.py @@ -29,7 +29,7 @@ def retry_grpc_request(func): def wrapper(self, *args, **kwargs): retry = kwargs.get("retry", 10) - exception = None + execption = None for i in range(retry): try: return func(self, *args, **kwargs) @@ -39,22 +39,15 @@ def wrapper(self, *args, **kwargs): logger.warning( f"Retry {i} to {class_name}.{func_name} with failure", ) - exception = e + execption = e time.sleep(5) - if exception: - logger.error(exception) - raise exception + if execption: + logger.error(execption) + raise execption return wrapper -def init_grpc_env(): - # At the cost of increased performance overhead, these provide greater - # stability in concurrent scenarios. - env_utils.set_env("GRPC_ENABLE_FORK_SUPPORT", "true") - env_utils.set_env("GRPC_POLL_STRATEGY", "poll") - - class MasterClient(Singleton): """MasterClient provides some APIs connect with the master service via gRPC call. @@ -82,7 +75,6 @@ def __init__(self, master_addr, node_id, node_type, timeout=5): f"Build master client with master_addr: {master_addr}, " f"node_id: {node_id}, node_type: {node_type}." ) - init_grpc_env() self._timeout = timeout self._master_addr = master_addr self._channel = grpc.build_channel(master_addr) diff --git a/dlrover/python/tests/test_master_client.py b/dlrover/python/tests/test_master_client.py index 7bebb8120..e24f296ff 100644 --- a/dlrover/python/tests/test_master_client.py +++ b/dlrover/python/tests/test_master_client.py @@ -15,7 +15,7 @@ import time import unittest -from dlrover.python.common import env_utils, grpc +from dlrover.python.common import grpc from dlrover.python.common.constants import ( NodeStatus, NodeType, @@ -35,9 +35,6 @@ def tearDown(self): self._master.stop() def test_open_channel(self): - self.assertEqual(env_utils.get_env("GRPC_ENABLE_FORK_SUPPORT"), "true") - self.assertEqual(env_utils.get_env("GRPC_POLL_STRATEGY"), "poll") - self.assertEqual(self._master_client._timeout, 0.5) self.assertEqual(self._master_client._timeout, 0.5) self._master_client.close_channel() self._master_client.open_channel()