feat: logging, tighter timeout control

.env

@@ -1,3 +1,4 @@
 WORKERS=1
 OMP_NUM_THREADS=1,1
 PORT=8051
+TIMEOUT=5

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

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:

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)

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
 

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}"
+)