diff --git a/comfy/__init__.py b/comfy/__init__.py index e69de29bb..f102a9cad 100644 --- a/comfy/__init__.py +++ b/comfy/__init__.py @@ -0,0 +1 @@ +__version__ = "0.0.1" diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 6704b5131..aad3ee0f3 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -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'): diff --git a/comfy/cli_args_types.py b/comfy/cli_args_types.py index 4275f0b7a..3a24f1031 100644 --- a/comfy/cli_args_types.py +++ b/comfy/cli_args_types.py @@ -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 diff --git a/comfy/cmd/main_pre.py b/comfy/cmd/main_pre.py index 51408b682..c3202765d 100644 --- a/comfy/cmd/main_pre.py +++ b/comfy/cmd/main_pre.py @@ -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"] diff --git a/comfy/tracing_compatibility.py b/comfy/tracing_compatibility.py new file mode 100644 index 000000000..7bcebb544 --- /dev/null +++ b/comfy/tracing_compatibility.py @@ -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) diff --git a/requirements.txt b/requirements.txt index dde0ec0cd..ae2cb825d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -36,4 +36,11 @@ lazy-object-proxy can_ada fsspec natsort -OpenEXR \ No newline at end of file +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 \ No newline at end of file