mirror of
https://github.com/wassname/ray.git
synced 2026-06-28 08:40:02 +08:00
Simon Mo
35 lines
718 B
Python
35 lines
718 B
Python
# coding: utf-8
|
|
import asyncio
|
|
import sys
|
|
|
|
import pytest
|
|
import numpy as np
|
|
|
|
import ray
|
|
from ray.cluster_utils import Cluster
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_asyncio_cluster_wait():
|
|
cluster = Cluster()
|
|
head_node = cluster.add_node()
|
|
cluster.add_node(resources={"OTHER_NODE": 100})
|
|
|
|
ray.init(address=head_node.address)
|
|
|
|
@ray.remote(num_cpus=0, resources={"OTHER_NODE": 1})
|
|
def get_array():
|
|
return np.random.random((192, 1080, 3)).astype(np.uint8) # ~ 0.5MB
|
|
|
|
object_ref = get_array.remote()
|
|
|
|
await asyncio.wait_for(object_ref, timeout=10)
|
|
|
|
ray.shutdown()
|
|
cluster.shutdown()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import pytest
|
|
sys.exit(pytest.main(["-v", __file__]))
|