def _prepare_test_graph(self, session_id, graph_key, mock_workers):
addr = f'127.0.0.1:{get_next_port()}'
a1 = mt.random.random((100,))
a2 = mt.random.random((100,))
s = a1 + a2
v1, v2 = mt.split(s, 2)
graph = TileableGraph([v1.data, v2.data])
builder = TileableGraphBuilder(graph)
next(iter(builder.build()))
with create_actor_pool(n_process=1, backend='gevent', address=addr) as pool:
pool.create_actor(SchedulerClusterInfoActor, [pool.cluster_info.address],
uid=SchedulerClusterInfoActor.default_uid())
resource_ref = pool.create_actor(ResourceActor, uid=ResourceActor.default_uid())
pool.create_actor(ChunkMetaActor, uid=ChunkMetaActor.default_uid())
pool.create_actor(AssignerActor, uid=AssignerActor.gen_uid(session_id))
graph_ref = pool.create_actor(GraphActor, session_id, graph_key, serialize_graph(graph),
uid=GraphActor.gen_uid(session_id, graph_key))
for w in mock_workers:
resource_ref.set_worker_meta(w, dict(hardware=dict(cpu=4, cpu_total=4, memory=1600)))
graph_ref.prepare_graph()
graph_ref.analyze_graph()
graph_ref.create_operand_actors(_start=False)
yield pool, graph_ref
async def test_run_task(actor_pool):
pool, session_id, meta_api, storage_api, manager = actor_pool
raw = np.random.RandomState(0).rand(10, 10)
a = mt.tensor(raw, chunk_size=5)
b = a + 1
graph = TileableGraph([b.data])
next(TileableGraphBuilder(graph).build())
task_id = await manager.submit_tileable_graph(graph, fuse_enabled=False)
assert isinstance(task_id, str)
await manager.wait_task(task_id)
task_result: TaskResult = await manager.get_task_result(task_id)
assert task_result.status == TaskStatus.terminated
assert task_result.error is None
result_tileables = (await manager.get_task_result_tileables(task_id))[0]
for i, chunk in enumerate(result_tileables.chunks):
result = await storage_api.get(chunk.key)
if i == 0:
expect = raw[:5, :5] + 1
elif i == 1:
expect = raw[:5, 5:] + 1
elif i == 2:
expect = raw[5:, :5] + 1
else:
expect = raw[5:, 5:] + 1
np.testing.assert_array_equal(result, expect)
def _gen_subtask(t, session_id):
graph = TileableGraph([t.data])
next(TileableGraphBuilder(graph).build())
chunk_graph = next(ChunkGraphBuilder(graph, fuse_enabled=False).build())
subtask = Subtask(new_task_id(), session_id, new_task_id(), chunk_graph)
return subtask
async def test_cancel_task(actor_pool):
pool, session_id, meta_api, storage_api, manager = actor_pool
def func():
time.sleep(20)
rs = [mr.spawn(func) for _ in range(10)]
graph = TileableGraph([r.data for r in rs])
next(TileableGraphBuilder(graph).build())
task_id = await manager.submit_tileable_graph(graph, fuse_enabled=False)
assert isinstance(task_id, str)
await asyncio.sleep(.5)
with Timer() as timer:
await manager.cancel_task(task_id)
result = await manager.get_task_result(task_id)
assert result.status == TaskStatus.terminated
assert timer.duration < 15