Input config.prasad (#21)

* #17 - Add file exporter support Signed-off-by: Prasad Mujumdar <[email protected]> * Added test Signed-off-by: Prasad Mujumdar <[email protected]> * Updated docs Signed-off-by: Prasad Mujumdar <[email protected]> * Added workflow name to default file name Signed-off-by: Prasad Mujumdar <[email protected]> * Updated tests per review feedback Signed-off-by: Prasad Mujumdar <[email protected]> --------- Signed-off-by: Prasad Mujumdar <[email protected]>
monocle2ai · Jul 31, 2024 · 3f79f15 · 3f79f15
1 parent bf249eb
commit 3f79f15
Show file tree

Hide file tree

Showing 5 changed files with 156 additions and 3 deletions.
diff --git a/Monocle_User_Guide.md b/Monocle_User_Guide.md
@@ -37,8 +37,7 @@ from langchain_openai import OpenAI
 from langchain.prompts import PromptTemplate
 
 # Call the setup Monocle telemetry method
-setup_monocle_telemetry(workflow_name = "simple_math_app",
-        span_processors=[BatchSpanProcessor(ConsoleSpanExporter())])
+setup_monocle_telemetry(workflow_name = "simple_math_app")
 
 llm = OpenAI()
 prompt = PromptTemplate.from_template("1 + {number} = ")
@@ -51,6 +50,19 @@ chain = LLMChain(llm=llm, prompt=prompt)
 chain.invoke({"number":2}, {"callbacks":[handler]})
 
 ```
+
+### Accessing monocle trace
+By default monocle generate traces in a json file created in the local directory where the application is running. The file name by default is monocle_trace_{workflow_name}\_{trace_id}\_{timestamp}.json where the trace_id is a unique number generated by monocle for every trace. Please refere to [Trace span json](Monocle_User_Guide.md#trace-span-json). The file path and format can be changed by setting those properties as argement to ```setup_monocle_telemetry()```. For example,
+```
+setup_monocle_telemetry(workflow_name = "simple_math_app",
+    span_processors=[BatchSpanProcessor(FileSpanExporter(
+        out_path = "/tmp",
+        file_prefix = "map_app_prod_trace_",
+        time_format = "%Y-%m-%d"))
+    ])
+```
+To print the trace on the console, use ```ConsoleSpanExporter()``` instead of ```FileSpanExporter()```
+
 ### Leveraging Monocle's extensibility to handle customization 
 When the out of box features from app frameworks are not sufficent, the app developers have to add custom code. For example, if you are extending a LLM class in LlamaIndex to use a model hosted in NVIDIA Triton. This new class is not know to Monocle. You can specify this new class method part of Monocle enabling API and it will be able to trace it.
 

diff --git a/src/monocle_apptrace/exporters/file_exporter.py b/src/monocle_apptrace/exporters/file_exporter.py
@@ -0,0 +1,60 @@
+from os import linesep, path
+from io import TextIOWrapper
+from datetime import datetime
+from typing import Optional, Callable, Sequence
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+from opentelemetry.sdk.resources import SERVICE_NAME, Resource
+
+class FileSpanExporter(SpanExporter):
+    DEFAULT_FILE_PREFIX:str = "monocle_trace_"
+    DEFAULT_TIME_FORMAT:str = "%Y-%m-%d_%H.%M.%S"
+    current_trace_id: int = None
+    current_file_path: str = None
+
+    def __init__(
+        self,
+        service_name: Optional[str] = None,
+        out_path:str = ".",
+        file_prefix = DEFAULT_FILE_PREFIX,
+        time_format = DEFAULT_TIME_FORMAT,
+        formatter: Callable[
+            [ReadableSpan], str
+        ] = lambda span: span.to_json()
+        + linesep,
+    ):
+        self.out_handle:TextIOWrapper = None
+        self.formatter = formatter
+        self.service_name = service_name
+        self.output_path = out_path
+        self.file_prefix = file_prefix
+        self.time_format = time_format
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        for span in spans:
+            if span.context.trace_id != self.current_trace_id:
+                self.rotate_file(span.resource.attributes[SERVICE_NAME],
+                                span.context.trace_id)
+            self.out_handle.write(self.formatter(span))
+        self.out_handle.flush()
+        return SpanExportResult.SUCCESS
+
+    def rotate_file(self, trace_name:str, trace_id:int) -> None:
+        self.reset_handle()
+        self.current_file_path = path.join(self.output_path,
+                        self.file_prefix + trace_name + "_" + hex(trace_id) + "_"
+                        + datetime.now().strftime(self.time_format) + ".json")
+        self.out_handle = open(self.current_file_path, "w")
+        self.current_trace_id = trace_id
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        self.out_handle.flush()
+        return True
+
+    def reset_handle(self) -> None:
+        if self.out_handle != None:
+            self.out_handle.close()
+            self.out_handle = None
+
+    def shutdown(self) -> None:
+        self.reset_handle()
diff --git a/src/monocle_apptrace/instrumentor.py b/src/monocle_apptrace/instrumentor.py
@@ -12,6 +12,7 @@
 from opentelemetry import trace
 from monocle_apptrace.wrap_common import CONTEXT_PROPERTIES_KEY
 from monocle_apptrace.wrapper import INBUILT_METHODS_LIST, WrapperMethod
+from monocle_apptrace.exporters.file_exporter import FileSpanExporter
 from opentelemetry.context import get_value, attach, set_value
 
 
@@ -88,7 +89,8 @@ def _uninstrument(self, **kwargs):
 
 def setup_monocle_telemetry(
         workflow_name: str,
-        span_processors: List[SpanProcessor] = [],
+        span_processors: List[SpanProcessor] =
+                [BatchSpanProcessor(FileSpanExporter())],
         wrapper_methods: List[WrapperMethod] = []):
     resource = Resource(attributes={
         SERVICE_NAME: workflow_name

diff --git a/tests/dummy_class.py b/tests/dummy_class.py
@@ -1,5 +1,23 @@
+from opentelemetry.trace import Tracer
+from monocle_apptrace.utils import with_tracer_wrapper
 
+@with_tracer_wrapper
+def dummy_wrapper(tracer: Tracer, to_wrap, wrapped, instance, args, kwargs):
+    if callable(to_wrap.get("span_name_getter")):
+        name = to_wrap.get("span_name_getter")(instance)
+    elif hasattr(instance, "name") and instance.name:
+        name = f"{to_wrap.get('span_name')}.{instance.name.lower()}"
+    elif to_wrap.get("span_name"):
+        name = to_wrap.get("span_name")
+    else:
+        name = f"dummy.{instance.__class__.__name__}"
+    kind = to_wrap.get("kind")
+    with tracer.start_as_current_span(name) as span:
+        return_value = wrapped(*args, **kwargs)
+
+    return return_value
 
 class DummyClass:
     def dummy_method(val: int):
         print("entering dummy_method: " + str(val))
+
diff --git a/tests/file_exporter_test.py b/tests/file_exporter_test.py
@@ -0,0 +1,61 @@
+import json
+import logging
+import os
+import unittest
+from dummy_class import DummyClass, dummy_wrapper
+
+from monocle_apptrace.instrumentor import setup_monocle_telemetry
+from monocle_apptrace.wrapper import WrapperMethod
+from monocle_apptrace.exporters.file_exporter import FileSpanExporter
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+
+logger = logging.getLogger()
+logger.setLevel(logging.DEBUG)
+fileHandler = logging.FileHandler('traces.txt','w')
+formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(name)s: %(message)s')
+fileHandler.setFormatter(formatter)
+logger.addHandler(fileHandler)
+
+
+class TestHandler(unittest.TestCase):
+    SPAN_NAME="dummy.span"
+    def test_file_exporter(self):
+        app_name = "file_test"
+        file_exporter = FileSpanExporter(time_format="%Y-%m-%d")
+        span_processor = BatchSpanProcessor(file_exporter)
+        setup_monocle_telemetry(
+            workflow_name=app_name,
+            span_processors=[
+                    span_processor
+            ],
+            wrapper_methods=[
+                WrapperMethod(
+                    package="dummy_class",
+                    object="DummyClass",
+                    method="dummy_method",
+                    span_name=self.SPAN_NAME,
+                    wrapper=dummy_wrapper)
+            ])
+        dummy_class_1 = DummyClass()
+
+        dummy_class_1.dummy_method()
+
+        span_processor.force_flush()
+        span_processor.shutdown()
+        trace_file_name = file_exporter.current_file_path
+
+        try:
+            with open(trace_file_name) as f:
+                trace_data = json.load(f)
+                trace_id_from_file = trace_data["context"]["trace_id"]
+                trace_id_from_exporter = hex(file_exporter.current_trace_id)
+                assert trace_id_from_file == trace_id_from_exporter
+
+                span_name = trace_data["name"]
+                assert self.SPAN_NAME == span_name
+
+            os.remove(trace_file_name)
+        except Exception as ex:
+            print("Got error " + str(ex))
+            assert false
+