mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-09 13:50:49 +08:00
80 lines
3.0 KiB
Python
80 lines
3.0 KiB
Python
from typing import Optional, Sequence
|
|
|
|
from aio_pika.abc import AbstractChannel
|
|
from opentelemetry.context import Context
|
|
from opentelemetry.sdk.trace.sampling import Sampler, SamplingResult, Decision
|
|
from opentelemetry.semconv.attributes.network_attributes import NETWORK_PEER_ADDRESS, NETWORK_PEER_PORT
|
|
from opentelemetry.trace import SpanKind, Link, TraceState
|
|
from opentelemetry.util.types import Attributes
|
|
|
|
|
|
def patch_spanbuilder_set_channel() -> None:
|
|
"""
|
|
The default SpanBuilder.set_channel does not work with aio_pika 9.1 and the refactored connection
|
|
attribute
|
|
"""
|
|
import opentelemetry.instrumentation.aio_pika.span_builder
|
|
from opentelemetry.instrumentation.aio_pika.span_builder import SpanBuilder
|
|
|
|
def set_channel(self: SpanBuilder, channel: AbstractChannel) -> None:
|
|
if hasattr(channel, "_connection"):
|
|
url = channel._connection.url
|
|
port = url.port or 5672
|
|
self._attributes.update(
|
|
{
|
|
NETWORK_PEER_ADDRESS: url.host,
|
|
NETWORK_PEER_PORT: port,
|
|
}
|
|
)
|
|
|
|
opentelemetry.instrumentation.aio_pika.span_builder.SpanBuilder.set_channel = set_channel # type: ignore[misc]
|
|
|
|
|
|
class ProgressSpanSampler(Sampler):
|
|
"""
|
|
Sampler which omits aio_pika messages destined/related to progress.
|
|
|
|
To debug which spans are being dropped, set environment variable:
|
|
OTEL_DEBUG_SAMPLER=1
|
|
"""
|
|
|
|
def __init__(self):
|
|
import logging
|
|
import os
|
|
self.logger = logging.getLogger(__name__)
|
|
self.debug = os.environ.get("OTEL_DEBUG_SAMPLER", "0") == "1"
|
|
|
|
def get_description(self) -> str:
|
|
return "Sampler which omits aio_pika messages destined/related to progress"
|
|
|
|
def should_sample(
|
|
self,
|
|
parent_context: Optional["Context"],
|
|
trace_id: int,
|
|
name: str,
|
|
kind: Optional[SpanKind] = None,
|
|
attributes: Attributes = None,
|
|
links: Optional[Sequence["Link"]] = None,
|
|
trace_state: Optional["TraceState"] = None,
|
|
) -> "SamplingResult":
|
|
should_drop = False
|
|
drop_reason = None
|
|
|
|
if attributes is not None and "messaging.destination" in attributes and attributes["messaging.destination"].endswith("progress"):
|
|
should_drop = True
|
|
drop_reason = f"messaging.destination={attributes['messaging.destination']}"
|
|
# the ephemeral reply channels are not required for correct span correlation
|
|
elif name.startswith(",amq_") or name.startswith("amq"):
|
|
should_drop = True
|
|
drop_reason = f"span name starts with amq: {name}"
|
|
|
|
if should_drop:
|
|
if self.debug:
|
|
self.logger.debug(f"[SAMPLER] DROPPED span '{name}' (kind={kind}): {drop_reason}")
|
|
return SamplingResult(Decision.DROP)
|
|
|
|
if self.debug:
|
|
self.logger.debug(f"[SAMPLER] RECORDING span '{name}' (kind={kind}, attributes={attributes})")
|
|
|
|
return SamplingResult(Decision.RECORD_AND_SAMPLE)
|