Repository: mars-project/mars

Issue:
[BUG] Applying resources under low memory may cause OOM
**Describe the bug**
When reallocating multiple batches of quota in MemQuotaActor, whether the allocation should apply depends on the memory size under processing before **the first** request is sent. Therefore the expected request order

* Apply request for quota 1
* Mark quota 1 as processing (total_processing += size1)
* Apply request for quota 2
* Mark quota 2 as processing (total_processing += size2)
* Apply request for quota 3 (total_processing now exceeds limitations)

could be

* Apply request for quota 1
* Apply request for quota 2
* Apply request for quota 3
* Mark quota 1 as processing (total_processing += size1)
* Mark quota 2 as processing (total_processing += size2)
* Mark quota 3 as processing (total_processing += size3)
* Possible OOM now occurs

Therefore when a resource request is sent, there should be an option to set the resource as processing once the request is accepted (i.e., request acceptance and marking the resource as processing should be an atomic operation).


Discussion:

Resolving patch:
diff --git a/mars/worker/execution.py b/mars/worker/execution.py
--- a/mars/worker/execution.py
+++ b/mars/worker/execution.py
@@ -713,7 +713,7 @@ def _start_calc(*_):
 
         # make sure that memory suffices before actually run execution
         logger.debug('Ensuring resource %r for graph %s', target_allocs, graph_key)
-        return self._mem_quota_ref.request_batch_quota(target_allocs, _promise=True) \
+        return self._mem_quota_ref.request_batch_quota(target_allocs, process_quota=True, _promise=True) \
             .then(lambda *_: self._deallocate_scheduler_resource(session_id, graph_key, delay=2)) \
             .then(_start_calc)
 
diff --git a/mars/worker/quota.py b/mars/worker/quota.py
--- a/mars/worker/quota.py
+++ b/mars/worker/quota.py
@@ -27,6 +27,18 @@
 QuotaDumpType = namedtuple('QuotaDumpType', 'allocations requests proc_sizes hold_sizes')
 
 
+class QuotaRequest(object):
+    __slots__ = 'req_size', 'delta', 'req_time', 'multiple', 'process_quota', 'callbacks'
+
+    def __init__(self, req_size, delta, req_time, multiple, process_quota, callbacks):
+        self.req_size = req_size
+        self.delta = delta
+        self.req_time = req_time
+        self.multiple = multiple
+        self.process_quota = process_quota
+        self.callbacks = callbacks
+
+
 class QuotaActor(WorkerActor):
     """
     Actor handling quota request and assignment
@@ -65,10 +77,11 @@ def _log_allocate(self, msg, *args, **kwargs):
 
     @promise.reject_on_exception
     @log_unhandled
-    def request_batch_quota(self, batch, callback=None):
+    def request_batch_quota(self, batch, process_quota=False, callback=None):
         """
         Request for resources in a batch
         :param batch: the request dict in form {request_key: request_size, ...}
+        :param process_quota: once handled, treat quota as processing
         :param callback: promise callback
         :return: if request is returned immediately, return True, otherwise False
         """
@@ -79,29 +92,33 @@ def request_batch_quota(self, batch, callback=None):
                 all_allocated = False
                 break
 
+        self._log_allocate('Receive batch quota request %r on %s.', batch, self.uid)
+        sorted_req = sorted(batch.items(), key=lambda tp: tp[0])
+        keys = tuple(tp[0] for tp in sorted_req)
+        values = tuple(tp[1] for tp in sorted_req)
+        delta = sum(v - self._allocations.get(k, 0) for k, v in batch.items())
+
         # if all requested and allocation can still be applied, apply directly
-        if all_allocated and self._has_space(0):
+        if all_allocated and self._has_space(delta):
             self._log_allocate('Quota request %r already allocated.', batch)
+            if process_quota:
+                self.process_quotas(list(batch.keys()))
             if callback is not None:
                 self.tell_promise(callback)
             return True
 
-        self._log_allocate('Receive batch quota request %r on %s.', batch, self.uid)
-        sorted_req = sorted(batch.items(), key=lambda tp: tp[0])
-        keys = tuple(tp[0] for tp in sorted_req)
-        values = tuple(tp[1] for tp in sorted_req)
-        delta = sum(v - self._allocations.get(k, 0) for k, v in batch.items())
         # make allocated requests the highest priority to be allocated
         return self._request_quota(keys, values, delta, callback, multiple=True,
-                                   make_first=all_allocated)
+                                   make_first=all_allocated, process_quota=process_quota)
 
     @promise.reject_on_exception
     @log_unhandled
-    def request_quota(self, key, quota_size, callback=None):
+    def request_quota(self, key, quota_size, process_quota=False, callback=None):
         """
         Request for resource
         :param key: request key
         :param quota_size: size of request quota
+        :param process_quota: once handled, treat quota as processing
         :param callback: promise callback
         :return: if request is returned immediately, return True, otherwise False
         """
@@ -112,7 +129,9 @@ def request_quota(self, key, quota_size, callback=None):
         if key in self._allocations:
             old_size = self._allocations[key]
             # if all requested and allocation can still be applied, apply directly
-            if old_size >= quota_size and self._has_space(0):
+            if old_size >= quota_size and self._has_space(quota_size - old_size):
+                if process_quota:
+                    self.process_quotas([key])
                 if callback is not None:
                     self.tell_promise(callback)
                 return True
@@ -122,10 +141,10 @@ def request_quota(self, key, quota_size, callback=None):
         else:
             old_size = 0
         return self._request_quota(key, quota_size, quota_size - old_size, callback,
-                                   make_first=make_first)
+                                   make_first=make_first, process_quota=process_quota)
 
     def _request_quota(self, keys, quota_sizes, delta, callback, multiple=False,
-                       make_first=False):
+                       make_first=False, process_quota=False):
         """
         Actually process requests
         :param keys: request keys
@@ -133,6 +152,7 @@ def _request_quota(self, keys, quota_sizes, delta, callback, multiple=False,
         :param delta: increase of allocate size
         :param callback: promise callback
         :param make_first: whether to move request keys to the highest priority
+        :param process_quota: once handled, treat quota as processing
         :return: if request is returned immediately, return True, otherwise False
         """
         if delta > self._total_size:
@@ -141,7 +161,7 @@ def _request_quota(self, keys, quota_sizes, delta, callback, multiple=False,
         if keys in self._requests:
             # already in request queue, store callback and quit
             if callback is not None:
-                self._requests[keys][-1].append(callback)
+                self._requests[keys].callbacks.append(callback)
             if make_first:
                 self._requests.move_to_end(keys, False)
             return False
@@ -154,7 +174,7 @@ def _request_quota(self, keys, quota_sizes, delta, callback, multiple=False,
                 self._log_allocate('Quota request met for key %r on %s.', keys, self.uid)
 
                 alter_allocation = self.alter_allocations if multiple else self.alter_allocation
-                alter_allocation(keys, quota_sizes, allocate=True)
+                alter_allocation(keys, quota_sizes, allocate=True, process_quota=process_quota)
 
                 if callback:
                     self.tell_promise(callback)
@@ -163,23 +183,27 @@ def _request_quota(self, keys, quota_sizes, delta, callback, multiple=False,
                 allocated = False
 
                 self._log_allocate('Quota request queued for key %r on %s.', keys, self.uid)
-                self._enqueue_request(keys, (quota_sizes, delta, time.time(), multiple, []),
-                                      callback=callback, make_first=make_first)
+                quota_request = QuotaRequest(quota_sizes, delta, time.time(), multiple, process_quota, [])
+                self._enqueue_request(keys, quota_request, callback=callback, make_first=make_first)
 
             self._process_requests()
             return allocated
         else:
             # current free space cannot satisfy the request, the request is queued
             self._log_allocate('Quota request unmet for key %r on %s.', keys, self.uid)
-            self._enqueue_request(keys, (quota_sizes, delta, time.time(), multiple, []),
-                                  callback=callback, make_first=make_first)
+            quota_request = QuotaRequest(quota_sizes, delta, time.time(), multiple, process_quota, [])
+            self._enqueue_request(keys, quota_request, callback=callback, make_first=make_first)
             return False
 
-    def _enqueue_request(self, keys, items, callback=None, make_first=False):
-        if keys not in self._requests:
-            self._requests[keys] = items
+    def _enqueue_request(self, keys, request, callback=None, make_first=False):
+        process_quota = request.process_quota
+        try:
+            request = self._requests[keys]
+            request.process_quota = process_quota
+        except KeyError:
+            self._requests[keys] = request
         if callback is not None:
-            self._requests[keys][-1].append(callback)
+            request.callbacks.append(callback)
 
         if make_first:
             self._requests.move_to_end(keys, False)
@@ -198,7 +222,7 @@ def cancel_requests(self, keys, reject_exc=None):
         for k in keys:
             try:
                 if reject_exc:
-                    for cb in self._requests[k][-1]:
+                    for cb in self._requests[k].callbacks:
                         self.tell_promise(cb, *reject_exc, **dict(_accept=False))
                 del self._requests[k]
                 logger.debug('Quota request %s cancelled', k)
@@ -343,16 +367,16 @@ def _process_requests(self):
         """
         removed = []
         for k, req in self._requests.items():
-            req_size, delta, req_time, multiple, callbacks = req
             try:
-                if self._has_space(delta):
-                    alter_allocation = self.alter_allocations if multiple else self.alter_allocation
-                    alter_allocation(k, req_size, handle_shrink=False, allocate=True)
-                    for cb in callbacks:
+                if self._has_space(req.delta):
+                    alter_allocation = self.alter_allocations if req.multiple else self.alter_allocation
+                    alter_allocation(k, req.req_size, handle_shrink=False, allocate=True,
+                                     process_quota=req.process_quota)
+                    for cb in req.callbacks:
                         self.tell_promise(cb)
                     if self._status_ref:
                         self._status_ref.update_mean_stats(
-                            'wait_time.' + self.uid.replace('Actor', ''), time.time() - req_time,
+                            'wait_time.' + self.uid.replace('Actor', ''), time.time() - req.req_time,
                             _tell=True, _wait=False)
                     removed.append(k)
                 else:
@@ -361,7 +385,7 @@ def _process_requests(self):
             except:  # noqa: E722
                 removed.append(k)
                 # just in case the quota is allocated
-                for cb in callbacks:
+                for cb in req.callbacks:
                     self.tell_promise(cb, *sys.exc_info(), **dict(_accept=False, _wait=False))
         for k in removed:
             self._requests.pop(k, None)


Test patch:
diff --git a/mars/worker/tests/test_quota.py b/mars/worker/tests/test_quota.py
--- a/mars/worker/tests/test_quota.py
+++ b/mars/worker/tests/test_quota.py
@@ -36,6 +36,7 @@ def _raiser(*_, **__):
 
             quota_ref = pool.create_actor(QuotaActor, 300, uid=QuotaActor.default_uid())
 
+            # test quota options with non-existing keys
             quota_ref.process_quotas(['non_exist'])
             quota_ref.hold_quotas(['non_exist'])
             quota_ref.release_quotas(['non_exist'])
@@ -43,13 +44,15 @@ def _raiser(*_, **__):
             with self.assertRaises(ValueError):
                 quota_ref.request_quota('ERROR', 1000)
 
+            # test quota request with immediate return
             self.assertTrue(quota_ref.request_quota('0', 100))
             self.assertTrue(quota_ref.request_quota('0', 50))
             self.assertTrue(quota_ref.request_quota('0', 200))
 
-            quota_ref.process_quotas(['0'])
+            # test request with process_quota=True
+            quota_ref.request_quota('0', 200, process_quota=True)
             self.assertIn('0', quota_ref.dump_data().proc_sizes)
-            quota_ref.alter_allocation('0', 190, new_key=('0', 0))
+            quota_ref.alter_allocation('0', 190, new_key=('0', 0), process_quota=True)
             self.assertEqual(quota_ref.dump_data().allocations[('0', 0)], 190)
 
             quota_ref.hold_quotas([('0', 0)])
@@ -89,9 +92,9 @@ def _raiser(*_, **__):
                     with self.assertRaises(ValueError):
                         self.get_result(5)
 
-            self.assertNotIn('1', quota_ref.dump_data().requests)
-            self.assertIn('2', quota_ref.dump_data().allocations)
-            self.assertNotIn('3', quota_ref.dump_data().allocations)
+                self.assertNotIn('1', quota_ref.dump_data().requests)
+                self.assertIn('2', quota_ref.dump_data().allocations)
+                self.assertNotIn('3', quota_ref.dump_data().allocations)
 
             quota_ref.release_quotas([('0', 1)])
             self.assertIn('3', quota_ref.dump_data().allocations)
@@ -147,22 +150,28 @@ def testBatchQuotaAllocation(self):
             quota_ref = pool.create_actor(QuotaActor, 300, uid=QuotaActor.default_uid())
 
             end_time = []
-            for idx in range(2):
-                x = str(idx)
-                with self.run_actor_test(pool) as test_actor:
+
+            with self.run_actor_test(pool) as test_actor:
+                for idx in (0, 1):
+                    x = str(idx)
                     ref = test_actor.promise_ref(QuotaActor.default_uid())
 
-                    def actual_exec(keys):
-                        for k in keys:
-                            ref.release_quotas([k])
+                    def actual_exec(b, set_result):
+                        self.assertTrue(ref.request_batch_quota(b, process_quota=True))
+                        self.assertEqual(set(b.keys()), set(quota_ref.dump_data().proc_sizes.keys()))
+                        ref.release_quotas(list(b.keys()))
                         end_time.append(time.time())
-                        test_actor.set_result(None)
+                        if set_result:
+                            test_actor.set_result(None)
 
                     keys = [x + '_0', x + '_1']
                     batch = dict((k, 100) for k in keys)
                     ref.request_batch_quota(batch, _promise=True) \
-                        .then(functools.partial(test_actor.run_later, actual_exec, keys, _delay=0.5))
-                    self.get_result(10)
+                        .then(functools.partial(test_actor.run_later, actual_exec, batch,
+                                                set_result=(idx == 1), _delay=0.5),
+                              lambda *exc: test_actor.set_result(exc, accept=False))
+
+                self.get_result(10)
 
             self.assertGreater(abs(end_time[0] - end_time[1]), 0.4)
             self.assertEqual(quota_ref.get_allocated_size(), 0)