auth log, bb4g fix

bb4g/__init__.py

@@ -1,12 +1,13 @@
+import enum
 import json
+import logging
 import math
 from datetime import datetime
 from typing import List
-import enum
+
 import requests
 from django.conf import settings as s
 from envipy_additional_information import EnviPyModel, UIConfig, WidgetType
-from envipy_additional_information import register
 
 from bridge.contracts import Classifier  # noqa: I001
 from bridge.dto import (
@@ -17,6 +18,9 @@ from bridge.dto import (
     TransformationProductPrediction,
 )  # noqa: I001
 
+logger = logging.getLogger("epdb")
+
+
 class SamplingAlgorithm(enum.Enum):
     EXACT = "exact"
 
@@ -88,7 +92,7 @@ class BB4G(Classifier):
         retries = 0
         while not started and retries < 5:
             res = requests.post(f"{self.url}/start", headers=header, data={}, proxies=s.PROXIES or None)
-
+            logger.info(f"Starting BB4G: {res.status_code}")
             if res.status_code == 200:
                 started = True
             elif res.status_code in [500, 502]:
@@ -166,18 +170,30 @@ class BB4G(Classifier):
                 "cutoff": self.config.cutoff,
             }
 
-            resp = requests.post(f"{self.url}/compute", headers=header, data=json.dumps(data), proxies=s.PROXIES or None)
+            retries = 0
+            while retries < 5:
+                resp = requests.post(f"{self.url}/compute", headers=header, data=json.dumps(data),
+                                     proxies=s.PROXIES or None)
 
-            resp.raise_for_status()
+                if resp.status_code == 418:
+                    retries += 1
+                    logger.info(f"BB4G predict hit a 418, retrying in 60 seconds")
+                    import time
+                    time.sleep(60)
+                    continue
 
-            for substrate, predictions in resp.json().items():
-                preds = {}
+                resp.raise_for_status()
 
-                for pred in predictions:
-                    prod = pred["prediction"]
-                    prob = math.exp(pred["log_likelihood"])
-                    preds[prod] = prob
+                for substrate, predictions in resp.json().items():
+                    preds = {}
 
-                result[substrate] = preds
+                    for pred in predictions:
+                        prod = pred["prediction"]
+                        prob = math.exp(pred["log_likelihood"])
+                        preds[prod] = prob
+
+                    result[substrate] = preds
+
+                break
 
         return result