scale-gp-beta 0.1.0a16__py3-none-any.whl → 0.1.0a18__py3-none-any.whl

scale_gp_beta/lib/tracing/trace_queue_manager.py

@@ -0,0 +1,181 @@
+import time
+import queue
+import atexit
+import logging
+import threading
+from typing import TYPE_CHECKING, Optional
+
+from scale_gp_beta import SGPClient, SGPClientError
+
+from .util import configure, is_disabled
+from .trace_exporter import TraceExporter
+
+if TYPE_CHECKING:
+    from .span import Span
+    from .trace import Trace
+
+# configurable by env vars?
+DEFAULT_MAX_QUEUE_SIZE = 4_000
+DEFAULT_TRIGGER_QUEUE_SIZE = 200
+DEFAULT_TRIGGER_CADENCE = 4.0
+DEFAULT_MAX_BATCH_SIZE = 50
+DEFAULT_RETRIES = 4
+
+WORKER_SLEEP_SECONDS = 0.1
+
+log: logging.Logger = logging.getLogger(__name__)
+
+
+class TraceQueueManager:
+    """Manage trace and spans queue
+    Store spans in-memory until the threshold has been reached then flush to server.
+
+    Optionally provide a client, if unprovided we will attempt to create a default client.
+    """
+
+    def __init__(
+        self,
+        client: Optional[SGPClient] = None,
+        max_queue_size: int = DEFAULT_MAX_QUEUE_SIZE,
+        max_batch_size: int = DEFAULT_MAX_BATCH_SIZE,
+        trigger_queue_size: int = DEFAULT_TRIGGER_QUEUE_SIZE,
+        trigger_cadence: float = DEFAULT_TRIGGER_CADENCE,
+        retries: int = DEFAULT_RETRIES,
+    ):
+        self._client = client
+        self._attempted_local_client_creation = False
+        self._trigger_queue_size = trigger_queue_size
+        self._trigger_cadence = trigger_cadence
+
+        self._reset_trigger_time()
+
+        self._exporter = TraceExporter(max_batch_size, retries)
+
+        self._shutdown_event = threading.Event()
+        self._queue: queue.Queue[Span] = queue.Queue(maxsize=max_queue_size)
+
+        if not is_disabled():
+            self._worker = threading.Thread(daemon=True, target=self._run)
+            self._worker.start()
+
+            # ensure the thread joins on exit
+            atexit.register(self.shutdown)
+
+    def register_client(self, client: SGPClient) -> None:
+        log.info("Registering client")
+        self._client = client
+
+    def shutdown(self, timeout: Optional[float] = None) -> None:
+        if is_disabled():
+            log.debug("No worker to shutdown")
+            return
+        log.info(f"Shutting down trace queue manager, joining worker thread with timeout {timeout}")
+        self._shutdown_event.set()
+        self._worker.join(timeout=timeout)
+        log.info("Shutdown complete")
+
+    def report_span_start(self, span: "Span") -> None:
+        # TODO: support making this optional. Current backend requires us to send span starts
+        try:
+            self._queue.put_nowait(span)
+        except queue.Full:
+            log.warning(f"Queue full, ignoring span {span.span_id}")
+
+    def report_span_end(self, span: "Span") -> None:
+        try:
+            self._queue.put_nowait(span)
+        except queue.Full:
+            log.warning(f"Queue full, ignoring span {span.span_id}")
+
+    def report_trace_start(self, trace: "Trace") -> None:
+        pass
+
+    def report_trace_end(self, trace: "Trace") -> None:
+        pass
+
+    def flush_queue(self) -> None:
+        self._export()
+
+    @property
+    def client(self) -> Optional[SGPClient]:
+        """
+        Use client provided client on init if available, otherwise attempt once to create a default one.
+        :return: SGPClient
+        """
+        if self._client is not None:
+            return self._client
+        if self._attempted_local_client_creation:
+            # Already tried and failed to create a client
+            return None
+
+        log.info("Tracing queue manager not initialized, attempting to create a default one")
+        try:
+            self.register_client(SGPClient())
+        except SGPClientError:
+            log.warning(
+                f"Failed to create SGPClient for tracing queue manager {self}, ignoring traces. Please initialize with a working client."
+            )
+        finally:
+            self._attempted_local_client_creation = True
+        return self._client
+
+    def _run(self) -> None:
+        # daemon worker loop
+        while not self._shutdown_event.is_set():
+            current_time = time.time()
+            queue_size = self._queue.qsize()
+
+            if queue_size >= self._trigger_queue_size or current_time >= self._next_trigger_time:
+                self._export()
+                self._reset_trigger_time()
+                continue
+
+            time.sleep(WORKER_SLEEP_SECONDS)
+
+        # flush all on shutdown
+        self._export()
+
+    def _export(self) -> None:
+        if self.client:
+            self._exporter.export(self.client, self._queue)
+
+    def _reset_trigger_time(self) -> None:
+        self._next_trigger_time = time.time() + self._trigger_cadence
+
+
+_global_tracing_queue_manager: Optional[TraceQueueManager] = None
+_queue_manager_lock = threading.Lock()
+
+
+def init(client: Optional[SGPClient] = None, disabled: Optional[bool] = None) -> None:
+    """Initialize the tracing backend
+
+    Good practice to always include this method call with a valid client at your program entrypoint.
+
+    Tracing will attempt to generate a default client if unprovided.
+
+    :param client: SGPClient
+    :param disabled: Set to True to disable tracing. Overrides environment variable ``DISABLE_SCALE_TRACING``
+    """
+    if disabled is not None:
+        configure(disabled=disabled)
+
+    global _global_tracing_queue_manager
+    if _global_tracing_queue_manager is not None:
+        return
+
+    with _queue_manager_lock:
+        if _global_tracing_queue_manager is None:
+            _global_tracing_queue_manager = TraceQueueManager(client)
+
+
+def tracing_queue_manager() -> TraceQueueManager:
+    global _global_tracing_queue_manager
+    if _global_tracing_queue_manager is None:
+        init(None)
+
+    if _global_tracing_queue_manager is None:
+        # should never happen... useful for linting
+        raise RuntimeError("Tracing queue manager failed to initialize.")
+
+    return _global_tracing_queue_manager

scale_gp_beta/lib/tracing/tracing.py

@@ -0,0 +1,180 @@
+import logging
+from typing import Any, Dict, Union, Optional
+
+from .span import Span, BaseSpan, NoOpSpan
+from .util import is_disabled
+from .scope import Scope
+from .trace import Trace, BaseTrace, NoOpTrace
+from .trace_queue_manager import tracing_queue_manager
+
+log: logging.Logger = logging.getLogger(__name__)
+
+
+def current_span() -> Optional[BaseSpan]:
+    """Retrieves the currently active span from the execution context.
+
+    This function relies on `contextvars` to manage the active span in
+    a context-local manner, making it safe for concurrent execution
+    environments (e.g., threads, asyncio tasks).
+
+    Returns:
+        Optional[BaseSpan]: The current BaseSpan instance if one is active,
+                            otherwise None. This could be a 'Span' or 'NoOpSpan'.
+    """
+    return Scope.get_current_span()
+
+
+def current_trace() -> Optional[BaseTrace]:
+    """Retrieves the currently active trace from the execution context.
+
+    Similarly, to `current_span()`, this uses `contextvars` for context-local
+    trace management.
+
+    Returns:
+        Optional[BaseTrace]: The current BaseTrace instance if one is active,
+                             otherwise None. This could be a 'Trace' or 'NoOpTrace'.
+    """
+    return Scope.get_current_trace()
+
+
+def flush_queue() -> None:
+    """
+    Blocking flush of all requests in the queue.
+
+    Useful for distributed applications to ensure spans have been committed before continuing.
+    :return:
+    """
+    queue_manager = tracing_queue_manager()
+    queue_manager.flush_queue()
+
+
+def create_trace(
+        name: str,
+        trace_id: Optional[str] = None,
+        root_span_id: Optional[str] = None,
+        metadata: Optional[Dict[str, Optional[str]]] = None,
+) -> BaseTrace:
+    """Creates a new trace and root span instance.
+
+    A trace represents a single, logical operation or workflow. It groups multiple
+    spans together. If tracing is disabled (via the 'DISABLE_SCALE_TRACING'
+    environment variable), a `NoOpTrace` instance is returned which performs no
+    actual tracing operations.
+
+    When a trace is started (e.g., by using it as a context manager or calling its
+    `start()` method), it becomes the `current_trace()` in the active scope.
+    Similarly, the root span instance becomes the `current_span()` in the active
+    scope.
+
+    Args:
+        name: The name of the trace.
+        trace_id (Optional[str]): An optional, user-defined ID for the trace.
+                                  If None, a unique trace ID will be generated.
+        root_span_id (Optional[str]): An optional, user-defined ID for the root span.
+        metadata (Optional[Dict[str, Optional[str]]]): An optional, user-defined metadata.
+
+    Returns:
+        BaseTrace: A `Trace` instance if tracing is enabled, or a `NoOpTrace`
+                   instance if tracing is disabled.
+    """
+    if is_disabled():
+        log.debug(f"Tracing is disabled. Not creating a new trace.")
+        return NoOpTrace(name=name, trace_id=trace_id, root_span_id=root_span_id, metadata=metadata)
+
+    active_trace = current_trace()
+    if active_trace is not None:
+        log.warning(f"Trace with id {active_trace.trace_id} is already active. Creating a new trace anyways.")
+
+    trace = Trace(name=name, trace_id=trace_id, queue_manager=tracing_queue_manager(), metadata=metadata)
+    log.debug(f"Created new trace: {trace.trace_id}")
+
+    return trace
+
+
+def create_span(
+    name: str,
+    span_id: Optional[str] = None,
+    input: Optional[Dict[str, Any]] = None,
+    output: Optional[Dict[str, Any]] = None,
+    metadata: Optional[Dict[str, Optional[str]]] = None,
+    trace: Optional[Union[BaseTrace, str]] = None,
+    parent_span: Optional[BaseSpan] = None,
+) -> BaseSpan:
+    """Creates a new span instance.
+
+    A span represents a single unit of work or operation within a trace. Spans
+    can be nested to form a hierarchy.
+
+    If tracing is disabled (via 'DISABLE_SCALE_TRACING' environment variable),
+    a `NoOpSpan` is returned. Additionally, if no explicit `parent` (Trace or Span)
+    is provided and there is no `current_trace()` active in the scope, a `NoOpSpan`
+    will also be returned to prevent orphaned spans.
+
+    When a span is started (e.g., via context manager or `start()`), it becomes
+    the `current_span()` in the active scope.
+
+    Args:
+        name (str): A descriptive name for the span (e.g., "database_query",
+                    "http_request").
+        span_id (Optional[str]): An optional, user-defined ID for the span.
+        input (Optional[dict[str, Any]], optional): A dictionary containing
+            input data or parameters relevant to this span's operation. Defaults to None.
+        output (Optional[dict[str, Any]], optional): A dictionary containing
+            output data or results from this span's operation. Defaults to None.
+        metadata (Optional[dict[str, Union[str, int, float, bool, None]]], optional):
+            A dictionary for arbitrary key-value pairs providing additional
+            context or annotations for the span. Values should be simple types.
+            Defaults to None.
+        trace (Optional[Union[BaseTrace, str]], optional): A `Trace` instance
+            or a trace ID. Used for explicit control. Default to trace fetched
+            from the active scope.
+        parent_span (Optional[BaseSpan], optional): A `Span` instance. Like
+            trace, used for explicit control. Defaults to span fetched from the
+            active scope.
+
+    Returns:
+        BaseSpan: A `Span` instance if tracing is enabled and a valid trace context
+                  exists, or a `NoOpSpan` otherwise.
+    """
+    queue_manager = tracing_queue_manager()
+    parent_span_id: Optional[str] = None
+
+    if parent_span is not None:
+        trace_id = parent_span.trace_id
+        parent_span_id = parent_span.span_id
+    elif trace is not None:
+        trace_id = trace if isinstance(trace, str) else trace.trace_id
+
+        parent_span = Scope.get_current_span()
+        parent_span_id = parent_span.span_id if parent_span else None
+    else:
+        trace = Scope.get_current_trace()
+        parent_span = Scope.get_current_span()
+
+        parent_span_id = parent_span.span_id if parent_span else None
+
+        if trace is None:
+            # need to think about default behavior here... do we create a trace, some other options?
+            # I am leaning towards setting it as an option as sometimes people might want to be succinct or when we
+            # build decorators we might want this functionality
+            log.debug(f"attempting to create a span with no trace")
+            return NoOpSpan(name=name, span_id=span_id, parent_span_id=parent_span_id)
+
+        trace_id = trace.trace_id
+
+    if is_disabled():
+        return NoOpSpan(name=name, span_id=span_id, parent_span_id=parent_span_id, trace_id=trace_id)
+
+    span = Span(
+        name=name,
+        span_id=span_id,
+        parent_span_id=parent_span_id,
+        trace_id=trace_id,
+        input=input or {},
+        output=output or {},
+        metadata=metadata or {},
+        queue_manager=queue_manager,
+    )
+    log.debug(f"Created new span: {span.span_id}")
+
+    return span

scale_gp_beta/lib/tracing/types.py

@@ -0,0 +1,40 @@
+"""
+This is necessary, unfortunately. Stainless does not provide these as enums, only as type annotations.
+
+For strict linting, we need to reference these enums.
+
+NOTE: These will have to be manually updated to support updated span_types and status.
+"""
+
+from typing_extensions import Literal
+
+SpanStatusLiterals = Literal["SUCCESS", "ERROR"]
+
+SpanTypeLiterals = Literal[
+        "TEXT_INPUT",
+        "TEXT_OUTPUT",
+        "COMPLETION_INPUT",
+        "COMPLETION",
+        "KB_RETRIEVAL",
+        "KB_INPUT",
+        "RERANKING",
+        "EXTERNAL_ENDPOINT",
+        "PROMPT_ENGINEERING",
+        "DOCUMENT_INPUT",
+        "MAP_REDUCE",
+        "DOCUMENT_SEARCH",
+        "DOCUMENT_PROMPT",
+        "CUSTOM",
+        "INPUT_GUARDRAIL",
+        "OUTPUT_GUARDRAIL",
+        "CODE_EXECUTION",
+        "DATA_MANIPULATION",
+        "EVALUATION",
+        "FILE_RETRIEVAL",
+        "KB_ADD_CHUNK",
+        "KB_MANAGEMENT",
+        "TRACER",
+        "AGENT_TRACER",
+        "AGENT_WORKFLOW",
+        "STANDALONE",
+]

scale_gp_beta/lib/tracing/util.py

@@ -0,0 +1,68 @@
+import os
+import uuid
+from typing import Optional
+from datetime import datetime, timezone
+from functools import cache
+
+_tracing_disabled: Optional[bool] = None
+
+
+def generate_trace_id() -> str:
+    """
+    Generate a unique trace ID.
+
+    Returns:
+        str: A unique trace ID.
+    """
+    return str(uuid.uuid4())
+
+
+def generate_span_id() -> str:
+    """
+    Generate a unique span ID.
+
+    Returns:
+        str: A unique span ID.
+    """
+    return str(uuid.uuid4())
+
+
+def iso_timestamp() -> str:
+    """
+    Generate an ISO 8601 timestamp.
+
+    Returns:
+        str: The current time in ISO 8601 format.
+    """
+    return datetime.now(timezone.utc).isoformat()
+
+
+def configure(disabled: bool) -> None:
+    """
+    Programmatically enable or disable tracing.
+
+    This setting takes precedence over the `DISABLE_SCALE_TRACING` environment
+    variable.
+
+    Args:
+        disabled (bool): Set to True to disable tracing, False to enable.
+    """
+    global _tracing_disabled
+    _tracing_disabled = disabled
+    is_disabled.cache_clear()
+
+
+@cache
+def is_disabled() -> bool:
+    """
+    Check if tracing is disabled, with programmatic control taking precedence.
+
+    Tracing is considered disabled if `configure(disabled=True)` has been called,
+    or if the `DISABLE_SCALE_TRACING` environment variable is set.
+
+    Returns:
+        bool: True if tracing is disabled, otherwise False.
+    """
+    if _tracing_disabled is not None:
+        return _tracing_disabled
+    return os.getenv("DISABLE_SCALE_TRACING") is not None