diff --git a/packages/syft/src/syft/util/telemetry.py b/packages/syft/src/syft/util/telemetry.py index d14b1235e2d..99a95d991d1 100644 --- a/packages/syft/src/syft/util/telemetry.py +++ b/packages/syft/src/syft/util/telemetry.py @@ -1,10 +1,11 @@ # stdlib +from collections.abc import Callable import logging import os from typing import Any +from typing import TypeVar # relative -from . import trace_decorator from .. import __version__ from .util import str_to_bool @@ -18,10 +19,22 @@ TRACING_ENABLED = str_to_bool(os.environ.get("TRACING", "False")) logger = logging.getLogger(__name__) +T = TypeVar("T", bound=Callable | type) + + +def no_instrument(__func_or_class: T | None = None, /, *args: Any, **kwargs: Any) -> T: + def noop_wrapper(__func_or_class: T) -> T: + return __func_or_class + + if __func_or_class is None: + return noop_wrapper # type: ignore + else: + return __func_or_class + def setup_instrumenter() -> Any: if not TRACING_ENABLED: - return trace_decorator.no_instrument + return no_instrument try: # third party @@ -35,6 +48,9 @@ def setup_instrumenter() -> Any: from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import BatchSpanProcessor + # relative + from .trace_decorator import instrument + # create a resource resource = Resource({"syft.version": __version__}) resource = resource.merge(OTELResourceDetector().detect()) @@ -53,10 +69,10 @@ def setup_instrumenter() -> Any: trace.set_tracer_provider(provider) logger.info("Added TracerProvider with BatchSpanProcessor") - return trace_decorator.instrument + return instrument except Exception as e: logger.error("Failed to import opentelemetry", exc_info=e) - return trace_decorator.no_instrument + return no_instrument def instrument_fastapi(app: Any) -> None: diff --git a/packages/syft/src/syft/util/trace_decorator.py b/packages/syft/src/syft/util/trace_decorator.py index 403ad188d6b..eaa259330cd 100644 --- a/packages/syft/src/syft/util/trace_decorator.py +++ b/packages/syft/src/syft/util/trace_decorator.py @@ -17,7 +17,7 @@ from opentelemetry.trace import Tracer from opentelemetry.trace.span import Span -__all__ = ["instrument", "no_instrument", "T"] +__all__ = ["instrument"] T = TypeVar("T", bound=Callable | type) @@ -172,13 +172,3 @@ async def wrap_with_span_async(*args: Any, **kwargs: Any) -> Callable: return span_decorator(_func_or_class) else: return span_decorator # type: ignore - - -def no_instrument(__func_or_class: T | None = None, /, *args: Any, **kwargs: Any) -> T: - def noop_wrapper(__func_or_class: T) -> T: - return __func_or_class - - if __func_or_class is None: - return noop_wrapper # type: ignore - else: - return __func_or_class