fix aoa engine time long

Author: xingmingyyj

python/paddle/distributed/flex_checkpoint/aoa/aoa_engine.py

@@ -92,15 +92,18 @@ def __init__(
         self.destination_state_shard_info = destination_state_shard_info
         self.left_var_to_right_var_mapping = {}
         self.right_var_from_left_var_mapping = {}
+        self.dst_state_keys = set()
+        self.init_dst_state_keys()
 
-    def get_all_dst_state_keys(self):
-        dst_state_keys = set()
+    def init_dst_state_keys(self):
         if self.destination_state_shard_info is None:
-            return dst_state_keys
+            return
         for k in self.destination_state_shard_info.keys():
             model_state_key, _ = split_optimizer_state_key(k)
-            dst_state_keys.add(model_state_key)
-        return dst_state_keys
+            self.dst_state_keys.add(model_state_key)
+
+    def get_all_dst_state_keys(self):
+        return self.dst_state_keys
 
     def get_all_src_state_keys(self):
         src_state_keys = set()