def testApi(self, *_):
timeout = 120 if 'CI' in os.environ else -1
with new_session(self._service_ep) as sess:
self.assertEqual(sess.count_workers(), 1)
a = mt.ones((100, 100), chunk_size=30)
b = mt.ones((100, 100), chunk_size=30)
c = a.dot(b)
result = sess.run(c, timeout=timeout)
np.testing.assert_array_equal(result, np.ones((100, 100)) * 100)
mock_resp = MockResponse(200, text='{"ERR_STR":')
with mock.patch('requests.Session.post',
side_effect=lambda *_, **__: mock_resp):
with self.assertRaises(ResponseMalformed):
sess.run(c, timeout=timeout)
mock_resp = MockResponse(
400,
json_data=dict(exc_info=base64.b64encode(
pickle.dumps(build_exc_info(ValueError))).decode()))
with mock.patch('requests.Session.post',
side_effect=lambda *_, **__: mock_resp):
with self.assertRaises(ValueError):
sess.run(c, timeout=timeout)
d = a * 100
self.assertIsNone(sess.run(d, fetch=False, timeout=120))
np.testing.assert_array_equal(sess.run(d, timeout=120),
np.ones((100, 100)) * 100)
async def test_mem_quota_allocation(actor_pool, enable_kill_slot):
from mars.utils import AttributeDict
mock_mem_stat = AttributeDict(dict(total=300, available=50, used=0, free=50))
mock_band_slot_manager_ref = await mo.create_actor(
MockBandSlotManagerActor, uid=BandSlotManagerActor.gen_uid('numa-0'),
address=actor_pool.external_address)
quota_ref = await mo.create_actor(
MemQuotaActor, (actor_pool.external_address, 'numa-0'), 300,
hard_limit=300, refresh_time=0.1,
enable_kill_slot=enable_kill_slot,
uid=MemQuotaActor.gen_uid('cpu-0'),
address=actor_pool.external_address) # type: Union[QuotaActorRef, mo.ActorRef]
with mock.patch('mars.resource.virtual_memory', new=lambda: mock_mem_stat):
time_recs = [time.time()]
async def task_fun():
await quota_ref.request_batch_quota({'req': 100})
await quota_ref.release_quotas(['req'])
time_recs.append(time.time())
task = asyncio.create_task(task_fun())
await asyncio.sleep(0.2)
assert 'req' not in (await quota_ref.dump_data()).allocations
mock_mem_stat['available'] = 150
mock_mem_stat['free'] = 150
await asyncio.wait_for(task, timeout=1)
assert 0.15 < abs(time_recs[0] - time_recs[1]) < 1
assert bool(await mock_band_slot_manager_ref.get_restart_record()) == enable_kill_slot
def testFailoverMessage(self):
mock_session_id = str(uuid.uuid4())
mock_graph_key = str(uuid.uuid4())
mock_chunk_key = str(uuid.uuid4())
addr = '127.0.0.1:%d' % get_next_port()
mock_worker_addr = '127.0.0.1:54132'
options.scheduler.worker_blacklist_time = 0.5
with create_actor_pool(n_process=1, backend='gevent',
address=addr) as pool:
cluster_info_ref = pool.create_actor(
SchedulerClusterInfoActor, [pool.cluster_info.address],
uid=SchedulerClusterInfoActor.default_uid())
session_manager_ref = pool.create_actor(
SessionManagerActor, uid=SessionManagerActor.default_uid())
resource_ref = pool.create_actor(ResourceActor,
uid=ResourceActor.default_uid())
pool.create_actor(ChunkMetaActor, uid=ChunkMetaActor.default_uid())
session_ref = pool.actor_ref(
session_manager_ref.create_session(mock_session_id))
chunk_meta_client = ChunkMetaClient(pool, cluster_info_ref)
chunk_meta_client.set_chunk_meta(mock_session_id,
mock_chunk_key,
size=80,
shape=(10, ),
workers=(mock_worker_addr, ))
with mock.patch(GraphActor.__module__ + '.' + GraphActor.__name__,
new=MockGraphActor):
session_ref.submit_tileable_graph(None, mock_graph_key)
graph_ref = pool.actor_ref(
GraphActor.gen_uid(mock_session_id, mock_graph_key))
expire_time = time.time(
) - options.scheduler.status_timeout - 1
resource_ref.set_worker_meta(mock_worker_addr,
dict(update_time=expire_time))
resource_ref.detect_dead_workers(_tell=True)
pool.sleep(0.2)
_, removes, lost_chunks = graph_ref.get_worker_change_args()
self.assertListEqual(removes, [mock_worker_addr])
self.assertListEqual(lost_chunks, [mock_chunk_key])
self.assertNotIn(mock_worker_addr,
resource_ref.get_workers_meta())
resource_ref.set_worker_meta(mock_worker_addr,
dict(update_time=time.time()))
self.assertNotIn(mock_worker_addr,
resource_ref.get_workers_meta())
pool.sleep(0.4)
resource_ref.set_worker_meta(mock_worker_addr,
dict(update_time=time.time()))
self.assertIn(mock_worker_addr,
resource_ref.get_workers_meta())