traceAI-transformers 0.1.0__py3-none-any.whl

traceai_transformers/__init__.py

@@ -0,0 +1,79 @@
+"""Trace local Transformers pipelines, independently of remote Hub clients."""
+from typing import Any, Collection
+
+from fi_instrumentation import FITracer, TraceConfig
+from opentelemetry import trace
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from wrapt import FunctionWrapper
+
+from ._wrappers import _PipelineWrapper
+from .version import __version__
+
+# Look up optional historical classes individually: Transformers 5 removed them.
+_TARGETS = (
+    ("TextGenerationPipeline", "text-generation", "LLM"),
+    ("FeatureExtractionPipeline", "feature-extraction", "EMBEDDING"),
+    ("Text2TextGenerationPipeline", "text2text-generation", "LLM"),
+    ("SummarizationPipeline", "summarization", "LLM"),
+    ("TranslationPipeline", "translation", "LLM"),
+    ("QuestionAnsweringPipeline", "question-answering", "CHAIN"),
+)
+
+
+class TransformersInstrumentor(BaseInstrumentor):
+    """Instrument supported pipeline calls; bare ``model.generate`` is not patched.
+
+    ``capture_content=False`` (default) records metadata only. Opt in to bounded
+    input/output previews with ``capture_content=True``; ``config=TraceConfig``
+    still controls privacy. Embedding vectors and lazy result payloads are never
+    collected. Supply ``tracer_provider`` to avoid changing the global provider.
+    """
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ("transformers>=4.42,<6",)
+
+    def _instrument(self, **kwargs: Any) -> None:
+        import transformers.pipelines as pipelines
+
+        config = kwargs.get("config")
+        if config is None:
+            config = TraceConfig()
+        if not isinstance(config, TraceConfig):
+            raise TypeError("config must be a TraceConfig")
+        capture = kwargs.get("capture_content", False)
+        if not isinstance(capture, bool):
+            raise TypeError("capture_content must be a bool")
+        tracer = FITracer(
+            trace.get_tracer(__name__, __version__, kwargs.get("tracer_provider")),
+            config=config,
+        )
+        self._patches = []
+        try:
+            for name, task, kind in _TARGETS:
+                cls = getattr(pipelines, name, None)
+                if cls is None:
+                    continue
+                own_method = "__call__" in vars(cls)
+                original = vars(cls)["__call__"] if own_method else cls.__call__
+                wrapper = _PipelineWrapper(tracer=tracer, config=config, capture=capture, task=task, kind=kind)
+                patched = FunctionWrapper(original, wrapper)
+                self._patches.append((cls, own_method, original, patched, wrapper))
+                cls.__call__ = patched
+        except BaseException:
+            self._uninstrument()
+            raise
+
+    def _uninstrument(self, **kwargs: Any) -> None:
+        for cls, own_method, original, patched, wrapper in reversed(self._patches):
+            # An outer, third-party patch may still hold ours: disable it without
+            # overwriting their patch. In-flight lazy iterators finish normally.
+            wrapper.enabled = False
+            if vars(cls).get("__call__") is patched:
+                if own_method:
+                    cls.__call__ = original
+                else:
+                    delattr(cls, "__call__")
+        self._patches = []
+
+
+__all__ = ["TransformersInstrumentor", "__version__"]

traceai_transformers/_wrappers.py

@@ -0,0 +1,336 @@
+"""Bounded metadata and iteration-scoped tracing without touching model outputs."""
+import json
+import logging
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from copy import copy
+from itertools import islice
+from time import time_ns
+from types import GeneratorType
+from typing import Any, Callable
+
+from fi_instrumentation import FITracer, TraceConfig
+from opentelemetry import context, trace
+from opentelemetry.context import _SUPPRESS_INSTRUMENTATION_KEY
+from opentelemetry.trace import Status, StatusCode
+from wrapt import ObjectProxy
+
+_ACTIVE: ContextVar[tuple[int, ...]] = ContextVar("traceai_transformers_active", default=())
+_MAX_PREVIEW = 4096
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.NullHandler())
+
+
+def _log_failure(action: str, error: Exception) -> None:
+    # Telemetry is fail-open, but not silent. Do not log exception messages or
+    # traceback source lines: they can bypass TraceConfig's content privacy.
+    logger.warning("Transformers instrumentation failed to %s (%s)", action, type(error).__name__)
+
+
+def _record_error(span, error: BaseException, config: TraceConfig) -> None:
+    if config.pii_redaction or config.hide_error_details:
+        # Supported FI core versions mask supplied attributes, but not the
+        # message/traceback generated by record_exception or status descriptions.
+        # Omit those surfaces entirely, including chained errors and notes;
+        # never stringify or mutate the exception being returned to the caller.
+        cls = type(error)
+        exception_type = cls.__qualname__
+        if cls.__module__ and cls.__module__ != "builtins":
+            exception_type = f"{cls.__module__}.{exception_type}"
+        span.add_event("exception", {
+            "exception.type": exception_type,
+            "exception.escaped": "False",
+        })
+        span.set_status(StatusCode.ERROR)
+    else:
+        span.record_exception(error)
+        span.set_status(Status(StatusCode.ERROR, str(error)))
+
+
+def _preview(value: object) -> str:
+    # No repr/default=str, tensor conversion, custom iteration or full JSON
+    # serialization of unbounded collections. A global node budget also bounds
+    # nested batches, not just each individual list.
+    budget = [64]
+
+    def trim(item, depth=0):
+        budget[0] -= 1
+        if budget[0] < 0 or depth > 5:
+            return "[truncated]"
+        if type(item) is str:
+            return item[:_MAX_PREVIEW]
+        if item is None or type(item) in (bool, int, float):
+            return item
+        if type(item) in (list, tuple):
+            return [trim(v, depth + 1) for v in item[:16] if budget[0] > 0]
+        if type(item) is dict:
+            result = {}
+            for key, val in islice(item.items(), 16):
+                if budget[0] <= 0 or len(result) >= 16:
+                    break
+                if type(key) is str:
+                    result[key[:128]] = trim(val, depth + 1)
+            return result
+        return "[not captured]"
+
+    return json.dumps(trim(value), ensure_ascii=True)[:_MAX_PREVIEW]
+
+
+class _Invocation:
+    def __init__(self, wrapper, instance, value, parent):
+        self.tracer = wrapper.tracer
+        self.config = wrapper.config
+        self.identity = id(instance)
+        self.parent = parent
+        self.name = "transformers." + wrapper.task
+        self.attributes = {
+            "gen_ai.span.kind": wrapper.kind,
+            "gen_ai.provider.name": "huggingface",
+            "gen_ai.operation.name": (
+                "embeddings" if wrapper.kind == "EMBEDDING" else
+                "text_completion" if wrapper.kind == "LLM" else "question_answering"
+            ),
+            "transformers.task": wrapper.task,
+        }
+        # Model architecture is safe metadata, not a local filesystem path or
+        # arbitrary user-supplied model identifier.
+        model_type = getattr(getattr(getattr(instance, "model", None), "config", None), "model_type", None)
+        if type(model_type) is str:
+            self.attributes["gen_ai.request.model"] = model_type[:256]
+        # Generation responses can contain the entire original prompt/chat.
+        # Never re-export hidden inputs through a full-text output preview.
+        self.capture_output = wrapper.capture and wrapper.kind != "EMBEDDING" and not (
+            wrapper.config.hide_outputs or wrapper.config.hide_output_text or
+            wrapper.config.hide_output_messages or wrapper.config.hide_inputs or
+            wrapper.config.hide_input_text or wrapper.config.hide_input_messages
+        )
+        if wrapper.capture and not (
+            wrapper.config.hide_inputs or wrapper.config.hide_input_text or
+            wrapper.config.hide_input_messages
+        ):
+            self.attributes["input.value"] = _preview(value)
+        self.span = None
+        self.started = False
+        self.finished = False
+
+    def start(self, start_time=None):
+        if not self.started:
+            self.started = True
+            try:
+                self.span = self.tracer.start_span(
+                    self.name, context=self.parent, attributes=self.attributes, start_time=start_time,
+                )
+            except Exception as error:
+                # An exporter/provider failure must not prevent local inference.
+                _log_failure("start span", error)
+            self.attributes = None
+
+    @contextmanager
+    def activate(self):
+        token = context.attach(self.parent)
+        active = _ACTIVE.set((*_ACTIVE.get(), self.identity))
+        try:
+            self.start()
+            if self.span is None:
+                yield
+            else:
+                with trace.use_span(self.span, end_on_exit=False, record_exception=False,
+                                    set_status_on_exception=False):
+                    yield
+        finally:
+            _ACTIVE.reset(active)
+            context.detach(token)
+
+    def finish(self, error=None, output=None):
+        if self.finished:
+            return
+        self.finished = True
+        if self.span is None:
+            return
+        try:
+            if error is not None:
+                _record_error(self.span, error, self.config)
+            else:
+                if self.capture_output and output is not None:
+                    self.span.set_attribute("output.value", _preview(output))
+                self.span.set_status(StatusCode.OK)
+        except Exception as telemetry_error:
+            _log_failure("record span result", telemetry_error)
+        finally:
+            try:
+                self.span.end()
+            except Exception as telemetry_error:
+                _log_failure("end span", telemetry_error)
+
+
+class _TracedIterator(Iterator):
+    """One span per lazy call, current only while the producer executes.
+
+    Payloads are never accumulated. Close an abandoned iterator explicitly;
+    finalization is only a best-effort fallback, as with Python generators.
+    """
+    def __init__(self, wrapped, invocation, *, iterator=None):
+        self._wrapped = wrapped
+        self._iterator = iterator
+        self._invocation = invocation
+        self._closed = False
+
+    def __next__(self):
+        if self._closed:
+            raise StopIteration
+        try:
+            with self._invocation.activate():
+                if self._iterator is None:
+                    self._iterator = iter(self._wrapped)
+                return next(self._iterator)
+        except StopIteration:
+            self._closed = True
+            self._invocation.finish()
+            self._wrapped = self._iterator = None
+            raise
+        except BaseException as error:
+            self._closed = True
+            self._invocation.finish(error)
+            self._wrapped = self._iterator = None
+            raise
+
+    def close(self):
+        if self._closed:
+            return
+        self._closed = True
+        try:
+            closer = getattr(self._wrapped, "close", None)
+            if closer is not None:
+                # Do not start an inference span for an unconsumed iterator.
+                if self._invocation.started:
+                    with self._invocation.activate():
+                        closer()
+                else:
+                    closer()
+        except BaseException as error:
+            self._invocation.finish(error)
+            raise
+        finally:
+            self._invocation.finish()
+            # Release DataLoader iterators/workers and their last batch promptly.
+            self._wrapped = self._iterator = None
+
+    def __getattr__(self, name):
+        return getattr(self._wrapped, name)
+
+    def __del__(self):
+        self._invocation.finish()
+
+
+class _TracedDataset(ObjectProxy):
+    """Keep upstream dataset protocols, with a separate span per traversal.
+
+    PipelineIterator is an IterableDataset whose __iter__ resets its producer,
+    including during a partial traversal. A one-shot Iterator wrapper loses that
+    behavior. ObjectProxy preserves __len__ and isinstance(..., IterableDataset),
+    which DataLoader uses to choose iterable rather than indexed access.
+    """
+    def __init__(self, wrapped, invocation):
+        super().__init__(wrapped)
+        # This template stays unstarted and retains only metadata and context,
+        # not source previews, outputs, or a pipeline/model reference.
+        self._self_template = invocation
+        self._self_cursor = _TracedIterator(wrapped, copy(invocation))
+
+    def __iter__(self):
+        self._self_cursor.close()
+        invocation = copy(self._self_template)
+        try:
+            iterator = iter(self.__wrapped__)
+        except BaseException as error:
+            invocation.start()
+            invocation.finish(error)
+            raise
+        self._self_cursor = _TracedIterator(self.__wrapped__, invocation, iterator=iterator)
+        return self
+
+    def __next__(self):
+        return next(self._self_cursor)
+
+    def __reduce_ex__(self, protocol):
+        # Pickle the native dataset, not this proxy: its invocation retains a
+        # tracer provider and parent context which cannot safely cross a spawn
+        # boundary (and may contain unpicklable locks or private metadata).
+        # Inference in a deserialized DataLoader worker is deliberately untraced.
+        return self.__wrapped__.__reduce_ex__(protocol)
+
+    def close(self):
+        # Explicit close ends the current traversal; a later iter() retains
+        # the underlying dataset's own ability (or inability) to start again.
+        self._self_cursor.close()
+
+
+def _input(args, kwargs):
+    if args:
+        return args[0]
+    for key in ("inputs", "text_inputs", "texts", "X"):
+        if key in kwargs:
+            return kwargs[key]
+    # Extractive QA uses named question/context rather than a single input.
+    return {key: kwargs[key] for key in ("question", "context") if key in kwargs}
+
+
+def _lazy_input(value, task):
+    if task == "question-answering":
+        # Upstream's QA argument handler materializes its inputs eagerly.
+        return False
+    from torch.utils.data import Dataset
+    return isinstance(value, (GeneratorType, Dataset))
+
+
+class _PipelineWrapper:
+    def __init__(self, *, tracer: FITracer, config: TraceConfig, capture: bool, task: str, kind: str) -> None:
+        self.tracer, self.config, self.capture = tracer, config, capture
+        self.task, self.kind = task, kind
+        self.enabled = True
+
+    def __call__(self, wrapped: Callable[..., Any], instance: Any,
+                 args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
+        if not self.enabled or id(instance) in _ACTIVE.get() or context.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+            return wrapped(*args, **kwargs)
+        try:
+            value = _input(args, kwargs)
+            invocation = _Invocation(self, instance, value, context.get_current())
+            lazy = _lazy_input(value, self.task)
+        except Exception as error:
+            _log_failure("prepare span metadata", error)
+            return wrapped(*args, **kwargs)
+        if lazy:
+            # The framework sets up the iterator here but does not run inference.
+            # Guard super().__call__ wrappers without creating a premature span.
+            active = _ACTIVE.set((*_ACTIVE.get(), id(instance)))
+            call_start = time_ns()
+            try:
+                result = wrapped(*args, **kwargs)
+            except BaseException as error:
+                invocation.start(start_time=call_start)
+                invocation.finish(error)
+                raise
+            finally:
+                _ACTIVE.reset(active)
+            if isinstance(result, Iterator):
+                # Do not keep even bounded source previews for lazy streams.
+                invocation.attributes.pop("input.value", None)
+                from torch.utils.data import IterableDataset
+                if isinstance(result, IterableDataset):
+                    return _TracedDataset(result, invocation)
+                return _TracedIterator(result, invocation)
+            # Some framework/backend paths materialize the result; avoid
+            # consuming or transforming it to manufacture lazy behavior.
+            invocation.start(start_time=call_start)
+            invocation.finish(output=result)
+            return result
+        try:
+            with invocation.activate():
+                result = wrapped(*args, **kwargs)
+        except BaseException as error:
+            invocation.finish(error)
+            raise
+        invocation.finish(output=result)
+        return result

traceai_transformers/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"  # x-release-please-version

traceai_transformers-0.1.0.dist-info/METADATA

@@ -0,0 +1,103 @@
+Metadata-Version: 2.5
+Name: traceAI-transformers
+Version: 0.1.0
+Summary: OpenTelemetry instrumentation for local Hugging Face Transformers pipelines
+Project-URL: Repository, https://github.com/future-agi/traceAI/tree/dev/python/frameworks/transformers
+Project-URL: Documentation, https://docs.futureagi.com/docs/tracing
+Author-email: Future AGI <no-reply@futureagi.com>
+License-Expression: Apache-2.0
+Requires-Python: <3.14,>=3.10
+Requires-Dist: fi-instrumentation-otel<2,>=1.2.2
+Requires-Dist: opentelemetry-instrumentation<1,>=0.50b0
+Requires-Dist: transformers<6,>=4.42
+Requires-Dist: wrapt<2,>=1.15
+Provides-Extra: dev
+Requires-Dist: mypy<2,>=1.15; extra == 'dev'
+Requires-Dist: pytest<9,>=8; extra == 'dev'
+Requires-Dist: torch<3,>=2.6; extra == 'dev'
+Requires-Dist: traceai-huggingface>=0.1.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# traceAI Transformers
+
+Trace local Hugging Face Transformers pipeline inference in Future AGI. For remote Hugging Face InferenceClient calls, use `traceAI-huggingface` instead; the two instrumentors can run together.
+
+## Install
+
+```bash
+pip install traceAI-transformers torch
+```
+
+Python 3.10–3.13 and Transformers 4.42–5.x are supported. Transformers 5 requires Python 3.11 or later. Use PyTorch pipelines; select the appropriate PyTorch build for your CPU or GPU.
+
+## Trace inference
+
+Set `FI_API_KEY` and `FI_SECRET_KEY` in your environment, then configure tracing before calling your pipeline:
+
+```python
+import os
+from fi_instrumentation import register
+from fi_instrumentation.fi_types import ProjectType
+from traceai_transformers import TransformersInstrumentor
+from transformers import pipeline
+
+provider = register(
+    project_name="local-transformers",
+    project_type=ProjectType.OBSERVE,
+    headers={
+        "FI_API_KEY": os.environ["FI_API_KEY"],
+        "FI_SECRET_KEY": os.environ["FI_SECRET_KEY"],
+    },
+    verbose=False,
+)
+TransformersInstrumentor().instrument(tracer_provider=provider)
+
+generate = pipeline(
+    "text-generation",
+    model="HuggingFaceTB/SmolLM2-135M-Instruct",
+    device=-1,
+)
+result = generate(
+    [{"role": "user", "content": "Explain what a tokenizer does in one sentence."}],
+    max_new_tokens=32,
+    do_sample=False,
+)
+print(result[0]["generated_text"][-1]["content"])
+provider.force_flush()
+provider.shutdown()
+```
+
+The first run downloads the model. `register` exports OTLP traces to `https://api.futureagi.com/tracer/v1/traces` with `FI_API_KEY` and `FI_SECRET_KEY` headers. This is a traces-only endpoint, not a destination for OTLP metrics or logs. Open the `local-transformers` project in Future AGI to inspect `transformers.text-generation` spans.
+
+## Content and privacy
+
+By default, only task, model architecture, timing, status and error information are recorded. Model paths, generation parameters, embedding vectors and token usage estimates are not collected.
+
+To include bounded input/output previews:
+
+```python
+from fi_instrumentation import TraceConfig
+
+TransformersInstrumentor().instrument(
+    tracer_provider=provider,
+    capture_content=True,
+    config=TraceConfig(hide_error_details=True),
+)
+```
+
+Choose this configuration on your first `instrument()` call, or call `uninstrument()` before changing it. Previews are limited to 4,096 characters and may be truncated. `TraceConfig` input/output privacy flags override content capture. Hiding inputs also suppresses output previews, because pipeline responses can contain the original prompt or conversation. Set `FI_HIDE_ERROR_DETAILS=true` to omit exception messages, stack traces and error descriptions. Error details can contain inputs even when content capture is off.
+
+## Supported calls
+
+| Pipeline | Span kind | Versions |
+| --- | --- | --- |
+| Text generation, including chat, batches and streamers | LLM | 4.42–5.x |
+| Feature extraction, including tensor results | EMBEDDING | 4.42–5.x |
+| Text-to-text generation, summarization and translation | LLM | 4.x only |
+| Extractive question answering | CHAIN | 4.x only |
+
+Transformers 5 removed the historical text-to-text, summarization, translation and extractive question-answering pipelines. Use a text-generation chat model on 5.x. Direct `model.generate()` calls, other tasks, TensorFlow pipelines, and model loading are not instrumented.
+
+Generator and dataset calls remain lazy. Each traversal gets a span that starts when the first result is requested and ends on exhaustion, error or explicit `iterator.close()`. Dataset results preserve their length and `DataLoader` compatibility. Starting another traversal preserves the dataset's reset behavior and ends any unfinished span from the previous traversal. The span is current only while producing the next result, not while your application handles it. Close partially consumed iterators to end their spans promptly. Lazy payloads are not collected. Multiprocessing `DataLoader` workers receive the native dataset through serialization, not the instrumentor's tracer provider or parent context; inference in those workers is not traced by this dataset wrapper. To trace inference inside a worker, initialize its own exporter there and add explicit spans around that worker's inference; the parent instrumentor does not trace a deserialized dataset. Streamer spans end when pipeline inference returns, independently of when queued tokens are consumed. When starting a worker thread, propagate OpenTelemetry context explicitly if it should inherit a parent span.
+
+Instrumentation is idempotent and supports `TransformersInstrumentor().uninstrument()` and re-instrumentation. It uses your supplied tracer provider without replacing the global provider.

traceai_transformers-0.1.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+traceai_transformers/__init__.py,sha256=A1tsx1bgfDu6tp1p-AMyCeIA-QQx6UcdmdMGdcIAn28,3362
+traceai_transformers/_wrappers.py,sha256=gbv0x_fiSmXWgDkEsEn8hoZpA-x--W5kQChHQX7Ow5E,13586
+traceai_transformers/version.py,sha256=IOEzhCbX916JIkvZr03GFY4bl9YqSkfjdqLO4qJ2FOs,50
+traceai_transformers-0.1.0.dist-info/METADATA,sha256=DDlSVwwhkX5S6tS_qQT3q2VSpVE45IAkFv21gW37bac,5658
+traceai_transformers-0.1.0.dist-info/WHEEL,sha256=2iZrJkDIkeDh0k9SikKMqFuRP-ua6mb4Ir1AsER4BEw,87
+traceai_transformers-0.1.0.dist-info/entry_points.txt,sha256=eMW-1VNv3gPxiIxD6HA4RuSx04B8D0_XUXrnhA6jVZQ,90
+traceai_transformers-0.1.0.dist-info/RECORD,,

traceai_transformers-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.33.0
+Root-Is-Purelib: true
+Tag: py3-none-any

traceai_transformers-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[opentelemetry_instrumentor]
+transformers = traceai_transformers:TransformersInstrumentor