def testTaskQueueActor(self):
with create_actor_pool(n_process=1, backend='gevent') as pool:
pool.create_actor(MockExecutionActor,
10,
uid=ExecutionActor.default_name())
quota_ref = pool.create_actor(QuotaActor,
30,
uid=MemQuotaActor.default_name())
pool.create_actor(TaskQueueActor,
4,
uid=TaskQueueActor.default_name())
session_id = str(uuid.uuid4())
chunk_keys = [str(uuid.uuid4()).replace('-', '') for _ in range(6)]
with self.run_actor_test(pool) as test_actor:
queue_ref = test_actor.promise_ref(
TaskQueueActor.default_name())
res_times = dict()
def callback_fun(key):
res_times[key] = time.time()
for idx, k in enumerate(chunk_keys):
depth = len(chunk_keys) - idx
queue_ref.enqueue_task(session_id, k, dict(depth=depth), _promise=True) \
.then(functools.partial(callback_fun, k))
gevent.sleep(1)
self.assertEqual(queue_ref.get_allocated_count(), 4)
queue_ref.update_priority(session_id, chunk_keys[-1],
dict(depth=len(chunk_keys)))
quota_ref.release_quota(chunk_keys[0])
queue_ref.release_task(session_id, chunk_keys[0])
quota_ref.release_quota(chunk_keys[1])
queue_ref.release_task(session_id, chunk_keys[1])
gevent.sleep(0.5)
self.assertIn(chunk_keys[-1], res_times)
for k in chunk_keys[:3]:
self.assertLessEqual(res_times[k],
res_times[chunk_keys[-1]] - 0.5)
self.assertIn(k, res_times)
def run_test(self, worker):
import mars.tensor as mt
from mars.worker import ExecutionActor
session_id = str(uuid.uuid4())
a = mt.ones((100, 50), chunk_size=30)
b = mt.ones((50, 200), chunk_size=30)
result = a.dot(b)
graph = result.build_graph(tiled=True)
executor_ref = self.promise_ref(ExecutionActor.default_name(), address=worker)
io_meta = dict(chunks=[c.key for c in result.chunks])
graph_key = str(id(graph))
executor_ref.enqueue_graph(session_id, graph_key, serialize_graph(graph),
io_meta, None, _promise=True) \
.then(lambda *_: executor_ref.start_execution(session_id, graph_key, _promise=True)) \
.then(lambda *_: setattr(self, '_replied', True))
