def test_failed_worker_pod(self):
"""
Start a pod running a python program destined to fail with
restart_policy="Never" to test failed_worker_count
"""
task_d = _TaskDispatcher({"f": (0, 10)}, {}, {}, 1, 1)
task_d.recover_tasks = MagicMock()
instance_manager = InstanceManager(
task_d,
job_name="test-failed-worker-pod-%d-%d" %
(int(time.time()), random.randint(1, 101)),
image_name="ubuntu:18.04",
worker_command=["/bin/bash"],
worker_args=["-c", "badcommand"],
namespace="default",
num_workers=3,
restart_policy="Never",
)
instance_manager.start_workers()
max_check_num = 20
for _ in range(max_check_num):
time.sleep(3)
counters = instance_manager.get_pod_counter(pod_type="worker")
if counters["Failed"] == 3:
break
instance_manager.stop_relaunch_and_remove_pods(pod_type="worker")
for _ in range(max_check_num):
time.sleep(3)
counters = instance_manager.get_pod_counter(pod_type="worker")
if not counters:
break
task_d.recover_tasks.assert_has_calls(
[call(0), call(1), call(2)], any_order=True)
def test_relaunch_ps_pod(self):
num_ps = 3
instance_manager = InstanceManager(
task_d=None,
job_name="test-relaunch-ps-pod-%d-%d" %
(int(time.time()), random.randint(1, 101)),
image_name="ubuntu:18.04",
ps_command=["/bin/bash"],
ps_args=["-c", "sleep 10 #"],
namespace="default",
num_ps=num_ps,
)
instance_manager.start_parameter_servers()
# Check we also have ps services started
for i in range(num_ps):
service = instance_manager._k8s_client.get_ps_service(i)
self.assertTrue(service.metadata.owner_references)
owner = service.metadata.owner_references[0]
self.assertEqual(owner.kind, "Pod")
self.assertEqual(owner.name,
instance_manager._k8s_client.get_ps_pod_name(i))
max_check_num = 60
for _ in range(max_check_num):
time.sleep(1)
counters = instance_manager.get_pod_counter(pod_type="ps")
if counters["Running"] + counters["Pending"] > 0:
break
# Note: There is a slight chance of race condition.
# Hack to find a ps to remove
all_current_ps = set()
all_live_ps = set()
with instance_manager._lock:
for k, (_, phase) in instance_manager._ps_pods_phase.items():
all_current_ps.add(k)
if phase in ["Running", "Pending"]:
all_live_ps.add(k)
self.assertTrue(all_live_ps)
ps_to_be_removed = all_live_ps.pop()
all_current_ps.remove(ps_to_be_removed)
instance_manager._remove_ps(ps_to_be_removed)
# Verify a new ps gets launched
found = False
for _ in range(max_check_num):
if found:
break
time.sleep(1)
with instance_manager._lock:
for k in instance_manager._ps_pods_phase:
if k not in all_current_ps:
found = True
else:
self.fail("Failed to find newly launched ps.")
instance_manager.stop_relaunch_and_remove_pods(pod_type="ps")
def test_relaunch_worker_pod(self):
num_workers = 3
task_d = _TaskDispatcher({"f": (0, 10)}, {}, {}, 1, 1)
instance_manager = InstanceManager(
task_d,
job_name="test-relaunch-worker-pod-%d-%d" %
(int(time.time()), random.randint(1, 101)),
image_name="ubuntu:18.04",
worker_command=["/bin/bash"],
worker_args=["-c", "sleep 10 #"],
namespace="default",
num_workers=num_workers,
)
instance_manager.start_workers()
max_check_num = 60
for _ in range(max_check_num):
time.sleep(1)
counters = instance_manager.get_pod_counter(pod_type="worker")
if counters["Running"] + counters["Pending"] > 0:
break
# Note: There is a slight chance of race condition.
# Hack to find a worker to remove
current_workers = set()
live_workers = set()
with instance_manager._lock:
for (
k,
(_, _, phase),
) in instance_manager._worker_pods_ip_phase.items():
current_workers.add(k)
if phase in ["Running", "Pending"]:
live_workers.add(k)
self.assertTrue(live_workers)
instance_manager._remove_worker(live_workers.pop())
# verify a new worker get launched
found = False
for _ in range(max_check_num):
if found:
break
time.sleep(1)
with instance_manager._lock:
for k in instance_manager._worker_pods_ip_phase:
if k not in range(num_workers, num_workers * 2):
found = True
else:
self.fail("Failed to find newly launched worker.")
instance_manager.stop_relaunch_and_remove_pods(pod_type="worker")
def test_create_delete_worker_pod(self):
task_d = _TaskDispatcher({"f": (0, 10)}, {}, {}, 1, 1)
task_d.recover_tasks = MagicMock()
instance_manager = InstanceManager(
task_d,
job_name="test-create-worker-pod-%d-%d" %
(int(time.time()), random.randint(1, 101)),
image_name="ubuntu:18.04",
worker_command=["/bin/bash"],
worker_args=["-c", "echo"],
namespace="default",
num_workers=2,
)
instance_manager.start_workers()
max_check_num = 20
for _ in range(max_check_num):
time.sleep(3)
counters = instance_manager.get_pod_counter(pod_type="worker")
if counters["Succeeded"] == 2:
break
instance_manager._not_created_worker_id = [2]
instance_manager._worker_pod_priority[2] = None
instance_manager._process_worker()
for _ in range(max_check_num):
time.sleep(3)
counters = instance_manager.get_pod_counter(pod_type="worker")
if counters["Succeeded"] == 3:
break
instance_manager.stop_relaunch_and_remove_pods(pod_type="worker")
for _ in range(max_check_num):
time.sleep(3)
counters = instance_manager.get_pod_counter(pod_type="worker")
if not counters:
break
self.assertFalse(counters)
def test_get_worker_addrs(self):
task_d = _TaskDispatcher({"f": (0, 10)}, {}, {}, 1, 1)
instance_manager = InstanceManager(
task_d,
job_name="test-create-worker-pod-%d-%d" %
(int(time.time()), random.randint(1, 101)),
image_name="ubuntu:18.04",
worker_command=["/bin/bash"],
worker_args=["-c", "sleep 5 #"],
namespace="default",
num_workers=3,
)
instance_manager.start_workers()
max_check_num = 20
for _ in range(max_check_num):
time.sleep(3)
counters = instance_manager.get_pod_counter(pod_type="worker")
if counters["Running"]:
worker_addrs = instance_manager._get_alive_worker_addr()
self.assertEqual(len(worker_addrs), counters["Running"])
instance_manager.stop_relaunch_and_remove_pods(pod_type="worker")