diff --git a/src/monocle_apptrace/exporters/monocle_exporters.py b/src/monocle_apptrace/exporters/monocle_exporters.py new file mode 100644 index 0000000..c3a87c3 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/src/monocle_apptrace/exporters/okahu/okahu_exporter.py b/src/monocle_apptrace/exporters/okahu/okahu_exporter.py index 4e002e4..7be6b58 100644 --- a/src/monocle_apptrace/exporters/okahu/okahu_exporter.py +++ b/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( diff --git a/src/monocle_apptrace/instrumentor.py b/src/monocle_apptrace/instrumentor.py index 8782133..503dcb5 100644 --- a/src/monocle_apptrace/instrumentor.py +++ b/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__ diff --git a/src/monocle_apptrace/metamodel/spans/span_types.json b/src/monocle_apptrace/metamodel/spans/span_types.json new file mode 100644 index 0000000..cefb4e4 --- /dev/null +++ b/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" + } + ] +} \ No newline at end of file diff --git a/tests/exporter_config_test.py b/tests/exporter_config_test.py new file mode 100644 index 0000000..dad945f --- /dev/null +++ b/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() \ No newline at end of file