Merge remote-tracking branch 'origin/edeutsch-kp-cache'

Author: edeutsch

code/ARAX/ARAXQuery/ARAX_background_tasker.py

@@ -11,6 +11,7 @@
 from importlib.metadata import version
 
 from ARAX_query_tracker import ARAXQueryTracker
+from Expand.trapi_query_cacher import KPQueryCacher
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../../")
 from RTXConfiguration import RTXConfiguration
@@ -128,6 +129,9 @@ def run_tasks(self):
                     eprint(result.stdout.decode('utf-8'))
         eprint("INFO: End listing databases area contents")
 
+        #### Set up the KP Cacher to be used for periodic refreshing
+        kp_cacher = KPQueryCacher()
+
         # Loop forever doing various things
         my_pid = os.getpid()
         while True:
@@ -198,14 +202,23 @@ def run_tasks(self):
                 n_clients += 1
                 n_ongoing_queries += n_queries
 
+            #### Refresh the KP cache
+            start_time = time.time()
+            kp_cacher.refresh_cache()
+            elapsed_time = time.time() - start_time
+            if elapsed_time < FREQ_CHECK_ONGOING_SEC - 1:
+                time_to_sleep = FREQ_CHECK_ONGOING_SEC - round(elapsed_time)
+            else:
+                time_to_sleep = 2
+
             load_tuple = psutil.getloadavg()
 
             timestamp = str(datetime.datetime.now().isoformat())
             eprint(f"{timestamp}: INFO: ARAXBackgroundTasker "
                    f"(PID {my_pid}) status: waiting. Current "
                    f"load is {load_tuple}, n_clients={n_clients}, "
                    f"n_ongoing_queries={n_ongoing_queries}")
-            time.sleep(FREQ_CHECK_ONGOING_SEC)
+            time.sleep(time_to_sleep)
 
 
 def main():

code/ARAX/ARAXQuery/ARAX_connect.py

@@ -10,15 +10,20 @@ def eprint(*args, **kwargs): print(*args, file=sys.stderr, **kwargs)
 import os
 from collections import Counter
 import copy
+import time
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 from Path_Finder.converter.EdgeExtractorFromPloverDB import EdgeExtractorFromPloverDB
 from Path_Finder.converter.ResultPerPathConverter import ResultPerPathConverter
 from Path_Finder.converter.Names import Names
 from Path_Finder.BidirectionalPathFinder import BidirectionalPathFinder
 
+from Expand.trapi_query_cacher import KPQueryCacher
+from ARAX_messenger import ARAXMessenger
+
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../../UI/OpenAPI/python-flask-server/")
 from openapi_server.models.knowledge_graph import KnowledgeGraph
+from openapi_server.models.pathfinder_analysis import PathfinderAnalysis
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../NodeSynonymizer/")
 from node_synonymizer import NodeSynonymizer
@@ -185,15 +190,64 @@ def apply(self, input_response, input_parameters):
         self.response.data['parameters'] = parameters
         self.parameters = parameters
 
-        getattr(self, '_' + self.__class__.__name__ + '__' + parameters[
-            'action'])()  # thank you https://stackoverflow.com/questions/11649848/call-methods-by-string
-
-        self.response.debug(f"Applying Connect to Message with parameters {parameters}")
+        #### Check the cache to see if we have this query cached already
+        start = time.time()
+        cacher = KPQueryCacher()
+        kp_curie = "PathFinder"
+        kp_url = "PathFinder"
+        response_envelope_as_dict = self.response.envelope.to_dict()
+        cleaned_parameters = self._clean_parameters(parameters)
+        pathfinder_input_data = { 'query_graph': response_envelope_as_dict['message']['query_graph'], 'parameters': cleaned_parameters }
+        self.response.info(f"Looking for a previously cached result from {kp_curie}")
+        response_data, response_code, elapsed_time, error = cacher.get_cached_result(kp_curie, pathfinder_input_data)
+        if response_code != -2: 
+            n_results = cacher._get_n_results(response_data)
+            self.response.info(f"Found a cached result with response_code={response_code}, n_results={n_results} from the cache in {elapsed_time:.3f} seconds")
+            self.response.envelope.message = ARAXMessenger().from_dict(response_data['message'])
+
+            # Hack to explicitly convert the analyses to PathfinderAnalysis objects because this doesn't work automatically. It should. Maybe move this into Messenger? FIXME
+            i_analysis = 0
+            for analysis_dict in response_data['message']['results'][0]['analyses']:
+                analysis_obj = PathfinderAnalysis.from_dict(analysis_dict)
+                self.response.envelope.message.results[0].analyses[i_analysis] = analysis_obj
+                i_analysis += 1
+
+        else:
+            self.response.debug(f"Applying Connect to Message with parameters {parameters}")
+
+            #### This will effectively call __connect_nodes() unless the user injects something else
+            getattr(self, '_' + self.__class__.__name__ + '__' + parameters[
+                'action'])()  # thank you https://stackoverflow.com/questions/11649848/call-methods-by-string
+
+            #### Store the result into the cache for next time
+            elapsed_time = time.time() - start
+            self.response.info(f"Got result from ARAX PathFinder Connect after {elapsed_time}. Converting to_dict()")
+            response_object = self.response.envelope.to_dict()
+            self.response.info(f"Storing resulting dict in the cache")
+            cacher.store_response(
+                kp_curie=kp_curie,
+                query_url=kp_url,
+                query_object=pathfinder_input_data,
+                response_object=response_object,
+                http_code=200,
+                elapsed_time=elapsed_time,
+                status="OK"
+            )
+            self.response.info(f"Stored result in the cache.")
 
         if self.report_stats:  # helper to report information in debug if class self.report_stats = True
             self.response = self.report_response_stats(self.response)
         return self.response
 
+
+    #### During processing, sometimes these parameters change from a string (of an integer) to an integer, so just force them all to strings for the purpose of cache comparison
+    def _clean_parameters(self, parameters):
+        cleaned_parameters = parameters.copy()
+        cleaned_parameters['max_path_length'] = str(cleaned_parameters['max_path_length'])
+        cleaned_parameters['max_pathfinder_paths'] = str(cleaned_parameters['max_pathfinder_paths'])
+        return cleaned_parameters
+
+
     def get_pinned_nodes(self):
         pinned_nodes = []
         for key, node in self.message.query_graph.nodes.items():

code/ARAX/ARAXQuery/ARAX_expander.py

@@ -30,6 +30,8 @@
 from openapi_server.models.attribute import Attribute
 from openapi_server.models.retrieval_source import RetrievalSource
 from Expand.trapi_querier import TRAPIQuerier
+from Expand.trapi_query_cacher import KPQueryCacher
+from ARAX_messenger import ARAXMessenger
 
 UNBOUND_NODES_KEY = "__UNBOUND__"
 
@@ -329,6 +331,8 @@ def apply(self, response, input_parameters, mode: str = "ARAX"):
             # Get any inferred results from ARAX Infer
             if inferred_qedge_keys:
                 response, overarching_kg = self.get_inferred_answers(inferred_qedge_keys, query_graph, response)
+                #### Update the local message with a potentially new message created in previous method call
+                message = response.envelope.message
                 if log.status != 'OK':
                     return response
                 # Now mark qedges as 'lookup' if this is an inferred query
@@ -549,6 +553,18 @@ def apply(self, response, input_parameters, mode: str = "ARAX"):
                                                  message.encountered_kryptonite_edges_info, response)
                 # Remove any paths that are now dead-ends
                 if inferred_qedge_keys and len(inferred_qedge_keys) == 1:
+
+                    #### Write some state information to files for debugging
+                    debug_filepath = os.path.dirname(os.path.abspath(__file__))
+                    if hasattr(response, 'dtd_from_cache') and response.dtd_from_cache is True:
+                        debug_filepath += "/zz_cache_"
+                    else:
+                        debug_filepath += "/zz_fresh_"
+                    with open(debug_filepath + "query_graph.json", 'w') as outfile:
+                        print(f"*******line 564: message.query_graph={message.query_graph}", file=outfile)
+                    with open(debug_filepath + "overarching_kg.json", 'w') as outfile:
+                        print(f"*******line 566: overarching_kg={overarching_kg}", file=outfile)
+
                     overarching_kg = self._remove_dead_end_paths(message.query_graph, overarching_kg, response)
                 else:
                     overarching_kg = self._remove_dead_end_paths(query_graph, overarching_kg, response)
@@ -668,13 +684,61 @@ def get_inferred_answers(inferred_qedge_keys: list[str],
                     infer_input_parameters = {"action": "drug_treatment_graph_expansion",
                                               'disease_curie': object_curie, 'qedge_id': inferred_qedge_key,
                                               'drug_curie': subject_curie}
-                    inferer = ARAXInfer()
-                    infer_response = inferer.apply(response, infer_input_parameters)
+
+                    #### Check the cache to see if we have this query cached already
+                    cacher = KPQueryCacher()
+                    enable_caching = False
+                    kp_curie = "xDTD"
+                    kp_url = "xDTD"
+                    if enable_caching:
+                        response.info(f"Looking for a previously cached result from {kp_curie}")
+                        response_data, response_code, elapsed_time, error = cacher.get_cached_result(kp_curie, infer_input_parameters)
+                    else:
+                        response.info(f"KP results caching for xDTD is currently disabled, pending further debugging")
+                    if enable_caching and response_code != -2: 
+                        n_results = cacher._get_n_results(response_data)
+                        response.info(f"Found a cached result with response_code={response_code}, n_results={n_results} from the cache in {elapsed_time:.3f} seconds")
+                        #### Transform the dict message into objects
+                        response.envelope.message = ARAXMessenger().from_dict(response_data['message'])
+                        response.envelope.message.encountered_kryptonite_edges_info = response_data['message']['encountered_kryptonite_edges_info']
+                        for node_key, node in response_data['message']['knowledge_graph']['nodes'].items():
+                            response.info(f"Copying qnode_keys for node {node_key}") 
+                            response.envelope.message.knowledge_graph.nodes[node_key].qnode_keys = node['qnode_keys']
+                        for edge_key, edge in response_data['message']['knowledge_graph']['edges'].items():
+                            response.info(f"Copying qedge_keys for edge {edge_key}") 
+                            response.envelope.message.knowledge_graph.edges[edge_key].qedge_keys = edge['qedge_keys']
+                        response.dtd_from_cache = True
+
+                    #### Else run the inferer to get the result and then cache it
+                    else:
+                        inferer = ARAXInfer()
+                        response.info(f"Launching ARAX inferer")
+                        infer_response = inferer.apply(response, infer_input_parameters)
+                        elapsed_time = time.time() - start
+                        response.info(f"Got result from ARAX inferer after {elapsed_time}. Converting to_dict()")
+                        response_object = response.envelope.to_dict()
+                        response_object['message']['encountered_kryptonite_edges_info'] = response.envelope.message.encountered_kryptonite_edges_info
+                        for node_key, node in response.envelope.message.knowledge_graph.nodes.items():
+                            response_object['message']['knowledge_graph']['nodes'][node_key]['qnode_keys'] = node.qnode_keys
+                        for edge_key, edge in response.envelope.message.knowledge_graph.edges.items():
+                            response_object['message']['knowledge_graph']['edges'][edge_key]['qedge_keys'] = edge.qedge_keys
+                        response.info(f"Storing result in the cache")
+                        cacher.store_response(
+                            kp_curie=kp_curie,
+                            query_url=kp_url,
+                            query_object=infer_input_parameters,
+                            response_object=response_object,
+                            http_code=200,
+                            elapsed_time=elapsed_time,
+                            status="OK"
+                        )
+                        response.info(f"Stored result in the cache.")
+
                     # return infer_response
-                    response = infer_response
+                    #response = infer_response  # these are already always the same object?
                     overarching_kg = eu.convert_standard_kg_to_qg_organized_kg(response.envelope.message.knowledge_graph)
 
-                    wait_time = round(time.time() - start)
+                    wait_time = round(time.time() - start, 2)
                     if response.status == "OK":
                         done_message = f"Returned {len(overarching_kg.edges_by_qg_id.get(inferred_qedge_key, dict()))} " \
                                        f"edges in {wait_time} seconds"