Skip to content

Commit

Permalink
prepend namespace
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfc-gh-gtokernliang
sfc-gh-gtokernliang committed Dec 21, 2024
1 parent 7dae6ff commit e1b685c
Show file tree
Hide file tree
Showing 4 changed files with 146 additions and 153 deletions.
181 changes: 57 additions & 124 deletions examples/experimental/otel_exporter.ipynb
View file
Original file line number Diff line number Diff line change
@@ -1,125 +1,14 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# !pip install opentelemetry-api\n",
"# !pip install opentelemetry-sdk"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from pathlib import Path\n",
"import sys\n",
"\n",
"# Add base dir to path to be able to access test folder.\n",
"base_dir = Path().cwd().parent.parent.resolve()\n",
"if str(base_dir) not in sys.path:\n",
" print(f\"Adding {base_dir} to sys.path\")\n",
" sys.path.append(str(base_dir))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from trulens.experimental.otel_tracing.core.instrument import instrument\n",
"\n",
"\n",
"class TestApp:\n",
" @instrument()\n",
" def respond_to_query(self, query: str) -> str:\n",
" return f\"answer: {self.nested(query)}\"\n",
"\n",
" @instrument(attributes={\"nested_attr1\": \"value1\"})\n",
" def nested(self, query: str) -> str:\n",
" return f\"nested: {self.nested2(query)}\"\n",
"\n",
" @instrument(\n",
" attributes=lambda ret, *args, **kwargs: {\n",
" \"nested2_ret\": ret,\n",
" \"nested2_args[0]\": args[0],\n",
" }\n",
" )\n",
" def nested2(self, query: str) -> str:\n",
" return f\"nested2: {query}\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"\n",
"import dotenv\n",
"from trulens.connectors.snowflake import SnowflakeConnector\n",
"from trulens.core.session import TruSession\n",
"from trulens.experimental.otel_tracing.core.init import init\n",
"\n",
"dotenv.load_dotenv()\n",
"\n",
"connection_params = {\n",
" \"account\": os.environ[\"SNOWFLAKE_ACCOUNT\"],\n",
" \"user\": os.environ[\"SNOWFLAKE_USER\"],\n",
" \"password\": os.environ[\"SNOWFLAKE_USER_PASSWORD\"],\n",
" \"database\": os.environ[\"SNOWFLAKE_DATABASE\"],\n",
" \"schema\": os.environ[\"SNOWFLAKE_SCHEMA\"],\n",
" \"warehouse\": os.environ[\"SNOWFLAKE_WAREHOUSE\"],\n",
" \"role\": os.environ[\"SNOWFLAKE_ROLE\"],\n",
"}\n",
"\n",
"connector = SnowflakeConnector(\n",
" **connection_params, database_redact_keys=True, database_args=None\n",
")\n",
"session = TruSession(connector=connector)\n",
"session.reset_database()\n",
"init(session, debug=True)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from trulens.apps.custom import TruCustomApp\n",
"\n",
"test_app = TestApp()\n",
"custom_app = TruCustomApp(test_app)\n",
"\n",
"with custom_app as recording:\n",
" test_app.respond_to_query(\"test\")"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "trulens",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3"
}
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# !pip install opentelemetry-api\n",
"# !pip install opentelemetry-sdk"
]
},
{
"cell_type": "code",
Expand All @@ -137,6 +26,26 @@
" sys.path.append(str(base_dir))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import logging\n",
"\n",
"root = logging.getLogger()\n",
"root.setLevel(logging.DEBUG)\n",
"handler = logging.StreamHandler(sys.stdout)\n",
"handler.setLevel(logging.DEBUG)\n",
"handler.addFilter(logging.Filter(\"trulens\"))\n",
"formatter = logging.Formatter(\n",
" \"%(asctime)s - %(name)s - %(levelname)s - %(message)s\"\n",
")\n",
"handler.setFormatter(formatter)\n",
"root.addHandler(handler)"
]
},
{
"cell_type": "code",
"execution_count": null,
Expand All @@ -156,13 +65,33 @@
" return f\"nested: {self.nested2(query)}\"\n",
"\n",
" @instrument(\n",
" attributes=lambda ret, *args, **kwargs: {\n",
" attributes=lambda ret, exception, *args, **kwargs: {\n",
" \"nested2_ret\": ret,\n",
" \"nested2_args[0]\": args[0],\n",
" }\n",
" )\n",
" def nested2(self, query: str) -> str:\n",
" return f\"nested2: {query}\""
" nested_result = \"\"\n",
"\n",
" try:\n",
" nested_result = self.nested3(query)\n",
" except Exception:\n",
" pass\n",
"\n",
" return f\"nested2: {nested_result}\"\n",
"\n",
" @instrument(\n",
" attributes=lambda ret, exception, *args, **kwargs: {\n",
" \"nested3_ex\": exception.args if exception else None,\n",
" \"nested3_ret\": ret,\n",
" \"selector_name\": \"special\",\n",
" \"cows\": \"moo\",\n",
" }\n",
" )\n",
" def nested3(self, query: str) -> str:\n",
" if query == \"throw\":\n",
" raise ValueError(\"nested3 exception\")\n",
" return \"nested3\""
]
},
{
Expand All @@ -178,6 +107,7 @@
"dotenv.load_dotenv()\n",
"\n",
"session = TruSession()\n",
"session.experimental_enable_feature(\"otel_tracing\")\n",
"session.reset_database()\n",
"init(session, debug=True)"
]
Expand All @@ -194,7 +124,10 @@
"custom_app = TruCustomApp(test_app)\n",
"\n",
"with custom_app as recording:\n",
" test_app.respond_to_query(\"test\")"
" test_app.respond_to_query(\"test\")\n",
"\n",
"with custom_app as recording:\n",
" test_app.respond_to_query(\"throw\")"
]
}
],
Expand Down
101 changes: 74 additions & 27 deletions src/core/trulens/experimental/otel_tracing/core/instrument.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,33 +4,55 @@

from opentelemetry import trace
from trulens.experimental.otel_tracing.core.init import TRULENS_SERVICE_NAME
from trulens.experimental.otel_tracing.core.semantic import (
TRULENS_SELECTOR_NAME,
)
from trulens.otel.semconv.trace import SpanAttributes

logger = logging.getLogger(__name__)


def instrument(
attributes: Optional[
Union[dict[str, Any], Callable[[Any, Any, Any], dict[str, Any]]]
Union[
dict[str, Any],
Callable[
[Optional[Any], Optional[Exception], Any, Any], dict[str, Any]
],
]
] = {},
):
"""
Decorator for marking functions to be instrumented in custom classes that are
wrapped by TruCustomApp, with OpenTelemetry tracing.
"""

def _validate_selector_name(final_attributes: dict[str, Any]):
if TRULENS_SELECTOR_NAME in final_attributes:
selector_name = final_attributes[TRULENS_SELECTOR_NAME]
def _validate_selector_name(attributes: dict[str, Any]) -> dict[str, Any]:
result = attributes.copy()

if (
SpanAttributes.SELECTOR_NAME_KEY in result
and SpanAttributes.SELECTOR_NAME in result
):
raise ValueError(
f"Both {SpanAttributes.SELECTOR_NAME_KEY} and {SpanAttributes.SELECTOR_NAME} cannot be set."
)

if SpanAttributes.SELECTOR_NAME in result:
# Transfer the trulens namespaced to the non-trulens namespaced key.
result[SpanAttributes.SELECTOR_NAME_KEY] = result[
SpanAttributes.SELECTOR_NAME
]
del result[SpanAttributes.SELECTOR_NAME]

if SpanAttributes.SELECTOR_NAME_KEY in result:
selector_name = result[SpanAttributes.SELECTOR_NAME_KEY]
if not isinstance(selector_name, str):
raise ValueError(
f"Selector name must be a string, not {type(selector_name)}"
)

def _validate_attributes(final_attributes: dict[str, Any]):
_validate_selector_name(final_attributes)
return result

def _validate_attributes(attributes: dict[str, Any]) -> dict[str, Any]:
return _validate_selector_name(attributes)
# TODO: validate OTEL attributes.
# TODO: validate span type attributes.

Expand All @@ -43,39 +65,64 @@ def wrapper(*args, **kwargs):
.start_as_current_span(
name=func.__name__,
)
) as parent_span:
span = trace.get_current_span()

span.set_attribute("name", func.__name__)
span.set_attribute("kind", "SPAN_KIND_TRULENS")
span.set_attribute(
"parent_span_id", parent_span.get_span_context().span_id
)
) as span:
ret = None
exception: Optional[Exception] = None

ret = func(*args, **kwargs)
try:
ret = func(*args, **kwargs)
except Exception as e:
# We want to get into the next clause to allow the users to still add attributes.
# It's on the user to deal with None as a return value.
exception = e

try:
attributes_to_add = {}

# Since we're decoratoring a method in a trulens app, the first argument is self,
# which we should ignore.
_self, *rest = args

# Set the user provider attributes.
if attributes:
if callable(attributes):
attributes_to_add = attributes(ret, *args, **kwargs)
attributes_to_add = attributes(
ret, exception, *rest, **kwargs
)
else:
attributes_to_add = attributes

logger.info(f"Attributes to add: {attributes_to_add}")

_validate_attributes(attributes_to_add)

for key, value in attributes_to_add.items():
span.set_attribute(key, value)

except Exception:
span.set_attribute("status", "STATUS_CODE_ERROR")
final_attributes = _validate_attributes(attributes_to_add)

prefix = "trulens."
if (
SpanAttributes.SPAN_TYPE in final_attributes
and final_attributes[SpanAttributes.SPAN_TYPE]
!= SpanAttributes.SpanType.UNKNOWN
):
prefix += (
final_attributes[SpanAttributes.SPAN_TYPE] + "."
)

for key, value in final_attributes.items():
span.set_attribute(prefix + key, value)

if (
key != SpanAttributes.SELECTOR_NAME_KEY
and SpanAttributes.SELECTOR_NAME_KEY
in final_attributes
):
span.set_attribute(
f"trulens.{final_attributes[SpanAttributes.SELECTOR_NAME_KEY]}.{key}",
value,
)

except Exception as e:
logger.error(f"Error setting attributes: {e}")
return None

span.set_attribute("status", "STATUS_CODE_UNSET")
return ret

return wrapper
Expand Down
1 change: 0 additions & 1 deletion src/core/trulens/experimental/otel_tracing/core/semantic.py
View file

This file was deleted.

16 changes: 15 additions & 1 deletion src/otel/semconv/trulens/otel/semconv/trace.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,21 @@ class SpanAttributes:
In some cases below, we also include span name or span name prefix.
"""

SPAN_TYPES = "trulens.span_types"
BASE = "trulens."
"""
Base prefix for the other keys.
"""

SPAN_TYPE = BASE + "span_type"

SPAN_TYPES = BASE + "span_types"

SELECTOR_NAME_KEY = "selector_name"

SELECTOR_NAME = BASE + "selector_name"
"""
User-defined selector name for the current span.
"""

class SpanType(str, Enum):
"""Span type attribute values.
Expand Down

0 comments on commit e1b685c

Please sign in to comment.