def testChunkMetaCache(self):
cache = ChunkMetaCache(9)
for idx in range(10):
cache['c%d' % idx] = WorkerMeta(idx, (idx, ), ('w0', ))
self.assertNotIn('c0', cache)
self.assertTrue(all('c%d' % idx in cache for idx in range(1, 10)))
self.assertListEqual(sorted(cache.get_worker_chunk_keys('w0')),
['c%d' % idx for idx in range(1, 10)])
dup_cache = copy.deepcopy(cache)
dup_cache.get('c1')
dup_cache['c10'] = WorkerMeta(10, (10, ), ('w0', ))
self.assertIsNone(dup_cache.get('c0'))
self.assertNotIn('c2', dup_cache)
self.assertIn('c1', dup_cache)
self.assertTrue(all('c%d' % idx in dup_cache for idx in range(3, 11)))
dup_cache = copy.deepcopy(cache)
_ = dup_cache['c1'] # noqa: F841
dup_cache['c10'] = WorkerMeta(10, (10, ), ('w0', ))
self.assertNotIn('c2', dup_cache)
self.assertIn('c1', dup_cache)
self.assertTrue(all('c%d' % idx in dup_cache for idx in range(3, 11)))
dup_cache = copy.deepcopy(cache)
dup_cache['c1'] = WorkerMeta(1, (1, ), ('w0', ))
dup_cache['c10'] = WorkerMeta(10, (10, ), ('w0', ))
self.assertNotIn('c2', dup_cache)
self.assertIn('c1', dup_cache)
self.assertTrue(all('c%d' % idx in dup_cache for idx in range(3, 11)))
def testChunkMetaStore(self):
store = ChunkMetaStore()
store['c0'] = WorkerMeta(0, (0, ), ('w0', ))
self.assertIn('c0', store)
self.assertEqual(store['c0'], WorkerMeta(0, (0, ), ('w0', )))
self.assertEqual(store.get('c0'), WorkerMeta(0, (0, ), ('w0', )))
self.assertIsNone(store.get('c1'))
self.assertSetEqual(store.get_worker_chunk_keys('w0'), {'c0'})
store['c0'] = WorkerMeta(0, (0, ), ('w1', ))
self.assertEqual(store.get_worker_chunk_keys('w0'), set())
self.assertSetEqual(store.get_worker_chunk_keys('w1'), {'c0'})
del store['c0']
self.assertNotIn('c0', store)
store['c1'] = WorkerMeta(1, (1, ), ('w0', 'w1'))
store['c2'] = WorkerMeta(2, (2, ), ('w1', ))
store['c3'] = WorkerMeta(3, (3, ), ('w0', ))
store['c4'] = WorkerMeta(4, (4, ), ('w0', ))
affected = store.remove_worker_keys('w0', lambda k: k[-1] < '4')
self.assertListEqual(affected, ['c3'])
self.assertEqual(store.get('c1'), WorkerMeta(1, (1, ), ('w1', )))
self.assertEqual(store.get('c2'), WorkerMeta(2, (2, ), ('w1', )))
self.assertSetEqual(store.get_worker_chunk_keys('w0'), {'c4'})
self.assertSetEqual(store.get_worker_chunk_keys('w1'), {'c1', 'c2'})
self.assertNotIn('c3', store)
self.assertIn('c4', store)
affected = store.remove_worker_keys('w0')
self.assertListEqual(affected, ['c4'])
self.assertNotIn('c4', store)
self.assertIsNone(store.get_worker_chunk_keys('w0'))
self.assertSetEqual(store.get_worker_chunk_keys('w1'), {'c1', 'c2'})
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 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 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()
def testChunkMetaActors(self, *_):
proc_count = 2
endpoints = [
'127.0.0.1:%d' % get_next_port() for _ in range(proc_count)
]
keys = []
def _mock_get_scheduler(key):
return endpoints[keys.index(key[1]) % len(endpoints)]
ChunkMetaClient.get_scheduler.side_effect = _mock_get_scheduler
session1 = str(uuid.uuid4())
session2 = str(uuid.uuid4())
with create_actor_pool(n_process=1,
backend='gevent',
address=endpoints[0]) as pool1:
cluster_info1 = pool1.create_actor(
SchedulerClusterInfoActor,
endpoints,
uid=SchedulerClusterInfoActor.default_name())
pool1.create_actor(ChunkMetaActor,
uid=ChunkMetaActor.default_name())
with create_actor_pool(n_process=1,
backend='gevent',
address=endpoints[1]) as pool2:
cluster_info2 = pool2.create_actor(
SchedulerClusterInfoActor,
endpoints,
uid=SchedulerClusterInfoActor.default_name())
pool2.create_actor(ChunkMetaActor,
uid=ChunkMetaActor.default_name())
actor_client = new_client()
client1 = ChunkMetaClient(
actor_client, actor_client.actor_ref(cluster_info1))
client2 = ChunkMetaClient(
actor_client, actor_client.actor_ref(cluster_info2))
loc_ref1 = actor_client.actor_ref(
ChunkMetaActor.default_name(), address=endpoints[0])
loc_ref2 = actor_client.actor_ref(
ChunkMetaActor.default_name(), address=endpoints[1])
key1 = (str(uuid.uuid4()), str(uuid.uuid4()))
key2 = str(uuid.uuid4())
key3 = str(uuid.uuid4())
key4 = (str(uuid.uuid4()), str(uuid.uuid4()))
key5 = str(uuid.uuid4())
key6 = str(uuid.uuid4())
keys = [key1, key2, key3, key4, key5, key6]
client1.set_chunk_size(session1, key1, 512)
client2.set_chunk_size(session1, key2, 1024)
client2.set_chunk_size(session2, key3, 1024)
self.assertEqual(client1.get_chunk_size(session1, key1), 512)
self.assertEqual(client2.get_chunk_size(session1, key2), 1024)
self.assertEqual(client1.get_chunk_size(session1, key2), 1024)
self.assertEqual(client2.get_chunk_size(session1, key1), 512)
self.assertListEqual(
client1.batch_get_chunk_size(session1, [key1, key2]),
[512, 1024])
self.assertListEqual(
client2.batch_get_chunk_size(session1, [key1, key2]),
[512, 1024])
client1.set_chunk_shape(session1, key1, (10, ))
client2.set_chunk_shape(session1, key2, (10, ) * 2)
client2.set_chunk_shape(session2, key3, (10, ) * 2)
self.assertEqual(client1.get_chunk_shape(session1, key1),
(10, ))
self.assertEqual(client2.get_chunk_shape(session1, key2),
(10, ) * 2)
self.assertEqual(client1.get_chunk_shape(session1, key2),
(10, ) * 2)
self.assertEqual(client2.get_chunk_shape(session1, key1),
(10, ))
self.assertListEqual(
client1.batch_get_chunk_shape(session1, [key1, key2]),
[(10, ), (10, ) * 2])
self.assertListEqual(
client2.batch_get_chunk_shape(session1, [key1, key2]),
[(10, ), (10, ) * 2])
mock_endpoint = '127.0.0.1:%d' % get_next_port()
with create_actor_pool(n_process=1,
backend='gevent',
address=mock_endpoint) as pool3:
cluster_info3 = pool3.create_actor(
SchedulerClusterInfoActor,
endpoints,
uid=SchedulerClusterInfoActor.default_name())
client3 = ChunkMetaClient(
actor_client, actor_client.actor_ref(cluster_info3))
self.assertListEqual(
client3.batch_get_chunk_shape(session1, [key1, key2]),
[(10, ), (10, ) * 2])
client1.add_worker(session1, key1, 'abc')
client1.add_worker(session1, key1, 'def')
client2.add_worker(session1, key2, 'ghi')
client1.add_worker(session2, key3, 'ghi')
self.assertEqual(sorted(client1.get_workers(session1, key1)),
sorted(('abc', 'def')))
self.assertEqual(sorted(client2.get_workers(session1, key2)),
sorted(('ghi', )))
batch_result = client1.batch_get_workers(
session1, [key1, key2])
self.assertEqual(sorted(batch_result[0]), sorted(
('abc', 'def')))
self.assertEqual(sorted(batch_result[1]), sorted(('ghi', )))
affected = []
for loc_ref in (loc_ref1, loc_ref2):
affected.extend(
loc_ref.remove_workers_in_session(session2, ['ghi']))
self.assertEqual(affected, [key3])
self.assertEqual(sorted(client1.get_workers(session1, key2)),
sorted(('ghi', )))
self.assertIsNone(client1.get_workers(session2, key3))
client1.delete_meta(session1, key1)
self.assertIsNone(client1.get_workers(session1, key1))
self.assertIsNone(
client1.batch_get_chunk_size(session1, [key1, key2])[0])
self.assertIsNone(
client1.batch_get_workers(session1, [key1, key2])[0])
client2.batch_delete_meta(session1, [key1, key2])
self.assertIsNone(client1.get_workers(session1, key2))
self.assertIsNone(
client1.batch_get_chunk_size(session1, [key1, key2])[1])
self.assertIsNone(
client1.batch_get_workers(session1, [key1, key2])[1])
meta4 = WorkerMeta(chunk_size=512,
chunk_shape=(10, ) * 2,
workers=(endpoints[0], ))
loc_ref2.batch_set_chunk_meta(session1, [key4], [meta4])
self.assertEqual(
loc_ref2.get_chunk_meta(session1, key4).chunk_size, 512)
self.assertEqual(
loc_ref2.get_chunk_meta(session1, key4).chunk_shape,
(10, ) * 2)
meta5 = WorkerMeta(chunk_size=512,
chunk_shape=(10, ) * 2,
workers=(endpoints[0], ))
meta6 = WorkerMeta(chunk_size=512,
chunk_shape=(10, ) * 2,
workers=(endpoints[0], ))
client1.batch_set_chunk_meta(session1, [key5, key6],
[meta5, meta6])
self.assertEqual(
loc_ref1.get_chunk_meta(session1, key5).chunk_size, 512)
self.assertEqual(
loc_ref2.get_chunk_meta(session1, key6).chunk_size, 512)
def testChunkBroadcast(self, *_):
proc_count = 2
endpoints = [
'127.0.0.1:%d' % get_next_port() for _ in range(proc_count)
]
keys = []
def _mock_get_scheduler(key):
return endpoints[keys.index(key[1]) % len(endpoints)]
ChunkMetaClient.get_scheduler.side_effect = _mock_get_scheduler
session_id = str(uuid.uuid4())
with create_actor_pool(n_process=1,
backend='gevent',
address=endpoints[0]) as pool1:
cluster_info1 = pool1.create_actor(
SchedulerClusterInfoActor,
endpoints,
uid=SchedulerClusterInfoActor.default_name())
pool1.create_actor(ChunkMetaActor,
uid=ChunkMetaActor.default_name())
with create_actor_pool(n_process=1,
backend='gevent',
address=endpoints[1]) as pool2:
cluster_info2 = pool2.create_actor(
SchedulerClusterInfoActor,
endpoints,
uid=SchedulerClusterInfoActor.default_name())
pool2.create_actor(ChunkMetaActor,
uid=ChunkMetaActor.default_name())
actor_client = new_client()
client1 = ChunkMetaClient(
actor_client, actor_client.actor_ref(cluster_info1))
client2 = ChunkMetaClient(
actor_client, actor_client.actor_ref(cluster_info2))
local_ref1 = actor_client.actor_ref(
ChunkMetaActor.default_name(), address=endpoints[0])
local_ref2 = actor_client.actor_ref(
ChunkMetaActor.default_name(), address=endpoints[1])
key1 = str(uuid.uuid4())
key2 = str(uuid.uuid4())
key3 = str(uuid.uuid4())
keys = [key1, key2, key3]
client1.set_chunk_broadcasts(session_id, key1, [endpoints[1]])
client1.set_chunk_size(session_id, key1, 512)
client1.set_chunk_shape(session_id, key1, (10, ) * 2)
client1.add_worker(session_id, key1, 'abc')
client2.set_chunk_broadcasts(session_id, key2, [endpoints[0]])
client2.set_chunk_size(session_id, key2, 512)
client1.set_chunk_shape(session_id, key2, (10, ) * 2)
client2.add_worker(session_id, key2, 'def')
pool2.sleep(0.1)
self.assertEqual(
local_ref1.get_chunk_meta(session_id, key1).chunk_size,
512)
self.assertEqual(
local_ref1.get_chunk_meta(session_id, key1).chunk_shape,
(10, ) * 2)
self.assertEqual(
local_ref1.get_chunk_broadcasts(session_id, key1),
[endpoints[1]])
self.assertEqual(
local_ref2.get_chunk_meta(session_id, key1).chunk_size,
512)
self.assertEqual(
local_ref2.get_chunk_meta(session_id, key1).chunk_shape,
(10, ) * 2)
self.assertEqual(
local_ref2.get_chunk_broadcasts(session_id, key2),
[endpoints[0]])
client1.batch_set_chunk_broadcasts(session_id, [key3],
[[endpoints[1]]])
meta3 = WorkerMeta(chunk_size=512,
chunk_shape=(10, ) * 2,
workers=(endpoints[0], ))
local_ref1.batch_set_chunk_meta(session_id, [key3], [meta3])
self.assertEqual(
local_ref2.get_chunk_meta(session_id, key3).chunk_size,
512)
self.assertEqual(
local_ref2.get_chunk_meta(session_id, key3).chunk_shape,
(10, ) * 2)
client1.delete_meta(session_id, key1)
pool2.sleep(0.1)
self.assertIsNone(local_ref1.get_chunk_meta(session_id, key1))
self.assertIsNone(local_ref2.get_chunk_meta(session_id, key1))
self.assertIsNone(
local_ref1.get_chunk_broadcasts(session_id, key1))
local_ref1.remove_workers_in_session(session_id, ['def'])
local_ref2.remove_workers_in_session(session_id, ['def'])
pool2.sleep(0.1)
self.assertIsNone(local_ref1.get_chunk_meta(session_id, key2))
self.assertIsNone(local_ref2.get_chunk_meta(session_id, key2))
self.assertIsNone(
local_ref2.get_chunk_broadcasts(session_id, key2))
def testAssignWithPreviousData(self):
r"""
Proper initial allocation should divide the graph like
0 1 | 2 3 | 4 5
\ / | \ / | \ /
6 | 7 | 8
"""
graph, _ = self._build_chunk_dag(
'0, 1, 2, 3, 4, 5, 6, 7, 8',
'0 -> 6, 1 -> 6, 2 -> 7, 3 -> 7, 4 -> 8, 5 -> 8')
# assign with partial mismatch
data_dist = {
'0':
dict(c00=WorkerMeta(chunk_size=5, workers=('w1', )),
c01=WorkerMeta(chunk_size=5, workers=('w2', ))),
'1':
dict(c10=WorkerMeta(chunk_size=10, workers=('w1', ))),
'2':
dict(c20=WorkerMeta(chunk_size=10, workers=('w3', ))),
'3':
dict(c30=WorkerMeta(chunk_size=10, workers=('w3', ))),
'4':
dict(c40=WorkerMeta(chunk_size=7, workers=('w3', ))),
}
analyzer = GraphAnalyzer(graph, dict(w1=24, w2=24, w3=24))
assignments = analyzer.calc_operand_assignments(
analyzer.get_initial_operand_keys(), input_chunk_metas=data_dist)
self.assertEqual(len(assignments), 6)
# explanation of the result:
# for '1', all data are in w1, hence assigned to w1
# '0' assigned to w1 according to connectivity
# '2' and '3' assigned to w3 according to connectivity
# '4' assigned to w2 because it has fewer data, and the slots of w3 is used up
self.assertEqual(assignments['0'], 'w1')
self.assertEqual(assignments['1'], 'w1')
self.assertEqual(assignments['2'], 'w3')
self.assertEqual(assignments['3'], 'w3')
self.assertEqual(assignments['4'], 'w2')
self.assertEqual(assignments['5'], 'w2')
# assign with full mismatch
data_dist = {
'0':
dict(c00=WorkerMeta(chunk_size=5, workers=('w1', )),
c01=WorkerMeta(chunk_size=5, workers=(
'w1',
'w2',
))),
'1':
dict(c10=WorkerMeta(chunk_size=10, workers=('w1', ))),
'2':
dict(c20=WorkerMeta(chunk_size=10, workers=('w3', ))),
'3':
dict(c30=WorkerMeta(chunk_size=10, workers=('w3', ))),
'4':
dict(c40=WorkerMeta(chunk_size=7, workers=('w2', ))),
'5':
dict(c50=WorkerMeta(chunk_size=7, workers=('w2', ))),
}
analyzer = GraphAnalyzer(graph, dict(w1=24, w2=24, w3=24))
assignments = analyzer.calc_operand_assignments(
analyzer.get_initial_operand_keys(), input_chunk_metas=data_dist)
self.assertEqual(len(assignments), 6)
self.assertEqual(assignments['0'], 'w1')
self.assertEqual(assignments['1'], 'w1')
self.assertEqual(assignments['2'], 'w3')
self.assertEqual(assignments['3'], 'w3')
self.assertEqual(assignments['4'], 'w2')
self.assertEqual(assignments['5'], 'w2')
def testAssignWithPreviousData(self):
import numpy as np
from mars.scheduler.chunkmeta import WorkerMeta
from mars.tensor.random import TensorRandint
from mars.tensor.arithmetic import TensorTreeAdd
graph = DAG()
r"""
Proper initial allocation should divide the graph like
U U | U U | U U
\ / | \ / | \ /
U | U | U
"""
inputs = [
tuple(
TensorRandint(_key=str(i * 2 +
j), dtype=np.float32()).new_chunk(
None, shape=(10, 10))
for j in range(2)) for i in range(3)
]
results = [
TensorTreeAdd(dtype=np.float32()).new_chunk(None, shape=(10, 10))
for _ in range(3)
]
for inp, r in zip(inputs, results):
r.op._inputs = list(inp)
graph.add_node(r)
for n in inp:
graph.add_node(n)
graph.add_edge(n, r)
# assign with partial mismatch
data_dist = {
'0':
dict(c00=WorkerMeta(chunk_size=5, workers=('w1', )),
c01=WorkerMeta(chunk_size=5, workers=('w2', ))),
'1':
dict(c10=WorkerMeta(chunk_size=10, workers=('w1', ))),
'2':
dict(c20=WorkerMeta(chunk_size=10, workers=('w3', ))),
'3':
dict(c30=WorkerMeta(chunk_size=10, workers=('w3', ))),
'4':
dict(c40=WorkerMeta(chunk_size=7, workers=('w3', ))),
}
analyzer = GraphAnalyzer(graph, dict(w1=24, w2=24, w3=24))
assignments = analyzer.calc_operand_assignments(
analyzer.get_initial_operand_keys(), input_chunk_metas=data_dist)
self.assertEqual(len(assignments), 6)
# explanation of the result:
# for '1', all data are in w1, hence assigned to w1
# '0' assigned to w1 according to connectivity
# '2' and '3' assigned to w3 according to connectivity
# '4' assigned to w2 because it has fewer data, and the slots of w3 is used up
self.assertEqual(assignments['0'], 'w1')
self.assertEqual(assignments['1'], 'w1')
self.assertEqual(assignments['2'], 'w3')
self.assertEqual(assignments['3'], 'w3')
self.assertEqual(assignments['4'], 'w2')
self.assertEqual(assignments['5'], 'w2')
# assign with full mismatch
data_dist = {
'0':
dict(c00=WorkerMeta(chunk_size=5, workers=('w1', )),
c01=WorkerMeta(chunk_size=5, workers=(
'w1',
'w2',
))),
'1':
dict(c10=WorkerMeta(chunk_size=10, workers=('w1', ))),
'2':
dict(c20=WorkerMeta(chunk_size=10, workers=('w3', ))),
'3':
dict(c30=WorkerMeta(chunk_size=10, workers=('w3', ))),
'4':
dict(c40=WorkerMeta(chunk_size=7, workers=('w2', ))),
'5':
dict(c50=WorkerMeta(chunk_size=7, workers=('w2', ))),
}
analyzer = GraphAnalyzer(graph, dict(w1=24, w2=24, w3=24))
assignments = analyzer.calc_operand_assignments(
analyzer.get_initial_operand_keys(), input_chunk_metas=data_dist)
self.assertEqual(len(assignments), 6)
self.assertEqual(assignments['0'], 'w1')
self.assertEqual(assignments['1'], 'w1')
self.assertEqual(assignments['2'], 'w3')
self.assertEqual(assignments['3'], 'w3')
self.assertEqual(assignments['4'], 'w2')
self.assertEqual(assignments['5'], 'w2')