Improved logging

* Change many print statements to logging statements, for better control
* Save log file to web app zip file for easier debugging

app/run_utils.py

@@ -44,7 +44,7 @@ def configure_logging(level=None):
     logging.basicConfig(
         level=level,
         format="[%(asctime)s] [%(filename)s:%(lineno)d] %(levelname)s - %(message)s",
-        datefmt="%Y-%m-%d %H:%M:%S",
+        datefmt="%Y-%m-%d %H:%M:%S %Z",
         handlers=[
             logging.StreamHandler(),  # Outputs logs to stderr by default
             # If you also want to log to a file, uncomment the following line:
@@ -93,8 +93,10 @@ def run_cli_command(
 
     assert len(args) == len(ARG_ORDER), f'Expected {len(ARG_ORDER)} arguments, got {len(args)}'
 
+    inference_log_level = os.environ.get("INFERENCE_LOG_LEVEL", os.environ.get("LOG_LEVEL", "WARNING"))
+
     all_arg_dict = {"protein_path": protein_path, "ligand": ligand, "config": config_path,
-                    "no_final_step_noise": True}
+                    "no_final_step_noise": True, "loglevel": inference_log_level}
     for arg_name, arg_val in zip(ARG_ORDER, args):
         all_arg_dict[arg_name] = arg_val
 
@@ -136,12 +138,18 @@ def run_cli_command(
                     capture_output=True,
                 )
                 logging.debug(f"Command output:\n{result.stdout}")
+                full_output = f"Standard out:\n{result.stdout}"
                 if result.stderr:
                     # Skip progress bar lines
                     stderr_lines = result.stderr.split("\n")
                     stderr_lines = filter(lambda x: "%|" not in x, stderr_lines)
                     stderr_text = "\n".join(stderr_lines)
                     logging.error(f"Command error:\n{stderr_text}")
+                    full_output += f"\nStandard error:\n{stderr_text}"
+
+                with open(f"{temp_dir_path}/output.log", "w") as log_file:
+                    log_file.write(full_output)
+
             else:
                 logging.debug("Skipping command execution")
                 artificial_output_dir = os.path.join(TEMP_DIR, "artificial_output")

datasets/process_mols.py

@@ -17,6 +17,7 @@
 from datasets.constants import aa_short2long, atom_order, three_to_one
 from datasets.parse_chi import get_chi_angles, get_coords, aa_idx2aa_short, get_onehot_sequence
 from utils.torsion import get_transformation_mask
+from utils.logging_utils import get_logger
 
 
 periodic_table = GetPeriodicTable()
@@ -305,11 +306,11 @@ def generate_conformer(mol):
     failures, id = 0, -1
     while failures < 3 and id == -1:
         if failures > 0:
-            print(f'rdkit coords could not be generated. trying again {failures}.')
+            get_logger().debug(f'rdkit coords could not be generated. trying again {failures}.')
         id = AllChem.EmbedMolecule(mol, ps)
         failures += 1
     if id == -1:
-        print('rdkit coords could not be generated without using random coords. using random coords now.')
+        get_logger().info('rdkit coords could not be generated without using random coords. using random coords now.')
         ps.useRandomCoords = True
         AllChem.EmbedMolecule(mol, ps)
         AllChem.MMFFOptimizeMolecule(mol, confId=0)
@@ -417,6 +418,7 @@ def write_mol_with_coords(mol, new_coords, path):
     w.write(mol)
     w.close()
 
+
 def read_molecule(molecule_file, sanitize=False, calc_charges=False, remove_hs=False):
     if molecule_file.endswith('.mol2'):
         mol = Chem.MolFromMol2File(molecule_file, sanitize=False, removeHs=False)
@@ -433,8 +435,8 @@ def read_molecule(molecule_file, sanitize=False, calc_charges=False, remove_hs=F
     elif molecule_file.endswith('.pdb'):
         mol = Chem.MolFromPDBFile(molecule_file, sanitize=False, removeHs=False)
     else:
-        return ValueError('Expect the format of the molecule_file to be '
-                          'one of .mol2, .sdf, .pdbqt and .pdb, got {}'.format(molecule_file))
+        raise ValueError('Expect the format of the molecule_file to be '
+                         'one of .mol2, .sdf, .pdbqt and .pdb, got {}'.format(molecule_file))
 
     try:
         if sanitize or calc_charges:
@@ -449,7 +451,12 @@ def read_molecule(molecule_file, sanitize=False, calc_charges=False, remove_hs=F
 
         if remove_hs:
             mol = Chem.RemoveHs(mol, sanitize=sanitize)
-    except:
+
+    except Exception as e:
+        # Print stacktrace
+        import traceback
+        msg = traceback.format_exc()
+        get_logger().warning(f"Failed to process molecule: {molecule_file}\n{msg}")
         return None
 
     return mol