wehub-resource-sync
75 lines
1.9 KiB
Python
75 lines
1.9 KiB
Python
import sys
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import (
|
|
PrometheusTimeseries,
|
|
wait_for_condition,
|
|
)
|
|
from ray.tests.test_task_metrics import METRIC_CONFIG, tasks_by_all
|
|
|
|
|
|
# Copied from similar test in test_reconstruction_2.py.
|
|
@pytest.mark.skipif(sys.platform == "win32", reason="No multi-node on Windows.")
|
|
def test_task_reconstruction(ray_start_cluster):
|
|
timeseries = PrometheusTimeseries()
|
|
cluster = ray_start_cluster
|
|
|
|
# Head node with no resources.
|
|
cluster.add_node(
|
|
num_cpus=0,
|
|
**METRIC_CONFIG,
|
|
)
|
|
info = ray.init(address=cluster.address)
|
|
|
|
node_to_kill = cluster.add_node(num_cpus=1, object_store_memory=10**8)
|
|
cluster.wait_for_nodes()
|
|
|
|
@ray.remote
|
|
def large_object():
|
|
print("RUN LARGE OBJECT")
|
|
return np.zeros(10**7, dtype=np.uint8)
|
|
|
|
@ray.remote
|
|
def dependent_task(x):
|
|
print("RUN DEP TASK")
|
|
return np.zeros(10**7, dtype=np.uint8)
|
|
|
|
obj = large_object.remote()
|
|
x = dependent_task.remote(obj)
|
|
|
|
expected = {
|
|
("large_object", "FINISHED", "0"): 1.0,
|
|
("dependent_task", "FINISHED", "0"): 1.0,
|
|
}
|
|
wait_for_condition(
|
|
lambda: tasks_by_all(info, timeseries) == expected,
|
|
timeout=20,
|
|
retry_interval_ms=500,
|
|
)
|
|
|
|
cluster.remove_node(node_to_kill, allow_graceful=False)
|
|
cluster.add_node(num_cpus=1, object_store_memory=10**8)
|
|
|
|
# Triggers reconstruction.
|
|
ray.get(x)
|
|
|
|
# No failures, yet IsRetry=1.
|
|
expected = {
|
|
("large_object", "FINISHED", "0"): 1.0,
|
|
("large_object", "FINISHED", "1"): 1.0,
|
|
("dependent_task", "FINISHED", "0"): 1.0,
|
|
("dependent_task", "FINISHED", "1"): 1.0,
|
|
}
|
|
wait_for_condition(
|
|
lambda: tasks_by_all(info, timeseries) == expected,
|
|
timeout=20,
|
|
retry_interval_ms=500,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|