Merge branch 'add-ml-framework-metrics' of github.com:newrelic/newrel…

…ic-python-agent into add-ml-framework-metrics

newrelic/core/otlp_utils.py

@@ -21,6 +21,7 @@
 
 import logging
 
+from newrelic.api.time_trace import get_service_linking_metadata
 from newrelic.common.encoding_utils import json_encode
 from newrelic.core.config import global_settings
 from newrelic.core.stats_engine import CountStats, TimeStats
@@ -124,8 +125,11 @@ def create_key_values_from_iterable(iterable):
     )
 
 
-def create_resource(attributes=None):
+def create_resource(attributes=None, attach_apm_entity=True):
     attributes = attributes or {"instrumentation.provider": "newrelic-opentelemetry-python-ml"}
+    if attach_apm_entity:
+        metadata = get_service_linking_metadata()
+        attributes.update(metadata)
     return Resource(attributes=create_key_values_from_iterable(attributes))
 
 
@@ -203,7 +207,7 @@ def stats_to_otlp_metrics(metric_data, start_time, end_time):
 
 
 def encode_metric_data(metric_data, start_time, end_time, resource=None, scope=None):
-    resource = resource or create_resource()
+    resource = resource or create_resource(attach_apm_entity=False)
     return MetricsData(
         resource_metrics=[
             ResourceMetrics(
@@ -220,24 +224,45 @@ def encode_metric_data(metric_data, start_time, end_time, resource=None, scope=N
 
 
 def encode_ml_event_data(custom_event_data, agent_run_id):
-    resource = create_resource()
-    ml_events = []
+    # An InferenceEvent is attached to a separate ML Model entity instead
+    # of the APM entity.
+    ml_inference_events = []
+    ml_apm_events = []
     for event in custom_event_data:
         event_info, event_attrs = event
+        event_type = event_info["type"]
         event_attrs.update(
             {
                 "real_agent_id": agent_run_id,
                 "event.domain": "newrelic.ml_events",
-                "event.name": event_info["type"],
+                "event.name": event_type,
             }
         )
         ml_attrs = create_key_values_from_iterable(event_attrs)
         unix_nano_timestamp = event_info["timestamp"] * 1e6
-        ml_events.append(
-            {
-                "time_unix_nano": int(unix_nano_timestamp),
-                "attributes": ml_attrs,
-            }
+        if event_type == "InferenceEvent":
+            ml_inference_events.append(
+                {
+                    "time_unix_nano": int(unix_nano_timestamp),
+                    "attributes": ml_attrs,
+                }
+            )
+        else:
+            ml_apm_events.append(
+                {
+                    "time_unix_nano": int(unix_nano_timestamp),
+                    "attributes": ml_attrs,
+                }
+            )
+
+    resource_logs = []
+    if ml_inference_events:
+        inference_resource = create_resource(attach_apm_entity=False)
+        resource_logs.append(
+            ResourceLogs(resource=inference_resource, scope_logs=[ScopeLogs(log_records=ml_inference_events)])
         )
+    if ml_apm_events:
+        apm_resource = create_resource()
+        resource_logs.append(ResourceLogs(resource=apm_resource, scope_logs=[ScopeLogs(log_records=ml_apm_events)]))
 
-    return LogsData(resource_logs=[ResourceLogs(resource=resource, scope_logs=[ScopeLogs(log_records=ml_events)])])
+    return LogsData(resource_logs=resource_logs)

newrelic/hooks/mlmodel_openai.py

@@ -15,12 +15,13 @@
 import uuid
 
 import openai
-from newrelic.common.package_version_utils import get_package_version
+
 from newrelic.api.function_trace import FunctionTrace
 from newrelic.api.time_trace import get_trace_linking_metadata
 from newrelic.api.transaction import current_transaction
 from newrelic.common.object_names import callable_name
 from newrelic.common.object_wrapper import wrap_function_wrapper
+from newrelic.common.package_version_utils import get_package_version
 from newrelic.core.config import global_settings
 
 

tests/agent_features/test_ml_events.py

@@ -58,23 +58,94 @@ def core_app(collector_agent_registration):
 
 
 @validate_ml_event_payload(
-    [{"foo": "bar", "real_agent_id": "1234567", "event.domain": "newrelic.ml_events", "event.name": "InferenceEvent"}]
+    {
+        "apm": [
+            {
+                "foo": "bar",
+                "real_agent_id": "1234567",
+                "event.domain": "newrelic.ml_events",
+                "event.name": "MyCustomEvent",
+            }
+        ]
+    }
 )
 @reset_core_stats_engine()
-def test_ml_event_payload_inside_transaction(core_app):
+def test_ml_event_payload_noninference_event_inside_transaction(core_app):
+    @background_task(name="test_ml_event_payload_inside_transaction")
+    def _test():
+        record_ml_event("MyCustomEvent", {"foo": "bar"})
+
+    _test()
+    core_app.harvest()
+
+
+@validate_ml_event_payload(
+    {
+        "inference": [
+            {
+                "foo": "bar",
+                "real_agent_id": "1234567",
+                "event.domain": "newrelic.ml_events",
+                "event.name": "InferenceEvent",
+            }
+        ]
+    }
+)
+@reset_core_stats_engine()
+def test_ml_event_payload_inference_event_inside_transaction(core_app):
+    @background_task(name="test_ml_event_payload_inside_transaction")
+    def _test():
+        record_ml_event("InferenceEvent", {"foo": "bar"})
+
+    _test()
+    core_app.harvest()
+
+
+@validate_ml_event_payload(
+    {
+        "apm": [
+            {
+                "foo": "bar",
+                "real_agent_id": "1234567",
+                "event.domain": "newrelic.ml_events",
+                "event.name": "MyCustomEvent",
+            }
+        ],
+        "inference": [
+            {
+                "foo": "bar",
+                "real_agent_id": "1234567",
+                "event.domain": "newrelic.ml_events",
+                "event.name": "InferenceEvent",
+            }
+        ],
+    }
+)
+@reset_core_stats_engine()
+def test_ml_event_payload_both_events_inside_transaction(core_app):
     @background_task(name="test_ml_event_payload_inside_transaction")
     def _test():
         record_ml_event("InferenceEvent", {"foo": "bar"})
+        record_ml_event("MyCustomEvent", {"foo": "bar"})
 
     _test()
     core_app.harvest()
 
 
 @validate_ml_event_payload(
-    [{"foo": "bar", "real_agent_id": "1234567", "event.domain": "newrelic.ml_events", "event.name": "InferenceEvent"}]
+    {
+        "inference": [
+            {
+                "foo": "bar",
+                "real_agent_id": "1234567",
+                "event.domain": "newrelic.ml_events",
+                "event.name": "InferenceEvent",
+            }
+        ]
+    }
 )
 @reset_core_stats_engine()
-def test_ml_event_payload_outside_transaction(core_app):
+def test_ml_event_payload_inference_event_outside_transaction(core_app):
     def _test():
         app = application()
         record_ml_event("InferenceEvent", {"foo": "bar"}, application=app)
@@ -83,6 +154,59 @@ def _test():
     core_app.harvest()
 
 
+@validate_ml_event_payload(
+    {
+        "apm": [
+            {
+                "foo": "bar",
+                "real_agent_id": "1234567",
+                "event.domain": "newrelic.ml_events",
+                "event.name": "MyCustomEvent",
+            }
+        ],
+        "inference": [
+            {
+                "foo": "bar",
+                "real_agent_id": "1234567",
+                "event.domain": "newrelic.ml_events",
+                "event.name": "InferenceEvent",
+            }
+        ],
+    }
+)
+@reset_core_stats_engine()
+def test_ml_event_payload_both_events_outside_transaction(core_app):
+    def _test():
+        app = application()
+        record_ml_event("InferenceEvent", {"foo": "bar"}, application=app)
+        record_ml_event("MyCustomEvent", {"foo": "bar"}, application=app)
+
+    _test()
+    core_app.harvest()
+
+
+@validate_ml_event_payload(
+    {
+        "apm": [
+            {
+                "foo": "bar",
+                "real_agent_id": "1234567",
+                "event.domain": "newrelic.ml_events",
+                "event.name": "MyCustomEvent",
+            }
+        ]
+    }
+)
+@reset_core_stats_engine()
+def test_ml_event_payload_noninference_event_outside_transaction(core_app):
+    def _test():
+        app = application()
+        record_ml_event("MyCustomEvent", {"foo": "bar"}, application=app)
+
+    _test()
+    core_app.harvest()
+
+
 @pytest.mark.parametrize(
     "params,expected",
     [
@@ -151,6 +275,7 @@ def test_record_ml_event_outside_transaction_params_not_a_dict():
 
 # Tests for ML Events configuration settings
 
+
 @override_application_settings({"ml_insights_events.enabled": False})
 @reset_core_stats_engine()
 @validate_ml_event_count(count=0)

tests/mlmodel_openai/test_chat_completion.py

@@ -19,7 +19,9 @@
 )
 from testing_support.validators.validate_ml_event_count import validate_ml_event_count
 from testing_support.validators.validate_ml_events import validate_ml_events
-from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
+from testing_support.validators.validate_transaction_metrics import (
+    validate_transaction_metrics,
+)
 
 from newrelic.api.background_task import background_task
 from newrelic.api.transaction import add_custom_attribute
@@ -127,8 +129,10 @@
 # One summary event, one system message, one user message, and one response message from the assistant
 @validate_ml_event_count(count=4)
 @validate_transaction_metrics(
-    name='test_chat_completion:test_openai_chat_completion_sync_in_txn_with_convo_id',
-    custom_metrics=[('Python/ML/OpenAI/%s' % openai.__version__, 1), ],
+    name="test_chat_completion:test_openai_chat_completion_sync_in_txn_with_convo_id",
+    custom_metrics=[
+        ("Python/ML/OpenAI/%s" % openai.__version__, 1),
+    ],
     background_task=True,
 )
 @background_task()
@@ -256,8 +260,10 @@ def test_openai_chat_completion_sync_outside_txn():
 @reset_core_stats_engine()
 @validate_ml_event_count(count=0)
 @validate_transaction_metrics(
-    name='test_chat_completion:test_openai_chat_completion_sync_ml_insights_disabled',
-    custom_metrics=[('Python/ML/OpenAI/%s' % openai.__version__, 1), ],
+    name="test_chat_completion:test_openai_chat_completion_sync_ml_insights_disabled",
+    custom_metrics=[
+        ("Python/ML/OpenAI/%s" % openai.__version__, 1),
+    ],
     background_task=True,
 )
 @background_task()
@@ -286,8 +292,10 @@ def test_openai_chat_completion_async_conversation_id_unset(loop, set_trace_info
 @validate_ml_events(chat_completion_recorded_events)
 @validate_ml_event_count(count=4)
 @validate_transaction_metrics(
-    name='test_chat_completion:test_openai_chat_completion_async_conversation_id_set',
-    custom_metrics=[('Python/ML/OpenAI/%s' % openai.__version__, 1), ],
+    name="test_chat_completion:test_openai_chat_completion_async_conversation_id_set",
+    custom_metrics=[
+        ("Python/ML/OpenAI/%s" % openai.__version__, 1),
+    ],
     background_task=True,
 )
 @background_task()
@@ -316,8 +324,10 @@ def test_openai_chat_completion_async_outside_transaction(loop):
 @reset_core_stats_engine()
 @validate_ml_event_count(count=0)
 @validate_transaction_metrics(
-    name='test_chat_completion:test_openai_chat_completion_async_disabled_ml_settings',
-    custom_metrics=[('Python/ML/OpenAI/%s' % openai.__version__, 1), ],
+    name="test_chat_completion:test_openai_chat_completion_async_disabled_ml_settings",
+    custom_metrics=[
+        ("Python/ML/OpenAI/%s" % openai.__version__, 1),
+    ],
     background_task=True,
 )
 @background_task()

tests/mlmodel_openai/test_embeddings.py

@@ -19,7 +19,9 @@
 )
 from testing_support.validators.validate_ml_event_count import validate_ml_event_count
 from testing_support.validators.validate_ml_events import validate_ml_events
-from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
+from testing_support.validators.validate_transaction_metrics import (
+    validate_transaction_metrics,
+)
 
 from newrelic.api.background_task import background_task
 
@@ -62,8 +64,10 @@
 @validate_ml_events(embedding_recorded_events)
 @validate_ml_event_count(count=1)
 @validate_transaction_metrics(
-    name='test_embeddings:test_openai_embedding_sync',
-    custom_metrics=[('Python/ML/OpenAI/%s' % openai.__version__, 1), ],
+    name="test_embeddings:test_openai_embedding_sync",
+    custom_metrics=[
+        ("Python/ML/OpenAI/%s" % openai.__version__, 1),
+    ],
     background_task=True,
 )
 @background_task()
@@ -82,8 +86,10 @@ def test_openai_embedding_sync_outside_txn():
 @reset_core_stats_engine()
 @validate_ml_event_count(count=0)
 @validate_transaction_metrics(
-    name='test_embeddings:test_openai_chat_completion_sync_disabled_settings',
-    custom_metrics=[('Python/ML/OpenAI/%s' % openai.__version__, 1), ],
+    name="test_embeddings:test_openai_chat_completion_sync_disabled_settings",
+    custom_metrics=[
+        ("Python/ML/OpenAI/%s" % openai.__version__, 1),
+    ],
     background_task=True,
 )
 @background_task()
@@ -96,8 +102,10 @@ def test_openai_chat_completion_sync_disabled_settings(set_trace_info):
 @validate_ml_events(embedding_recorded_events)
 @validate_ml_event_count(count=1)
 @validate_transaction_metrics(
-    name='test_embeddings:test_openai_embedding_async',
-    custom_metrics=[('Python/ML/OpenAI/%s' % openai.__version__, 1), ],
+    name="test_embeddings:test_openai_embedding_async",
+    custom_metrics=[
+        ("Python/ML/OpenAI/%s" % openai.__version__, 1),
+    ],
     background_task=True,
 )
 @background_task()
@@ -121,8 +129,10 @@ def test_openai_embedding_async_outside_transaction(loop):
 @reset_core_stats_engine()
 @validate_ml_event_count(count=0)
 @validate_transaction_metrics(
-    name='test_embeddings:test_openai_embedding_async_disabled_ml_insights_events',
-    custom_metrics=[('Python/ML/OpenAI/%s' % openai.__version__, 1), ],
+    name="test_embeddings:test_openai_embedding_async_disabled_ml_insights_events",
+    custom_metrics=[
+        ("Python/ML/OpenAI/%s" % openai.__version__, 1),
+    ],
     background_task=True,
 )
 @background_task()