Add sklearn instrumentation and ML model feature support

MANIFEST.in

@@ -9,3 +9,4 @@ include newrelic/packages/wrapt/LICENSE
 include newrelic/packages/wrapt/README
 include newrelic/packages/urllib3/LICENSE.txt
 include newrelic/packages/isort/LICENSE
+include newrelic/packages/opentelemetry_proto/LICENSE.txt

THIRD_PARTY_NOTICES.md

@@ -26,6 +26,15 @@ Distributed under the following license(s):
 * [The MIT License](http://opensource.org/licenses/MIT)
 
 
+## [opentelemetry-proto](https://pypi.org/project/opentelemetry-proto)
+
+Copyright (c) The OpenTelemetry Authors
+
+Distributed under the following license(s):
+
+* [The Apache License, Version 2.0 License](https://opensource.org/license/apache-2-0/)
+
+
 ## [six](https://pypi.org/project/six)
 
 Copyright (c) 2010-2013 Benjamin Peterson

newrelic/agent.py

@@ -59,6 +59,7 @@
 from newrelic.api.transaction import record_custom_metric as __record_custom_metric
 from newrelic.api.transaction import record_custom_metrics as __record_custom_metrics
 from newrelic.api.transaction import record_log_event as __record_log_event
+from newrelic.api.transaction import record_ml_event as __record_ml_event
 from newrelic.api.transaction import set_background_task as __set_background_task
 from newrelic.api.transaction import set_transaction_name as __set_transaction_name
 from newrelic.api.transaction import suppress_apdex_metric as __suppress_apdex_metric
@@ -152,6 +153,7 @@ def __asgi_application(*args, **kwargs):
 from newrelic.api.message_transaction import (
     wrap_message_transaction as __wrap_message_transaction,
 )
+from newrelic.api.ml_model import wrap_mlmodel as __wrap_mlmodel
 from newrelic.api.profile_trace import ProfileTraceWrapper as __ProfileTraceWrapper
 from newrelic.api.profile_trace import profile_trace as __profile_trace
 from newrelic.api.profile_trace import wrap_profile_trace as __wrap_profile_trace
@@ -206,11 +208,6 @@ def __asgi_application(*args, **kwargs):
 # EXPERIMENTAL - Generator traces are currently experimental and may not
 # exist in this form in future versions of the agent.
 
-
-# EXPERIMENTAL - Profile traces are currently experimental and may not
-# exist in this form in future versions of the agent.
-
-
 initialize = __initialize
 extra_settings = __wrap_api_call(__extra_settings, "extra_settings")
 global_settings = __wrap_api_call(__global_settings, "global_settings")
@@ -248,6 +245,7 @@ def __asgi_application(*args, **kwargs):
 record_custom_metrics = __wrap_api_call(__record_custom_metrics, "record_custom_metrics")
 record_custom_event = __wrap_api_call(__record_custom_event, "record_custom_event")
 record_log_event = __wrap_api_call(__record_log_event, "record_log_event")
+record_ml_event = __wrap_api_call(__record_ml_event, "record_ml_event")
 accept_distributed_trace_payload = __wrap_api_call(
     __accept_distributed_trace_payload, "accept_distributed_trace_payload"
 )
@@ -341,3 +339,4 @@ def __asgi_application(*args, **kwargs):
 wrap_out_function = __wrap_api_call(__wrap_out_function, "wrap_out_function")
 insert_html_snippet = __wrap_api_call(__insert_html_snippet, "insert_html_snippet")
 verify_body_exists = __wrap_api_call(__verify_body_exists, "verify_body_exists")
+wrap_mlmodel = __wrap_api_call(__wrap_mlmodel, "wrap_mlmodel")

newrelic/api/application.py

@@ -142,10 +142,22 @@ def record_custom_metrics(self, metrics):
         if self.active and metrics:
             self._agent.record_custom_metrics(self._name, metrics)
 
+    def record_dimensional_metric(self, name, value, tags=None):
+        if self.active:
+            self._agent.record_dimensional_metric(self._name, name, value, tags)
+
+    def record_dimensional_metrics(self, metrics):
+        if self.active and metrics:
+            self._agent.record_dimensional_metrics(self._name, metrics)
+
     def record_custom_event(self, event_type, params):
         if self.active:
             self._agent.record_custom_event(self._name, event_type, params)
 
+    def record_ml_event(self, event_type, params):
+        if self.active:
+            self._agent.record_ml_event(self._name, event_type, params)
+
     def record_transaction(self, data):
         if self.active:
             self._agent.record_transaction(self._name, data)

newrelic/api/ml_model.py

@@ -0,0 +1,35 @@
+# Copyright 2010 New Relic, Inc.
+#
+# Licensed 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.
+
+import sys
+
+from newrelic.common.object_names import callable_name
+from newrelic.hooks.mlmodel_sklearn import _nr_instrument_model
+
+
+def wrap_mlmodel(model, name=None, version=None, feature_names=None, label_names=None, metadata=None):
+    model_callable_name = callable_name(model)
+    _class = model.__class__.__name__
+    module = sys.modules[model_callable_name.split(":")[0]]
+    _nr_instrument_model(module, _class)
+    if name:
+        model._nr_wrapped_name = name
+    if version:
+        model._nr_wrapped_version = version
+    if feature_names:
+        model._nr_wrapped_feature_names = feature_names
+    if label_names:
+        model._nr_wrapped_label_names = label_names
+    if metadata:
+        model._nr_wrapped_metadata = metadata

newrelic/api/transaction.py

@@ -60,11 +60,15 @@
     DST_NONE,
     DST_TRANSACTION_TRACER,
 )
-from newrelic.core.config import CUSTOM_EVENT_RESERVOIR_SIZE, LOG_EVENT_RESERVOIR_SIZE
+from newrelic.core.config import (
+    CUSTOM_EVENT_RESERVOIR_SIZE,
+    LOG_EVENT_RESERVOIR_SIZE,
+    ML_EVENT_RESERVOIR_SIZE,
+)
 from newrelic.core.custom_event import create_custom_event
 from newrelic.core.log_event_node import LogEventNode
 from newrelic.core.stack_trace import exception_stack
-from newrelic.core.stats_engine import CustomMetrics, SampledDataSet
+from newrelic.core.stats_engine import CustomMetrics, DimensionalMetrics, SampledDataSet
 from newrelic.core.thread_utilization import utilization_tracker
 from newrelic.core.trace_cache import (
     TraceCacheActiveTraceError,
@@ -305,6 +309,7 @@ def __init__(self, application, enabled=None, source=None):
         self.synthetics_header = None
 
         self._custom_metrics = CustomMetrics()
+        self._dimensional_metrics = DimensionalMetrics()
 
         global_settings = application.global_settings
 
@@ -328,12 +333,14 @@ def __init__(self, application, enabled=None, source=None):
             self._custom_events = SampledDataSet(
                 capacity=self._settings.event_harvest_config.harvest_limits.custom_event_data
             )
+            self._ml_events = SampledDataSet(capacity=self._settings.event_harvest_config.harvest_limits.ml_event_data)
             self._log_events = SampledDataSet(
                 capacity=self._settings.event_harvest_config.harvest_limits.log_event_data
             )
         else:
             self._custom_events = SampledDataSet(capacity=CUSTOM_EVENT_RESERVOIR_SIZE)
             self._log_events = SampledDataSet(capacity=LOG_EVENT_RESERVOIR_SIZE)
+            self._ml_events = SampledDataSet(capacity=ML_EVENT_RESERVOIR_SIZE)
 
     def __del__(self):
         self._dead = True
@@ -580,10 +587,12 @@ def __exit__(self, exc, value, tb):
             errors=tuple(self._errors),
             slow_sql=tuple(self._slow_sql),
             custom_events=self._custom_events,
+            ml_events=self._ml_events,
             log_events=self._log_events,
             apdex_t=self.apdex,
             suppress_apdex=self.suppress_apdex,
             custom_metrics=self._custom_metrics,
+            dimensional_metrics=self._dimensional_metrics,
             guid=self.guid,
             cpu_time=self._cpu_user_time_value,
             suppress_transaction_trace=self.suppress_transaction_trace,
@@ -1607,6 +1616,16 @@ def record_custom_metrics(self, metrics):
         for name, value in metrics:
             self._custom_metrics.record_custom_metric(name, value)
 
+    def record_dimensional_metric(self, name, value, tags=None):
+        self._dimensional_metrics.record_dimensional_metric(name, value, tags)
+
+    def record_dimensional_metrics(self, metrics):
+        for metric in metrics:
+            name, value = metric[:2]
+            tags = metric[2] if len(metric) >= 3 else None
+
+            self._dimensional_metrics.record_dimensional_metric(name, value, tags)
+
     def record_custom_event(self, event_type, params):
         settings = self._settings
 
@@ -1620,6 +1639,19 @@ def record_custom_event(self, event_type, params):
         if event:
             self._custom_events.add(event, priority=self.priority)
 
+    def record_ml_event(self, event_type, params):
+        settings = self._settings
+
+        if not settings:
+            return
+
+        if not settings.ml_insights_events.enabled:
+            return
+
+        event = create_custom_event(event_type, params)
+        if event:
+            self._ml_events.add(event, priority=self.priority)
+
     def _intern_string(self, value):
         return self._string_cache.setdefault(value, value)
 
@@ -1913,6 +1945,44 @@ def record_custom_metrics(metrics, application=None):
         application.record_custom_metrics(metrics)
 
 
+def record_dimensional_metric(name, value, tags=None, application=None):
+    if application is None:
+        transaction = current_transaction()
+        if transaction:
+            transaction.record_dimensional_metric(name, value, tags)
+        else:
+            _logger.debug(
+                "record_dimensional_metric has been called but no "
+                "transaction was running. As a result, the following metric "
+                "has not been recorded. Name: %r Value: %r Tags: %r. To correct this "
+                "problem, supply an application object as a parameter to this "
+                "record_dimensional_metrics call.",
+                name,
+                value,
+                tags,
+            )
+    elif application.enabled:
+        application.record_dimensional_metric(name, value, tags)
+
+
+def record_dimensional_metrics(metrics, application=None):
+    if application is None:
+        transaction = current_transaction()
+        if transaction:
+            transaction.record_dimensional_metrics(metrics)
+        else:
+            _logger.debug(
+                "record_dimensional_metrics has been called but no "
+                "transaction was running. As a result, the following metrics "
+                "have not been recorded: %r. To correct this problem, "
+                "supply an application object as a parameter to this "
+                "record_dimensional_metric call.",
+                list(metrics),
+            )
+    elif application.enabled:
+        application.record_dimensional_metrics(metrics)
+
+
 def record_custom_event(event_type, params, application=None):
     """Record a custom event.
 
@@ -1941,6 +2011,34 @@ def record_custom_event(event_type, params, application=None):
         application.record_custom_event(event_type, params)
 
 
+def record_ml_event(event_type, params, application=None):
+    """Record a machine learning custom event.
+
+    Args:
+        event_type (str): The type (name) of the ml event.
+        params (dict): Attributes to add to the event.
+        application (newrelic.api.Application): Application instance.
+
+    """
+
+    if application is None:
+        transaction = current_transaction()
+        if transaction:
+            transaction.record_ml_event(event_type, params)
+        else:
+            _logger.debug(
+                "record_ml_event has been called but no "
+                "transaction was running. As a result, the following event "
+                "has not been recorded. event_type: %r params: %r. To correct "
+                "this problem, supply an application object as a parameter to "
+                "this record_ml_event call.",
+                event_type,
+                params,
+            )
+    elif application.enabled:
+        application.record_ml_event(event_type, params)
+
+
 def record_log_event(message, level=None, timestamp=None, application=None, priority=None):
     """Record a log event.