Add watcher test. (#146)

Author: gongweibao

python/edl/tests/unittests/test_cluster_watcher.py

@@ -0,0 +1,77 @@
+# Copyright (c) 2020 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 time
+import unittest
+from edl.tests.unittests import etcd_test_base
+from edl.utils import cluster as edl_cluster
+from edl.utils import constants
+from edl.utils import exceptions
+from edl.utils import cluster_watcher
+from edl.utils import pod as edl_pod
+
+
+class TestWatcher(etcd_test_base.EtcdTestBase):
+ def setUp(self):
+ super(TestWatcher, self).setUp("test_watcher")
+
+ def test_watcher_stage_changed(self):
+ cluster = edl_cluster.Cluster()
+ cluster._stage = "0"
+ print("cluster 0 ids:", cluster.to_json(), cluster.get_pods_ids_list())
+ self._etcd.set_server_permanent(constants.ETCD_CLUSTER,
+ constants.ETCD_CLUSTER,
+ cluster.to_json())
+ watcher = cluster_watcher.Watcher(self._job_env, cluster)
+
+ cluster._stage = "1"
+ print("cluster 1 ids:", cluster.to_json(), cluster.get_pods_ids_list())
+ self._etcd.set_server_permanent(constants.ETCD_CLUSTER,
+ constants.ETCD_CLUSTER,
+ cluster.to_json())
+ time.sleep(constants.ETCD_TTL)
+ self.assertTrue(watcher.changed)
+
+ def test_watch_valid(self):
+ try:
+ cluster = edl_cluster.Cluster()
+ self._etcd.set_server_permanent(constants.ETCD_CLUSTER,
+ constants.ETCD_CLUSTER,
+ cluster.to_json())
+ watcher = cluster_watcher.Watcher(self._job_env, cluster)
+ self._etcd.remove_server(constants.ETCD_CLUSTER,
+ constants.ETCD_CLUSTER)
+ time.sleep(constants.ETCD_TTL)
+ except exceptions.EdlTableError as e:
+ pass
+
+ def test_watcher_ids_changed(self):
+ cluster = edl_cluster.Cluster()
+ print("cluster 0 ids:", cluster.to_json(), cluster.get_pods_ids_list())
+ self._etcd.set_server_permanent(constants.ETCD_CLUSTER,
+ constants.ETCD_CLUSTER,
+ cluster.to_json())
+ watcher = cluster_watcher.Watcher(self._job_env, cluster)
+
+ pod = edl_pod.Pod()
+ cluster._pods.append(pod)
+ print("cluster 1 ids:", cluster.to_json(), cluster.get_pods_ids_list())
+ self._etcd.set_server_permanent(constants.ETCD_CLUSTER,
+ constants.ETCD_CLUSTER,
+ cluster.to_json())
+ time.sleep(constants.ETCD_TTL)
+ self.assertTrue(watcher.changed)
+
+
+if __name__ == '__main__':
+ unittest.main()

python/edl/utils/cluster.py

@@ -20,17 +20,13 @@
 import time
 import traceback
 import uuid
-from edl.discovery import etcd_client
+
 from edl.utils import constants
 from edl.utils import error_utils
 from edl.utils import exceptions
 from edl.utils import json_serializable
-from edl.utils import leader_pod
 from edl.utils import pod as edl_pod
-from edl.utils import resource_pods as edl_resource_pods
 from edl.utils import status as edl_status
-from edl.utils import train_status as edl_train_status
-from edl.utils.log_utils import logger
 
 
 class Cluster(json_serializable.Serializable):
@@ -163,3 +159,15 @@ def load_from_etcd(etcd, timeout=60):
  cluster = Cluster()
  cluster.from_json(value)
  return cluster
+
+
+@error_utils.handle_errors_until_timeout
+def wait_to_load_from_etcd(etcd, timeout=60):
+ cluster = load_from_etcd(etcd, timeout=60)
+ if cluster is None:
+ raise exceptions.EdlTableError(
+ "can't load cluster from etcd path:{}".format(
+ etcd.get_full_path(constants.ETCD_CLUSTER,
+ constants.ETCD_CLUSTER)))
+
+ return cluster

python/edl/utils/watcher.py renamed to python/edl/utils/cluster_watcher.py

@@ -13,64 +13,44 @@
 # limitations under the License.
 
 import copy
+import threading
 import time
-from edl.discovery.etcd_client import EtcdClient
+from edl.discovery import etcd_client
 from edl.utils import cluster as edl_cluster
-from edl.utils import constants
 from edl.utils.log_utils import logger
-from threading import Lock, Thread, Event
 
 
 class Watcher(object):
- def __init__(self, job_env, cluster, pod):
- self._etcd = None
-
+ def __init__(self, job_env, cluster):
  self._job_id = job_env.job_id
 
  # current context
  self._cluster = copy.copy(cluster)
- self._leader_id = cluster.get_pod_leader_id()
- self._current_pod = pod
 
- self._new_cluster = None
- self._new_leader_id = None
+ self._new_cluster = cluster
  self._changed = False
  logger.info("watcher gets the init cluster:{}".format(self._cluster))
 
- self._lock = Lock()
- self._stop = Event()
+ self._lock = threading.Lock()
+ self._stop = threading.Event()
 
+ self._etcd = None
  self._t_watcher = None
+ self._job_env = job_env
 
  # assign value
- self._etcd = EtcdClient(self._job_env.etcd_endpoints, root=job_id)
+ self._etcd = etcd_client.EtcdClient(
+ self._job_env.etcd_endpoints, root=self._job_id)
  self._etcd.init()
 
- self._t_watcher = Thread(target=self._watcher)
+ self._t_watcher = threading.Thread(target=self._watcher)
  self._t_watcher.start()
 
  def _watcher(self):
- begin = time.time()
  while not self._stop.is_set():
- # if leader_id changed?
- servers = self._etcd.get_service(constants.ETCD_POD_RANK)
- assert len(servers) <= 1
- if len(servers) == 0:
- time.sleep(1)
- continue
-
- with self._lock:
- self._new_leader_id = s.info
-
  # if cluster changed?
- value, _, _, _, _, = etcd._get_server(constants.ETCD_CLUSTER,
- self._new_leader_id)
- if value is None:
- time.sleep(1)
- continue
- new_cluster = edl_cluster.Cluster()
- new_cluster.from_json(value)
-
+ new_cluster = edl_cluster.wait_to_load_from_etcd(
+ self._etcd, timeout=60)
  with self._lock:
  self._new_cluster = new_cluster