FEAT-modin-project#7445: Add metrics interface so third-parties can collect metrics from the modin frontend (modin-project#7444)

Adds an interface for collecting frontend API statistics by third party systems. Specifically this will be used by Snowflake's pandas engine for Modin to collect data on interactive use cases.

Signed-off-by: John Kew <[email protected]>
Co-authored-by: Jonathan Shi <[email protected]>
Co-authored-by: Mahesh Vashishtha <[email protected]>
Co-authored-by: Devin Petersohn <[email protected]>

Author: 4 people

.github/workflows/ci.yml

@@ -156,6 +156,7 @@ jobs:
       - run: python -m pytest modin/tests/interchange/dataframe_protocol/base
       - run: python -m pytest modin/tests/test_dataframe_api_standard.py
       - run: python -m pytest modin/tests/test_logging.py
+      - run: python -m pytest modin/tests/test_metrics.py
       - uses: ./.github/actions/upload-coverage
 
   test-defaults:

docs/usage_guide/advanced_usage/index.rst

@@ -10,6 +10,7 @@ Advanced Usage
    progress_bar
    modin_xgboost
    modin_logging
+   modin_metrics
    batch
    modin_engines
 

docs/usage_guide/advanced_usage/modin_metrics.rst

@@ -0,0 +1,47 @@
+Modin Metrics
+=============
+
+Modin allows for third-party systems to register a metrics handler to collect specific API statistics.
+Metrics have a name and a value, can be aggregated, discarded, or emitted without impact to the program.
+
+CPU load, memory usage, and disk usage are all typical metrics; but modin currently only emits metrics on API timings which can be used to optimize end-user interactive performance. New metrics may 
+be added in the future.
+
+It is the responsibility of the handler to process or forward these metrics. The name of the metric will 
+be in "dot format" and all lowercase, similar to graphite or rrd. The value is an integer or float.
+
+Example metric names include:
+
+* 'modin.core-dataframe.pandasdataframe.copy_index_cache'
+* 'modin.core-dataframe.pandasdataframe.transpose'
+* 'modin.query-compiler.pandasquerycompiler.transpose'
+* 'modin.query-compiler.basequerycompiler.columnarize'
+* 'modin.pandas-api.series.__init__'
+* 'modin.pandas-api.dataframe._reduce_dimension'
+* 'modin.pandas-api.dataframe.sum'
+
+Handlers are functions of the form: `fn(str, int|float)` and can be registered with:
+
+.. code-block:: python
+
+  import modin.pandas as pd
+  from modin.logging.metrics import add_metric_handler
+
+  def func(name: str, value: int | float):
+    print(f"Got metric {name} value {value}")
+
+  add_metric_handler(func)
+
+.. warning:: 
+  A metric handler should be non-blocking, returning within 100ms, although this is not enforced. It must not throw exceptions or it will
+  be deregistered. These restrictions are to help guard against the implementation of a metrics collector which would impact
+  interactice performance significantly. The data from metrics should generally be offloaded to another system for processing
+  and not involve any blocking network calls.
+
+Metrics are enabled by default. Modin metrics can be disabled like so:
+
+.. code-block:: python
+
+  import modin.pandas as pd
+  from modin.config import MetricsMode
+  MetricsMode.disable()

modin/config/__init__.py

@@ -36,6 +36,7 @@
     LogMemoryInterval,
     LogMode,
     Memory,
+    MetricsMode,
     MinColumnPartitionSize,
     MinPartitionSize,
     MinRowPartitionSize,
@@ -108,6 +109,7 @@
     "LogMode",
     "LogMemoryInterval",
     "LogFileSize",
+    "MetricsMode",
     # Plugin settings
     "DocModule",
 ]

modin/config/envvars.py

@@ -624,6 +624,31 @@ def get(cls) -> int:
         return log_file_size
 
 
+class MetricsMode(EnvironmentVariable, type=ExactStr):
+    """
+    Set ``MetricsMode`` value to disable/enable metrics collection.
+
+    Metric handlers are registered through `add_metric_handler` and can
+    be used to record graphite-style timings or values. It is the
+    responsibility of the handler to define how those emitted metrics
+    are handled.
+    """
+
+    varname = "MODIN_METRICS_MODE"
+    choices = ("enable", "disable")
+    default = "enable"
+
+    @classmethod
+    def enable(cls) -> None:
+        """Enable all metric collection."""
+        cls.put("enable")
+
+    @classmethod
+    def disable(cls) -> None:
+        """Disable all metric collection."""
+        cls.put("disable")
+
+
 class PersistentPickle(EnvironmentVariable, type=bool):
     """Whether serialization should be persistent."""
 

modin/logging/__init__.py

@@ -14,10 +14,14 @@
 from .class_logger import ClassLogger  # noqa: F401
 from .config import get_logger  # noqa: F401
 from .logger_decorator import disable_logging, enable_logging  # noqa: F401
+from .metrics import add_metric_handler, clear_metric_handler, emit_metric
 
 __all__ = [
     "ClassLogger",
     "get_logger",
     "enable_logging",
     "disable_logging",
+    "emit_metric",
+    "add_metric_handler",
+    "clear_metric_handler",
 ]

modin/logging/logger_decorator.py

@@ -16,14 +16,15 @@
 
 ``enable_logging`` is used for decorating individual Modin functions or classes.
 """
-
 from __future__ import annotations
 
 from functools import wraps
+from time import perf_counter
 from types import FunctionType, MethodType
 from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, overload
 
 from modin.config import LogMode
+from modin.logging.metrics import emit_metric
 
 from .config import LogLevel, get_logger
 
@@ -121,8 +122,11 @@ def decorator(obj: Fn) -> Fn:
 
         assert isinstance(modin_layer, str), "modin_layer is somehow not a string!"
 
-        start_line = f"START::{modin_layer.upper()}::{name or obj.__name__}"
-        stop_line = f"STOP::{modin_layer.upper()}::{name or obj.__name__}"
+        api_call_name = f"{name or obj.__name__}"
+        log_line = f"{modin_layer.upper()}::{api_call_name}"
+        metric_name = f"{modin_layer.lower()}.{api_call_name.lower()}"
+        start_line = f"START::{log_line}"
+        stop_line = f"STOP::{log_line}"
 
         @wraps(obj)
         def run_and_log(*args: Tuple, **kwargs: Dict) -> Any:
@@ -140,13 +144,17 @@ def run_and_log(*args: Tuple, **kwargs: Dict) -> Any:
             -------
             Any
             """
+            start_time = perf_counter()
             if LogMode.get() == "disable":
-                return obj(*args, **kwargs)
+                result = obj(*args, **kwargs)
+                emit_metric(metric_name, perf_counter() - start_time)
+                return result
 
             logger = get_logger()
             logger.log(log_level, start_line)
             try:
                 result = obj(*args, **kwargs)
+                emit_metric(metric_name, perf_counter() - start_time)
             except BaseException as e:
                 # Only log the exception if a deeper layer of the modin stack has not
                 # already logged it.

modin/logging/metrics.py

@@ -0,0 +1,79 @@
+# Licensed to Modin Development Team under one or more contributor license agreements.
+# See the NOTICE file distributed with this work for additional information regarding
+# copyright ownership.  The Modin Development Team licenses this file to you under the
+# Apache License, Version 2.0 (the "License"); you may not use this file except in
+# compliance with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software distributed under
+# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific language
+# governing permissions and limitations under the License.
+
+"""
+Module contains metrics handler functions.
+
+Allows for the registration of functions to collect
+API metrics.
+"""
+
+import re
+from typing import Callable, Union
+
+from modin.config.envvars import MetricsMode
+
+metric_name_pattern = r"[a-zA-Z\._\-0-9]+$"
+_metric_handlers: list[Callable[[str, Union[int, float]], None]] = []
+
+
+# Metric/Telemetry hooks can be implemented by plugin engines
+# to collect discrete data on how modin is performing at the
+# high level modin layer.
+def emit_metric(name: str, value: Union[int, float]) -> None:
+    """
+    Emit a metric using the set of registered handlers.
+
+    Parameters
+    ----------
+    name : str, required
+            Name of the metric, in dot-format.
+    value : int or float required
+            Value of the metric.
+    """
+    if MetricsMode.get() == "disable":
+        return
+    if not re.fullmatch(metric_name_pattern, name):
+        raise KeyError(
+            f"Metrics name is not in metric-name dot format, (eg. modin.dataframe.hist.duration ): {name}"
+        )
+
+    handlers = _metric_handlers.copy()
+    for fn in handlers:
+        try:
+            fn(f"modin.{name}", value)
+        except Exception:
+            clear_metric_handler(fn)
+
+
+def add_metric_handler(handler: Callable[[str, Union[int, float]], None]) -> None:
+    """
+    Add a metric handler to Modin which can collect metrics.
+
+    Parameters
+    ----------
+    handler : Callable, required
+    """
+    _metric_handlers.append(handler)
+
+
+def clear_metric_handler(handler: Callable[[str, Union[int, float]], None]) -> None:
+    """
+    Remove a metric handler from Modin.
+
+    Parameters
+    ----------
+    handler : Callable, required
+    """
+    if handler in _metric_handlers:
+        _metric_handlers.remove(handler)

modin/tests/test_metrics.py

@@ -0,0 +1,105 @@
+# Licensed to Modin Development Team under one or more contributor license agreements.
+# See the NOTICE file distributed with this work for additional information regarding
+# copyright ownership.  The Modin Development Team licenses this file to you under the
+# Apache License, Version 2.0 (the "License"); you may not use this file except in
+# compliance with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software distributed under
+# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific language
+# governing permissions and limitations under the License.
+
+from typing import Union
+
+import pytest
+
+import modin.logging
+import modin.pandas as pd
+from modin.config import MetricsMode
+from modin.logging.metrics import (
+    _metric_handlers,
+    add_metric_handler,
+    clear_metric_handler,
+    emit_metric,
+)
+
+
+class FakeTelemetryClient:
+
+    def __init__(self):
+        self._metrics = {}
+        self._metric_handler = None
+
+    def metric_handler_fail(self, name: str, value: Union[int, float]):
+        raise KeyError("Poorly implemented metric handler")
+
+    def metric_handler_pass(self, name: str, value: Union[int, float]):
+        self._metrics[name] = value
+
+
+@modin.logging.enable_logging
+def func(do_raise):
+    if do_raise:
+        raise ValueError()
+
+
+@pytest.fixture()
+def metric_client():
+    MetricsMode.enable()
+    client = FakeTelemetryClient()
+    yield client
+    clear_metric_handler(client._metric_handler)
+    MetricsMode.disable()
+
+
+def test_metrics_api_timings(metric_client):
+    assert len(_metric_handlers) == 0
+    metric_client._metric_handler = metric_client.metric_handler_pass
+    add_metric_handler(metric_client._metric_handler)
+    assert len(_metric_handlers) == 1
+    assert _metric_handlers[0] == metric_client._metric_handler
+    func(do_raise=False)
+    assert len(metric_client._metrics) == 1
+    assert metric_client._metrics["modin.pandas-api.func"] is not None
+    assert metric_client._metrics["modin.pandas-api.func"] > 0.0
+
+
+def test_df_metrics(metric_client):
+    metric_client._metric_handler = metric_client.metric_handler_pass
+    add_metric_handler(metric_client._metric_handler)
+    df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
+    df.sum()
+    assert len(metric_client._metrics) == 54
+    assert metric_client._metrics["modin.pandas-api.dataframe.sum"] is not None
+    assert metric_client._metrics["modin.pandas-api.dataframe.sum"] > 0.0
+
+
+def test_metrics_handler_fails(metric_client):
+    assert len(metric_client._metrics) == 0
+    metric_client._metric_handler = metric_client.metric_handler_fail
+    add_metric_handler(metric_client._metric_handler)
+    assert len(_metric_handlers) == 1
+    func(do_raise=False)
+    assert len(_metric_handlers) == 0
+    assert len(metric_client._metrics) == 0
+
+
+def test_emit_name_enforced():
+    MetricsMode.enable()
+    with pytest.raises(KeyError):
+        emit_metric("Not::A::Valid::Metric::Name", 1.0)
+
+
+def test_metrics_can_be_opt_out(metric_client):
+    MetricsMode.enable()
+    assert len(metric_client._metrics) == 0
+    metric_client._metric_handler = metric_client.metric_handler_pass
+    add_metric_handler(metric_client._metric_handler)
+    # If Metrics are disabled after the addition of a handler
+    # no metrics are emitted
+    MetricsMode.disable()
+    assert len(_metric_handlers) == 1
+    func(do_raise=False)
+    assert len(metric_client._metrics) == 0