chore: import upstream snapshot with attribution

python/paddle/distributed/launch/utils/etcd_client.py

@@ -0,0 +1,180 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import time
+
+import etcd3
+
+
+class ETCDClient:
+    def __init__(self, host, port, retry_times=20):
+        self.retry_times = retry_times
+        times = 0
+        while times < self.retry_times:
+            try:
+                self.client = etcd3.client(host=host, port=port)
+                break
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Initialize etcd client failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(
+                f"Initialize etcd client failed failed after {self.retry_times} times."
+            )
+
+    def put(self, key, value, lease=None, prev_kv=False):
+        times = 0
+        while times < self.retry_times:
+            try:
+                return self.client.put(key, value, lease, prev_kv)
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Put failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(f"Put failed after {self.retry_times} times.")
+
+    def get(self, key):
+        times = 0
+        while times < self.retry_times:
+            try:
+                return self.client.get(key)
+
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Get {key} failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(
+                f"Get {key} failed after {self.retry_times} times."
+            )
+
+    def delete(self, key, prev_kv=False, return_response=False):
+        times = 0
+        while times < self.retry_times:
+            try:
+                return self.client.delete(key, prev_kv, return_response)
+                break
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Delete {key} failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(
+                f"Delete {key} failed after {self.retry_times} times."
+            )
+
+    def get_prefix(self, key_prefix, sort_order=None, sort_target='key'):
+        times = 0
+        while times < self.retry_times:
+            try:
+                return self.client.get_prefix(key_prefix)
+                break
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Get prefix {key_prefix} failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(
+                f"Get prefix {key_prefix} failed after {self.retry_times} times."
+            )
+
+    def delete_prefix(self, prefix):
+        times = 0
+        while times < self.retry_times:
+            try:
+                return self.client.delete_prefix(prefix)
+                break
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Delete prefix {prefix} failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(
+                f"Delete prefix {prefix} failed after {self.retry_times} times."
+            )
+
+    def lease(self, ttl, lease_id=None):
+        times = 0
+        while times < self.retry_times:
+            try:
+                return self.client.lease(ttl, lease_id)
+                break
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Lease failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(f"Lease failed after {self.retry_times} times.")
+
+    def add_watch_prefix_callback(self, key_prefix, callback, **kwargs):
+        times = 0
+        while times < self.retry_times:
+            try:
+                return self.client.add_watch_prefix_callback(
+                    key_prefix, callback, **kwargs
+                )
+                break
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Add watch prefix callback failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(
+                f"Add watch prefix callback failed after {self.retry_times} times."
+            )
+
+    def cancel_watch(self, watch_id):
+        times = 0
+        while times < self.retry_times:
+            try:
+                return self.client.cancel_watch(watch_id)
+                break
+            except Exception as e:
+                times += 1
+                logging.info(
+                    f"Cancel watch failed with exception {e}, retry after 1 second."
+                )
+                time.sleep(1)
+
+        if times >= self.retry_times:
+            raise ValueError(
+                f"Cancel watch failed after {self.retry_times} times."
+            )