Okahu and log exporter (#60)

Signed-off-by: Prasad Mujumdar <[email protected]>

src/monocle_apptrace/exporters/monocle_exporters.py

@@ -0,0 +1,27 @@
+from typing import Dict, Any
+import os, warnings
+from importlib import import_module
+from opentelemetry.sdk.trace.export import SpanExporter, ConsoleSpanExporter
+from monocle_apptrace.exporters.file_exporter import FileSpanExporter
+
+monocle_exporters:Dict[str, Any] = {
+    "s3": {"module": "monocle_apptrace.exporters.aws.s3_exporter", "class": "S3SpanExporter"},
+    "blob": {"module":"monocle_apptrace.exporters.azure.blob_exporter", "class": "AzureBlobSpanExporter"},
+    "okahu": {"module":"monocle_apptrace.exporters.okahu.okahu_exporter", "class": "OkahuSpanExporter"},
+    "file": {"module":"monocle_apptrace.exporters.file_exporter", "class": "FileSpanExporter"}
+}
+
+def get_monocle_exporter() -> SpanExporter:
+    exporter_name = os.environ.get("MONOCLE_EXPORTER", "file")
+    try:
+        exporter_class_path  = monocle_exporters[exporter_name]
+    except Exception as ex:
+        warnings.warn(f"Unsupported Monocle span exporter setting {exporter_name}, using default FileSpanExporter.")
+        return FileSpanExporter()
+    try:
+        exporter_module = import_module(exporter_class_path.get("module"))
+        exporter_class = getattr(exporter_module, exporter_class_path.get("class"))
+        return exporter_class()
+    except Exception as ex:
+        warnings.warn(f"Unable to set Monocle span exporter to {exporter_name}, error {ex}. Using ConsoleSpanExporter")
+        return ConsoleSpanExporter()

src/monocle_apptrace/exporters/okahu/okahu_exporter.py

@@ -29,9 +29,7 @@ def __init__(
         api_key: str = os.environ.get("OKAHU_API_KEY")
         self._closed = False
         if not api_key:
-            logger.warning("OKAHU_API_KEY not set. Using ConsoleSpanExporter instead.")
-            self.exporter = ConsoleSpanExporter()
-            return
+            raise ValueError("OKAHU_API_KEY not set.")
         self.timeout = timeout or 15
         self.session = session or requests.Session()
         self.session.headers.update(

src/monocle_apptrace/instrumentor.py

@@ -1,4 +1,4 @@
-import logging
+import logging, os
 from typing import Collection, List
 from wrapt import wrap_function_wrapper
 from opentelemetry.trace import get_tracer
@@ -12,7 +12,7 @@
 from monocle_apptrace.utils import process_wrapper_method_config
 from monocle_apptrace.wrap_common import SESSION_PROPERTIES_KEY
 from monocle_apptrace.wrapper import INBUILT_METHODS_LIST, WrapperMethod
-from monocle_apptrace.exporters.file_exporter import FileSpanExporter
+from monocle_apptrace.exporters.monocle_exporters import get_monocle_exporter
 
 logger = logging.getLogger(__name__)
 
@@ -90,7 +90,7 @@ def setup_monocle_telemetry(
     resource = Resource(attributes={
         SERVICE_NAME: workflow_name
     })
-    span_processors = span_processors or [BatchSpanProcessor(FileSpanExporter())]
+    span_processors = span_processors or [BatchSpanProcessor(get_monocle_exporter())]
     trace_provider = TracerProvider(resource=resource)
     tracer_provider_default = trace.get_tracer_provider()
     provider_type = type(tracer_provider_default).__name__

src/monocle_apptrace/metamodel/spans/span_types.json

@@ -0,0 +1,16 @@
+{
+    "span_types" : [
+        {
+            "type": "inference",
+            "description": "Model inference span"
+        },
+        {
+            "type": "retrieval",
+            "description": "vector embedding retrieval"
+        },
+        {
+            "type": "workflow",
+            "description": "workflow orchetraction at top level"
+        }
+    ]
+}

tests/exporter_config_test.py

@@ -0,0 +1,27 @@
+import os
+import unittest
+
+from monocle_apptrace.exporters.monocle_exporters import get_monocle_exporter
+
+class TestHandler(unittest.TestCase):
+    def test_default_exporter(self):
+        default_exporter = get_monocle_exporter()
+        assert default_exporter.__class__.__name__ == "FileSpanExporter"
+
+    def test_fallback_exporter(self):
+        """ No Okahu API key, it should fall back to console exporter"""
+        os.environ["MONOCLE_EXPORTER"] = "okahu"
+        default_exporter = get_monocle_exporter()
+        assert default_exporter.__class__.__name__ == "ConsoleSpanExporter"
+
+    def test_set_exporter(self):
+        os.environ["MONOCLE_EXPORTER"] = "okahu"
+        os.environ["OKAHU_API_KEY"] = "foo"
+        default_exporter = get_monocle_exporter()
+        assert default_exporter.__class__.__name__ == "OkahuSpanExporter"
+
+if __name__ == "__main__":
+    handler = TestHandler()
+    handler.test_default_exporter()
+    handler.test_fallback_exporter()
+    handler.test_set_exporter()