refine

fed/_private/global_context.py

@@ -19,9 +19,6 @@ def __init__(self) -> None:
         self._seq_count = 0
         self._cleanup_manager = CleanupManager()
 
-    def set_exit_on_failure_sending(self, func):
-        self._cleanup_manager.set_exit_on_failure_sending(func)
-
     def next_seq_id(self):
         self._seq_count += 1
         return self._seq_count
@@ -39,4 +36,5 @@ def get_global_context():
 
 def clear_global_context():
     global _global_context
+    _global_context._cleanup_manager._stop_gracefully()
     _global_context = None

fed/api.py

@@ -34,7 +34,6 @@
     start_recv_proxy,
     start_send_proxy,
 )
-# from fed.cleanup import set_exit_on_failure_sending, wait_sending
 from fed.fed_object import FedObject
 from fed.utils import is_ray_object_refs, setup_logger
 
@@ -251,9 +250,7 @@ def shutdown():
     """
     Shutdown a RayFed client.
     """
-    # wait_sending()
     compatible_utils._clear_internal_kv()
-    get_global_context()._cleanup_manager._stop_gracefully()
     clear_global_context()
     logger.info('Shutdowned rayfed.')