참고소스 수정본
This commit is contained in:
LASTA_DEV01\lasta
@@ -0,0 +1,3 @@
|
||||
from guardrails.integrations.databricks.ml_flow_instrumentor import MlFlowInstrumentor
|
||||
|
||||
__all__ = ["MlFlowInstrumentor"]
|
||||
@@ -0,0 +1,504 @@
|
||||
from functools import wraps
|
||||
import inspect
|
||||
import sys
|
||||
from typing import (
|
||||
Any,
|
||||
AsyncIterator,
|
||||
Awaitable,
|
||||
Callable,
|
||||
Coroutine,
|
||||
Iterator,
|
||||
Union,
|
||||
)
|
||||
|
||||
from guardrails import Guard, AsyncGuard, settings
|
||||
from guardrails_ai.types import ValidationResult
|
||||
from guardrails.run import Runner, StreamRunner, AsyncRunner, AsyncStreamRunner
|
||||
from guardrails.validator_base import Validator
|
||||
from guardrails.version import GUARDRAILS_VERSION
|
||||
from guardrails.telemetry.guard_tracing import (
|
||||
add_guard_attributes,
|
||||
trace_stream_guard,
|
||||
trace_async_stream_guard,
|
||||
)
|
||||
from guardrails.telemetry.runner_tracing import add_step_attributes, add_call_attributes
|
||||
from guardrails.telemetry.validator_tracing import add_validator_attributes
|
||||
from guardrails.classes.generic.stack import Stack
|
||||
from guardrails.classes.llm.llm_response import LLMResponse
|
||||
from guardrails.classes.history.iteration import Iteration
|
||||
from guardrails.classes.output_type import OT
|
||||
from guardrails.classes.validation_outcome import ValidationOutcome
|
||||
from guardrails.utils.safe_get import safe_get
|
||||
|
||||
try:
|
||||
import mlflow
|
||||
import mlflow.tracing
|
||||
import mlflow.tracing.provider
|
||||
from mlflow.entities.span_status import SpanStatusCode
|
||||
except ImportError:
|
||||
raise ImportError("Please install mlflow to use this instrumentor")
|
||||
|
||||
|
||||
if sys.version_info.minor < 10:
|
||||
from guardrails.utils.polyfills import anext
|
||||
|
||||
|
||||
# TODO: Abstract these methods and common logic into a base class
|
||||
# that can be extended by other instrumentors
|
||||
class MlFlowInstrumentor:
|
||||
"""Instruments Guardrails to send traces to MLFlow."""
|
||||
|
||||
def __init__(self, experiment_name: str):
|
||||
self.experiment_name = experiment_name
|
||||
# Disable legacy OTEL tracing to avoid duplicate spans
|
||||
settings.disable_tracing = True
|
||||
|
||||
def instrument(self):
|
||||
mlflow.tracing.enable()
|
||||
mlflow.set_experiment(self.experiment_name)
|
||||
|
||||
wrapped_guard_execute = self._instrument_guard(Guard._execute)
|
||||
setattr(Guard, "_execute", wrapped_guard_execute)
|
||||
|
||||
wrapped_async_guard_execute = self._instrument_async_guard(AsyncGuard._execute)
|
||||
setattr(AsyncGuard, "_execute", wrapped_async_guard_execute)
|
||||
|
||||
wrapped_runner_step = self._instrument_runner_step(Runner.step)
|
||||
setattr(Runner, "step", wrapped_runner_step)
|
||||
|
||||
wrapped_stream_runner_step = self._instrument_stream_runner_step(
|
||||
StreamRunner.step
|
||||
)
|
||||
setattr(StreamRunner, "step", wrapped_stream_runner_step)
|
||||
|
||||
wrapped_async_runner_step = self._instrument_async_runner_step(
|
||||
AsyncRunner.async_step
|
||||
)
|
||||
setattr(AsyncRunner, "async_step", wrapped_async_runner_step)
|
||||
|
||||
wrapped_async_stream_runner_step = self._instrument_async_stream_runner_step(
|
||||
AsyncStreamRunner.async_step # type: ignore
|
||||
)
|
||||
setattr(AsyncStreamRunner, "async_step", wrapped_async_stream_runner_step)
|
||||
|
||||
wrapped_runner_call = self._instrument_runner_call(Runner.call)
|
||||
setattr(Runner, "call", wrapped_runner_call)
|
||||
|
||||
wrapped_async_runner_call = self._instrument_async_runner_call(
|
||||
AsyncRunner.async_call
|
||||
)
|
||||
setattr(AsyncRunner, "async_call", wrapped_async_runner_call)
|
||||
|
||||
import guardrails
|
||||
|
||||
validators = guardrails.hub.__dir__() # type: ignore
|
||||
|
||||
for validator_name in validators:
|
||||
export = getattr(guardrails.hub, validator_name) # type: ignore
|
||||
if isinstance(export, type) and issubclass(export, Validator):
|
||||
wrapped_validator_validate = self._instrument_validator_validate(
|
||||
export.validate
|
||||
)
|
||||
setattr(export, "validate", wrapped_validator_validate)
|
||||
|
||||
wrapped_validator_async_validate = (
|
||||
self._instrument_validator_async_validate(export.async_validate)
|
||||
)
|
||||
setattr(export, "async_validate", wrapped_validator_async_validate)
|
||||
|
||||
setattr(guardrails.hub, validator_name, export) # type: ignore
|
||||
|
||||
def _instrument_guard(
|
||||
self,
|
||||
guard_execute: Callable[
|
||||
..., Union[ValidationOutcome[OT], Iterator[ValidationOutcome[OT]]]
|
||||
],
|
||||
):
|
||||
@wraps(guard_execute)
|
||||
def _guard_execute_wrapper(
|
||||
*args, **kwargs
|
||||
) -> Union[ValidationOutcome[OT], Iterator[ValidationOutcome[OT]]]:
|
||||
with mlflow.start_span(
|
||||
name="guardrails/guard",
|
||||
span_type="guard",
|
||||
attributes={
|
||||
"guardrails.version": GUARDRAILS_VERSION,
|
||||
"type": "guardrails/guard",
|
||||
},
|
||||
) as guard_span:
|
||||
guard_self = args[0]
|
||||
# Use default bc we don't know if the Guard instance exists yet or not
|
||||
history = Stack(max_length=10)
|
||||
|
||||
if guard_self is not None and isinstance(guard_self, Guard):
|
||||
guard_span.set_attribute("guard.name", guard_self.name)
|
||||
history = guard_self.history
|
||||
|
||||
try:
|
||||
result = guard_execute(*args, **kwargs)
|
||||
if isinstance(result, Iterator) and not isinstance(
|
||||
result, ValidationOutcome
|
||||
):
|
||||
return trace_stream_guard(guard_span, result, history) # type: ignore
|
||||
add_guard_attributes(guard_span, history, result) # type: ignore
|
||||
return result
|
||||
except Exception as e:
|
||||
guard_span.set_status(status=SpanStatusCode.ERROR)
|
||||
raise e
|
||||
|
||||
return _guard_execute_wrapper
|
||||
|
||||
def _instrument_async_guard(
|
||||
self,
|
||||
guard_execute: Callable[
|
||||
...,
|
||||
Coroutine[
|
||||
Any,
|
||||
Any,
|
||||
Union[
|
||||
ValidationOutcome[OT],
|
||||
Awaitable[ValidationOutcome[OT]],
|
||||
AsyncIterator[ValidationOutcome[OT]],
|
||||
],
|
||||
],
|
||||
],
|
||||
):
|
||||
@wraps(guard_execute)
|
||||
async def _async_guard_execute_wrapper(
|
||||
*args, **kwargs
|
||||
) -> Union[
|
||||
ValidationOutcome[OT],
|
||||
Awaitable[ValidationOutcome[OT]],
|
||||
AsyncIterator[ValidationOutcome[OT]],
|
||||
]:
|
||||
with mlflow.start_span(
|
||||
name="guardrails/guard",
|
||||
span_type="guard",
|
||||
attributes={
|
||||
"guardrails.version": GUARDRAILS_VERSION,
|
||||
"type": "guardrails/guard",
|
||||
"async": True,
|
||||
},
|
||||
) as guard_span:
|
||||
guard_self = args[0]
|
||||
# Use default bc we don't know if the Guard instance exists yet or not
|
||||
history = Stack(max_length=10)
|
||||
|
||||
if guard_self is not None and isinstance(guard_self, Guard):
|
||||
guard_span.set_attribute("guard.name", guard_self.name)
|
||||
history = guard_self.history
|
||||
|
||||
try:
|
||||
result = await guard_execute(*args, **kwargs)
|
||||
if isinstance(result, AsyncIterator):
|
||||
return trace_async_stream_guard(guard_span, result, history) # type: ignore
|
||||
res = result
|
||||
if inspect.isawaitable(result):
|
||||
res = await result
|
||||
add_guard_attributes(guard_span, history, res) # type: ignore
|
||||
return res
|
||||
except Exception as e:
|
||||
guard_span.set_status(status=SpanStatusCode.ERROR)
|
||||
raise e
|
||||
|
||||
return _async_guard_execute_wrapper
|
||||
|
||||
def _instrument_runner_step(self, runner_step: Callable[..., Iteration]):
|
||||
@wraps(runner_step)
|
||||
def trace_step_wrapper(*args, **kwargs) -> Iteration:
|
||||
with mlflow.start_span(
|
||||
name="guardrails/guard/step",
|
||||
span_type="step",
|
||||
attributes={
|
||||
"guardrails.version": GUARDRAILS_VERSION,
|
||||
"type": "guardrails/guard/step",
|
||||
},
|
||||
) as step_span:
|
||||
try:
|
||||
response = runner_step(*args, **kwargs)
|
||||
add_step_attributes(step_span, response, *args, **kwargs) # type: ignore
|
||||
return response
|
||||
except Exception as e:
|
||||
step_span.set_status(status=SpanStatusCode.ERROR)
|
||||
add_step_attributes(step_span, None, *args, **kwargs) # type: ignore
|
||||
raise e
|
||||
|
||||
return trace_step_wrapper
|
||||
|
||||
def _instrument_stream_runner_step(
|
||||
self, runner_step: Callable[..., Iterator[ValidationOutcome[OT]]]
|
||||
):
|
||||
@wraps(runner_step)
|
||||
def trace_stream_step_wrapper(
|
||||
*args, **kwargs
|
||||
) -> Iterator[ValidationOutcome[OT]]:
|
||||
with mlflow.start_span(
|
||||
name="guardrails/guard/step",
|
||||
span_type="step",
|
||||
attributes={
|
||||
"guardrails.version": GUARDRAILS_VERSION,
|
||||
"type": "guardrails/guard/step",
|
||||
"stream": True,
|
||||
},
|
||||
) as step_span:
|
||||
exception = None
|
||||
try:
|
||||
gen = runner_step(*args, **kwargs)
|
||||
next_exists = True
|
||||
while next_exists:
|
||||
try:
|
||||
res = next(gen)
|
||||
yield res
|
||||
except StopIteration:
|
||||
next_exists = False
|
||||
except Exception as e:
|
||||
step_span.set_status(status=SpanStatusCode.ERROR)
|
||||
exception = e
|
||||
finally:
|
||||
call = safe_get(args, 8, kwargs.get("call_log", None))
|
||||
iteration = call.iterations.last if call else None
|
||||
add_step_attributes(step_span, iteration, *args, **kwargs) # type: ignore
|
||||
if exception:
|
||||
raise exception
|
||||
|
||||
return trace_stream_step_wrapper
|
||||
|
||||
def _instrument_async_runner_step(
|
||||
self, runner_step: Callable[..., Awaitable[Iteration]]
|
||||
):
|
||||
@wraps(runner_step)
|
||||
async def trace_async_step_wrapper(*args, **kwargs) -> Iteration:
|
||||
with mlflow.start_span(
|
||||
name="guardrails/guard/step",
|
||||
span_type="step",
|
||||
attributes={
|
||||
"guardrails.version": GUARDRAILS_VERSION,
|
||||
"type": "guardrails/guard/step",
|
||||
"async": True,
|
||||
},
|
||||
) as step_span:
|
||||
try:
|
||||
response = await runner_step(*args, **kwargs)
|
||||
add_step_attributes(step_span, response, *args, **kwargs) # type: ignore
|
||||
return response
|
||||
except Exception as e:
|
||||
step_span.set_status(status=SpanStatusCode.ERROR)
|
||||
add_step_attributes(step_span, None, *args, **kwargs) # type: ignore
|
||||
raise e
|
||||
|
||||
return trace_async_step_wrapper
|
||||
|
||||
def _instrument_async_stream_runner_step(
|
||||
self, runner_step: Callable[..., AsyncIterator[ValidationOutcome[OT]]]
|
||||
) -> Callable[..., AsyncIterator[ValidationOutcome[OT]]]:
|
||||
@wraps(runner_step)
|
||||
async def trace_async_stream_step_wrapper(
|
||||
*args, **kwargs
|
||||
) -> AsyncIterator[ValidationOutcome[OT]]:
|
||||
with mlflow.start_span(
|
||||
name="guardrails/guard/step",
|
||||
span_type="step",
|
||||
attributes={
|
||||
"guardrails.version": GUARDRAILS_VERSION,
|
||||
"type": "guardrails/guard/step",
|
||||
"async": True,
|
||||
"stream": True,
|
||||
},
|
||||
) as step_span:
|
||||
exception = None
|
||||
try:
|
||||
gen = runner_step(*args, **kwargs)
|
||||
next_exists = True
|
||||
while next_exists:
|
||||
try:
|
||||
res = await anext(gen)
|
||||
yield res
|
||||
except StopIteration:
|
||||
next_exists = False
|
||||
except StopAsyncIteration:
|
||||
next_exists = False
|
||||
except Exception as e:
|
||||
step_span.set_status(status=SpanStatusCode.ERROR)
|
||||
exception = e
|
||||
finally:
|
||||
call = safe_get(args, 3, kwargs.get("call_log", None))
|
||||
iteration = call.iterations.last if call else None
|
||||
add_step_attributes(step_span, iteration, *args, **kwargs) # type: ignore
|
||||
if exception:
|
||||
raise exception
|
||||
|
||||
return trace_async_stream_step_wrapper
|
||||
|
||||
def _instrument_runner_call(self, runner_call: Callable[..., LLMResponse]):
|
||||
@wraps(runner_call)
|
||||
def trace_call_wrapper(*args, **kwargs):
|
||||
with mlflow.start_span(
|
||||
name="guardrails/guard/step/call",
|
||||
span_type="LLM",
|
||||
attributes={
|
||||
"guardrails.version": GUARDRAILS_VERSION,
|
||||
"type": "guardrails/guard/step/call",
|
||||
},
|
||||
) as call_span:
|
||||
try:
|
||||
response = runner_call(*args, **kwargs)
|
||||
add_call_attributes(call_span, response, *args, **kwargs) # type: ignore
|
||||
return response
|
||||
except Exception as e:
|
||||
call_span.set_status(status=SpanStatusCode.ERROR)
|
||||
add_call_attributes(call_span, None, *args, **kwargs) # type: ignore
|
||||
raise e
|
||||
|
||||
return trace_call_wrapper
|
||||
|
||||
def _instrument_async_runner_call(
|
||||
self, runner_call: Callable[..., Awaitable[LLMResponse]]
|
||||
):
|
||||
@wraps(runner_call)
|
||||
async def trace_async_call_wrapper(*args, **kwargs):
|
||||
with mlflow.start_span(
|
||||
name="guardrails/guard/step/call",
|
||||
span_type="LLM",
|
||||
attributes={
|
||||
"guardrails.version": GUARDRAILS_VERSION,
|
||||
"type": "guardrails/guard/step/call",
|
||||
"async": True,
|
||||
},
|
||||
) as call_span:
|
||||
try:
|
||||
response = await runner_call(*args, **kwargs)
|
||||
add_call_attributes(call_span, response, *args, **kwargs) # type: ignore
|
||||
return response
|
||||
except Exception as e:
|
||||
call_span.set_status(status=SpanStatusCode.ERROR)
|
||||
add_call_attributes(call_span, None, *args, **kwargs) # type: ignore
|
||||
raise e
|
||||
|
||||
return trace_async_call_wrapper
|
||||
|
||||
def _instrument_validator_validate(
|
||||
self, validator_validate: Callable[..., ValidationResult]
|
||||
):
|
||||
@wraps(validator_validate)
|
||||
def trace_validator_wrapper(*args, **kwargs):
|
||||
validator_name = "validator"
|
||||
obj_id = id(validator_validate)
|
||||
on_fail_descriptor = "unknown"
|
||||
init_kwargs = {}
|
||||
validation_session_id = "unknown"
|
||||
|
||||
validator_self = args[0]
|
||||
if validator_self is not None and isinstance(validator_self, Validator):
|
||||
validator_name = validator_self.rail_alias
|
||||
obj_id = id(validator_self)
|
||||
on_fail_descriptor = validator_self.on_fail_descriptor
|
||||
init_kwargs = validator_self._kwargs
|
||||
|
||||
validator_span_name = f"{validator_name}.validate"
|
||||
|
||||
# Skip this instrumentation in the case of async
|
||||
# when the parent span cannot be fetched from the current context
|
||||
# because Validator.validate is running in a ThreadPoolExecutor
|
||||
parent_span = mlflow.get_current_active_span()
|
||||
if not parent_span:
|
||||
return validator_validate(*args, **kwargs)
|
||||
|
||||
with mlflow.start_span(
|
||||
name=validator_span_name,
|
||||
span_type="validator",
|
||||
attributes={
|
||||
"guardrails.version": GUARDRAILS_VERSION,
|
||||
"type": "guardrails/guard/step/validator",
|
||||
},
|
||||
) as validator_span:
|
||||
try:
|
||||
resp = validator_validate(*args, **kwargs)
|
||||
add_validator_attributes(
|
||||
*args,
|
||||
validator_span=validator_span, # type: ignore
|
||||
validator_name=validator_name,
|
||||
obj_id=obj_id,
|
||||
on_fail_descriptor=on_fail_descriptor,
|
||||
result=resp,
|
||||
init_kwargs=init_kwargs,
|
||||
validation_session_id=validation_session_id,
|
||||
**kwargs,
|
||||
)
|
||||
return resp
|
||||
except Exception as e:
|
||||
validator_span.set_status(status=SpanStatusCode.ERROR)
|
||||
add_validator_attributes(
|
||||
*args,
|
||||
validator_span=validator_span, # type: ignore
|
||||
validator_name=validator_name,
|
||||
obj_id=obj_id,
|
||||
on_fail_descriptor=on_fail_descriptor,
|
||||
result=None,
|
||||
init_kwargs=init_kwargs,
|
||||
validation_session_id=validation_session_id,
|
||||
**kwargs,
|
||||
)
|
||||
raise e
|
||||
|
||||
return trace_validator_wrapper
|
||||
|
||||
def _instrument_validator_async_validate(
|
||||
self,
|
||||
validator_async_validate: Callable[..., Coroutine[Any, Any, ValidationResult]],
|
||||
):
|
||||
@wraps(validator_async_validate)
|
||||
async def trace_async_validator_wrapper(*args, **kwargs):
|
||||
validator_name = "validator"
|
||||
obj_id = id(validator_async_validate)
|
||||
on_fail_descriptor = "unknown"
|
||||
init_kwargs = {}
|
||||
validation_session_id = "unknown"
|
||||
|
||||
validator_self = args[0]
|
||||
if validator_self is not None and isinstance(validator_self, Validator):
|
||||
validator_name = validator_self.rail_alias
|
||||
obj_id = id(validator_self)
|
||||
on_fail_descriptor = validator_self.on_fail_descriptor
|
||||
init_kwargs = validator_self._kwargs
|
||||
|
||||
validator_span_name = f"{validator_name}.validate"
|
||||
|
||||
with mlflow.start_span(
|
||||
name=validator_span_name,
|
||||
span_type="validator",
|
||||
attributes={
|
||||
"guardrails.version": GUARDRAILS_VERSION,
|
||||
"type": "guardrails/guard/step/validator",
|
||||
"async": True,
|
||||
},
|
||||
) as validator_span:
|
||||
try:
|
||||
resp = await validator_async_validate(*args, **kwargs)
|
||||
add_validator_attributes(
|
||||
*args,
|
||||
validator_span=validator_span, # type: ignore
|
||||
validator_name=validator_name,
|
||||
obj_id=obj_id,
|
||||
on_fail_descriptor=on_fail_descriptor,
|
||||
result=resp,
|
||||
init_kwargs=init_kwargs,
|
||||
validation_session_id=validation_session_id,
|
||||
**kwargs,
|
||||
)
|
||||
return resp
|
||||
except Exception as e:
|
||||
validator_span.set_status(status=SpanStatusCode.ERROR)
|
||||
add_validator_attributes(
|
||||
*args,
|
||||
validator_span=validator_span, # type: ignore
|
||||
validator_name=validator_name,
|
||||
obj_id=obj_id,
|
||||
on_fail_descriptor=on_fail_descriptor,
|
||||
result=None,
|
||||
init_kwargs=init_kwargs,
|
||||
validation_session_id=validation_session_id,
|
||||
**kwargs,
|
||||
)
|
||||
raise e
|
||||
|
||||
return trace_async_validator_wrapper
|
||||
Reference in New Issue
Block a user