def testDistributedContext(self):
self.start_processes(etcd=False)
sess = new_session(self.session_manager_ref.address)
rs = np.random.RandomState(0)
context = DistributedContext(
scheduler_address=self.session_manager_ref.address,
session_id=sess.session_id)
raw1 = rs.rand(10, 10)
a = mt.tensor(raw1, chunk_size=4)
a.execute(session=sess, timeout=self.timeout, name='test')
tileable_infos = context.get_named_tileable_infos('test')
self.assertEqual(a.key, tileable_infos.tileable_key)
self.assertEqual(a.shape, tileable_infos.tileable_shape)
nsplits = context.get_tileable_metas([a.key],
filter_fields=['nsplits'])[0][0]
self.assertEqual(((4, 4, 2), (4, 4, 2)), nsplits)
r = context.get_tileable_data(a.key)
np.testing.assert_array_equal(raw1, r)
indexes = [slice(3, 9), slice(0, 7)]
r = context.get_tileable_data(a.key, indexes)
np.testing.assert_array_equal(raw1[tuple(indexes)], r)
indexes = [[1, 4, 2, 4, 5], slice(None, None, None)]
r = context.get_tileable_data(a.key, indexes)
np.testing.assert_array_equal(raw1[tuple(indexes)], r)
indexes = ([9, 1, 2, 0], [0, 0, 4, 4])
r = context.get_tileable_data(a.key, indexes)
np.testing.assert_array_equal(raw1[[9, 1, 2, 0], [0, 0, 4, 4]], r)
def testDistributedContext(self):
self.start_processes(etcd=False)
session_id = uuid.uuid1()
actor_client = new_client()
rs = np.random.RandomState(0)
context = DistributedContext(
scheduler_address=self.scheduler_endpoints[0],
session_id=session_id)
session_ref = actor_client.actor_ref(
self.session_manager_ref.create_session(session_id))
raw1 = rs.rand(10, 10)
a = mt.tensor(raw1, chunk_size=4)
graph = a.build_graph()
targets = [a.key]
graph_key = uuid.uuid1()
session_ref.submit_tileable_graph(json.dumps(graph.to_json()),
graph_key,
target_tileables=targets,
names=['test'])
state = self.wait_for_termination(actor_client, session_ref, graph_key)
self.assertEqual(state, GraphState.SUCCEEDED)
tileable_infos = context.get_named_tileable_infos('test')
self.assertEqual(a.key, tileable_infos.tileable_key)
self.assertEqual(a.shape, tileable_infos.tileable_shape)
nsplits = context.get_tileable_metas([a.key],
filter_fields=['nsplits'])[0][0]
self.assertEqual(((4, 4, 2), (4, 4, 2)), nsplits)
r = context.get_tileable_data(a.key)
np.testing.assert_array_equal(raw1, r)
indexes = [slice(3, 9), slice(0, 7)]
r = context.get_tileable_data(a.key, indexes)
np.testing.assert_array_equal(raw1[tuple(indexes)], r)
indexes = [[1, 4, 2, 4, 5], slice(None, None, None)]
r = context.get_tileable_data(a.key, indexes)
np.testing.assert_array_equal(raw1[tuple(indexes)], r)
indexes = ([9, 1, 2, 0], [0, 0, 4, 4])
r = context.get_tileable_data(a.key, indexes)
np.testing.assert_array_equal(raw1[[9, 1, 2, 0], [0, 0, 4, 4]], r)