From 69589f36eefcef97da42900065bf7c382e544090 Mon Sep 17 00:00:00 2001
From: "Kevin M. Jablonka" <kevin.jablonka@epfl.ch>
Date: Fri, 12 Nov 2021 08:58:21 +0100
Subject: [PATCH] feat: logging, tighter timeout control

---
 .env                   |  1 +
 Dockerfile             |  2 +-
 docker-compose.yml     |  2 ++
 xtbservice/ir.py       | 21 ++++++++++++++++++---
 xtbservice/optimize.py |  3 +++
 xtbservice/settings.py | 12 +++++++++---
 6 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/.env b/.env
index 480eca3..b27d544 100644
--- a/.env
+++ b/.env
@@ -1,3 +1,4 @@
 WORKERS=1
 OMP_NUM_THREADS=1,1
 PORT=8051
+TIMEOUT=5
diff --git a/Dockerfile b/Dockerfile
index adba4cc..4eea84e 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -12,4 +12,4 @@ RUN pip install --no-cache-dir -r requirements.txt
 
 COPY README.md .
 
-CMD gunicorn -w $WORKERS xtbservice.xtbservice:app -b 0.0.0.0:$PORT -k uvicorn.workers.UvicornWorker
+CMD gunicorn -w $WORKERS xtbservice.xtbservice:app -b 0.0.0.0:$PORT -k uvicorn.workers.UvicornWorker -t $TIMEOUT --keep-alive $TIMEOUT
diff --git a/docker-compose.yml b/docker-compose.yml
index 6979c23..de895a7 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -10,6 +10,8 @@ services:
       - WORKERS=${WORKERS}
       - OMP_NUM_THREADS=${OMP_NUM_THREADS}
       - MAX_ATOMS=50
+      - TIMEOUT=${TIMEOUT}
+      - LOG_LEVEL=debug
     ports:
       - ${PORT}:${PORT}
     volumes:
diff --git a/xtbservice/ir.py b/xtbservice/ir.py
index 55fbefa..561bf9b 100644
--- a/xtbservice/ir.py
+++ b/xtbservice/ir.py
@@ -7,6 +7,7 @@
 import wrapt_timeout_decorator
 from ase import Atoms
 from ase.vibrations import Infrared
+from fastapi.logger import logger
 from rdkit import Chem
 from scipy import spatial
 from xtb.ase.calculator import XTB
@@ -27,13 +28,15 @@ def run_xtb_ir(
 ) -> IRResult:
     if mol is None:
         raise Exception
-    
+
     this_hash = ir_hash(atoms, method)
+    logger.debug(f"Running IR for {this_hash}")
     moi = atoms.get_moments_of_inertia()
     linear = sum(moi > 0.01) == 2
     result = ir_cache.get(this_hash)
 
     if result is None:
+        logger.debug(f"IR not in cache for {this_hash}, running")
         atoms.pbc = False
         atoms.calc = XTB(method=method)
 
@@ -68,7 +71,11 @@ def run_xtb_ir(
         ]
 
         mode_info, has_imaginary, has_large_imaginary = compile_modes_info(
-            ir, linear, displacement_alignments, bond_displacements, bonds,
+            ir,
+            linear,
+            displacement_alignments,
+            bond_displacements,
+            bonds,
         )
         result = IRResult(
             wavenumbers=list(spectrum[0]),
@@ -271,7 +278,15 @@ def get_displacement_xyz_for_mode(ir, frequencies, symbols, n):
     for i, pos in enumerate(ir.atoms.positions):
         xyz_file.append(
             "%2s %12.5f %12.5f %12.5f %12.5f %12.5f %12.5f\n"
-            % (symbols[i], pos[0], pos[1], pos[2], mode[i, 0], mode[i, 1], mode[i, 2],)
+            % (
+                symbols[i],
+                pos[0],
+                pos[1],
+                pos[2],
+                mode[i, 0],
+                mode[i, 1],
+                mode[i, 2],
+            )
         )
 
     xyz_file_string = "".join(xyz_file)
diff --git a/xtbservice/optimize.py b/xtbservice/optimize.py
index 00316eb..460fa6c 100644
--- a/xtbservice/optimize.py
+++ b/xtbservice/optimize.py
@@ -5,6 +5,7 @@
 
 from ase import Atoms
 from ase.optimize.lbfgs import LBFGS
+from fastapi.logger import logger
 from xtb.ase.calculator import XTB
 
 from .cache import opt_cache
@@ -24,11 +25,13 @@ def run_xtb_opt(
     maxiter: int = 100,
 ) -> OptimizationResult:
     this_hash = opt_hash(atoms, method)
+    logger.debug(f"Running optimization with hash {this_hash}")
     try:
         result = opt_cache.get(this_hash)
     except KeyError:
         pass
     if result is None:
+        logger.debug(f"Optimization not found in cache, running")
         mol = deepcopy(atoms)
         mol.pbc = False
 
diff --git a/xtbservice/settings.py b/xtbservice/settings.py
index 7dbd3b4..ff51c71 100644
--- a/xtbservice/settings.py
+++ b/xtbservice/settings.py
@@ -1,6 +1,12 @@
 # -*- coding: utf-8 -*-
 import os
 
-IMAGINARY_FREQ_THRESHOLD = os.getenv("IMAGINARY_FREQ_THRESHOLD", 10)
-MAX_ATOMS = os.getenv("MAX_ATOMS", 50)
-TIMEOUT = os.getenv("TIMEOUT", 100)
+from fastapi.logger import logger
+
+IMAGINARY_FREQ_THRESHOLD = int(os.getenv("IMAGINARY_FREQ_THRESHOLD", 10))
+MAX_ATOMS = int(os.getenv("MAX_ATOMS", 50))
+TIMEOUT = int(os.getenv("TIMEOUT", 100))
+
+logger.info(
+    f"Settings: IMAGINARY_FREQ_THRESHOLD: {IMAGINARY_FREQ_THRESHOLD}, MAX_ATOMS: {MAX_ATOMS}, TIMEOUT: {TIMEOUT}"
+)