Merge pull request #107 from BioPack-team/fix-timeout-overrun

Properly cancel Tier 0 execution at timeout (fix #74)

Author: tokebe

src/retriever/data_tiers/tier_0/base_query.py

@@ -12,7 +12,6 @@
     Infores,
     KnowledgeGraphDict,
     QueryGraphDict,
-    ResultDict,
 )
 from retriever.utils.logs import TRAPILogger
 from retriever.utils.trapi import (
@@ -43,21 +42,12 @@ async def execute(self) -> LookupArtifacts:
             start_time = time.time()
             self.job_log.info("Starting lookup against Tier 0...")
 
-            try:
-                timeout = None if self.ctx.timeout[0] < 0 else self.ctx.timeout[0]
-                self.job_log.debug(
-                    f"Tier 0 timeout is {'disabled' if timeout is None else f'{timeout}s'}."
-                )
-                async with asyncio.timeout(timeout):
-                    backend_results = await self.get_results(self.qgraph)
-
-            except TimeoutError:
-                self.job_log.error("Tier 0 operation timed out.")
-                backend_results = BackendResult(
-                    results=list[ResultDict](),
-                    knowledge_graph=KnowledgeGraphDict(nodes={}, edges={}),
-                    auxiliary_graphs=dict[AuxGraphID, AuxiliaryGraphDict](),
-                )
+            timeout = None if self.ctx.timeout[0] < 0 else self.ctx.timeout[0]
+            self.job_log.debug(
+                f"Tier 0 timeout is {'disabled' if timeout is None else f'{timeout}s'}."
+            )
+            async with asyncio.timeout(timeout):
+                backend_results = await self.get_results(self.qgraph)
 
             with tracer.start_as_current_span("update_kg"):
                 normalize_kgraph(
@@ -90,10 +80,13 @@ async def execute(self) -> LookupArtifacts:
                 self.aux_graphs,
                 self.job_log.get_logs(),
             )
-        except Exception:
-            self.job_log.exception(
-                "Unhandled exception occurred while processing Tier 0. See logs for details."
-            )
+        except Exception as e:
+            if isinstance(e, TimeoutError):
+                self.job_log.error("Tier 0 operation timed out.")
+            elif not isinstance(e, asyncio.exceptions.CancelledError):
+                self.job_log.exception(
+                    "Unhandled exception occurred while processing Tier 0 query. See logs for details."
+                )
             return LookupArtifacts(
                 [], self.kgraph, self.aux_graphs, self.job_log.get_logs(), error=True
             )

src/retriever/utils/biolink.py

@@ -86,7 +86,7 @@ def reverse_qualifier_constraints(
     """Reverse a given list of qualifier constraints."""
     new = list[QualifierConstraintDict]()
     for constraint in qualifier_constraints:
-        new_qualifier_set = set[QualifierDict]()
+        new_qualifier_set = list[QualifierDict]()
         for qualifier in constraint["qualifier_set"]:
             new_qualifier = QualifierDict(**qualifier)
             if "object" in qualifier["qualifier_type_id"]:
@@ -105,7 +105,7 @@ def reverse_qualifier_constraints(
                 # BUG: Technically invalid if we can't reverse the predicate
                 # but this is vanishingly rare and not worth addressing right now
                 new_qualifier["qualifier_value"] = inverse
-            new_qualifier_set.add(new_qualifier)
+            new_qualifier_set.append(new_qualifier)
         new.append(constraint)
     return new
 

src/retriever/utils/logs.py

@@ -42,12 +42,14 @@ def trapi_level_to_int(level: TRAPILogLevel) -> int:
 def format_trapi_log(
     level: LogLevel,
     message: str,
-    timestamp: datetime | None = None,
+    timestamp: str | datetime | None = None,
     trace: str | None = None,
 ) -> LogEntryDict:
     """Format a loguru message into a TRAPI-spec LogEntry."""
     if timestamp is None:
         timestamp = datetime.now().astimezone()
+    if isinstance(timestamp, str):
+        timestamp = datetime.fromisoformat(timestamp)
     log_entry = LogEntryDict(
         level=log_level_to_trapi(level),
         message=message,