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))