Fix cancellation of trainings

If a user performed a training -> cancel -> training -> cancel, the last
cancel did not terminate properly and the API blocked, as the process
did not join.

Since we are using a process pool with only 1 slot, we can get the
running processes in the pool and kill them.

Author: alvarolopez

deepaas/model/v2/wrapper.py

@@ -22,6 +22,7 @@
 import multiprocessing
 import multiprocessing.pool
 import os
+import signal
 import tempfile
 
 from aiohttp import web
@@ -80,6 +81,8 @@ def __init__(self, name, model_obj, app):
         self._train_workers = CONF.train_workers
         self._train_executor = self._init_train_executor()
 
+        self._setup_cleanup()
+
         schema = getattr(self.model_obj, "schema", None)
 
         if isinstance(schema, dict):
@@ -109,6 +112,13 @@ def __init__(self, name, model_obj, app):
 
         self.response_schema = schema
 
+    def _setup_cleanup(self):
+        self._app.on_cleanup.append(self._close_executors)
+
+    async def _close_executors(self, app):
+        self._train_executor.shutdown()
+        self._predict_executor.shutdown()
+
     def _init_predict_executor(self):
         n = self._predict_workers
         executor = concurrent.futures.ThreadPoolExecutor(max_workers=n)
@@ -119,22 +129,6 @@ def _init_train_executor(self):
         executor = CancellablePool(max_workers=n)
         return executor
 
-#        run = sconcurrent.futures.elf.loop.run_in_executor
-
-#        fs = [run(executor, self.warm, path) for i in range(0, n)]
-#        await asyncio.gather(*fs)
-#
-        async def close_executor():
-            self._executor.shutdown()
-
-#        async def close_executor():
-#            fs = [run(executor, self.clean) for i in range(0, n)]
-#            await asyncio.shield(asyncio.gather(*fs))
-#            executor.shutdown(wait=True)
-
-        self._app.on_cleanup.append(close_executor)
-#        app['executor'] = executor
-
     @contextlib.contextmanager
     def _catch_error(self):
         name = self.name
@@ -408,6 +402,13 @@ def _on_err(err):
         try:
             return await fut
         except asyncio.CancelledError:
+            # This is ugly, but since our pools only have one slot we can
+            # kill the process before termination
+            try:
+                pool._pool[0].kill()
+            except AttributeError:
+                os.kill(pool._pool[0].pid,
+                        signal.SIGKILL)
             pool.terminate()
             usable_pool = self._new_pool()
         finally:
@@ -416,6 +417,6 @@ def _on_err(err):
             self._change.set()
 
     def shutdown(self):
-        for p in self._working | self._free:
+        for p in self._working:
             p.terminate()
         self._free.clear()

deepaas/tests/test_v2_models.py

@@ -69,7 +69,8 @@ def get_train_args(self):
         self.assertRaises(NotImplementedError, m.train)
         self.assertRaises(NotImplementedError, m.get_train_args)
 
-    def test_bad_schema(self):
+    @mock.patch("deepaas.model.v2.wrapper.ModelWrapper._setup_cleanup")
+    def test_bad_schema(self, m_clean):
         class Model(object):
             schema = []
 
@@ -80,7 +81,8 @@ class Model(object):
             self.app
         )
 
-    def test_validate_no_schema(self):
+    @mock.patch("deepaas.model.v2.wrapper.ModelWrapper._setup_cleanup")
+    def test_validate_no_schema(self, m_clean):
         class Model(object):
             schema = None
 
@@ -91,7 +93,8 @@ class Model(object):
             None
         )
 
-    def test_invalid_schema(self):
+    @mock.patch("deepaas.model.v2.wrapper.ModelWrapper._setup_cleanup")
+    def test_invalid_schema(self, m_clean):
         class Model(object):
             schema = object()
 
@@ -102,7 +105,8 @@ class Model(object):
             self.app
         )
 
-    def test_marshmallow_schema(self):
+    @mock.patch("deepaas.model.v2.wrapper.ModelWrapper._setup_cleanup")
+    def test_marshmallow_schema(self, m_clean):
         class Schema(marshmallow.Schema):
             foo = m_fields.Str()
 
@@ -118,7 +122,8 @@ class Model(object):
             {"foo": 1.0}
         )
 
-    def test_dict_schema(self):
+    @mock.patch("deepaas.model.v2.wrapper.ModelWrapper._setup_cleanup")
+    def test_dict_schema(self, m_clean):
         class Model(object):
             schema = {
                 "foo": m_fields.Str()
@@ -151,8 +156,9 @@ def test_dummy_model(self):
         for arg, val in itertools.chain(pargs.items(), targs.items()):
             self.assertIsInstance(val, fields.Field)
 
+    @mock.patch("deepaas.model.v2.wrapper.ModelWrapper._setup_cleanup")
     @test_utils.unittest_run_loop
-    async def test_dummy_model_with_wrapper(self):
+    async def test_dummy_model_with_wrapper(self, m_clean):
         w = v2_wrapper.ModelWrapper("foo", v2_test.TestModel(), self.app)
         task = w.predict()
         await task
@@ -175,8 +181,10 @@ async def test_dummy_model_with_wrapper(self):
         for arg, val in itertools.chain(pargs.items(), targs.items()):
             self.assertIsInstance(val, fields.Field)
 
+    @mock.patch("deepaas.model.v2.wrapper.ModelWrapper._setup_cleanup")
     @test_utils.unittest_run_loop
-    async def test_model_with_not_implemented_attributes_and_wrapper(self):
+    async def test_model_with_not_implemented_attributes_and_wrapper(self,
+                                                                     m_clean):
         w = v2_wrapper.ModelWrapper("foo", object(), self.app)
 
         # NOTE(aloga): Cannot use assertRaises here directly, as testtools
@@ -203,25 +211,28 @@ async def test_model_with_not_implemented_attributes_and_wrapper(self):
         for arg, val in itertools.chain(pargs.items(), targs.items()):
             self.assertIsInstance(val, fields.Field)
 
+    @mock.patch("deepaas.model.v2.wrapper.ModelWrapper._setup_cleanup")
     @mock.patch('deepaas.model.loading.get_available_models')
-    def test_loading_ok(self, mock_loading):
+    def test_loading_ok(self, mock_loading, m_clean):
         mock_loading.return_value = {uuid.uuid4().hex: "bar"}
         deepaas.model.v2.register_models(self.app)
         mock_loading.assert_called()
         for m in deepaas.model.v2.MODELS.values():
             self.assertIsInstance(m, v2_wrapper.ModelWrapper)
 
+    @mock.patch("deepaas.model.v2.wrapper.ModelWrapper._setup_cleanup")
     @mock.patch('deepaas.model.loading.get_available_models')
-    def test_loading_ok_singleton(self, mock_loading):
+    def test_loading_ok_singleton(self, mock_loading, m_clean):
         mock_loading.return_value = {uuid.uuid4().hex: "bar"}
         deepaas.model.v2.register_models(self.app)
         deepaas.model.v2.register_models(self.app)
         mock_loading.assert_called_once()
         for m in deepaas.model.v2.MODELS.values():
             self.assertIsInstance(m, v2_wrapper.ModelWrapper)
 
+    @mock.patch("deepaas.model.v2.wrapper.ModelWrapper._setup_cleanup")
     @mock.patch('deepaas.model.loading.get_available_models')
-    def test_loading_error(self, mock_loading):
+    def test_loading_error(self, mock_loading, m_clean):
         mock_loading.return_value = {}
         deepaas.model.v2.register_models(self.app)
         mock_loading.assert_called()