wehub-resource-sync
90 lines
3.3 KiB
Python
90 lines
3.3 KiB
Python
import os
|
|
import sys
|
|
import time
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
|
|
from ray.util.client.ray_client_helpers import ray_start_client_server
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info >= (3, 12), reason="flaky test, remove when fixed")
|
|
def test_dataclient_disconnect_on_request():
|
|
# Client can't signal graceful shutdown to server after unrecoverable
|
|
# error. Lower grace period so we don't have to sleep as long before
|
|
# checking new connection data.
|
|
with patch.dict(
|
|
os.environ, {"RAY_CLIENT_RECONNECT_GRACE_PERIOD": "5"}
|
|
), ray_start_client_server() as ray:
|
|
assert ray.is_connected()
|
|
|
|
@ray.remote
|
|
def f():
|
|
return 42
|
|
|
|
assert ray.get(f.remote()) == 42
|
|
# Force grpc to error by sending garbage request
|
|
with pytest.raises(ConnectionError):
|
|
ray.worker.data_client._blocking_send(Mock())
|
|
|
|
# Client should be disconnected
|
|
assert not ray.is_connected()
|
|
|
|
# Test that a new connection can be made
|
|
time.sleep(5) # Give server time to clean up old connection
|
|
connection_data = ray.connect("localhost:50051")
|
|
assert connection_data["num_clients"] == 1
|
|
assert ray.get(f.remote()) == 42
|
|
|
|
|
|
@pytest.mark.skipif(sys.version_info >= (3, 12), reason="flaky test, remove when fixed")
|
|
def test_dataclient_disconnect_before_request():
|
|
# Client can't signal graceful shutdown to server after unrecoverable
|
|
# error. Lower grace period so we don't have to sleep as long before
|
|
# checking new connection data.
|
|
with patch.dict(
|
|
os.environ, {"RAY_CLIENT_RECONNECT_GRACE_PERIOD": "5"}
|
|
), ray_start_client_server() as ray:
|
|
assert ray.is_connected()
|
|
|
|
@ray.remote
|
|
def f():
|
|
return 42
|
|
|
|
assert ray.get(f.remote()) == 42
|
|
# Force grpc to error by queueing garbage request. This simulates
|
|
# the data channel shutting down for connection issues between
|
|
# different remote calls.
|
|
ray.worker.data_client.request_queue.put(Mock())
|
|
|
|
# The following two assertions are relatively brittle. Consider a more
|
|
# robust mechanism if they fail with code changes or become flaky.
|
|
|
|
# The next remote call should error since the data channel will shut
|
|
# down because of the invalid input above. Two cases can happen:
|
|
# (1) Data channel shuts down after `f.remote()` finishes.
|
|
# error is raised to `ray.get()`. The next background operation
|
|
# will disconnect Ray client.
|
|
# (2) Data channel shuts down before `f.remote()` is called.
|
|
# `f.remote()` will raise the error and disconnect the client.
|
|
with pytest.raises(ConnectionError):
|
|
ray.get(f.remote())
|
|
|
|
with pytest.raises(
|
|
ConnectionError, match="Ray client has already been disconnected"
|
|
):
|
|
ray.get(f.remote())
|
|
|
|
# Client should be disconnected
|
|
assert not ray.is_connected()
|
|
|
|
# Test that a new connection can be made
|
|
time.sleep(5) # Give server time to clean up old connection
|
|
connection_data = ray.connect("localhost:50051")
|
|
assert connection_data["num_clients"] == 1
|
|
assert ray.get(f.remote()) == 42
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|