Run training and prediction as tasks

This change implements the execution of the training and prediction
calls as tasks, that are spawned in different executors using asyncio.
This means that the ModelWrapper returns tasks for the predict and train
methods, that are not anymore async methods. The callee must then handle
the task (e.g. await for it to be completed) before rendering the
result.

The approach we follow is different for each call, as follows:

- We use separate executor pools for trainings and predictions. The
  number of workers for each can be configured through the command line,
  although in most cases it does not make too much sense to provide a
  larger worker pool for trainings.

- Prediction calls are spawned using a ThreadPoolExecutor. We have found
  several difficulties in spawning them as a processes using the
  ProcessPoolExecutor. Until we find the root cause of this we stick to
  using threads.

- The prediction pool can be "warmed" meaning that the model can be
  initialized before the API starts. This is usefult for avoiding
  uncessary waits for the clients interacting with the API.

- Even though Predictions are tasks, we do not allow to cancel them,
  since the we cannot cancel threads. However we use tasks to allow for
  async calls, that will be implemented in a separate patch

- Trainings use a special CancellablePool, that uses processes under the
  hood. We use this special pool in order to be able to cancel the
  tasks, as the multiprocessing.futures.ProcessPoolExecutor does not
  allow to cancel tasks spawned there.

- We implement a full fledged /train endpoint, where we can POST,
  DELETE, GET, etc. in order to create, delete and get information about
  a running training.

Sem-Ver: api-break

Author: alvarolopez

deepaas/api/__init__.py

@@ -56,7 +56,7 @@ async def get_app(doc="/docs"):
 
         LOG.info("Serving loaded V1 models: %s", list(model.V1_MODELS.keys()))
 
-    model.register_v2_models()
+    model.register_v2_models(APP)
 
     v2app = v2.get_app()
     APP.add_subapp("/v2", v2app)

deepaas/api/v2/predict.py

@@ -57,7 +57,10 @@ def __init__(self, model_name, model_obj):
             @aiohttp_apispec.response_schema(responses.Failure(), 400)
             @aiohttpparser.parser.use_args(args)
             async def post(self, request, args):
-                ret = await self.model_obj.predict(**args)
+                task = self.model_obj.predict(**args)
+                await task
+
+                ret = task.result()
 
                 if self.model_obj.has_schema:
                     self.model_obj.validate_response(ret)

deepaas/api/v2/responses.py

@@ -16,6 +16,7 @@
 
 import marshmallow
 from marshmallow import fields
+from marshmallow import validate
 
 
 class Location(marshmallow.Schema):
@@ -56,3 +57,19 @@ class ModelMeta(marshmallow.Schema):
     version = fields.Str(required=False, description='Model version')
     url = fields.Str(required=False, description='Model url')
     links = fields.List(fields.Nested(Location))
+
+
+class Training(marshmallow.Schema):
+    uuid = fields.UUID(required=True, description='Training identifier')
+    date = fields.DateTime(required=True, description='Training start time')
+    status = fields.Str(
+        required=True,
+        description='Training status',
+        enum=["running", "error", "completed", "cancelled"],
+        validate=validate.OneOf(["running", "error", "completed", "cancelled"])
+    )
+    message = fields.Str(description="Optional message explaining status")
+
+
+class TrainingList(marshmallow.Schema):
+    trainings = fields.List(fields.Nested(Training))

deepaas/api/v2/train.py

@@ -14,13 +14,21 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
+import asyncio
+import datetime
+import uuid
+
 from aiohttp import web
 import aiohttp_apispec
+from oslo_log import log
 from webargs import aiohttpparser
 import webargs.core
 
+from deepaas.api.v2 import responses
 from deepaas import model
 
+LOG = log.getLogger("deepaas.api.v2.train")
+
 
 def setup_routes(app):
     # In the next lines we iterate over the loaded models and create the
@@ -36,6 +44,29 @@ class Handler(object):
             def __init__(self, model_name, model_obj):
                 self.model_name = model_name
                 self.model_obj = model_obj
+                self._trainings = {}
+
+            def build_train_response(self, uuid_):
+                training = self._trainings.get(uuid_, None)
+                if training:
+                    ret = {}
+                    ret["date"] = training["date"]
+                    ret["uuid"] = uuid_
+
+                    if training["task"].cancelled():
+                        ret["status"] = "cancelled"
+                    elif training["task"].done():
+                        exc = training["task"].exception()
+                        if exc:
+                            ret["status"] = "error"
+                            ret["message"] = "%s" % exc
+                        else:
+                            ret["status"] = "done"
+                    else:
+                        ret["status"] = "running"
+                    return ret
+                else:
+                    return None
 
             @aiohttp_apispec.docs(
                 tags=["models"],
@@ -44,10 +75,66 @@ def __init__(self, model_name, model_obj):
             @aiohttp_apispec.querystring_schema(args)
             @aiohttpparser.parser.use_args(args)
             async def post(self, request, args):
-                ret = await self.model_obj.train(**args)
-                # FIXME(aloga): what are we returning here? We need to take
-                # care of these responses as well.
+                uuid_ = uuid.uuid4().hex
+                train_task = self.model_obj.train(**args)
+                self._trainings[uuid_] = {
+                    "date": str(datetime.datetime.now()),
+                    "task": train_task,
+                }
+                ret = self.build_train_response(uuid_)
                 return web.json_response(ret)
 
+            @aiohttp_apispec.docs(
+                tags=["models"],
+                summary="Cancel a running training"
+            )
+            async def delete(self, request):
+                uuid_ = request.match_info["uuid"]
+                training = self._trainings.get(uuid_, None)
+                if training:
+                    training["task"].cancel()
+                    try:
+                        await asyncio.wait_for(training["task"], 5)
+                    except asyncio.TimeoutError:
+                        pass
+                    LOG.info("Training %s has been cancelled" % uuid_)
+                    ret = self.build_train_response(uuid_)
+                    return web.json_response(ret)
+                else:
+                    raise web.HTTPNotFound()
+
+            @aiohttp_apispec.docs(
+                tags=["models"],
+                summary="Get a list of trainings (running or completed)"
+            )
+            @aiohttp_apispec.response_schema(responses.TrainingList(), 200)
+            async def index(self, request):
+
+                ret = []
+                for uuid_, training in self._trainings.items():
+                    aux = self.build_train_response(uuid_)
+                    ret.append(aux)
+
+                return web.json_response(ret)
+
+            @aiohttp_apispec.docs(
+                tags=["models"],
+                summary="Get status of a training"
+            )
+            @aiohttp_apispec.response_schema(responses.Training(), 200)
+            async def get(self, request):
+                uuid_ = request.match_info["uuid"]
+                ret = self.build_train_response(uuid_)
+                if ret:
+                    return web.json_response(ret)
+                else:
+                    raise web.HTTPNotFound()
+
         hdlr = Handler(model_name, model_obj)
         app.router.add_post("/models/%s/train" % model_name, hdlr.post)
+        app.router.add_get("/models/%s/train" % model_name, hdlr.index)
+        app.router.add_get("/models/%s/train/{uuid}" % model_name, hdlr.get)
+        app.router.add_delete(
+            "/models/%s/train/{uuid}" % model_name,
+            hdlr.delete
+        )

deepaas/config.py

@@ -42,13 +42,19 @@
 "/debug" endpoint. Default is to not provide this information. This will not
 provide logging information about the API itself.
 """),
-    cfg.IntOpt('model-workers',
-               short='n',
+    cfg.IntOpt('predict-workers',
+               short='p',
                default=1,
                help="""
-Specify the number of workers *per model* that we will initialize. If using a
-CPU you probably want to increase this number, if using a GPU probably you want
-to leave it to 1. (defaults to 1)
+Specify the number of workers to spawn for prediction tasks. If using a CPU you
+probably want to increase this number, if using a GPU probably you want to
+leave it to 1. (defaults to 1)
+"""),
+    cfg.IntOpt('train-workers',
+               default=1,
+               help="""
+Specify the number of workers to spawn for training tasks. Unless you know what
+you are doing you should leave this number to 1. (defaults to 1)
 """),
 ]
 

deepaas/model/__init__.py

@@ -31,11 +31,11 @@ def register_v1_models():
     return v1.register_models()
 
 
-def register_v2_models():
+def register_v2_models(app):
     """Register V2 models.
 
     This method has to be called before the API is spawned, so that we
     can look up the correct entry points and load the defined models.
     """
 
-    return v2.register_models()
+    return v2.register_models(app)

deepaas/model/v2/__init__.py

@@ -27,7 +27,7 @@
 MODELS_LOADED = False
 
 
-def register_models():
+def register_models(app):
     global MODELS
     global MODELS_LOADED
 
@@ -36,7 +36,7 @@ def register_models():
 
     try:
         for name, model in loading.get_available_models("v2").items():
-            MODELS[name] = wrapper.ModelWrapper(name, model)
+            MODELS[name] = wrapper.ModelWrapper(name, model, app)
     except Exception as e:
         LOG.warning("Error loading models: %s", e)
 
@@ -52,14 +52,15 @@ def register_models():
 
     try:
         for name, model in loading.get_available_models("v1").items():
-            MODELS[name] = wrapper.ModelWrapper(name, model)
+            MODELS[name] = wrapper.ModelWrapper(name, model, app)
     except Exception as e:
         LOG.warning("Error loading models: %s", e)
 
     if not MODELS:
         LOG.info("No models found with V2 or V1 namespace, loading test model")
         MODELS["deepaas-test"] = wrapper.ModelWrapper(
             "deepaas-test",
-            test.TestModel()
+            test.TestModel(),
+            app
         )
     MODELS_LOADED = True

deepaas/model/v2/test.py

@@ -14,6 +14,8 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
+import time
+
 from oslo_log import log
 from webargs import fields
 from webargs import validate
@@ -56,8 +58,11 @@ def predict(self, **kwargs):
         }
 
     def train(self, *args, **kwargs):
+        sleep = kwargs.get("sleep", 1)
         LOG.debug("Got the following arguments: %s", args)
         LOG.debug("Got the following kw arguments: %s", kwargs)
+        LOG.debug("Starting training, ending in %is" % sleep)
+        time.sleep(sleep)
 
     def get_predict_args(self):
         return {
@@ -81,7 +86,7 @@ def get_predict_args(self):
 
     def get_train_args(self):
         return {
-            "parameter_one": fields.Int(
+            "sleep": fields.Int(
                 required=True,
                 descripton='This is a integer parameter, and it is '
                            'a required one.'