matchbox-db 0.7.14.dev43__tar.gz → 0.7.14.dev46__tar.gz

{matchbox_db-0.7.14.dev43/src/matchbox_db.egg-info → matchbox_db-0.7.14.dev46}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: matchbox-db
-Version: 0.7.14.dev43
+Version: 0.7.14.dev46
 Summary: A framework for orchestrating and comparing data linking and deduplication methodologies.
 Author: Department for Business and Trade
 Project-URL: Documentation, https://uktrade.github.io/matchbox/
@@ -29,6 +29,7 @@ Requires-Dist: sqlglot[rs]>=26.12.1
 Requires-Dist: tenacity>=9.1.2
 Requires-Dist: textual>=0.87.1
 Requires-Dist: typer>=0.15.0
+Requires-Dist: psutil>=7.1.3
 Provides-Extra: server
 Requires-Dist: adbc-driver-postgresql>=1.5.0; extra == "server"
 Requires-Dist: alembic>=1.15.2; extra == "server"

{matchbox_db-0.7.14.dev43 → matchbox_db-0.7.14.dev46}/pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
     "tenacity>=9.1.2",
     "textual>=0.87.1",
     "typer>=0.15.0",
+    "psutil>=7.1.3"
 ]
 
 [project.optional-dependencies]

{matchbox_db-0.7.14.dev43 → matchbox_db-0.7.14.dev46}/src/matchbox/client/_handler.py

@@ -65,7 +65,7 @@ from matchbox.common.exceptions import (
     MatchboxUserNotFoundError,
 )
 from matchbox.common.hash import hash_to_base64
-from matchbox.common.logging import logger, profile
+from matchbox.common.logging import logger, profile_time
 
 URLEncodeHandledType = str | int | float | bytes | StrEnum
 
@@ -422,7 +422,7 @@ def update_resolution(
     return ResourceOperationStatus.model_validate(res.json())
 
 
-@profile(kwarg="path")
+@profile_time(kwarg="path")
 @http_retry
 def get_resolution(path: ResolutionPath) -> Resolution | None:
     """Get a resolution from Matchbox."""
@@ -435,7 +435,7 @@ def get_resolution(path: ResolutionPath) -> Resolution | None:
     return Resolution.model_validate(res.json())
 
 
-@profile(kwarg="path")
+@profile_time(kwarg="path")
 @http_retry
 def set_data(path: ResolutionPath, data: pl.DataFrame | Table) -> None:
     """Upload source hashes or model results to server."""
@@ -482,7 +482,7 @@ def get_resolution_stage(path: ResolutionPath) -> UploadStage:
     return upload_info.stage
 
 
-@profile(kwarg="path")
+@profile_time(kwarg="path")
 @http_retry
 def get_results(path: ModelResolutionPath) -> Table:
     """Get model results from Matchbox."""

{matchbox_db-0.7.14.dev43 → matchbox_db-0.7.14.dev46}/src/matchbox/client/dags.py

@@ -28,7 +28,7 @@ from matchbox.common.exceptions import (
     MatchboxCollectionNotFoundError,
     MatchboxResolutionNotFoundError,
 )
-from matchbox.common.logging import logger, profile
+from matchbox.common.logging import logger, profile_time
 from matchbox.common.transform import truth_int_to_float
 
 
@@ -561,7 +561,7 @@ class DAG:
         # If no matches, _handler will raise
         return {from_source: list(matches[0].source_id), **to_sources_results}
 
-    @profile(kwarg="node")
+    @profile_time(kwarg="node")
     def resolve(
         self,
         node: ResolutionName | None = None,

{matchbox_db-0.7.14.dev43 → matchbox_db-0.7.14.dev46}/src/matchbox/client/models/models.py

@@ -24,7 +24,7 @@ from matchbox.common.dtos import (
 )
 from matchbox.common.exceptions import MatchboxResolutionNotFoundError
 from matchbox.common.hash import hash_model_results
-from matchbox.common.logging import logger, profile
+from matchbox.common.logging import logger, profile_mem, profile_time
 from matchbox.common.transform import truth_float_to_int, truth_int_to_float
 
 if TYPE_CHECKING:
@@ -234,7 +234,8 @@ class Model:
         result = _handler.delete_resolution(path=self.resolution_path, certain=certain)
         return result.success
 
-    @profile(attr="name")
+    @profile_mem()
+    @profile_time(attr="name")
     def run(
         self, for_validation: bool = False, cache_queries: bool = False
     ) -> ModelResults:
@@ -284,7 +285,8 @@ class Model:
         return self.results
 
     @post_run
-    @profile(attr="name")
+    @profile_mem()
+    @profile_time(attr="name")
     def sync(self) -> None:
         """Send the model config and results to the server.
 

{matchbox_db-0.7.14.dev43 → matchbox_db-0.7.14.dev46}/src/matchbox/client/sources.py

@@ -24,7 +24,7 @@ from matchbox.common.dtos import (
 )
 from matchbox.common.exceptions import MatchboxResolutionNotFoundError
 from matchbox.common.hash import HashMethod, hash_arrow_table, hash_rows
-from matchbox.common.logging import logger, profile
+from matchbox.common.logging import logger, profile_mem, profile_time
 
 if TYPE_CHECKING:
     from matchbox.client.dags import DAG
@@ -223,7 +223,7 @@ class Source:
             return False
         return self.config == other.config
 
-    @profile(attr="name")
+    @profile_time(attr="name")
     def fetch(
         self,
         qualify_names: bool = False,
@@ -270,7 +270,8 @@ class Source:
                 return_type=return_type,
             )
 
-    @profile(attr="name")
+    @profile_mem()
+    @profile_time(attr="name")
     def run(self, batch_size: int | None = None) -> ArrowTable:
         """Hash a dataset from its warehouse, ready to be inserted, and cache hashes.
 
@@ -375,7 +376,8 @@ class Source:
         return self.config.f(self.name, fields)
 
     @post_run
-    @profile(attr="name")
+    @profile_mem()
+    @profile_time(attr="name")
     def sync(self) -> None:
         """Send the source config and hashes to the server.
 

{matchbox_db-0.7.14.dev43 → matchbox_db-0.7.14.dev46}/src/matchbox/common/logging.py

@@ -4,9 +4,11 @@ import functools
 import importlib.metadata
 import logging
 import time
+import tracemalloc
 from collections.abc import Callable
 from typing import Any, Final, Literal, ParamSpec, TypeVar
 
+import psutil
 from rich.console import Console
 from rich.progress import (
     BarColumn,
@@ -115,14 +117,14 @@ T = TypeVar("T")
 P = ParamSpec("P")
 
 
-def profile(
+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 functions and methods using logger.
+    """Decorator to profile running time of functions and methods using logger.
 
     Args:
         logger: The logger to use.
@@ -170,3 +172,45 @@ def profile(
         return wrapper
 
     return decorator
+
+
+def profile_mem(
+    logger: PrefixedLoggerAdapter = logger,
+    level: int = logging.DEBUG,
+    prefix: str | None = "Memory",
+) -> Callable[[Callable[P, T]], Callable[P, T]]:
+    """Decorator to profile memory use 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.
+    """
+
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
+        @functools.wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+            tracemalloc.start()
+
+            try:
+                return func(*args, **kwargs)
+            finally:
+                # Final and peak mem use of decorated function
+                curr, peak = tracemalloc.get_traced_memory()
+                tracemalloc.stop()
+                # OS memory availability after function run
+                mem = psutil.virtual_memory()
+
+                msg = (
+                    f"Memory usage of {func.__name__}: "
+                    f"current {round(curr / 10**6, 2)} MB, "
+                    f"peak {round(peak / 10**6, 2)} MB. "
+                    f"Total memory: {round(mem[0] / 10**6, 2)} MB, "
+                    f"available memory: {round(mem[1] / 10**6, 2)} MB."
+                )
+
+                logger.log(level, msg, prefix=prefix)
+
+        return wrapper
+
+    return decorator

{matchbox_db-0.7.14.dev43 → matchbox_db-0.7.14.dev46/src/matchbox_db.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: matchbox-db
-Version: 0.7.14.dev43
+Version: 0.7.14.dev46
 Summary: A framework for orchestrating and comparing data linking and deduplication methodologies.
 Author: Department for Business and Trade
 Project-URL: Documentation, https://uktrade.github.io/matchbox/
@@ -29,6 +29,7 @@ Requires-Dist: sqlglot[rs]>=26.12.1
 Requires-Dist: tenacity>=9.1.2
 Requires-Dist: textual>=0.87.1
 Requires-Dist: typer>=0.15.0
+Requires-Dist: psutil>=7.1.3
 Provides-Extra: server
 Requires-Dist: adbc-driver-postgresql>=1.5.0; extra == "server"
 Requires-Dist: alembic>=1.15.2; extra == "server"

{matchbox_db-0.7.14.dev43 → matchbox_db-0.7.14.dev46}/src/matchbox_db.egg-info/requires.txt

@@ -19,6 +19,7 @@ sqlglot[rs]>=26.12.1
 tenacity>=9.1.2
 textual>=0.87.1
 typer>=0.15.0
+psutil>=7.1.3
 
 [server]
 adbc-driver-postgresql>=1.5.0

{matchbox_db-0.7.14.dev43 → matchbox_db-0.7.14.dev46}/uv.lock

@@ -1283,6 +1283,7 @@ dependencies = [
     { name = "pandas" },
     { name = "polars" },
     { name = "polars-hash" },
+    { name = "psutil" },
     { name = "psycopg", extra = ["binary", "pool"] },
     { name = "pyarrow" },
     { name = "pydantic" },
@@ -1351,6 +1352,7 @@ requires-dist = [
     { name = "pandas", specifier = ">=2.2.3" },
     { name = "polars", specifier = ">=1.32.3" },
     { name = "polars-hash", specifier = ">=0.5.3" },
+    { name = "psutil", specifier = ">=7.1.3" },
     { name = "psycopg", extras = ["binary", "pool"], specifier = ">=3.2.6" },
     { name = "pyarrow", specifier = ">=17.0.0" },
     { name = "pydantic", specifier = ">=2.9.2" },