Add common stats collection objects

.gitignore

@@ -26,6 +26,7 @@ site/
 /CLAUDE.md
 /GEMINI.md
 /AGENTS.md
+.codex
 
 # .agents dir
 .agents
@@ -209,4 +210,4 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+#.idea/

src/matchbox/client/_handler/collections.py

@@ -32,7 +32,7 @@
 from matchbox.common.exceptions import (
     MatchboxServerFileError,
 )
-from matchbox.common.logging import logger, profile_time
+from matchbox.common.logging import logger
 
 # Collection management
 
@@ -241,7 +241,6 @@ def update_step(
     return ResourceOperationStatus.model_validate(res.json())
 
 
-@profile_time(kwarg="path")
 @http_retry
 def get_step(path: StepPath) -> Step | None:
     """Get a step from Matchbox."""
@@ -254,7 +253,6 @@ def get_step(path: StepPath) -> Step | None:
     return Step.model_validate(res.json())
 
 
-@profile_time(kwarg="path")
 @http_retry
 def set_data(path: StepPath, data: pl.DataFrame | Table) -> str:
     """Upload any step data to server."""
@@ -293,7 +291,6 @@ def set_data(path: StepPath, data: pl.DataFrame | Table) -> str:
     return upload_id
 
 
-@profile_time(kwarg="path")
 @http_retry
 def get_data(path: ModelStepPath | ResolverStepPath) -> Table:
     """Download step data from Matchbox."""

src/matchbox/client/dags.py

@@ -1,6 +1,7 @@
 """Objects to define a DAG which indexes, deduplicates and links data."""
 
 import tempfile
+import time
 from collections import deque
 from enum import StrEnum
 from pathlib import Path
@@ -37,7 +38,8 @@
     MatchboxCollectionNotFoundError,
     MatchboxStepNotFoundError,
 )
-from matchbox.common.logging import log_mem_usage, logger, profile_time
+from matchbox.common.logging import logger
+from matchbox.common.stats import DAGStats
 
 
 class DAGNodeExecutionStatus(StrEnum):
@@ -72,6 +74,7 @@ def __init__(
         self._run: RunID | None = None
         self.nodes: dict[StepName, Source | Model | Resolver] = {}
         self.graph: dict[StepName, list[StepName]] = {}
+        self.stats = DAGStats()
 
         CACHE_DIR.mkdir(parents=True, exist_ok=True)
         self._cache_dir = tempfile.TemporaryDirectory(dir=str(CACHE_DIR))
@@ -624,6 +627,9 @@ def run_and_sync(
         if batch_size is None:
             batch_size = settings.batch_size
 
+        self.stats.reset()
+        wall_start = time.perf_counter()
+
         sequence: list[StepName] = self.sequence
 
         # Identify skipped nodes
@@ -652,6 +658,8 @@ def run_and_sync(
         status = {
             step_name: DAGNodeExecutionStatus.SKIPPED for step_name in skipped_nodes
         }
+        for step_name in skipped_nodes:
+            self.stats.ensure_step(step_name)
 
         for step_name in sequence:
             node = self.nodes[step_name]
@@ -666,7 +674,7 @@ def run_and_sync(
                     node.run()
                 node.sync()
                 if profile:
-                    log_mem_usage()
+                    self.stats.record_mem(step_name)
             except Exception as e:
                 logger.error(f"❌ {node.name} failed: {e}")
                 raise e
@@ -676,9 +684,11 @@ def run_and_sync(
                 node.clear_data()
                 logger.info("Cleared node data")
                 if profile:
-                    log_mem_usage()
+                    self.stats.record_mem(step_name)
         logger.info("\n" + self.draw(status=status))
 
+        self.stats.dag_run_seconds = time.perf_counter() - wall_start
+
     def set_default(self) -> None:
         """Set the current run as the default for the collection.
 
@@ -740,7 +750,6 @@ def lookup_key(
         return {from_source: list(matches[0].source_id), **to_sources_results}
 
     @validate_call
-    @profile_time(kwarg="node")
     def get_matches(
         self,
         resolver: ResolverStepName | None = None,
@@ -759,6 +768,8 @@ def get_matches(
         if not isinstance(resolver, Resolver):
             raise ValueError("get_matches can only query from resolver nodes")
 
+        self.stats.reset()
+
         available_sources = {
             node_name: self.get_source(node_name) for node_name in resolver.sources
         }
@@ -784,14 +795,15 @@ def get_matches(
         query_results: list[pl.DataFrame] = []
         for source_name in filtered_source_names:
             resolved_sources.append(available_sources[source_name])
-            query_results.append(
-                pl.from_arrow(
-                    _handler.query(
-                        source=available_sources[source_name].path,
-                        resolver=resolver.path,
-                        return_leaf_id=True,
+            with DAGStats.time("query", name=source_name, stats=self.stats):
+                query_results.append(
+                    pl.from_arrow(
+                        _handler.query(
+                            source=available_sources[source_name].path,
+                            resolver=resolver.path,
+                            return_leaf_id=True,
+                        )
                     )
                 )
-            )
 
         return ResolverMatches(sources=resolved_sources, query_results=query_results)

src/matchbox/client/models/models.py

@@ -24,7 +24,8 @@
     StepType,
 )
 from matchbox.common.hash import hash_arrow_table
-from matchbox.common.logging import logger, profile_time
+from matchbox.common.logging import logger
+from matchbox.common.stats import DAGStats
 
 if TYPE_CHECKING:
     from matchbox.client.dags import DAG
@@ -206,7 +207,7 @@ def path(self) -> ModelStepPath:
             name=self.name,
         )
 
-    @profile_time(attr="name")
+    @DAGStats.time("compute")
     def compute_scores(
         self, left_df: DataFrame, right_df: DataFrame | None = None
     ) -> DataFrame:
@@ -238,20 +239,21 @@ def run(
         log_prefix = f"Run {self.name}"
         logger.info("Executing left query", prefix=log_prefix)
 
-        left_df = (
-            left_data
-            if left_data is not None
-            else self.left_query.data(cache_leaf_ids=(not low_memory))
-        )
-        right_df = None
-
-        if self.config.type == ModelType.LINKER:
-            logger.info("Executing right query", prefix=log_prefix)
-            right_df = (
-                right_data
-                if right_data is not None
-                else self.right_query.data(cache_leaf_ids=(not low_memory))
+        with DAGStats.time("query", name=self.name, stats=self._stats):
+            left_df = (
+                left_data
+                if left_data is not None
+                else self.left_query.data(cache_leaf_ids=(not low_memory))
             )
+            right_df = None
+
+            if self.config.type == ModelType.LINKER:
+                logger.info("Executing right query", prefix=log_prefix)
+                right_df = (
+                    right_data
+                    if right_data is not None
+                    else self.right_query.data(cache_leaf_ids=(not low_memory))
+                )
 
         logger.info("Running model logic", prefix=log_prefix)
         scores = self.compute_scores(left_df, right_df)

src/matchbox/client/queries.py

@@ -19,7 +19,6 @@
 from matchbox.client.models.linkers.base import Linker, LinkerSettings
 from matchbox.common.db import QueryReturnClass, QueryReturnType
 from matchbox.common.dtos import QueryCombineType, QueryConfig
-from matchbox.common.logging import profile_time
 
 if TYPE_CHECKING:
     from matchbox.client.dags import DAG
@@ -207,7 +206,6 @@ def data_raw(
 
             return _convert_df(raw_data.collect(), return_type=return_type)
 
-    @profile_time()
     def data(
         self,
         raw_data: pl.DataFrame | None = None,

src/matchbox/client/resolvers/resolvers.py

@@ -23,7 +23,8 @@
 )
 from matchbox.common.exceptions import MatchboxStepTypeError
 from matchbox.common.hash import hash_clusters
-from matchbox.common.logging import logger, profile_time
+from matchbox.common.logging import logger
+from matchbox.common.stats import DAGStats
 
 if TYPE_CHECKING:
     from matchbox.client.dags import DAG
@@ -159,14 +160,13 @@ def path(self) -> ResolverStepPath:
             name=self.name,
         )
 
-    @profile_time(attr="name")
+    @DAGStats.time("compute")
     def compute_clusters(
         self, model_edges: Mapping[StepName, pl.DataFrame]
     ) -> pl.DataFrame:
         """Delegate cluster computation to the configured resolver instance."""
         return self.resolver_instance.compute_clusters(model_edges=model_edges)
 
-    @profile_time(attr="name")
     def run(self) -> pl.DataFrame:
         """Run the resolver and materialise cluster assignments."""
         model_edges: dict[StepName, pl.DataFrame] = {}

src/matchbox/client/sources.py

@@ -24,7 +24,8 @@
     StepType,
 )
 from matchbox.common.hash import HashMethod, hash_rows
-from matchbox.common.logging import logger, profile_time
+from matchbox.common.logging import logger
+from matchbox.common.stats import DAGStats
 
 if TYPE_CHECKING:
     from matchbox.client.dags import DAG
@@ -280,7 +281,7 @@ def sample(
         """Peek at the top n entries in a source."""
         return next(self.fetch(batch_size=n, return_type=return_type))
 
-    @profile_time(attr="name")
+    @DAGStats.time("hash")
     def run(self, batch_size: int | None = None) -> pl.DataFrame:
         """Hash a dataset from its warehouse, ready to be inserted, and cache hashes.
 

src/matchbox/client/steps.py

@@ -19,7 +19,8 @@
 )
 from matchbox.common.exceptions import MatchboxStepNotFoundError
 from matchbox.common.hash import hash_arrow_table
-from matchbox.common.logging import logger, profile_time
+from matchbox.common.logging import logger
+from matchbox.common.stats import DAGStats
 
 if TYPE_CHECKING:
     from matchbox.client.dags import DAG
@@ -79,6 +80,7 @@ def __init__(
         self.name = name
         self.description = description
         self._local_data: pl.DataFrame | None = None
+        self._stats: DAGStats = dag.stats
 
     # Local data access
 
@@ -172,7 +174,7 @@ def download(self) -> pl.DataFrame:
         return self._local_data
 
     @post_run
-    @profile_time(attr="name")
+    @DAGStats.time("sync")
     def sync(self) -> None:
         """Send step config and local data to the server.
 

src/matchbox/common/logging.py

@@ -1,14 +1,9 @@
 """Logging utilities."""
 
-import functools
 import importlib.metadata
 import logging
-import os
-import time
-from collections.abc import Callable
-from typing import Any, Final, Literal, ParamSpec, TypeVar
+from typing import Any, Final, Literal
 
-import psutil
 from rich.console import Console
 from rich.progress import (
     BarColumn,
@@ -95,79 +90,6 @@ def build_progress_bar(console_: Console | None = None) -> Progress:
     )
 
 
-T = TypeVar("T")
-P = ParamSpec("P")
-
-
-def profile_time(
-    logger: PrefixedLoggerAdapter = logger,
-    level: int = logging.INFO,
-    prefix: str | None = "Profiling",
-    attr: str | None = None,
-    kwarg: str | None = None,
-) -> Callable[[Callable[P, T]], Callable[P, T]]:
-    """Decorator to profile running time of functions and methods using logger.
-
-    Args:
-        logger: The logger to use.
-        level: The level to use to log the profiling information. It defaults to INFO.
-        prefix: Prefix to pass to the logged message.
-        attr: Attribute name to extract from instantiated class.
-        kwarg: Argument name to extract from function call.
-
-    `attr` should be used when we want to include some class atribute in the log, e.g.
-    node name in Source.
-    `kwarg` should be used when we want to include the value passed to some function
-    argument, e.g. path in `set_data`. This will only work if the argument is passed
-    to the function as a kwarg, and not a positional argument.
-    """
-
-    def decorator(func: Callable[P, T]) -> Callable[P, T]:
-        @functools.wraps(func)
-        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
-            # If attr, will try to get its value from class instance
-            if attr is not None:
-                # If class, first argument will be self
-                self = args[0] if args else None
-                node = getattr(self, attr) if self and hasattr(self, attr) else None
-            # If kwarg, will try to get value passed to function
-            if kwarg is not None:
-                value = kwargs.get(kwarg)
-
-            start = time.perf_counter()
-            try:
-                return func(*args, **kwargs)
-            finally:
-                duration = time.perf_counter() - start
-
-                if attr:
-                    msg = f"`{func.__name__}` in node `{node}` took {duration:.3f}s"
-                elif kwarg:
-                    msg = (
-                        f"`{func.__name__}` with {kwarg} `{value}` took {duration:.3f}s"
-                    )
-                else:
-                    msg = f"`{func.__name__}` took {duration:.3f}s"
-
-                logger.log(level, msg, prefix=prefix)
-
-        return wrapper
-
-    return decorator
-
-
-def log_mem_usage(
-    logger: PrefixedLoggerAdapter = logger,
-    level: int = logging.INFO,
-    prefix: str | None = "Profiling",
-) -> None:
-    """Log current memory usage for this process."""
-    proc = psutil.Process(os.getpid())
-    usage = proc.memory_info().rss / (1024**2)
-    msg = f"Current memory used by process (MiB): {usage}"
-    logger.log(level, msg, prefix=prefix)
-
-
 def get_formatter() -> logging.Formatter:
     """Retrieve plugin registered in 'matchbox.logging' entry point, or fallback."""
     global _PLUGINS