wehub-resource-sync
55 lines
1.8 KiB
Python
55 lines
1.8 KiB
Python
import sys
|
|
from time import sleep
|
|
|
|
import pytest
|
|
|
|
from ray._common.test_utils import wait_for_condition
|
|
from ray.tests.conftest_docker import * # noqa
|
|
|
|
|
|
# TODO(sang): Also check temp dir
|
|
@pytest.mark.skipif(sys.platform != "linux", reason="Only works on linux.")
|
|
def test_ray_session_name_preserved(docker_cluster):
|
|
get_nodes_script = """
|
|
import ray
|
|
ray.init("auto")
|
|
print(ray._private.worker._global_node.session_name)
|
|
"""
|
|
head, worker = docker_cluster
|
|
|
|
def get_session_name(to_head=True):
|
|
if to_head:
|
|
output = head.exec_run(cmd=f"python -c '{get_nodes_script}'")
|
|
else:
|
|
output = worker.exec_run(cmd=f"python -c '{get_nodes_script}'")
|
|
session_name = output.output.decode().strip().split("\n")[-1]
|
|
print("Output: ", output.output.decode().strip().split("\n"))
|
|
assert output.exit_code == 0
|
|
return session_name
|
|
|
|
# Make sure two nodes are alive
|
|
wait_for_condition(get_session_name, to_head=True)
|
|
session_name_head = get_session_name(to_head=True)
|
|
wait_for_condition(get_session_name, to_head=False)
|
|
session_name_worker = get_session_name(to_head=False)
|
|
assert session_name_head == session_name_worker
|
|
print("head killed")
|
|
head.kill()
|
|
|
|
sleep(2)
|
|
|
|
head.restart()
|
|
|
|
wait_for_condition(get_session_name, timeout=30, to_head=True)
|
|
session_name_head_after_restart = get_session_name(to_head=True)
|
|
wait_for_condition(get_session_name, to_head=False)
|
|
session_name_worker_after_restart = get_session_name(to_head=False)
|
|
assert session_name_worker_after_restart == session_name_head_after_restart
|
|
assert session_name_head == session_name_head_after_restart
|
|
assert session_name_worker_after_restart == session_name_worker
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
sys.exit(pytest.main(["-sv", __file__]))
|