float timeout for sub-second jobs

Author: tokebe

src/retriever/data_tiers/tier_0/base_query.py

@@ -45,7 +45,7 @@ async def execute(self) -> LookupArtifacts:
             self.job_log.info("Starting lookup against Tier 0...")
 
             try:
-                timeout = None if self.ctx.timeout < 0 else self.ctx.timeout - 0.5
+                timeout = None if self.ctx.timeout < 0 else self.ctx.timeout
                 async with asyncio.timeout(timeout):
                     backend_results = await self.get_results(
                         QueryGraphDict(**self.qgraph.model_dump())

src/retriever/lookup/qgx.py

@@ -235,7 +235,7 @@ async def execute_branch(
             while parallel_tasks:
                 timeout = (
                     max(
-                        self.ctx.timeout - (time.time() - self.start_time) - 0.5,
+                        self.ctx.timeout - (time.time() - self.start_time),
                         0,
                     )
                     if self.ctx.timeout >= 0
@@ -561,7 +561,7 @@ async def timeout(self) -> None:
         try:
             if self.ctx.timeout < 0:
                 return
-            await asyncio.sleep(max(self.ctx.timeout - 0.5, 0))
+            await asyncio.sleep(self.ctx.timeout)
             self.job_log.error("QGX hit timeout, attempting wrapup...")
             self.terminate = True
         except asyncio.CancelledError:

src/retriever/query.py

@@ -60,7 +60,7 @@ async def make_query(
     Unhandled errors are handled by middleware.
     """
     job_id = uuid.uuid4().hex
-    timeout: int = {
+    timeout: float = {
         "lookup": CONFIG.job.lookup.timeout,
         "metakg": CONFIG.job.metakg.timeout,
     }[func]

src/retriever/types/general.py

@@ -54,7 +54,7 @@ class QueryInfo(NamedTuple):
     body: Query | AsyncQuery | None
     job_id: str
     tiers: set[TierNumber]
-    timeout: int
+    timeout: float
 
 
 class BackendResults(TypedDict):

src/retriever/types/trapi.py

@@ -25,7 +25,7 @@ class ParametersDict(TypedDict):
     """Query Parameters."""
 
     tiers: list[int]
-    timeout: int
+    timeout: float
 
 
 class QueryDict(TypedDict):

src/retriever/types/trapi_pydantic.py

@@ -15,7 +15,7 @@ class Parameters(BaseModel):
     """Parameters that govern some elements of query execution behavior."""
 
     timeout: Annotated[
-        int | None,
+        float | None,
         Field(
             description="Custom query timeout in seconds. Defaults to server default if not set. Set to -1 to disable timeout entirely."
         ),