def testRegister(self):
from mars.core.graph import DAG
fake_result = np.random.rand(10, 10)
fake_size = (fake_result.nbytes * 2, fake_result.nbytes * 2)
def fake_execute(ctx, op):
ctx[op.outputs[0].key] = fake_result
def fake_estimate(ctx, op):
ctx[op.outputs[0].key] = fake_size
register(FakeOperand, fake_execute, fake_estimate)
graph = DAG()
chunk = FakeOperand().new_chunk(None, shape=(10, 10))
graph.add_node(chunk.data)
executor = Executor()
res = executor.execute_graph(graph, keys=[chunk.key])[0]
np.testing.assert_array_equal(res, fake_result)
size = executor.execute_graph(graph, keys=[chunk.key], mock=True)[0]
self.assertEqual(size, fake_size)
graph = DAG()
chunk = SubFakeOperand().new_chunk(None, shape=(10, 10))
graph.add_node(chunk.data)
executor = Executor()
res = executor.execute_graph(graph, keys=[chunk.key])[0]
np.testing.assert_array_equal(res, fake_result)
def testEmptyGraph(self, *_):
session_id = str(uuid.uuid4())
addr = f'127.0.0.1:{get_next_port()}'
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))
resource_ref.set_worker_meta('localhost:12345',
dict(hardware=dict(cpu_total=4)))
resource_ref.set_worker_meta('localhost:23456',
dict(hardware=dict(cpu_total=4)))
graph_key = str(uuid.uuid4())
serialized_graph = serialize_graph(DAG())
graph_ref = pool.create_actor(GraphActor,
session_id,
graph_key,
serialized_graph,
uid=GraphActor.gen_uid(
session_id, graph_key))
graph_ref.execute_graph()
self.assertEqual(graph_ref.get_state(), GraphState.SUCCEEDED)
def testInitialAssignsWithInputs(self):
import numpy as np
from mars.tensor.random import TensorRandint
from mars.tensor.arithmetic import TensorTreeAdd
n1 = TensorRandint(state=np.random.RandomState(0),
dtype=np.dtype(np.float32())).new_chunk(None,
shape=(10,
10))
n2 = TensorRandint(state=np.random.RandomState(1),
dtype=np.dtype(np.float32())).new_chunk(None,
shape=(10,
10))
n3 = TensorTreeAdd(args=[],
dtype=np.dtype(np.float32())).new_chunk(None,
shape=(10,
10))
n3.op._inputs = [n1.data, n2.data]
n4 = TensorTreeAdd(args=[],
dtype=np.dtype(np.float32())).new_chunk(None,
shape=(10,
10))
n4.op._inputs = [n3.data]
graph = DAG()
graph.add_node(n1)
graph.add_node(n3)
graph.add_node(n4)
graph.add_edge(n1, n3)
graph.add_edge(n3, n4)
analyzer = GraphAnalyzer(graph, {})
ext_chunks = analyzer.collect_external_input_chunks(initial=False)
self.assertListEqual(ext_chunks[n3.op.key], [n2.key])
self.assertEqual(
len(analyzer.collect_external_input_chunks(initial=True)), 0)
def test_dag():
r"""
1 --- 4
2 --- 6
\ /
5
/
3
"""
dag = DAG()
[dag.add_node(i) for i in range(1, 7)]
dag.add_edge(1, 4)
dag.add_edge(2, 6)
dag.add_edge(2, 5)
dag.add_edge(5, 6)
dag.add_edge(3, 5)
with pytest.raises(KeyError):
dag.add_edge(1, 10)
with pytest.raises(KeyError):
dag.add_edge(10, 1)
assert set(dag[2]) == {5, 6}
assert list(dag.topological_iter()) == [3, 2, 5, 6, 1, 4]
assert list(dag.dfs()) == [3, 2, 5, 6, 1, 4]
assert list(dag.bfs()) == [1, 2, 3, 4, 5, 6]
dag.add_edge(6, 1)
dag.add_edge(1, 2)
with pytest.raises(KeyError):
for _ in dag.iter_predecessors(-1):
pass
with pytest.raises(KeyError):
for _ in dag.iter_successors(-1):
pass
with pytest.raises(GraphContainsCycleError):
_ = list(dag.topological_iter())
dag.remove_edge(2, 5)
assert dag.has_successor(2, 5) is False
with pytest.raises(KeyError):
dag.remove_edge(2, 5)
rev_dag = dag.build_reversed()
for n in dag:
assert n in rev_dag
assert all(
rev_dag.has_successor(n, pred)
for pred in dag.predecessors(n)) is True
undigraph = dag.build_undirected()
for n in dag:
assert n in undigraph
assert all(
undigraph.has_predecessor(pred, n)
for pred in dag.predecessors(n)) is True
assert all(
undigraph.has_successor(n, pred)
for pred in dag.predecessors(n)) is True
dag_copy = dag.copy()
for n in dag:
assert n in dag_copy
assert all(
dag_copy.has_successor(pred, n)
for pred in dag_copy.predecessors(n)) is True
def _build_chunk_dag(node_str, edge_str):
from mars.tensor.random import TensorRandint
from mars.tensor.arithmetic import TensorTreeAdd
char_dag = DAG()
for s in node_str.split(','):
char_dag.add_node(s.strip())
for s in edge_str.split(','):
l, r = s.split('->')
char_dag.add_edge(l.strip(), r.strip())
chunk_dag = DAG()
str_to_chunk = dict()
for s in char_dag.topological_iter():
if char_dag.count_predecessors(s):
c = TensorTreeAdd(args=[],
_key=s,
dtype=np.dtype(np.float32())).new_chunk(
None, shape=(10, 10)).data
inputs = c.op._inputs = [
str_to_chunk[ps] for ps in char_dag.predecessors(s)
]
else:
c = TensorRandint(_key=s, dtype=np.dtype(
np.float32())).new_chunk(None, shape=(10, 10)).data
inputs = []
str_to_chunk[s] = c
chunk_dag.add_node(c)
for inp in inputs:
chunk_dag.add_edge(inp, c)
return chunk_dag, str_to_chunk
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)