Merge pull request #29 from monocle2ai/kshitiz/add_aws_lambda_infra

Added AWS lambda infra attribute in traces

src/monocle_apptrace/constants.py

@@ -1,16 +1,22 @@
-# Azure environment constants 
+# Azure environment constants
 AZURE_ML_ENDPOINT_ENV_NAME = "AZUREML_ENTRY_SCRIPT"
 AZURE_FUNCTION_WORKER_ENV_NAME = "FUNCTIONS_WORKER_RUNTIME"
 AZURE_APP_SERVICE_ENV_NAME = "WEBSITE_SITE_NAME"
+AWS_LAMBDA_ENV_NAME = "AWS_LAMBDA_RUNTIME_API"
 
 # Azure naming reference can be found here
 # https://learn.microsoft.com/en-us/azure/cloud-adoption-framework/ready/azure-best-practices/resource-abbreviations
 AZURE_FUNCTION_NAME = "azure.func"
 AZURE_APP_SERVICE_NAME = "azure.asp"
 AZURE_ML_SERVICE_NAME = "azure.mlw"
+AWS_LAMBDA_SERVICE_NAME = "aws.lambda"
 
 azure_service_map = {
     AZURE_ML_ENDPOINT_ENV_NAME: AZURE_ML_SERVICE_NAME,
     AZURE_APP_SERVICE_ENV_NAME: AZURE_APP_SERVICE_NAME,
     AZURE_FUNCTION_WORKER_ENV_NAME: AZURE_FUNCTION_NAME
 }
+
+aws_service_map = {
+    AWS_LAMBDA_ENV_NAME: AWS_LAMBDA_SERVICE_NAME
+}

src/monocle_apptrace/utils.py

@@ -3,7 +3,7 @@
 from importlib import import_module
 import os
 from opentelemetry.trace import Span
-from monocle_apptrace.constants import AZURE_APP_SERVICE_ENV_NAME, AZURE_APP_SERVICE_NAME, AZURE_FUNCTION_NAME, AZURE_FUNCTION_WORKER_ENV_NAME, AZURE_ML_ENDPOINT_ENV_NAME, AZURE_ML_SERVICE_NAME
+from monocle_apptrace.constants import azure_service_map, aws_service_map
 
 def set_span_attribute(span, name, value):
     if value is not None:
@@ -65,9 +65,9 @@ def get_wrapper_method(package_name: str, method_name: str):
     return getattr(wrapper_module, method_name)
 
 def update_span_with_infra_name(span: Span, span_key: str):
-    if AZURE_FUNCTION_WORKER_ENV_NAME in os.environ:
-        span.set_attribute(span_key, AZURE_FUNCTION_NAME)
-    elif AZURE_APP_SERVICE_ENV_NAME in os.environ:
-        span.set_attribute(span_key, AZURE_APP_SERVICE_NAME)
-    elif AZURE_ML_ENDPOINT_ENV_NAME in os.environ:
-        span.set_attribute(span_key, AZURE_ML_SERVICE_NAME)
+    for key,val  in azure_service_map.items():
+        if key in os.environ:
+            span.set_attribute(span_key, val)
+    for key,val  in aws_service_map.items():
+        if key in os.environ:
+            span.set_attribute(span_key, val)

tests/file_exporter_test.py

@@ -57,5 +57,5 @@ def test_file_exporter(self):
             os.remove(trace_file_name)
         except Exception as ex:
             print("Got error " + str(ex))
-            assert false
+            assert False
 

tests/langchain_test.py

@@ -18,7 +18,16 @@
 from langchain_community.vectorstores import faiss
 from langchain_core.messages.ai import AIMessage
 from langchain_core.runnables import RunnablePassthrough
-from monocle_apptrace.constants import AZURE_APP_SERVICE_ENV_NAME, AZURE_APP_SERVICE_NAME, AZURE_FUNCTION_NAME, AZURE_FUNCTION_WORKER_ENV_NAME, AZURE_ML_ENDPOINT_ENV_NAME, AZURE_ML_SERVICE_NAME
+from monocle_apptrace.constants import (
+    AZURE_APP_SERVICE_ENV_NAME,
+    AZURE_APP_SERVICE_NAME,
+    AZURE_FUNCTION_NAME,
+    AZURE_FUNCTION_WORKER_ENV_NAME,
+    AZURE_ML_ENDPOINT_ENV_NAME,
+    AZURE_ML_SERVICE_NAME,
+    AWS_LAMBDA_ENV_NAME,
+    AWS_LAMBDA_SERVICE_NAME
+)
 from monocle_apptrace.instrumentor import (
     MonocleInstrumentor,
     set_context_properties,
@@ -112,6 +121,7 @@ def tearDown(self) -> None:
        ("1", AZURE_ML_ENDPOINT_ENV_NAME, AZURE_ML_SERVICE_NAME),
        ("2", AZURE_FUNCTION_WORKER_ENV_NAME, AZURE_FUNCTION_NAME),
        ("3", AZURE_APP_SERVICE_ENV_NAME, AZURE_APP_SERVICE_NAME),
+       ("4", AWS_LAMBDA_ENV_NAME, AWS_LAMBDA_SERVICE_NAME),
     ])
     @patch.object(requests.Session, 'post')
     def test_llm_chain(self, test_name, test_input_infra, test_output_infra, mock_post):