async def test_assign_gpu_tasks(actor_pool): pool, session_id, assigner_ref, cluster_api, meta_api = actor_pool input1 = TensorFetch(key='a', source_key='a', dtype=np.dtype(int)).new_chunk([]) input2 = TensorFetch(key='b', source_key='b', dtype=np.dtype(int)).new_chunk([]) result_chunk = TensorTreeAdd(args=[input1, input2], gpu=True) \ .new_chunk([input1, input2]) chunk_graph = ChunkGraph([result_chunk]) chunk_graph.add_node(input1) chunk_graph.add_node(input2) chunk_graph.add_node(result_chunk) chunk_graph.add_edge(input1, result_chunk) chunk_graph.add_edge(input2, result_chunk) await meta_api.set_chunk_meta(input1, memory_size=200, store_size=200, bands=[('address0', 'numa-0')]) await meta_api.set_chunk_meta(input2, memory_size=200, store_size=200, bands=[('address0', 'numa-0')]) subtask = Subtask('test_task', session_id, chunk_graph=chunk_graph) [result] = await assigner_ref.assign_subtasks([subtask]) assert result[1].startswith('gpu')
def testMockExecuteSize(self): import mars.tensor as mt from mars.core.graph import DAG from mars.tensor.fetch import TensorFetch from mars.tensor.arithmetic import TensorTreeAdd graph_add = DAG() input_chunks = [] for _ in range(2): fetch_op = TensorFetch(dtype=np.dtype('int64')) inp_chunk = fetch_op.new_chunk(None, shape=(100, 100)).data input_chunks.append(inp_chunk) add_op = TensorTreeAdd(args=input_chunks, dtype=np.dtype('int64')) add_chunk = add_op.new_chunk(input_chunks, shape=(100, 100), dtype=np.dtype('int64')).data graph_add.add_node(add_chunk) for inp_chunk in input_chunks: graph_add.add_node(inp_chunk) graph_add.add_edge(inp_chunk, add_chunk) executor = Executor() res = executor.execute_graph(graph_add, [add_chunk.key], compose=False, mock=True)[0] self.assertEqual(res, (80000, 80000)) self.assertEqual(executor.mock_max_memory, 80000) for _ in range(3): new_add_op = TensorTreeAdd(args=[add_chunk], dtype=np.dtype('int64')) new_add_chunk = new_add_op.new_chunk([add_chunk], shape=(100, 100), dtype=np.dtype('int64')).data graph_add.add_node(new_add_chunk) graph_add.add_edge(add_chunk, new_add_chunk) add_chunk = new_add_chunk executor = Executor() res = executor.execute_graph(graph_add, [add_chunk.key], compose=False, mock=True)[0] self.assertEqual(res, (80000, 80000)) self.assertEqual(executor.mock_max_memory, 160000) a = mt.random.rand(10, 10, chunk_size=10) b = a[:, mt.newaxis, :] - a r = mt.triu(mt.sqrt(b**2).sum(axis=2)) executor = Executor() res = executor.execute_tensor(r, concat=False, mock=True) # larger than maximal memory size in calc procedure self.assertGreaterEqual(res[0][0], 800) self.assertGreaterEqual(executor.mock_max_memory, 8000)
async def test_execute_tensor(actor_pool): pool, session_id, meta_api, storage_api, execution_ref = actor_pool data1 = np.random.rand(10, 10) data2 = np.random.rand(10, 10) input1 = TensorFetch(key='input1', source_key='input2', dtype=np.dtype(int)).new_chunk([]) input2 = TensorFetch(key='input2', source_key='input2', dtype=np.dtype(int)).new_chunk([]) result_chunk = TensorTreeAdd(args=[input1, input2]) \ .new_chunk([input1, input2], shape=data1.shape, dtype=data1.dtype) await meta_api.set_chunk_meta(input1, memory_size=data1.nbytes, store_size=data1.nbytes, bands=[(pool.external_address, 'numa-0')]) await meta_api.set_chunk_meta(input2, memory_size=data1.nbytes, store_size=data2.nbytes, bands=[(pool.external_address, 'numa-0')]) # todo use different storage level when storage ready await storage_api.put(input1.key, data1) await storage_api.put(input2.key, data2) chunk_graph = ChunkGraph([result_chunk]) chunk_graph.add_node(input1) chunk_graph.add_node(input2) chunk_graph.add_node(result_chunk) chunk_graph.add_edge(input1, result_chunk) chunk_graph.add_edge(input2, result_chunk) subtask = Subtask('test_task', session_id=session_id, chunk_graph=chunk_graph) await execution_ref.run_subtask(subtask, 'numa-0', pool.external_address) # check if results are correct result = await storage_api.get(result_chunk.key) np.testing.assert_array_equal(data1 + data2, result) # check if quota computations are correct quota_ref = await mo.actor_ref(QuotaActor.gen_uid('numa-0'), address=pool.external_address) [quota] = await quota_ref.get_batch_quota_reqs() assert quota[(subtask.subtask_id, subtask.subtask_id)] == data1.nbytes # check if metas are correct result_meta = await meta_api.get_chunk_meta(result_chunk.key) assert result_meta['object_id'] == result_chunk.key assert result_meta['shape'] == result.shape
def testPrepareSpilled(self): from mars.worker.spill import write_spill_file pool_address = '127.0.0.1:%d' % get_next_port() session_id = str(uuid.uuid4()) mock_data = np.array([1, 2, 3, 4]) options.worker.spill_directory = tempfile.mkdtemp(prefix='mars_worker_prep_spilled-') with create_actor_pool(n_process=1, backend='gevent', address=pool_address) as pool: self.create_standard_actors(pool, pool_address, with_daemon=False, with_status=False) pool.create_actor(SpillActor) pool.create_actor(CpuCalcActor) cluster_info_ref = pool.actor_ref(WorkerClusterInfoActor.default_uid()) chunk_meta_client = ChunkMetaClient(pool, cluster_info_ref) pool.actor_ref(ChunkHolderActor.default_uid()) import mars.tensor as mt from mars.tensor.fetch import TensorFetch arr = mt.ones((4,), chunk_size=4) arr_add = mt.array(mock_data) result_tensor = arr + arr_add graph = result_tensor.build_graph(compose=False, tiled=True) modified_chunk = arr_add.chunks[0] arr_add.chunks[0]._op = TensorFetch( dtype=modified_chunk.dtype, _outputs=[weakref.ref(o) for o in modified_chunk.op.outputs], _key=modified_chunk.op.key) # test meta missing with self.run_actor_test(pool) as test_actor: graph_key = str(uuid.uuid4()) execution_ref = test_actor.promise_ref(ExecutionActor.default_uid()) execution_ref.execute_graph(session_id, graph_key, serialize_graph(graph), dict(chunks=[result_tensor.chunks[0].key]), None, _promise=True) \ .then(lambda *_: test_actor.set_result(None)) \ .catch(lambda *exc: test_actor.set_result(exc, False)) with self.assertRaises(DependencyMissing): self.get_result() chunk_meta_client.set_chunk_meta(session_id, modified_chunk.key, size=mock_data.nbytes, shape=mock_data.shape, workers=('0.0.0.0:1234', pool_address)) write_spill_file(modified_chunk.key, mock_data) # test read from spilled file with self.run_actor_test(pool) as test_actor: def _validate(_): data = test_actor._chunk_store.get(session_id, result_tensor.chunks[0].key) assert_array_equal(data, mock_data + np.ones((4,))) graph_key = str(uuid.uuid4()) execution_ref = test_actor.promise_ref(ExecutionActor.default_uid()) execution_ref.execute_graph(session_id, graph_key, serialize_graph(graph), dict(chunks=[result_tensor.chunks[0].key]), None, _promise=True) \ .then(_validate) \ .then(lambda *_: test_actor.set_result(None)) \ .catch(lambda *exc: test_actor.set_result(exc, False)) self.get_result()
def testPrepareQuota(self, *_): pinned = [True] def _mock_pin(_graph_key, chunk_keys): from mars.errors import PinChunkFailed if pinned[0]: raise PinChunkFailed return chunk_keys ChunkHolderActor.pin_chunks.side_effect = _mock_pin pool_address = '127.0.0.1:%d' % get_next_port() session_id = str(uuid.uuid4()) mock_data = np.array([1, 2, 3, 4]) with create_actor_pool(n_process=1, backend='gevent', address=pool_address) as pool: self.create_standard_actors(pool, pool_address, with_daemon=False, with_status=False) pool.create_actor(MockSenderActor, mock_data, 'in', uid='w:mock_sender') pool.create_actor(CpuCalcActor) cluster_info_ref = pool.actor_ref(WorkerClusterInfoActor.default_uid()) chunk_meta_client = ChunkMetaClient(pool, cluster_info_ref) import mars.tensor as mt from mars.tensor.fetch import TensorFetch arr = mt.ones((4,), chunk_size=4) arr_add = mt.array(mock_data) result_tensor = arr + arr_add graph = result_tensor.build_graph(compose=False, tiled=True) modified_chunk = arr_add.chunks[0] arr_add.chunks[0]._op = TensorFetch( dtype=modified_chunk.dtype, _outputs=[weakref.ref(o) for o in modified_chunk.op.outputs], _key=modified_chunk.op.key) chunk_meta_client.set_chunk_meta(session_id, modified_chunk.key, size=mock_data.nbytes, shape=mock_data.shape, workers=('0.0.0.0:1234', pool_address)) with self.run_actor_test(pool) as test_actor: graph_key = str(uuid.uuid4()) execution_ref = test_actor.promise_ref(ExecutionActor.default_uid()) start_time = time.time() execution_ref.execute_graph( session_id, graph_key, serialize_graph(graph), dict(chunks=[result_tensor.chunks[0].key]), None, _tell=True) execution_ref.add_finish_callback(session_id, graph_key, _promise=True) \ .then(lambda *_: test_actor.set_result(time.time())) \ .catch(lambda *exc: test_actor.set_result(exc, False)) def _delay_fun(): time.sleep(1) pinned[0] = False threading.Thread(target=_delay_fun).start() finish_time = self.get_result() self.assertGreaterEqual(finish_time, start_time + 1)
def testCopyAdd(self): t1 = ones((3, 4), chunk_size=2) t2 = ones(4, chunk_size=2) t3 = t1 + t2 t3.tiles() c = t3.chunks[0] inputs = c.op.lhs, TensorFetch().new_chunk( c.op.rhs.inputs, shape=c.op.rhs.shape, index=c.op.rhs.index, _key=c.op.rhs.key) new_c = c.op.copy().reset_key().new_chunk(inputs, shape=c.shape, _key='new_key') self.assertEqual(new_c.key, 'new_key') self.assertIs(new_c.inputs[1], new_c.op.rhs) self.assertIsInstance(new_c.inputs[1].op, TensorFetch)
async def test_assigner(actor_pool): pool, session_id, assigner_ref, meta_api = actor_pool input1 = TensorFetch(key='a', source_key='a', dtype=np.dtype(int)).new_chunk([]) input2 = TensorFetch(key='b', source_key='b', dtype=np.dtype(int)).new_chunk([]) input3 = TensorFetch(key='c', source_key='c', dtype=np.dtype(int)).new_chunk([]) result_chunk = TensorTreeAdd(args=[input1, input2, input3]) \ .new_chunk([input1, input2, input3]) chunk_graph = ChunkGraph([result_chunk]) chunk_graph.add_node(input1) chunk_graph.add_node(input2) chunk_graph.add_node(input3) chunk_graph.add_node(result_chunk) chunk_graph.add_edge(input1, result_chunk) chunk_graph.add_edge(input2, result_chunk) chunk_graph.add_edge(input3, result_chunk) await meta_api.set_chunk_meta(input1, memory_size=200, store_size=200, bands=[('address0', 'numa-0')]) await meta_api.set_chunk_meta(input2, memory_size=400, store_size=400, bands=[('address1', 'numa-0')]) await meta_api.set_chunk_meta(input3, memory_size=400, store_size=400, bands=[('address2', 'numa-0')]) subtask = Subtask('test_task', session_id, chunk_graph=chunk_graph) [result] = await assigner_ref.assign_subtasks([subtask]) assert result in (('address1', 'numa-0'), ('address2', 'numa-0'))
def test_copy_add(): t1 = ones((3, 4), chunk_size=2) t2 = ones(4, chunk_size=2) t3 = t1 + t2 t3 = tile(t3) c = t3.chunks[0] inputs = c.op.lhs, TensorFetch().new_chunk(c.op.rhs.inputs, shape=c.op.rhs.shape, index=c.op.rhs.index, _key=c.op.rhs.key) new_c = c.op.copy().reset_key().new_chunk(inputs, shape=c.shape, _key='new_key') assert new_c.key == 'new_key' assert new_c.inputs[1] is new_c.op.rhs assert isinstance(new_c.inputs[1].op, TensorFetch)
def _build_test_graph(data_list): from mars.tensor.fetch import TensorFetch from mars.tensor.arithmetic import TensorTreeAdd inputs = [] for idx, d in enumerate(data_list): chunk_key = 'chunk-%d' % idx fetch_chunk = TensorFetch(to_fetch_key=chunk_key, dtype=d.dtype) \ .new_chunk([], shape=d.shape, _key=chunk_key) inputs.append(fetch_chunk) add_chunk = TensorTreeAdd(data_list[0].dtype).new_chunk(inputs, shape=data_list[0].shape) exec_graph = DAG() exec_graph.add_node(add_chunk) for input_chunk in inputs: exec_graph.add_node(input_chunk) exec_graph.add_edge(input_chunk, add_chunk) return exec_graph, inputs, add_chunk
def _build_test_graph(data_list): from mars.tensor.fetch import TensorFetch from mars.tensor.arithmetic import TensorTreeAdd inputs = [] for idx, d in enumerate(data_list): chunk_key = f'chunk-{random.randint(0, 999)}-{idx}' fetch_chunk = TensorFetch(to_fetch_key=chunk_key, dtype=d.dtype) \ .new_chunk([], shape=d.shape, _key=chunk_key) inputs.append(fetch_chunk) add_chunk = TensorTreeAdd(args=inputs, dtype=data_list[0].dtype) \ .new_chunk(inputs, shape=data_list[0].shape) exec_graph = ChunkGraph([add_chunk.data]) exec_graph.add_node(add_chunk.data) for input_chunk in inputs: exec_graph.add_node(input_chunk.data) exec_graph.add_edge(input_chunk.data, add_chunk.data) return exec_graph, inputs, add_chunk
def testPrepareSpilled(self): pool_address = '127.0.0.1:%d' % get_next_port() session_id = str(uuid.uuid4()) mock_data = np.array([1, 2, 3, 4]) options.worker.spill_directory = tempfile.mkdtemp(prefix='mars_worker_prep_spilled-') with create_actor_pool(n_process=1, backend='gevent', address=pool_address) as pool: self.create_standard_actors(pool, pool_address, with_daemon=False, with_status=False) pool.create_actor(IORunnerActor) pool.create_actor(CpuCalcActor) pool.create_actor(InProcHolderActor) import mars.tensor as mt from mars.tensor.fetch import TensorFetch arr = mt.ones((4,), chunk_size=4) arr_add = mt.array(mock_data) result_tensor = arr + arr_add graph = result_tensor.build_graph(compose=False, tiled=True) arr_add = get_tiled(arr_add) result_tensor = get_tiled(result_tensor) modified_chunk = arr_add.chunks[0] arr_add.chunks[0]._op = TensorFetch( dtype=modified_chunk.dtype, _outputs=[weakref.ref(o) for o in modified_chunk.op.outputs], _key=modified_chunk.op.key) # test meta missing with self.run_actor_test(pool) as test_actor: graph_key = str(uuid.uuid4()) execution_ref = test_actor.promise_ref(ExecutionActor.default_uid()) execution_ref.execute_graph(session_id, graph_key, serialize_graph(graph), dict(chunks=[result_tensor.chunks[0].key]), None, _promise=True) \ .then(lambda *_: test_actor.set_result(None)) \ .catch(lambda *exc: test_actor.set_result(exc, False)) with self.assertRaises(DependencyMissing): self.get_result() metas = {modified_chunk.key: WorkerMeta( mock_data.nbytes, mock_data.shape, ('0.0.0.0:1234', pool_address))} # test read from spilled file with self.run_actor_test(pool) as test_actor: self.waitp( test_actor.storage_client.put_objects( session_id, [modified_chunk.key], [mock_data], [DataStorageDevice.PROC_MEMORY]) .then(lambda *_: test_actor.storage_client.copy_to( session_id, [modified_chunk.key], [DataStorageDevice.DISK])) ) test_actor.storage_client.delete(session_id, [modified_chunk.key], [DataStorageDevice.PROC_MEMORY]) def _validate(_): data = test_actor.shared_store.get(session_id, result_tensor.chunks[0].key) assert_array_equal(data, mock_data + np.ones((4,))) graph_key = str(uuid.uuid4()) execution_ref = test_actor.promise_ref(ExecutionActor.default_uid()) execution_ref.execute_graph(session_id, graph_key, serialize_graph(graph), dict(chunks=[result_tensor.chunks[0].key]), metas, _promise=True) \ .then(_validate) \ .then(lambda *_: test_actor.set_result(None)) \ .catch(lambda *exc: test_actor.set_result(exc, False)) self.get_result()
def testFetchRemoteData(self): pool_address = '127.0.0.1:%d' % get_next_port() session_id = str(uuid.uuid4()) mock_data = np.array([1, 2, 3, 4]) with create_actor_pool(n_process=1, backend='gevent', address=pool_address, distributor=MarsDistributor(2, 'w:0:')) as pool: self.create_standard_actors(pool, pool_address, with_daemon=False, with_status=False, with_resource=True) pool.create_actor(CpuCalcActor) pool.create_actor(InProcHolderActor) pool.create_actor(MockSenderActor, [mock_data], 'in', uid='w:mock_sender') import mars.tensor as mt from mars.tensor.fetch import TensorFetch arr = mt.ones((4,), chunk_size=4) arr_add = mt.array(mock_data) result_tensor = arr + arr_add graph = result_tensor.build_graph(compose=False, tiled=True) arr_add = get_tiled(arr_add) result_tensor = get_tiled(result_tensor) modified_chunk = arr_add.chunks[0] arr_add.chunks[0]._op = TensorFetch( dtype=modified_chunk.dtype, _outputs=[weakref.ref(o) for o in modified_chunk.op.outputs], _key=modified_chunk.op.key) with self.run_actor_test(pool) as test_actor: graph_key = str(uuid.uuid4()) execution_ref = test_actor.promise_ref(ExecutionActor.default_uid()) execution_ref.execute_graph(session_id, graph_key, serialize_graph(graph), dict(chunks=[result_tensor.chunks[0].key]), None, _tell=True) execution_ref.add_finish_callback(session_id, graph_key, _promise=True) \ .then(lambda *_: test_actor.set_result(None)) \ .catch(lambda *exc: test_actor.set_result(exc, False)) with self.assertRaises(DependencyMissing): self.get_result() metas = {modified_chunk.key: WorkerMeta(mock_data.nbytes, mock_data.shape, ('0.0.0.0:1234',))} with self.run_actor_test(pool) as test_actor: graph_key = str(uuid.uuid4()) execution_ref = test_actor.promise_ref(ExecutionActor.default_uid()) execution_ref.execute_graph(session_id, graph_key, serialize_graph(graph), dict(chunks=[result_tensor.chunks[0].key]), metas, _tell=True) execution_ref.add_finish_callback(session_id, graph_key, _promise=True) \ .then(lambda *_: test_actor.set_result(None)) \ .catch(lambda *exc: test_actor.set_result(exc, False)) with self.assertRaises(DependencyMissing): self.get_result() metas[modified_chunk.key] = WorkerMeta( mock_data.nbytes, mock_data.shape, ('0.0.0.0:1234', pool_address.replace('127.0.0.1', 'localhost'))) with self.run_actor_test(pool) as test_actor: def _validate(_): data = test_actor.shared_store.get(session_id, result_tensor.chunks[0].key) assert_array_equal(data, mock_data + np.ones((4,))) graph_key = str(uuid.uuid4()) execution_ref = test_actor.promise_ref(ExecutionActor.default_uid()) execution_ref.execute_graph(session_id, graph_key, serialize_graph(graph), dict(chunks=[result_tensor.chunks[0].key]), metas, _tell=True) execution_ref.add_finish_callback(session_id, graph_key, _promise=True) \ .then(_validate) \ .then(lambda *_: test_actor.set_result(None)) \ .catch(lambda *exc: test_actor.set_result(exc, False)) self.get_result()
def testPrepareQuota(self, *_): pinned = True orig_pin = SharedHolderActor.pin_data_keys def _mock_pin(self, session_id, chunk_keys, token): from mars.errors import PinDataKeyFailed if pinned: raise PinDataKeyFailed return orig_pin(self, session_id, chunk_keys, token) pool_address = '127.0.0.1:%d' % get_next_port() session_id = str(uuid.uuid4()) mock_data = np.array([1, 2, 3, 4]) with patch_method(SharedHolderActor.pin_data_keys, new=_mock_pin), \ create_actor_pool(n_process=1, backend='gevent', address=pool_address) as pool: self.create_standard_actors(pool, pool_address, with_daemon=False, with_status=False) pool.create_actor(MockSenderActor, [mock_data], 'in', uid='w:mock_sender') pool.create_actor(CpuCalcActor) pool.create_actor(InProcHolderActor) pool.actor_ref(WorkerClusterInfoActor.default_uid()) import mars.tensor as mt from mars.tensor.fetch import TensorFetch arr = mt.ones((4,), chunk_size=4) arr_add = mt.array(mock_data) result_tensor = arr + arr_add graph = result_tensor.build_graph(compose=False, tiled=True) arr_add = get_tiled(arr_add) result_tensor = get_tiled(result_tensor) modified_chunk = arr_add.chunks[0] arr_add.chunks[0]._op = TensorFetch( dtype=modified_chunk.dtype, _outputs=[weakref.ref(o) for o in modified_chunk.op.outputs], _key=modified_chunk.op.key) metas = {modified_chunk.key: WorkerMeta( mock_data.nbytes, mock_data.shape, ('0.0.0.0:1234', pool_address.replace('127.0.0.1', 'localhost')))} with self.run_actor_test(pool) as test_actor: graph_key = str(uuid.uuid4()) execution_ref = test_actor.promise_ref(ExecutionActor.default_uid()) start_time = time.time() execution_ref.execute_graph( session_id, graph_key, serialize_graph(graph), dict(chunks=[result_tensor.chunks[0].key]), metas, _tell=True) execution_ref.add_finish_callback(session_id, graph_key, _promise=True) \ .then(lambda *_: test_actor.set_result(time.time())) \ .catch(lambda *exc: test_actor.set_result(exc, False)) def _delay_fun(): nonlocal pinned time.sleep(0.5) pinned = False threading.Thread(target=_delay_fun).start() finish_time = self.get_result() self.assertGreaterEqual(finish_time, start_time + 0.5)
def testSimpleExecution(self): pool_address = '127.0.0.1:%d' % get_next_port() with create_actor_pool(n_process=1, backend='gevent', address=pool_address) as pool: self.create_standard_actors(pool, pool_address, with_daemon=False) pool.create_actor(CpuCalcActor, uid='w:1:calc-a') pool.create_actor(InProcHolderActor) import mars.tensor as mt from mars.tensor.datasource import TensorOnes from mars.tensor.fetch import TensorFetch arr = mt.ones((10, 8), chunk_size=10) arr_add = mt.ones((10, 8), chunk_size=10) arr2 = arr + arr_add graph = arr2.build_graph(compose=False, tiled=True) arr = get_tiled(arr) arr2 = get_tiled(arr2) metas = dict() for chunk in graph: if isinstance(chunk.op, TensorOnes): chunk._op = TensorFetch( dtype=chunk.dtype, _outputs=[weakref.ref(o) for o in chunk.op.outputs], _key=chunk.op.key) metas[chunk.key] = WorkerMeta(chunk.nbytes, chunk.shape, pool_address) with self.run_actor_test(pool) as test_actor: session_id = str(uuid.uuid4()) storage_client = test_actor.storage_client self.waitp( storage_client.put_objects(session_id, [arr.chunks[0].key], [np.ones((10, 8), dtype=np.int16)], [DataStorageDevice.SHARED_MEMORY]), ) execution_ref = test_actor.promise_ref(ExecutionActor.default_uid()) def _validate(_): data = test_actor.shared_store.get(session_id, arr2.chunks[0].key) assert_array_equal(data, 2 * np.ones((10, 8))) graph_key = str(uuid.uuid4()) execution_ref.execute_graph(session_id, graph_key, serialize_graph(graph), dict(chunks=[arr2.chunks[0].key]), metas, _promise=True) \ .then(_validate) \ .then(lambda *_: test_actor.set_result(None)) \ .catch(lambda *exc: test_actor.set_result(exc, False)) self.get_result() with self.run_actor_test(pool) as test_actor: execution_ref = test_actor.promise_ref(ExecutionActor.default_uid()) def _validate(_): data = test_actor.shared_store.get(session_id, arr2.chunks[0].key) assert_array_equal(data, 2 * np.ones((10, 8))) execution_ref.add_finish_callback(session_id, graph_key, _promise=True) \ .then(_validate) \ .then(lambda *_: test_actor.set_result(None)) \ .catch(lambda *exc: test_actor.set_result(exc, False)) self.get_result()
async def test_execute_with_cancel(actor_pool, cancel_phase): pool, session_id, meta_api, storage_api, execution_ref = actor_pool # config for different phases ref_to_delay = None if cancel_phase == 'prepare': ref_to_delay = await mo.actor_ref(StorageManagerActor.default_uid(), address=pool.external_address) elif cancel_phase == 'quota': ref_to_delay = await mo.actor_ref(QuotaActor.gen_uid('numa-0'), address=pool.external_address) elif cancel_phase == 'slot': ref_to_delay = await mo.actor_ref( BandSlotManagerActor.gen_uid('numa-0'), address=pool.external_address) if ref_to_delay: await ref_to_delay.set_delay_fetch_time(100) def delay_fun(delay, _inp1): time.sleep(delay) return delay input1 = TensorFetch(key='input1', source_key='input1', dtype=np.dtype(int)).new_chunk([]) remote_result = RemoteFunction(function=delay_fun, function_args=[100, input1], function_kwargs={}, n_output=1) \ .new_chunk([input1]) data1 = np.random.rand(10, 10) await meta_api.set_chunk_meta(input1, memory_size=data1.nbytes, store_size=data1.nbytes, bands=[(pool.external_address, 'numa-0')]) await storage_api.put(input1.key, data1) chunk_graph = ChunkGraph([remote_result]) chunk_graph.add_node(input1) chunk_graph.add_node(remote_result) chunk_graph.add_edge(input1, remote_result) subtask = Subtask(f'test_task_{uuid.uuid4()}', session_id=session_id, chunk_graph=chunk_graph) aiotask = asyncio.create_task( execution_ref.run_subtask(subtask, 'numa-0', pool.external_address)) await asyncio.sleep(1) with Timer() as timer: await execution_ref.cancel_subtask(subtask.subtask_id, kill_timeout=1) with pytest.raises(asyncio.CancelledError): await asyncio.wait_for(aiotask, timeout=30) assert timer.duration < 6 # check for different phases if ref_to_delay is not None: assert await ref_to_delay.get_is_cancelled() await ref_to_delay.set_delay_fetch_time(0) # test if slot is restored remote_tileable = mr.spawn(delay_fun, args=(0.5, None)) graph = TileableGraph([remote_tileable.data]) next(TileableGraphBuilder(graph).build()) chunk_graph = next(ChunkGraphBuilder(graph, fuse_enabled=False).build()) subtask = Subtask(f'test_task2_{uuid.uuid4()}', session_id=session_id, chunk_graph=chunk_graph) await asyncio.wait_for(execution_ref.run_subtask(subtask, 'numa-0', pool.external_address), timeout=30)
async def test_assign_cpu_tasks(actor_pool): pool, session_id, assigner_ref, cluster_api, meta_api = actor_pool input1 = TensorFetch(key='a', source_key='a', dtype=np.dtype(int)).new_chunk([]) input2 = TensorFetch(key='b', source_key='b', dtype=np.dtype(int)).new_chunk([]) input3 = TensorFetch(key='c', source_key='c', dtype=np.dtype(int)).new_chunk([]) result_chunk = TensorTreeAdd(args=[input1, input2, input3]) \ .new_chunk([input1, input2, input3]) chunk_graph = ChunkGraph([result_chunk]) chunk_graph.add_node(input1) chunk_graph.add_node(input2) chunk_graph.add_node(input3) chunk_graph.add_node(result_chunk) chunk_graph.add_edge(input1, result_chunk) chunk_graph.add_edge(input2, result_chunk) chunk_graph.add_edge(input3, result_chunk) await meta_api.set_chunk_meta(input1, memory_size=200, store_size=200, bands=[('address0', 'numa-0')]) await meta_api.set_chunk_meta(input2, memory_size=400, store_size=400, bands=[('address1', 'numa-0')]) await meta_api.set_chunk_meta(input3, memory_size=400, store_size=400, bands=[('address2', 'numa-0')]) await cluster_api.set_node_status(node='address1', role=NodeRole.WORKER, status=NodeStatus.STOPPING) await cluster_api.set_node_status(node='address3', role=NodeRole.WORKER, status=NodeStatus.STOPPING) subtask = Subtask('test_task', session_id, chunk_graph=chunk_graph) [result] = await assigner_ref.assign_subtasks([subtask]) assert result in (('address0', 'numa-0'), ('address2', 'numa-0')) subtask.expect_bands = [('address0', 'numa-0')] [result] = await assigner_ref.assign_subtasks([subtask]) assert result == ('address0', 'numa-0') subtask.expect_bands = [('address0', 'numa-0'), ('address1', 'numa-0')] [result] = await assigner_ref.assign_subtasks([subtask]) assert result == ('address0', 'numa-0') subtask.expect_bands = [('address1', 'numa-0')] [result] = await assigner_ref.assign_subtasks([subtask]) assert result in (('address0', 'numa-0'), ('address2', 'numa-0')) result_chunk.op.gpu = True subtask = Subtask('test_task', session_id, chunk_graph=chunk_graph) with pytest.raises(NoMatchingSlots) as err: await assigner_ref.assign_subtasks([subtask]) assert 'gpu' in str(err.value)
def testSimpleExecution(self): pool_address = '127.0.0.1:%d' % get_next_port() with create_actor_pool(n_process=1, backend='gevent', address=pool_address) as pool: self.create_standard_actors(pool, pool_address, with_daemon=False) pool.create_actor(CpuCalcActor, uid='w:1:calc-a') import mars.tensor as mt from mars.tensor.datasource import TensorOnes from mars.tensor.fetch import TensorFetch arr = mt.ones((10, 8), chunk_size=10) arr_add = mt.ones((10, 8), chunk_size=10) arr2 = arr + arr_add graph = arr2.build_graph(compose=False, tiled=True) for chunk in graph: if isinstance(chunk.op, TensorOnes): chunk._op = TensorFetch( dtype=chunk.dtype, _outputs=[weakref.ref(o) for o in chunk.op.outputs], _key=chunk.op.key) with self.run_actor_test(pool) as test_actor: session_id = str(uuid.uuid4()) chunk_holder_ref = test_actor.promise_ref(ChunkHolderActor.default_uid()) refs = test_actor._chunk_store.put(session_id, arr.chunks[0].key, np.ones((10, 8), dtype=np.int16)) chunk_holder_ref.register_chunk(session_id, arr.chunks[0].key) del refs refs = test_actor._chunk_store.put(session_id, arr_add.chunks[0].key, np.ones((10, 8), dtype=np.int16)) chunk_holder_ref.register_chunk(session_id, arr_add.chunks[0].key) del refs execution_ref = test_actor.promise_ref(ExecutionActor.default_uid()) def _validate(_): data = test_actor._chunk_store.get(session_id, arr2.chunks[0].key) assert_array_equal(data, 2 * np.ones((10, 8))) graph_key = str(uuid.uuid4()) execution_ref.execute_graph(session_id, graph_key, serialize_graph(graph), dict(chunks=[arr2.chunks[0].key]), None, _promise=True) \ .then(_validate) \ .then(lambda *_: test_actor.set_result(None)) \ .catch(lambda *exc: test_actor.set_result(exc, False)) self.get_result() with self.run_actor_test(pool) as test_actor: execution_ref = test_actor.promise_ref(ExecutionActor.default_uid()) def _validate(_): data = test_actor._chunk_store.get(session_id, arr2.chunks[0].key) assert_array_equal(data, 2 * np.ones((10, 8))) execution_ref.add_finish_callback(session_id, graph_key, _promise=True) \ .then(_validate) \ .then(lambda *_: test_actor.set_result(None)) \ .catch(lambda *exc: test_actor.set_result(exc, False)) self.get_result()