[Feature] Basic logging of Jobs, Model Evaluation (#169)

Co-authored-by: Tim Lorsbach <tim@lorsba.ch>
Reviewed-on: enviPath/enviPy#169

epdb/tasks.py

@@ -1,10 +1,13 @@
 import logging
-from typing import Optional
-from celery.utils.functional import LRUCache
-from celery import shared_task
-from epdb.models import Pathway, Node, EPModel, Setting
-from epdb.logic import SPathway
+from datetime import datetime
+from typing import Callable, Optional
+from uuid import uuid4
 
+from celery import shared_task
+from celery.utils.functional import LRUCache
+
+from epdb.logic import SPathway
+from epdb.models import EPModel, JobLog, Node, Package, Pathway, Setting, User
 
 logger = logging.getLogger(__name__)
 ML_CACHE = LRUCache(3)  # Cache the three most recent ML models to reduce load times.
@@ -16,6 +19,40 @@ def get_ml_model(model_pk: int):
     return ML_CACHE[model_pk]
 
 
+def dispatch_eager(user: "User", job: Callable, *args, **kwargs):
+    try:
+        x = job(*args, **kwargs)
+        log = JobLog()
+        log.user = user
+        log.task_id = uuid4()
+        log.job_name = job.__name__
+        log.status = "SUCCESS"
+        log.done_at = datetime.now()
+        log.task_result = str(x) if x else None
+        log.save()
+
+        return x
+    except Exception as e:
+        logger.exception(e)
+        raise e
+
+
+def dispatch(user: "User", job: Callable, *args, **kwargs):
+    try:
+        x = job.delay(*args, **kwargs)
+        log = JobLog()
+        log.user = user
+        log.task_id = x.task_id
+        log.job_name = job.__name__
+        log.status = "INITIAL"
+        log.save()
+
+        return x.result
+    except Exception as e:
+        logger.exception(e)
+        raise e
+
+
 @shared_task(queue="background")
 def mul(a, b):
     return a * b
@@ -33,17 +70,55 @@ def send_registration_mail(user_pk: int):
     pass
 
 
-@shared_task(queue="model")
-def build_model(model_pk: int):
+@shared_task(bind=True, queue="model")
+def build_model(self, model_pk: int):
     mod = EPModel.objects.get(id=model_pk)
-    mod.build_dataset()
-    mod.build_model()
+
+    if JobLog.objects.filter(task_id=self.request.id).exists():
+        JobLog.objects.filter(task_id=self.request.id).update(status="RUNNING", task_result=mod.url)
+
+    try:
+        mod.build_dataset()
+        mod.build_model()
+    except Exception as e:
+        if JobLog.objects.filter(task_id=self.request.id).exists():
+            JobLog.objects.filter(task_id=self.request.id).update(
+                status="FAILED", task_result=mod.url
+            )
+
+        raise e
+
+    if JobLog.objects.filter(task_id=self.request.id).exists():
+        JobLog.objects.filter(task_id=self.request.id).update(status="SUCCESS", task_result=mod.url)
+
+    return mod.url
 
 
-@shared_task(queue="model")
-def evaluate_model(model_pk: int):
+@shared_task(bind=True, queue="model")
+def evaluate_model(self, model_pk: int, multigen: bool, package_pks: Optional[list] = None):
+    packages = None
+
+    if package_pks:
+        packages = Package.objects.filter(pk__in=package_pks)
+
     mod = EPModel.objects.get(id=model_pk)
-    mod.evaluate_model()
+    if JobLog.objects.filter(task_id=self.request.id).exists():
+        JobLog.objects.filter(task_id=self.request.id).update(status="RUNNING", task_result=mod.url)
+
+    try:
+        mod.evaluate_model(multigen, eval_packages=packages)
+    except Exception as e:
+        if JobLog.objects.filter(task_id=self.request.id).exists():
+            JobLog.objects.filter(task_id=self.request.id).update(
+                status="FAILED", task_result=mod.url
+            )
+
+        raise e
+
+    if JobLog.objects.filter(task_id=self.request.id).exists():
+        JobLog.objects.filter(task_id=self.request.id).update(status="SUCCESS", task_result=mod.url)
+
+    return mod.url
 
 
 @shared_task(queue="model")
@@ -52,9 +127,13 @@ def retrain(model_pk: int):
     mod.retrain()
 
 
-@shared_task(queue="predict")
+@shared_task(bind=True, queue="predict")
 def predict(
-    pw_pk: int, pred_setting_pk: int, limit: Optional[int] = None, node_pk: Optional[int] = None
+    self,
+    pw_pk: int,
+    pred_setting_pk: int,
+    limit: Optional[int] = None,
+    node_pk: Optional[int] = None,
 ) -> Pathway:
     pw = Pathway.objects.get(id=pw_pk)
     setting = Setting.objects.get(id=pred_setting_pk)
@@ -65,6 +144,9 @@ def predict(
     pw.kv.update(**{"status": "running"})
     pw.save()
 
+    if JobLog.objects.filter(task_id=self.request.id).exists():
+        JobLog.objects.filter(task_id=self.request.id).update(status="RUNNING", task_result=pw.url)
+
     try:
         # regular prediction
         if limit is not None:
@@ -89,7 +171,18 @@ def predict(
     except Exception as e:
         pw.kv.update({"status": "failed"})
         pw.save()
+
+        if JobLog.objects.filter(task_id=self.request.id).exists():
+            JobLog.objects.filter(task_id=self.request.id).update(
+                status="FAILED", task_result=pw.url
+            )
+
         raise e
 
     pw.kv.update(**{"status": "completed"})
     pw.save()
+
+    if JobLog.objects.filter(task_id=self.request.id).exists():
+        JobLog.objects.filter(task_id=self.request.id).update(status="SUCCESS", task_result=pw.url)
+
+    return pw.url