test: analyze cache hits/misses

Author: akhileshh

pychunkedgraph/__init__.py

@@ -1 +1,56 @@
 __version__ = "3.1.6"
+
+import sys
+import warnings
+import logging as stdlib_logging  # Use alias to avoid conflict with pychunkedgraph.logging
+
+# Suppress annoying warning from python_jsonschema_objects dependency
+warnings.filterwarnings("ignore", message="Schema id not specified", module="python_jsonschema_objects")
+
+# Export logging levels for convenience
+DEBUG = stdlib_logging.DEBUG
+INFO = stdlib_logging.INFO
+WARNING = stdlib_logging.WARNING
+ERROR = stdlib_logging.ERROR
+
+# Set up library-level logger with NullHandler (Python logging best practice)
+stdlib_logging.getLogger(__name__).addHandler(stdlib_logging.NullHandler())
+
+
+def configure_logging(level=stdlib_logging.INFO, format_str=None, stream=None):
+    """
+    Configure logging for pychunkedgraph. Call this to enable log output.
+
+    Works in Jupyter notebooks and scripts.
+
+    Args:
+        level: Logging level (default: INFO). Use pychunkedgraph.DEBUG, .INFO, .WARNING, .ERROR
+        format_str: Custom format string (optional)
+        stream: Output stream (default: sys.stdout for Jupyter compatibility)
+
+    Example:
+        import pychunkedgraph
+        pychunkedgraph.configure_logging()  # Enable INFO level logging
+        pychunkedgraph.configure_logging(pychunkedgraph.DEBUG)  # Enable DEBUG level
+    """
+    if format_str is None:
+        format_str = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    if stream is None:
+        stream = sys.stdout
+
+    # Get root logger for pychunkedgraph
+    logger = stdlib_logging.getLogger(__name__)
+    logger.setLevel(level)
+
+    # Remove existing handlers and add fresh StreamHandler
+    # This allows reconfiguring with different levels/formats
+    for h in logger.handlers[:]:
+        if isinstance(h, stdlib_logging.StreamHandler) and not isinstance(h, stdlib_logging.NullHandler):
+            logger.removeHandler(h)
+
+    handler = stdlib_logging.StreamHandler(stream)
+    handler.setLevel(level)
+    handler.setFormatter(stdlib_logging.Formatter(format_str))
+    logger.addHandler(handler)
+
+    return logger

pychunkedgraph/graph/cache.py

@@ -2,6 +2,8 @@
 """
 Cache nodes, parents, children and cross edges.
 """
+import traceback
+from collections import defaultdict
 from sys import maxsize
 from datetime import datetime
 
@@ -40,6 +42,30 @@ def __init__(self, cg):
         self.children_cache = LRUCache(maxsize=maxsize)
         self.cross_chunk_edges_cache = LRUCache(maxsize=maxsize)
 
+        # Stats tracking for cache hits/misses
+        self.stats = {
+            "parents": {"hits": 0, "misses": 0, "calls": 0},
+            "children": {"hits": 0, "misses": 0, "calls": 0},
+            "cross_chunk_edges": {"hits": 0, "misses": 0, "calls": 0},
+        }
+        # Track where calls/misses come from
+        self.call_sources = defaultdict(lambda: defaultdict(lambda: {"calls": 0, "misses": 0}))
+
+    def _get_caller(self, skip_frames=2):
+        """Get caller info (filename:line:function)."""
+        stack = traceback.extract_stack()
+        # Skip frames: _get_caller, the cache method, and go to actual caller
+        if len(stack) > skip_frames:
+            frame = stack[-(skip_frames + 1)]
+            return f"{frame.filename.split('/')[-1]}:{frame.lineno}:{frame.name}"
+        return "unknown"
+
+    def _record_call(self, cache_type, misses=0):
+        """Record a call and its source."""
+        caller = self._get_caller(skip_frames=3)
+        self.call_sources[cache_type][caller]["calls"] += 1
+        self.call_sources[cache_type][caller]["misses"] += misses
+
     def __len__(self):
         return (
             len(self.parents_cache)
@@ -52,14 +78,53 @@ def clear(self):
         self.children_cache.clear()
         self.cross_chunk_edges_cache.clear()
 
+    def get_stats(self):
+        """Return stats with hit rates calculated."""
+        result = {}
+        for name, s in self.stats.items():
+            total = s["hits"] + s["misses"]
+            hit_rate = s["hits"] / total if total > 0 else 0
+            result[name] = {
+                **s,
+                "total": total,
+                "hit_rate": f"{hit_rate:.1%}",
+                "sources": dict(self.call_sources[name]),
+            }
+        return result
+
+    def reset_stats(self):
+        for s in self.stats.values():
+            s["hits"] = 0
+            s["misses"] = 0
+            s["calls"] = 0
+        self.call_sources.clear()
+
     def parent(self, node_id: np.uint64, *, time_stamp: datetime = None):
+        self.stats["parents"]["calls"] += 1
+        is_cached = node_id in self.parents_cache
+        miss_count = 0 if is_cached else 1
+        if is_cached:
+            self.stats["parents"]["hits"] += 1
+        else:
+            self.stats["parents"]["misses"] += 1
+        self._record_call("parents", misses=miss_count)
+
         @cached(cache=self.parents_cache, key=lambda node_id: node_id)
         def parent_decorated(node_id):
             return self._cg.get_parent(node_id, raw_only=True, time_stamp=time_stamp)
 
         return parent_decorated(node_id)
 
     def children(self, node_id):
+        self.stats["children"]["calls"] += 1
+        is_cached = node_id in self.children_cache
+        miss_count = 0 if is_cached else 1
+        if is_cached:
+            self.stats["children"]["hits"] += 1
+        else:
+            self.stats["children"]["misses"] += 1
+        self._record_call("children", misses=miss_count)
+
         @cached(cache=self.children_cache, key=lambda node_id: node_id)
         def children_decorated(node_id):
             children = self._cg.get_children(node_id, raw_only=True)
@@ -69,6 +134,15 @@ def children_decorated(node_id):
         return children_decorated(node_id)
 
     def cross_chunk_edges(self, node_id, *, time_stamp: datetime = None):
+        self.stats["cross_chunk_edges"]["calls"] += 1
+        is_cached = node_id in self.cross_chunk_edges_cache
+        miss_count = 0 if is_cached else 1
+        if is_cached:
+            self.stats["cross_chunk_edges"]["hits"] += 1
+        else:
+            self.stats["cross_chunk_edges"]["misses"] += 1
+        self._record_call("cross_chunk_edges", misses=miss_count)
+
         @cached(cache=self.cross_chunk_edges_cache, key=lambda node_id: node_id)
         def cross_edges_decorated(node_id):
             edges = self._cg.get_cross_chunk_edges(
@@ -82,7 +156,13 @@ def parents_multiple(self, node_ids: np.ndarray, *, time_stamp: datetime = None)
         node_ids = np.array(node_ids, dtype=NODE_ID, copy=False)
         if not node_ids.size:
             return node_ids
+        self.stats["parents"]["calls"] += 1
         mask = np.in1d(node_ids, np.fromiter(self.parents_cache.keys(), dtype=NODE_ID))
+        hits = int(np.sum(mask))
+        misses = len(node_ids) - hits
+        self.stats["parents"]["hits"] += hits
+        self.stats["parents"]["misses"] += misses
+        self._record_call("parents", misses=misses)
         parents = node_ids.copy()
         parents[mask] = self._parent_vec(node_ids[mask])
         parents[~mask] = self._cg.get_parents(
@@ -96,7 +176,13 @@ def children_multiple(self, node_ids: np.ndarray, *, flatten=False):
         node_ids = np.array(node_ids, dtype=NODE_ID, copy=False)
         if not node_ids.size:
             return result
+        self.stats["children"]["calls"] += 1
         mask = np.in1d(node_ids, np.fromiter(self.children_cache.keys(), dtype=NODE_ID))
+        hits = int(np.sum(mask))
+        misses = len(node_ids) - hits
+        self.stats["children"]["hits"] += hits
+        self.stats["children"]["misses"] += misses
+        self._record_call("children", misses=misses)
         cached_children_ = self._children_vec(node_ids[mask])
         result.update({id_: c_ for id_, c_ in zip(node_ids[mask], cached_children_)})
         result.update(self._cg.get_children(node_ids[~mask], raw_only=True))
@@ -114,9 +200,15 @@ def cross_chunk_edges_multiple(
         node_ids = np.array(node_ids, dtype=NODE_ID, copy=False)
         if not node_ids.size:
             return result
+        self.stats["cross_chunk_edges"]["calls"] += 1
         mask = np.in1d(
             node_ids, np.fromiter(self.cross_chunk_edges_cache.keys(), dtype=NODE_ID)
         )
+        hits = int(np.sum(mask))
+        misses = len(node_ids) - hits
+        self.stats["cross_chunk_edges"]["hits"] += hits
+        self.stats["cross_chunk_edges"]["misses"] += misses
+        self._record_call("cross_chunk_edges", misses=misses)
         cached_edges_ = self._cross_chunk_edges_vec(node_ids[mask])
         result.update(
             {id_: edges_ for id_, edges_ in zip(node_ids[mask], cached_edges_)}

pychunkedgraph/graph/operation.py

@@ -1,5 +1,6 @@
 # pylint: disable=invalid-name, missing-docstring, too-many-lines, protected-access, broad-exception-raised
 
+import logging
 from abc import ABC, abstractmethod
 from collections import namedtuple
 from datetime import datetime
@@ -16,6 +17,8 @@
 import numpy as np
 from google.cloud import bigtable
 
+logger = logging.getLogger(__name__)
+
 from . import locks
 from . import edits
 from . import types
@@ -444,6 +447,19 @@ def execute(
                         operation_id=lock.operation_id,
                         timestamp=override_ts if override_ts else timestamp,
                     )
+                # Log cache stats
+                if self.cg.cache:
+                    stats = self.cg.cache.get_stats()
+                    lines = [f"[Op {lock.operation_id}] Cache:"]
+                    for name, s in stats.items():
+                        lines.append(f"  {name}: {s['hit_rate']} hit ({s['hits']}/{s['total']}) calls={s['calls']}")
+                        # Show top miss sources if any
+                        if s.get("sources"):
+                            top_sources = sorted(s["sources"].items(), key=lambda x: -x[1]["misses"])[:3]
+                            if top_sources and any(src[1]["misses"] > 0 for src in top_sources):
+                                src_str = ", ".join(f"{k}({v['misses']})" for k, v in top_sources if v["misses"] > 0)
+                                lines.append(f"    miss sources: {src_str}")
+                    logger.info("\n".join(lines))
                 if self.cg.meta.READ_ONLY:
                     # return without persisting changes
                     return GraphEditOperation.Result(