def setUp(self):
endpoint = '127.0.0.1:%d' % get_next_port()
self.endpoint = endpoint
self.pool = create_actor_pool(n_process=1, backend='gevent', address=endpoint)
self.pool.create_actor(SchedulerClusterInfoActor, [endpoint],
uid=SchedulerClusterInfoActor.default_name())
self.pool.create_actor(SessionManagerActor, uid=SessionManagerActor.default_name())
self.pool.create_actor(ResourceActor, uid=ResourceActor.default_name())
self.api = MarsAPI(endpoint)
class Test(unittest.TestCase):
def setUp(self):
endpoint = '127.0.0.1:%d' % get_next_port()
self.endpoint = endpoint
self.pool = create_actor_pool(n_process=1, backend='gevent', address=endpoint)
self.pool.create_actor(SchedulerClusterInfoActor, [endpoint],
uid=SchedulerClusterInfoActor.default_name())
self.pool.create_actor(SessionManagerActor, uid=SessionManagerActor.default_name())
self.pool.create_actor(ResourceActor, uid=ResourceActor.default_name())
self.api = MarsAPI(endpoint)
def tearDown(self):
self.pool.stop()
@patch_method(GraphActor.execute_graph)
def testApi(self, *_):
self.assertEqual(0, self.api.count_workers())
session_id = 'mock_session_id'
self.api.create_session(session_id)
self.assertEqual(1, len(self.api.session_manager.get_sessions()))
self.api.delete_session(session_id)
self.assertEqual(0, len(self.api.session_manager.get_sessions()))
self.api.create_session(session_id)
serialized_graph = 'mock_serialized_graph'
graph_key = 'mock_graph_key'
targets = 'mock_targets'
self.api.submit_graph(session_id, serialized_graph, graph_key, targets)
graph_uid = GraphActor.gen_uid(session_id, graph_key)
graph_ref = self.api.get_actor_ref(graph_uid)
self.assertTrue(self.pool.has_actor(graph_ref))
state = self.api.get_graph_state(session_id, graph_key)
self.assertEqual(GraphState('preparing'), state)
self.api.stop_graph(session_id, graph_key)
state = self.api.get_graph_state(session_id, graph_key)
self.assertEqual(GraphState('cancelled'), state)
self.api.delete_graph(session_id, graph_key)
self.assertFalse(self.pool.has_actor(graph_ref))
@patch_method(GraphActor.get_tileable_chunk_indexes)
@patch_method(ChunkMetaClient.batch_get_chunk_shape)
def testGetTensorNsplits(self, *_):
session_id = 'mock_session_id'
graph_key = 'mock_graph_key'
tensor_key = 'mock_tensor_key'
serialized_graph = 'mock_serialized_graph'
graph_uid = GraphActor.gen_uid(session_id, graph_key)
self.pool.create_actor(GraphActor, session_id, serialized_graph, graph_key, uid=graph_uid)
self.pool.create_actor(ChunkMetaActor, uid=ChunkMetaActor.default_name())
mock_indexes = [
OrderedDict(zip(['chunk_key1', 'chunk_key2', 'chunk_key3', 'chunk_key4'],
[(0, ), (1,), (2,), (3,)])),
OrderedDict(zip(['chunk_key1', 'chunk_key2', 'chunk_key3', 'chunk_key4'],
[(0, 0), (0, 1), (1, 0), (1, 1)]))
]
mock_shapes = [
[(3,), (4,), (5,), (6,)],
[(3, 4), (3, 2), (2, 4), (2, 2)]
]
GraphActor.get_tileable_chunk_indexes.side_effect = mock_indexes
ChunkMetaClient.batch_get_chunk_shape.side_effect = mock_shapes
nsplits = self.api.get_tileable_nsplits(session_id, graph_key, tensor_key)
self.assertEqual(((3, 4, 5, 6),), nsplits)
nsplits = self.api.get_tileable_nsplits(session_id, graph_key, tensor_key)
self.assertEqual(((3, 2), (4, 2)), nsplits)