wehub-resource-sync
66 lines
2.0 KiB
Python
66 lines
2.0 KiB
Python
import json
|
|
import sys
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import wait_for_condition
|
|
from ray._private.internal_api import get_memory_info_reply, get_state_from_address
|
|
from ray._private.test_utils import (
|
|
RPC_FAILURE_MAP,
|
|
RPC_FAILURE_TYPES,
|
|
)
|
|
from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
|
|
|
|
|
|
@pytest.mark.parametrize("deterministic_failure", RPC_FAILURE_TYPES)
|
|
def test_free_objects_idempotent(
|
|
monkeypatch, shutdown_only, deterministic_failure, ray_start_cluster
|
|
):
|
|
failure = RPC_FAILURE_MAP[deterministic_failure].copy()
|
|
failure["num_failures"] = 1
|
|
monkeypatch.setenv(
|
|
"RAY_testing_rpc_failure",
|
|
json.dumps({"ObjectManagerService.grpc_client.FreeObjects": failure}),
|
|
)
|
|
|
|
@ray.remote
|
|
def simple_task(big_object_ref_list):
|
|
ray.get(big_object_ref_list[0])
|
|
return "ok"
|
|
|
|
cluster = ray_start_cluster
|
|
remote_node_1 = cluster.add_node(num_cpus=1)
|
|
remote_node_2 = cluster.add_node(num_cpus=1)
|
|
ray.init(address=cluster.address)
|
|
|
|
big_object_ref = ray.put(np.zeros(100 * 1024 * 1024))
|
|
|
|
# Propagate the big object to the remote nodes' plasma stores
|
|
result_ref_1 = simple_task.options(
|
|
scheduling_strategy=NodeAffinitySchedulingStrategy(
|
|
node_id=remote_node_1.node_id, soft=False
|
|
)
|
|
).remote([big_object_ref])
|
|
result_ref_2 = simple_task.options(
|
|
scheduling_strategy=NodeAffinitySchedulingStrategy(
|
|
node_id=remote_node_2.node_id, soft=False
|
|
)
|
|
).remote([big_object_ref])
|
|
|
|
assert ray.get([result_ref_1, result_ref_2]) == ["ok", "ok"]
|
|
|
|
del big_object_ref
|
|
|
|
def get_cluster_memory_usage():
|
|
state = get_state_from_address(ray.get_runtime_context().gcs_address)
|
|
reply = get_memory_info_reply(state)
|
|
return reply.store_stats.object_store_bytes_used
|
|
|
|
wait_for_condition(lambda: get_cluster_memory_usage() == 0, timeout=30)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|