diff --git a/.gitignore b/.gitignore index 12e43b18..baea668f 100644 --- a/.gitignore +++ b/.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/ \ No newline at end of file +#.idea/ diff --git a/src/matchbox/client/_handler/collections.py b/src/matchbox/client/_handler/collections.py index d93a5d92..6badcb3e 100644 --- a/src/matchbox/client/_handler/collections.py +++ b/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.""" diff --git a/src/matchbox/client/dags.py b/src/matchbox/client/dags.py index 0f95abc8..e621e137 100644 --- a/src/matchbox/client/dags.py +++ b/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) diff --git a/src/matchbox/client/models/models.py b/src/matchbox/client/models/models.py index 935607bb..167713f2 100644 --- a/src/matchbox/client/models/models.py +++ b/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) diff --git a/src/matchbox/client/queries.py b/src/matchbox/client/queries.py index 1d753ce7..3f2a75d6 100644 --- a/src/matchbox/client/queries.py +++ b/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, diff --git a/src/matchbox/client/resolvers/resolvers.py b/src/matchbox/client/resolvers/resolvers.py index 3f63bb95..a32d2c84 100644 --- a/src/matchbox/client/resolvers/resolvers.py +++ b/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] = {} diff --git a/src/matchbox/client/sources.py b/src/matchbox/client/sources.py index f3d7be64..14fb02e3 100644 --- a/src/matchbox/client/sources.py +++ b/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. diff --git a/src/matchbox/client/steps.py b/src/matchbox/client/steps.py index 2a529b2f..05171cf3 100644 --- a/src/matchbox/client/steps.py +++ b/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. diff --git a/src/matchbox/common/logging.py b/src/matchbox/common/logging.py index 80da656e..49d90729 100644 --- a/src/matchbox/common/logging.py +++ b/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 diff --git a/src/matchbox/common/stats.py b/src/matchbox/common/stats.py new file mode 100644 index 00000000..086c80df --- /dev/null +++ b/src/matchbox/common/stats.py @@ -0,0 +1,159 @@ +"""Query and processing time statistics for DAGs.""" + +import functools +import os +import time +from collections.abc import Callable +from contextlib import ContextDecorator +from typing import Any, Self + +import psutil + +from matchbox.common.dtos import StepName +from matchbox.common.logging import logger + + +def log_mem_usage(name: str) -> float: + """Log current process memory usage for name and return the value in MiB.""" + usage = psutil.Process(os.getpid()).memory_info().rss / (1024**2) + logger.info(f"Memory in `{name}`: {usage:.1f} MiB", prefix="Stats") + return usage + + +class _StatLogger(ContextDecorator): + """Context manager / decorator that measures a metric's change and records it. + + Generic over any zero-argument callable returning a float, e.g. + time.perf_counter for elapsed time, or a memory-reading function. + + As decorator: resolves name and stats from the decorated instance at call + time. As context manager: name and stats must be passed explicitly. + """ + + def __init__( + self, + operation: str, + *, + metric: str, + measure: Callable[[], float], + name: str | None = None, + stats: "DAGStats | None" = None, + ) -> None: + self.operation = operation + self.metric = metric + self.measure = measure + self.name = name + self.stats = stats + self._start: float | None = None + self._instance: Any = None + + def _recreate_cm(self) -> "_StatLogger": + """Return a fresh instance to avoid mutable state leakage across calls.""" + return _StatLogger( + self.operation, + metric=self.metric, + measure=self.measure, + name=self.name, + stats=self.stats, + ) + + def __enter__(self) -> Self: + self._start = self.measure() + return self + + def __exit__(self, *exc_info: object) -> None: + if self._start is None: + return + delta = self.measure() - self._start + + stats = self.stats + name = self.name + + if stats is not None and name is not None: + stats.record_metric(name, self.metric, delta, operation=self.operation) + + logger.info( + f"`{self.operation}` in `{name}`: {self.metric} changed by {delta:.3f}", + prefix="Stats", + ) + + def __call__(self, func: Callable[..., Any]) -> Callable[..., Any]: + """Wrap func so each call is measured against the bound instance.""" + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: # noqa: ANN401 + instance = args[0] if args else None + stat_logger = self._recreate_cm() + stat_logger._instance = instance + stat_logger.name = getattr(instance, "name", None) + stat_logger.stats = getattr(instance, "_stats", None) + with stat_logger: + return func(*args, **kwargs) + + return wrapper + + +class DAGStats: + """Collects per-step timing and memory statistics for a DAG run.""" + + def __init__(self) -> None: + """Initialise empty metric storage.""" + self._metrics: dict[StepName, dict[str, dict[str | None, float]]] = {} + self.dag_run_seconds: float | None = None + + def reset(self) -> None: + """Clear all collected statistics.""" + self._metrics.clear() + self.dag_run_seconds = None + + def ensure_step(self, name: StepName) -> None: + """Register a step even if no metric has been recorded for it yet.""" + self._metrics.setdefault(name, {}) + + def record_metric( + self, name: StepName, metric: str, value: float, *, operation: str | None = None + ) -> None: + """Record a metric value for a step, optionally scoped to an operation.""" + self._metrics.setdefault(name, {}).setdefault(metric, {})[operation] = value + + def record_mem(self, name: StepName) -> None: + """Record current process memory usage as a point estimate for name.""" + self.record_metric(name, "mem", log_mem_usage(name)) + + @staticmethod + def time( + operation: str, + *, + name: str | None = None, + stats: "DAGStats | None" = None, + ) -> _StatLogger: + """Return a _StatLogger that times wall-clock duration for operation.""" + return _StatLogger( + operation, metric="time", measure=time.perf_counter, name=name, stats=stats + ) + + @property + def timings(self) -> dict[StepName, dict[str, float]]: + """Per-step operation timings.""" + return { + name: dict(metrics.get("time", {})) + for name, metrics in self._metrics.items() + } + + @property + def mem(self) -> dict[StepName, float]: + """Latest point-estimate memory reading per step, in MiB.""" + return { + name: metrics["mem"][None] + for name, metrics in self._metrics.items() + if "mem" in metrics and None in metrics["mem"] + } + + @property + def total_run_seconds(self) -> float: + """Sum of all recorded operation durations, excluding dag_run_seconds.""" + return sum( + value + for metrics in self._metrics.values() + for value in metrics.get("time", {}).values() + ) diff --git a/src/matchbox/server/uploads.py b/src/matchbox/server/uploads.py index cecfe13b..d4697239 100644 --- a/src/matchbox/server/uploads.py +++ b/src/matchbox/server/uploads.py @@ -24,7 +24,8 @@ StepType, ) from matchbox.common.exceptions import MatchboxServerFileError -from matchbox.common.logging import configure_celery_logging, log_mem_usage, logger +from matchbox.common.logging import configure_celery_logging, logger +from matchbox.common.stats import log_mem_usage from matchbox.server.base import ( MatchboxBackends, MatchboxDBAdapter, @@ -280,7 +281,7 @@ def process_upload_celery( MBDB._disconnect_adbc() celery_logger.info("Uploading data for step %s, ID %s", str(step_path), upload_id) - log_mem_usage() + log_mem_usage(str(step_path)) upload_function = partial( process_upload, @@ -311,7 +312,7 @@ def process_upload_celery( raise finally: - log_mem_usage() + log_mem_usage(str(step_path)) # Cleanup connections if CELERY_SETTINGS.backend_type == MatchboxBackends.POSTGRES: diff --git a/test/client/test_dags.py b/test/client/test_dags.py index 56711b24..6ed5fd10 100644 --- a/test/client/test_dags.py +++ b/test/client/test_dags.py @@ -1613,3 +1613,163 @@ def test_dag_set_default_unreachable_nodes(sqla_sqlite_warehouse: Engine) -> Non with pytest.raises(ValueError, match="unreachable"): dag.set_default() + + +def test_run_and_sync_records_stats( + sqla_sqlite_warehouse: Engine, +) -> None: + """run_and_sync resets stats and records dag_run_seconds. + + When run/sync are mocked the decorated method bodies are bypassed, so + individual operation timings may be empty. We assert the wall-clock + field and that every node gets an entry (empty or not). + """ + foo_tkit = source_factory( + name="foo", engine=sqla_sqlite_warehouse + ).write_to_location() + bar_tkit = source_factory( + name="bar", engine=sqla_sqlite_warehouse + ).write_to_location() + + dag = TestkitDAG().dag + foo = dag.source(**foo_tkit.into_dag()) + dag.source(**bar_tkit.into_dag()) + foo_dedupe = foo.query().deduper( + name="foo_dedupe", + model_class=NaiveDeduper, + model_settings={"unique_fields": []}, + ) + foo_dedupe.resolver(name="root", resolver_class=Components) + + with ( + patch.object(Source, "run"), + patch.object(Source, "sync"), + patch.object(Model, "run"), + patch.object(Model, "sync"), + patch.object(Resolver, "run"), + patch.object(Resolver, "sync"), + ): + dag.run_and_sync() + + # Wall clock is always recorded + assert dag.stats.dag_run_seconds is not None + assert dag.stats.dag_run_seconds > 0 + + +def test_get_matches_records_stats( + matchbox_api: MockRouter, +) -> None: + """get_matches records 'query' timing per source.""" + foo = source_factory(name="foo", location_name="sqlite") + bar = source_factory(name="bar", location_name="postgres") + + foo_data = pa.Table.from_pylist( + [{"id": 1, "leaf_id": 1, "key": "1"}], + schema=SCHEMA_QUERY_WITH_LEAVES, + ) + bar_data = pa.Table.from_pylist( + [{"id": 2, "leaf_id": 2, "key": "a"}], + schema=SCHEMA_QUERY_WITH_LEAVES, + ) + + dag = DAG("companies") + + matchbox_api.get(f"/collections/{dag.name}").mock( + return_value=Response( + 200, + json=Collection(name=dag.name, runs=[], default_run=None).model_dump(), + ) + ) + matchbox_api.post(f"/collections/{dag.name}/runs").mock( + return_value=Response(200, json=Run(run_id=1, steps={}).model_dump()), + ) + + foo_source = dag.source(**foo.into_dag()) + bar_source = dag.source(**bar.into_dag()) + foo_dedupe = foo_source.query().deduper( + name="foo_dedupe", + model_class=NaiveDeduper, + model_settings={"unique_fields": []}, + ) + foo_bar = foo_source.query().linker( + bar_source.query(), + name="foo_bar", + model_class=DeterministicLinker, + model_settings={"comparisons": "l.field=r.field"}, + ) + foo_bar_resolver = foo_bar.resolver( + foo_dedupe, name="foo_bar_resolver", resolver_class=Components + ) + dag.new_run() + + matchbox_api.get( + "/query", + params={ + "source": "foo", + "run_id": 1, + "collection": dag.name, + "resolver": foo_bar_resolver.name, + "return_leaf_id": "True", + }, + ).mock(return_value=Response(200, content=table_to_buffer(foo_data).read())) + + matchbox_api.get( + "/query", + params={ + "source": "bar", + "run_id": 1, + "collection": dag.name, + "resolver": foo_bar_resolver.name, + "return_leaf_id": "True", + }, + ).mock(return_value=Response(200, content=table_to_buffer(bar_data).read())) + + dag.get_matches() + + assert "foo" in dag.stats.timings + assert "query" in dag.stats.timings["foo"] + assert "bar" in dag.stats.timings + assert "query" in dag.stats.timings["bar"] + + +def test_stats_reset_between_calls( + sqla_sqlite_warehouse: Engine, +) -> None: + """Stats are reset between run_and_sync calls.""" + foo_tkit = source_factory( + name="foo", engine=sqla_sqlite_warehouse + ).write_to_location() + + dag = TestkitDAG().dag + foo = dag.source(**foo_tkit.into_dag()) + foo.query().deduper( + name="foo_dedupe", + model_class=NaiveDeduper, + model_settings={"unique_fields": []}, + ).resolver(name="root", resolver_class=Components) + + with ( + patch.object(Source, "run"), + patch.object(Source, "sync"), + patch.object(Model, "run"), + patch.object(Model, "sync"), + patch.object(Resolver, "run"), + patch.object(Resolver, "sync"), + ): + dag.run_and_sync() + assert dag.stats.dag_run_seconds is not None + first_run_seconds = dag.stats.dag_run_seconds + + with ( + patch.object(Source, "run"), + patch.object(Source, "sync"), + patch.object(Model, "run"), + patch.object(Model, "sync"), + patch.object(Resolver, "run"), + patch.object(Resolver, "sync"), + ): + dag.run_and_sync() + + # dag_run_seconds should reflect only the second run (reset worked) + assert dag.stats.dag_run_seconds is not None + assert dag.stats.dag_run_seconds != first_run_seconds diff --git a/test/common/test_stats.py b/test/common/test_stats.py new file mode 100644 index 00000000..5722d301 --- /dev/null +++ b/test/common/test_stats.py @@ -0,0 +1,115 @@ +"""Tests for the DAGStats statistics class.""" + +import logging +import time + +import pytest + +from matchbox.common.stats import DAGStats + + +class _DummyStep: + """Minimal stand-in for a step node with name and _stats.""" + + def __init__(self, name: str, stats: DAGStats) -> None: + self.name = name + self._stats = stats + + @DAGStats.time("op") + def do_something(self) -> int: + return 42 + + +class _StepWithoutStats: + """A class without _stats attribute — decorator should no-op.""" + + name = "no_stats" + + @DAGStats.time("op") + def do_something(self) -> int: + return 99 + + +def test_time_as_decorator(caplog: pytest.LogCaptureFixture) -> None: + """Decorator mode records timing and logs.""" + stats = DAGStats() + step = _DummyStep(name="step_a", stats=stats) + + with caplog.at_level(logging.INFO, logger="matchbox"): + result = step.do_something() + + assert result == 42 + assert "step_a" in stats.timings + assert "op" in stats.timings["step_a"] + assert stats.timings["step_a"]["op"] >= 0 + assert any("Stats" in r.message for r in caplog.records) + + +def test_time_as_context_manager(caplog: pytest.LogCaptureFixture) -> None: + """Context manager mode records timing and logs.""" + stats = DAGStats() + + with ( + caplog.at_level(logging.INFO, logger="matchbox"), + DAGStats.time("query", name="source_x", stats=stats), + ): + time.sleep(0.001) + + assert "source_x" in stats.timings + assert "query" in stats.timings["source_x"] + assert stats.timings["source_x"]["query"] >= 0 + + +def test_record_mem(caplog: pytest.LogCaptureFixture) -> None: + """record_mem populates mem dict and logs.""" + stats = DAGStats() + + with caplog.at_level(logging.INFO, logger="matchbox"): + stats.record_mem("step_b") + + assert "step_b" in stats.mem + assert stats.mem["step_b"] > 0 + assert any("Memory" in r.message for r in caplog.records) + + +def test_reset() -> None: + """reset clears all collected statistics.""" + stats = DAGStats() + stats.record_metric("step_a", "time", 1.0, operation="op") + stats.record_mem("step_a") + stats.dag_run_seconds = 5.0 + + stats.reset() + + assert stats.timings == {} + assert stats.mem == {} + assert stats.dag_run_seconds is None + + +def test_total_run_seconds() -> None: + """total_run_seconds sums all operation durations.""" + stats = DAGStats() + stats.record_metric("step_a", "time", 1.0, operation="hash") + stats.record_metric("step_a", "time", 2.0, operation="sync") + stats.record_metric("step_b", "time", 3.5, operation="compute") + + assert stats.total_run_seconds == 6.5 + + +def test_dag_run_seconds_excluded_from_total() -> None: + """dag_run_seconds is not included in total_run_seconds.""" + stats = DAGStats() + stats.record_metric("step_a", "time", 1.0, operation="op") + stats.dag_run_seconds = 100.0 + + assert stats.total_run_seconds == 1.0 + + +def test_timer_no_stats_no_crash(caplog: pytest.LogCaptureFixture) -> None: + """Decorator on a class without _stats does not crash.""" + step = _StepWithoutStats() + + with caplog.at_level(logging.INFO, logger="matchbox"): + result = step.do_something() + + assert result == 99 diff --git a/test/e2e/test_e2e_dag.py b/test/e2e/test_e2e_dag.py index 44e903e3..72919496 100644 --- a/test/e2e/test_e2e_dag.py +++ b/test/e2e/test_e2e_dag.py @@ -285,6 +285,11 @@ def test_dag_pipeline_creation_and_rerun(self) -> None: logging.info("Running DAG for the first time") dag.run_and_sync() + # Stats should cover every node + assert len(dag.stats.timings) == len(dag.nodes) + assert dag.stats.dag_run_seconds is not None + assert dag.stats.dag_run_seconds > 0 + assert DAG.list_all() == [dag.name] # Update metadata of one node, will check later @@ -320,6 +325,11 @@ def test_dag_pipeline_creation_and_rerun(self) -> None: # Can retrieve whole lookup dag1_lookup = dag.get_matches().as_lookup() + # get_matches should record a "query" entry per queried source + for source_name in dag.default_resolver.sources: + assert source_name in dag.stats.timings + assert "query" in dag.stats.timings[source_name] + # Set as new default dag.set_default() @@ -356,6 +366,11 @@ def test_dag_pipeline_creation_and_rerun(self) -> None: assert rerun_dag.run != dag.run rerun_dag.run_and_sync() + # Re-run should have fresh stats with all nodes + assert len(rerun_dag.stats.timings) == len(rerun_dag.nodes) + assert rerun_dag.stats.dag_run_seconds is not None + assert rerun_dag.stats.dag_run_seconds > 0 + # The lookup is identical assert_frame_equal( rerun_dag.get_matches().as_lookup(),