This commit is contained in:
doctorpangloss 2024-04-17 08:20:07 -07:00
parent e49c662c7f
commit 2643730acc
6 changed files with 162 additions and 10 deletions

View File

@ -0,0 +1 @@
__version__ = "0.0.1"

View File

@ -1,16 +1,17 @@
from __future__ import annotations
import enum
import logging
import sys
from importlib.metadata import entry_points
from types import ModuleType
from typing import Optional, Any, Callable
from typing import Optional, Any
import configargparse
import configargparse as argparse
import enum
from . import __version__
from . import options
from .cli_args_types import LatentPreviewMethod, Configuration, ConfigurationExtender
import sys
class EnumAction(argparse.Action):
@ -179,6 +180,24 @@ def create_parser() -> argparse.ArgumentParser:
parser.add_argument("--verbose", action="store_true", help="Enables more debug prints.")
parser.add_argument("--disable-known-models", action="store_true", help="Disables automatic downloads of known models and prevents them from appearing in the UI.")
parser.add_argument("--max-queue-size", type=int, default=65536, help="The API will reject prompt requests if the queue's size exceeds this value.")
# tracing
parser.add_argument("--otel-service-name", type=str, default="comfyui", env_var="OTEL_SERVICE_NAME", help="The name of the service or application that is generating telemetry data.")
parser.add_argument("--otel-service-version", type=str, default=__version__, env_var="OTEL_SERVICE_VERSION", help="The version of the service or application that is generating telemetry data.")
parser.add_argument("--otel-exporter-otlp-endpoint", type=str, default=None, env_var="OTEL_EXPORTER_OTLP_ENDPOINT", help="A base endpoint URL for any signal type, with an optionally-specified port number. Helpful for when you're sending more than one signal to the same endpoint and want one environment variable to control the endpoint.")
parser.add_argument("--otel-exporter-otlp-traces-endpoint", type=str, default=None, env_var="OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", help="Endpoint URL for trace data only, with an optionally-specified port number. Typically ends with v1/traces when using OTLP/HTTP.")
parser.add_argument("--otel-exporter-otlp-metrics-endpoint", type=str, default=None, env_var="OTEL_EXPORTER_OTLP_METRICS_ENDPOINT", help="Endpoint URL for metric data only, with an optionally-specified port number. Typically ends with v1/metrics when using OTLP/HTTP.")
parser.add_argument("--otel-exporter-otlp-logs-endpoint", type=str, default=None, env_var="OTEL_EXPORTER_OTLP_LOGS_ENDPOINT", help="Endpoint URL for log data only, with an optionally-specified port number. Typically ends with v1/logs when using OTLP/HTTP.")
parser.add_argument("--otel-exporter-otlp-headers", type=str, default=None, env_var="OTEL_EXPORTER_OTLP_HEADERS", help="A list of headers to apply to all outgoing data (traces, metrics, and logs). To add a username and password to your headers, set this field to ``")
parser.add_argument("--otel-exporter-otlp-traces-headers", type=str, default=None, env_var="OTEL_EXPORTER_OTLP_TRACES_HEADERS", help="A list of headers to apply to all outgoing traces.")
parser.add_argument("--otel-exporter-otlp-metrics-headers", type=str, default=None, env_var="OTEL_EXPORTER_OTLP_METRICS_HEADERS", help="A list of headers to apply to all outgoing metrics.")
parser.add_argument("--otel-exporter-otlp-logs-headers", type=str, default=None, env_var="OTEL_EXPORTER_OTLP_LOGS_HEADERS", help="A list of headers to apply to all outgoing logs.")
parser.add_argument("--otel-exporter-otlp-timeout", type=str, default=None, env_var="OTEL_EXPORTER_OTLP_TIMEOUT", help="The timeout value for all outgoing data (traces, metrics, and logs) in milliseconds.")
parser.add_argument("--otel-exporter-otlp-traces-timeout", type=str, default=None, env_var="OTEL_EXPORTER_OTLP_TRACES_TIMEOUT", help="The timeout value for all outgoing traces in milliseconds.")
parser.add_argument("--otel-exporter-otlp-metrics-timeout", type=str, default=None, env_var="OTEL_EXPORTER_OTLP_METRICS_TIMEOUT", help="The timeout value for all outgoing metrics in milliseconds.")
parser.add_argument("--otel-exporter-otlp-logs-timeout", type=str, default=None, env_var="OTEL_EXPORTER_OTLP_LOGS_TIMEOUT", help="The timeout value for all outgoing logs in milliseconds.")
parser.add_argument("--otel-exporter-otlp-protocol", type=str, default=None, env_var="OTEL_EXPORTER_OTLP_PROTOCOL", help="Specifies the OTLP transport protocol to be used for all telemetry data.")
parser.add_argument("--otel-exporter-otlp-traces-protocol", type=str, default=None, env_var="OTEL_EXPORTER_OTLP_TRACES_PROTOCOL", help="Specifies the OTLP transport protocol to be used for trace data.")
parser.add_argument("--otel-exporter-otlp-metrics-protocol", type=str, default=None, env_var="OTEL_EXPORTER_OTLP_METRICS_PROTOCOL", help="Specifies the OTLP transport protocol to be used for metrics data.")
# now give plugins a chance to add configuration
for entry_point in entry_points().select(group='comfyui.custom_config'):

View File

@ -1,8 +1,11 @@
# Define a class for your command-line arguments
import enum
from typing import Optional, List, Callable, Literal
from typing import Optional, List, Callable
import configargparse as argparse
from . import __version__
ConfigurationExtender = Callable[[argparse.ArgParser], Optional[argparse.ArgParser]]
@ -81,6 +84,23 @@ class Configuration(dict):
verbose (bool): Shows extra output for debugging purposes such as import errors of custom nodes.
disable_known_models (bool): Disables automatic downloads of known models and prevents them from appearing in the UI.
max_queue_size (int): The API will reject prompt requests if the queue's size exceeds this value.
otel_service_name (str): The name of the service or application that is generating telemetry data. Default: "comfyui".
otel_service_version (str): The version of the service or application that is generating telemetry data. Default: "0.0.1".
otel_exporter_otlp_endpoint (Optional[str]): A base endpoint URL for any signal type, with an optionally-specified port number. Helpful for when you're sending more than one signal to the same endpoint and want one environment variable to control the endpoint.
otel_exporter_otlp_traces_endpoint (Optional[str]): Endpoint URL for trace data only, with an optionally-specified port number. Typically ends with v1/traces when using OTLP/HTTP.
otel_exporter_otlp_metrics_endpoint (Optional[str]): Endpoint URL for metric data only, with an optionally-specified port number. Typically ends with v1/metrics when using OTLP/HTTP.
otel_exporter_otlp_logs_endpoint (Optional[str]): Endpoint URL for log data only, with an optionally-specified port number. Typically ends with v1/logs when using OTLP/HTTP.
otel_exporter_otlp_headers (Optional[str]): A list of headers to apply to all outgoing data (traces, metrics, and logs).
otel_exporter_otlp_traces_headers (Optional[str]): A list of headers to apply to all outgoing traces.
otel_exporter_otlp_metrics_headers (Optional[str]): A list of headers to apply to all outgoing metrics.
otel_exporter_otlp_logs_headers (Optional[str]): A list of headers to apply to all outgoing logs.
otel_exporter_otlp_timeout (Optional[str]): The timeout value for all outgoing data (traces, metrics, and logs) in milliseconds.
otel_exporter_otlp_traces_timeout (Optional[str]): The timeout value for all outgoing traces in milliseconds.
otel_exporter_otlp_metrics_timeout (Optional[str]): The timeout value for all outgoing metrics in milliseconds.
otel_exporter_otlp_logs_timeout (Optional[str]): The timeout value for all outgoing logs in milliseconds.
otel_exporter_otlp_protocol (Optional[str]): Specifies the OTLP transport protocol to be used for all telemetry data.
otel_exporter_otlp_traces_protocol (Optional[str]): Specifies the OTLP transport protocol to be used for trace data.
otel_exporter_otlp_metrics_protocol (Optional[str]): Specifies the OTLP transport protocol to be used for metrics data.
"""
def __init__(self, **kwargs):
@ -146,6 +166,25 @@ class Configuration(dict):
self.external_address: Optional[str] = None
self.disable_known_models: bool = False
self.max_queue_size: int = 65536
# from opentracing docs
self.otel_service_name: str = "comfyui"
self.otel_service_version: str = "0.0.1"
self.otel_exporter_otlp_endpoint: Optional[str] = None
self.otel_exporter_otlp_traces_endpoint: Optional[str] = None
self.otel_exporter_otlp_metrics_endpoint: Optional[str] = None
self.otel_exporter_otlp_logs_endpoint: Optional[str] = None
self.otel_exporter_otlp_headers: Optional[str] = None
self.otel_exporter_otlp_traces_headers: Optional[str] = None
self.otel_exporter_otlp_metrics_headers: Optional[str] = None
self.otel_exporter_otlp_logs_headers: Optional[str] = None
self.otel_exporter_otlp_timeout: Optional[str] = None
self.otel_exporter_otlp_traces_timeout: Optional[str] = None
self.otel_exporter_otlp_metrics_timeout: Optional[str] = None
self.otel_exporter_otlp_logs_timeout: Optional[str] = None
self.otel_exporter_otlp_protocol: Optional[str] = None
self.otel_exporter_otlp_traces_protocol: Optional[str] = None
self.otel_exporter_otlp_metrics_protocol: Optional[str] = None
for key, value in kwargs.items():
self[key] = value

View File

@ -7,12 +7,22 @@ Use this instead of cli_args to import the args:
It will enable command line argument parsing. If this isn't desired, you must author your own implementation of these fixes.
"""
import logging
import os
import warnings
from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.aio_pika import AioPikaInstrumentor
from opentelemetry.instrumentation.aiohttp_server import AioHttpServerInstrumentor
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
from opentelemetry.semconv.resource import ResourceAttributes as ResAttrs
from .. import options
import warnings
import logging
from ..tracing_compatibility import ProgressSpanSampler
from ..tracing_compatibility import patch_spanbuilder_set_channel
options.enable_args_parsing()
if os.name == "nt":
@ -32,4 +42,31 @@ if args.deterministic:
os.environ['CUBLAS_WORKSPACE_CONFIG'] = ":4096:8"
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
__all__ = ["args"]
def _create_tracer():
resource = Resource.create({
ResAttrs.SERVICE_NAME: args.otel_service_name,
ResAttrs.SERVICE_VERSION: args.otel_service_version,
})
# omit progress spans from aio pika
sampler = ProgressSpanSampler()
provider = TracerProvider(resource=resource, sampler=sampler)
otlp_exporter = OTLPSpanExporter() if args.otel_exporter_otlp_endpoint is not None else ConsoleSpanExporter()
processor = BatchSpanProcessor(otlp_exporter)
provider.add_span_processor(processor)
trace.set_tracer_provider(provider)
# enable instrumentation
patch_spanbuilder_set_channel()
AioPikaInstrumentor().instrument()
AioHttpServerInstrumentor().instrument()
return trace.get_tracer(args.otel_service_name)
tracer = _create_tracer()
__all__ = ["args", "tracer"]

View File

@ -0,0 +1,49 @@
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.trace import SpanAttributes
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(
{
SpanAttributes.NET_PEER_NAME: url.host,
SpanAttributes.NET_PEER_PORT: port,
}
)
opentelemetry.instrumentation.aio_pika.span_builder.SpanBuilder.set_channel = set_channel # type: ignore[misc]
class ProgressSpanSampler(Sampler):
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":
if attributes is not None and "messaging.destination" in attributes and attributes["messaging.destination"].endswith("progress"):
return SamplingResult(Decision.DROP)
return SamplingResult(Decision.RECORD_AND_SAMPLE)

View File

@ -36,4 +36,11 @@ lazy-object-proxy
can_ada
fsspec
natsort
OpenEXR
OpenEXR
opentelemetry-distro
opentelemetry-exporter-otlp
opentelemetry-exporter-otlp-proto-http
opentelemetry-exporter-otlp-proto-grpc
opentelemetry-propagator-jaeger
opentelemetry-instrumentation-aio-pika
opentelemetry-instrumentation-aiohttp-server @ git+https://github.com/open-telemetry/opentelemetry-python-contrib.git@v0.45b0#subdirectory=instrumentation/opentelemetry-instrumentation-aiohttp-server