scale-gp-beta 0.1.0a17__py3-none-any.whl → 0.1.0a19__py3-none-any.whl

scale_gp_beta/_version.py

@@ -1,4 +1,4 @@
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
 __title__ = "scale_gp_beta"
-__version__ = "0.1.0-alpha.17"  # x-release-please-version
+__version__ = "0.1.0-alpha.19"  # x-release-please-version

scale_gp_beta/lib/tracing/__init__.py

@@ -0,0 +1,4 @@
+from .tracing import create_span, flush_queue, create_trace, current_span, current_trace
+from .trace_queue_manager import init
+
+__all__ = ["init", "create_span", "create_trace", "current_trace", "current_span", "flush_queue"]

scale_gp_beta/lib/tracing/exceptions.py

@@ -0,0 +1,5 @@
+from scale_gp_beta._exceptions import SGPClientError
+
+
+class ParamsCreationError(SGPClientError):
+    pass

scale_gp_beta/lib/tracing/scope.py

@@ -0,0 +1,104 @@
+import contextvars
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from .span import BaseSpan
+    from .trace import BaseTrace
+
+
+_current_span: contextvars.ContextVar["BaseSpan | None"] = contextvars.ContextVar("current_span", default=None)
+
+_current_trace: contextvars.ContextVar["BaseTrace | None"] = contextvars.ContextVar("current_trace", default=None)
+
+
+class Scope:
+    """
+    Manages the currently active span and trace within a context.
+
+    This class provides methods to get, set, and reset the current `BaseSpan`
+    and `BaseTrace` using `contextvars`. This allows for context-local
+    storage of the active span and trace.
+
+    Both traces and spans are managed in a way that allows for nesting.
+    While traces are not typically expected to be nested, this class handles
+    such scenarios gracefully by managing them with context variables, similar
+    to how spans are managed.
+    """
+
+    @classmethod
+    def get_current_span(cls) -> Optional["BaseSpan"]:
+        """
+        Retrieves the currently active span from the context.
+
+        Returns:
+            Optional["BaseSpan"]: The currently active span, or None if no
+                                 span is active in the current context.
+        """
+        return _current_span.get()
+
+    @classmethod
+    def set_current_span(cls, span: Optional["BaseSpan"]) -> contextvars.Token[Optional["BaseSpan"]]:
+        """
+        Sets the currently active span in the context.
+
+        Args:
+            span: The span to set as the current active span. Can be None
+                  to indicate no active span.
+
+        Returns:
+            contextvars.Token[Optional["BaseSpan"]]: A token that can be used
+                                                    to reset the context variable
+                                                    to its previous state.
+        """
+        return _current_span.set(span)
+
+    @classmethod
+    def reset_current_span(cls, token: contextvars.Token[Optional["BaseSpan"]]) -> None:
+        """
+        Resets the current span in the context to its previous state.
+
+        Args:
+            token: The token returned by a previous call to `set_current_span`.
+                   This token is used to restore the context variable to the
+                   value it had before the `set` call that D_GENERATED the token.
+        """
+        _current_span.reset(token)
+
+    @classmethod
+    def get_current_trace(cls) -> Optional["BaseTrace"]:
+        """
+        Retrieves the currently active trace from the context.
+
+        Returns:
+            Optional["BaseTrace"]: The currently active trace, or None if no
+                                  trace is active in the current context.
+        """
+        return _current_trace.get()
+
+    @classmethod
+    def set_current_trace(cls, trace: Optional["BaseTrace"]) -> contextvars.Token[Optional["BaseTrace"]]:
+        """
+        Sets the currently active trace in the context.
+
+        Args:
+            trace: The trace to set as the current active trace. Can be None
+                   to indicate no active trace.
+
+        Returns:
+            contextvars.Token[Optional["BaseTrace"]]: A token that can be used
+                                                     to reset the context variable
+                                                     to its previous state.
+        """
+        return _current_trace.set(trace)
+
+    @classmethod
+    def reset_current_trace(cls, token: contextvars.Token[Optional["BaseTrace"]]) -> None:
+        """
+        Resets the current trace in the context to its previous state.
+
+        Args:
+            token: The token returned by a previous call to `set_current_trace`.
+                   This token is used to restore the context variable to the
+                   value it had before the `set` call that D_GENERATED the token.
+        """
+        _current_trace.reset(token)

scale_gp_beta/lib/tracing/span.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Type, Optional
+from typing_extensions import override
+
+from scale_gp_beta.types.span_upsert_batch_params import Item as SpanCreateRequest
+
+from .util import iso_timestamp, generate_span_id
+from .scope import Scope
+from .types import SpanTypeLiterals, SpanStatusLiterals
+from .exceptions import ParamsCreationError
+
+if TYPE_CHECKING:
+    import contextvars
+    from types import TracebackType
+
+    from .trace_queue_manager import TraceQueueManager
+
+log: logging.Logger = logging.getLogger(__name__)
+
+
+class BaseSpan:
+    """Base class for all span types, providing common attributes and context management.
+
+    A span represents a single unit of work or operation within a trace. This base
+    class defines the core interface and properties for spans, such as name, IDs,
+    timestamps, and methods for starting and ending the span's lifecycle.
+
+    It is intended to be subclassed by concrete span implementations like `Span`
+    (for active tracing) and `NoOpSpan` (for when tracing is disabled).
+
+    Attributes:
+        name (str): The human-readable name of the span.
+        trace_id (str): The ID of the trace this span belongs to.
+        span_id (str): The unique ID of this span.
+        parent_span_id (Optional[str]): The ID of the parent span, if this is a child span.
+        start_time (Optional[str]): ISO 8601 timestamp of when the span started.
+                                     Set by the `start()` method.
+        end_time (Optional[str]): ISO 8601 timestamp of when the span ended.
+                                   Set by the `end()` method.
+        input (Optional[dict[str, Any]]): Input data or parameters for the span's operation.
+        output (Optional[dict[str, Any]]): Output data or results from the span's operation.
+        metadata (Optional[dict[str, Any]]): Additional arbitrary key-value metadata for the span.
+        _contextvar_token (Optional[contextvars.Token]): Token for managing the span's presence
+                                                        in the current execution context.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        trace_id: Optional[str] = None,
+        queue_manager: Optional[TraceQueueManager] = None,
+        span_id: Optional[str] = None,
+        parent_span_id: Optional[str] = None,
+        input: Optional[dict[str, Any]] = None,
+        output: Optional[dict[str, Any]] = None,
+        metadata: Optional[dict[str, Any]] = None,
+        span_type: SpanTypeLiterals = "STANDALONE"
+    ):
+        self.name = name
+        self.trace_id = trace_id or "no_trace_id"
+        self.span_id: str = span_id or generate_span_id()
+        self.parent_span_id = parent_span_id
+        self.start_time: Optional[str] = None
+        self.end_time: Optional[str] = None
+        self.input = input
+        self.output = output
+        self.metadata = metadata
+        self.span_type: SpanTypeLiterals = span_type
+        self.status: SpanStatusLiterals = "SUCCESS"
+        self._queue_manager = queue_manager
+
+        self._contextvar_token: Optional[contextvars.Token[Optional[BaseSpan]]] = None
+
+    def start(self) -> None:
+        pass
+
+    def end(self) -> None:
+        pass
+
+    def __enter__(self) -> BaseSpan:
+        self.start()
+        return self
+
+    def __exit__(
+            self,
+            exc_type: Optional[Type[BaseException]],
+            exc_val: Optional[BaseException],
+            exc_tb: Optional[TracebackType]
+    ) -> None:
+        # Naively record details in metadata for now, note that error capture only supported in context manager
+        # TODO: support error observations when using direct span.start() and span.end()
+        if exc_type is not None:
+            if self.metadata is None:
+                self.metadata = {}
+            self.metadata["error"] = True
+            self.metadata["error.type"] = exc_type.__name__
+            self.metadata["error.message"] = str(exc_val)
+            self.status = "ERROR"
+        self.end()
+
+    def to_request_params(self) -> SpanCreateRequest:
+        if self.start_time is None:
+            raise ParamsCreationError("No start time specified")
+
+        request_data = SpanCreateRequest(
+            name=self.name,
+            id=self.span_id,
+            trace_id=self.trace_id,
+            start_timestamp=self.start_time,
+            input=self.input or {},
+            output=self.output or {},
+            metadata=self.metadata or {},
+            status=self.status,
+            type=self.span_type
+        )
+
+        if self.end_time is not None:
+            request_data["end_timestamp"] = self.end_time
+
+        # parent_span_id is optional (root spans)
+        if self.parent_span_id is not None:
+            request_data["parent_id"] = self.parent_span_id
+
+        return request_data
+
+
+class NoOpSpan(BaseSpan):
+    @override
+    def start(self) -> None:
+        if self.start_time is not None:
+            log.warning(f"Span {self.name}: {self.span_id} has already started at {self.start_time}")
+            return
+
+        self.start_time = iso_timestamp()
+        self._contextvar_token = Scope.set_current_span(self)
+
+    @override
+    def end(self) -> None:
+        if self.end_time is not None:
+            log.warning(f"Span {self.name}: {self.span_id} has already ended at {self.end_time}")
+            return
+
+        if self._contextvar_token is None:
+            log.warning(f"Span {self.name}: {self.span_id} has not started yet.")
+            return
+
+        self.end_time = iso_timestamp()
+        Scope.reset_current_span(self._contextvar_token)
+        self._contextvar_token = None
+
+
+class Span(BaseSpan):
+    """An operational span implementation that records and reports tracing data.
+
+    `Span` instances represent actual units of work that are part of an active trace.
+    They record timestamps, manage their context in `Scope`, and interact with the
+    `TraceQueueManager` to report their start and end events for later export.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        trace_id: str,
+        queue_manager: TraceQueueManager,
+        span_id: Optional[str] = None,
+        parent_span_id: Optional[str] = None,
+        input: Optional[dict[str, Any]] = None,
+        output: Optional[dict[str, Any]] = None,
+        metadata: Optional[dict[str, Any]] = None,
+        span_type: SpanTypeLiterals = "STANDALONE",
+    ):
+        super().__init__(name, trace_id, queue_manager, span_id, parent_span_id, input, output, metadata, span_type)
+        self._queue_manager: TraceQueueManager = queue_manager
+        self.trace_id: str = trace_id
+
+    @override
+    def start(self) -> None:
+        """Starts the operational Span.
+
+        Sets the `start_time`, reports the span start to the `TraceQueueManager`
+        , and registers this span as the current span.
+        """
+        if self.start_time is not None:
+            log.warning(f"Span {self.name}: {self.span_id} has already started at {self.start_time}")
+            return
+
+        self.start_time = iso_timestamp()
+        self._queue_manager.report_span_start(self)
+        self._contextvar_token = Scope.set_current_span(self)
+
+    @override
+    def end(self) -> None:
+        """Ends the operational Span.
+
+        Sets the `end_time`, reports the span end (with its complete data) to the
+        `TraceQueueManager` for queuing and export, and resets this span from the
+        `Scope`.
+        """
+        if self.end_time is not None:
+            log.warning(f"Span {self.name}: {self.span_id} has already ended at {self.end_time}")
+            return
+        if self._contextvar_token is None:
+            log.warning(
+                (
+                    f"Span {self.name}: {self.span_id} attempting to end without a valid context token. "
+                    "Was start() called and completed successfully?"
+                )
+            )
+            return
+
+        self.end_time = iso_timestamp()
+        self._queue_manager.report_span_end(self)
+        Scope.reset_current_span(self._contextvar_token)
+        self._contextvar_token = None

scale_gp_beta/lib/tracing/trace.py

@@ -0,0 +1,118 @@
+import logging
+import contextvars
+from types import TracebackType
+from typing import Dict, Type, Optional
+from typing_extensions import override
+
+from .span import Span, NoOpSpan
+from .util import generate_trace_id
+from .scope import Scope
+from .trace_queue_manager import TraceQueueManager
+
+log: logging.Logger = logging.getLogger(__name__)
+
+
+class BaseTrace:
+    def __init__(self, queue_manager: Optional[TraceQueueManager], trace_id: Optional[str] = None):
+        self.trace_id = trace_id or generate_trace_id()
+        self.queue_manager = queue_manager
+
+        self._in_progress = False
+        self._contextvar_token: Optional[contextvars.Token[Optional[BaseTrace]]] = None
+
+    def start(self) -> None:
+        pass
+
+    def end(self) -> None:
+        pass
+
+    def __enter__(self) -> "BaseTrace":
+        self.start()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]] = None,
+        exc_value: Optional[BaseException] = None,
+        traceback: Optional[TracebackType] = None,
+    ) -> None:
+        self.end()
+
+
+class NoOpTrace(BaseTrace):
+    def __init__(
+            self,
+            name: str,
+            queue_manager: Optional[TraceQueueManager] = None,
+            trace_id: Optional[str] = None,
+            root_span_id: Optional[str] = None,
+            metadata: Optional[Dict[str, Optional[str]]] = None,
+    ):
+        super().__init__(queue_manager, trace_id)
+
+        self.root_span = NoOpSpan(
+            name=name,
+            span_id=root_span_id,
+            trace_id=self.trace_id,
+            queue_manager=queue_manager,
+            metadata=metadata,
+            span_type="TRACER"
+        )
+
+    @override
+    def start(self) -> None:
+        self.root_span.start()
+
+    @override
+    def end(self) -> None:
+        self.root_span.end()
+
+
+class Trace(BaseTrace):
+    def __init__(
+            self,
+            name: str,
+            queue_manager: TraceQueueManager,
+            trace_id: Optional[str] = None,
+            root_span_id: Optional[str] = None,
+            metadata: Optional[Dict[str, Optional[str]]] = None,
+    ):
+        super().__init__(queue_manager, trace_id)
+        self.queue_manager: TraceQueueManager = queue_manager
+
+        self.root_span = Span(
+            name=name,
+            span_id=root_span_id,
+            trace_id=self.trace_id,
+            queue_manager=queue_manager,
+            metadata=metadata,
+            span_type="TRACER"
+        )
+
+    @override
+    def start(self) -> None:
+        if self._in_progress:
+            log.warning(f"Trace already started: {self.trace_id}")
+            return
+
+        self._in_progress = True
+        self.queue_manager.report_trace_start(self)  # no-op
+        self._contextvar_token = Scope.set_current_trace(self)
+
+        self.root_span.start()
+
+    @override
+    def end(self) -> None:
+        if not self._in_progress:
+            log.warning(f"Ending trace which is not active: {self.trace_id}")
+            return
+        if self._contextvar_token is None:
+            log.warning(f"Ending trace which is not active: {self.trace_id}, contextvar_token not set")
+            return
+
+        self._in_progress = False
+        self.queue_manager.report_trace_end(self)  # no-op
+        Scope.reset_current_trace(self._contextvar_token)
+        self._contextvar_token = None
+
+        self.root_span.end()

scale_gp_beta/lib/tracing/trace_exporter.py

@@ -0,0 +1,87 @@
+import time
+import logging
+from queue import Empty, Queue
+from typing import TYPE_CHECKING, List
+from typing_extensions import TypeAlias
+
+from scale_gp_beta import SGPClient
+from scale_gp_beta._exceptions import APIError
+from scale_gp_beta.types.span_upsert_batch_params import Item as SpanCreateRequest
+
+from .exceptions import ParamsCreationError
+
+if TYPE_CHECKING:
+    from .tracing import Span
+
+INITIAL_BACKOFF = 0.4
+MAX_BACKOFF = 20
+
+log: logging.Logger = logging.getLogger(__name__)
+
+SpanRequestBatch: TypeAlias = List[SpanCreateRequest]
+
+
+class TraceExporter:
+    def __init__(
+        self,
+        max_batch_size: int,
+        max_retries: int,
+        backoff: float = INITIAL_BACKOFF,
+        max_backoff: float = MAX_BACKOFF,
+    ):
+        self.max_batch_size = max_batch_size
+        self.max_retries = max_retries
+        self.backoff = backoff
+        self.max_backoff = max_backoff
+
+    def export(self, client: SGPClient, queue: "Queue[Span]") -> None:
+        # this is also thread safe, two threads can call this method with the same queue at once
+        # the ordering of the requests might be randomly split between the two threads, but they should all be picked up
+        batches: list[SpanRequestBatch] = self._create_batches(queue)
+
+        log.info(f"Exporting {len(batches)} span batches")
+
+        for batch in batches:
+            self._export_batch(batch, client)
+
+    def _export_batch(self, batch: SpanRequestBatch, client: SGPClient) -> None:
+        attempts_remaining = self.max_retries
+        backoff_delay = self.backoff
+        while attempts_remaining > 0:
+            attempts_remaining -= 1
+            try:
+                client.spans.upsert_batch(items=batch)
+                return
+            except APIError as e:
+                log.warning(
+                    f"API error occurred while exporting batch: {e.message}, attempts remaining: {attempts_remaining}"
+                )
+                if attempts_remaining == 0:
+                    continue
+                time.sleep(backoff_delay)
+                backoff_delay = min(backoff_delay * 2, self.max_backoff)
+
+        log.error(f"Failed to export span batch after {self.max_retries} attempts, dropping...")
+
+    def _create_batches(self, queue: "Queue[Span]") -> "list[SpanRequestBatch]":
+        """Drain the queue and return a list of batches"""
+        batches: list[SpanRequestBatch] = []
+
+        while True:
+            span_batch: SpanRequestBatch = []
+
+            while len(span_batch) < self.max_batch_size and queue.qsize() > 0:
+                try:
+                    span: "Span" = queue.get_nowait()
+                    span_request = span.to_request_params()
+                    span_batch.append(span_request)
+                except Empty:
+                    break
+                except ParamsCreationError as e:
+                    log.warning(f"ParamsCreationError: {e}\ndropping...")
+
+            if not span_batch:
+                break
+            batches.append(span_batch)
+
+        return batches

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

{scale_gp_beta-0.1.0a17.dist-info → scale_gp_beta-0.1.0a19.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: scale-gp-beta
-Version: 0.1.0a17
+Version: 0.1.0a19
 Summary: The official Python library for the Scale GP API
 Project-URL: Homepage, https://github.com/scaleapi/sgp-python-beta
 Project-URL: Repository, https://github.com/scaleapi/sgp-python-beta

{scale_gp_beta-0.1.0a17.dist-info → scale_gp_beta-0.1.0a19.dist-info}/RECORD

@@ -11,7 +11,7 @@ scale_gp_beta/_resource.py,sha256=siZly_U6D0AOVLAzaOsqUdEFFzVMbWRj-ml30nvRp7E,11
 scale_gp_beta/_response.py,sha256=GemuybPk0uemovTlGHyHkj-ScYTTDJA0jqH5FQqIPwQ,28852
 scale_gp_beta/_streaming.py,sha256=fcCSGXslmi2SmmkM05g2SACXHk2Mj7k1X5uMBu6U5s8,10112
 scale_gp_beta/_types.py,sha256=0wSs40TefKMPBj2wQKenEeZ0lzedoHClNJeqrpAgkII,6204
-scale_gp_beta/_version.py,sha256=8X-aROY-qc-hXo3LtgtE__YX1vTzDaeOZIxZVj1_jfE,174
+scale_gp_beta/_version.py,sha256=QpUK6D2Mrqt-7NO3RPbtzHQemQKxo2IcY3na8iBT4mE,174
 scale_gp_beta/pagination.py,sha256=t-U86PYxl20VRsz8VXOMJJDe7HxkX7ISFMvRNbBNy9s,4054
 scale_gp_beta/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 scale_gp_beta/_utils/__init__.py,sha256=PNZ_QJuzZEgyYXqkO1HVhGkj5IU9bglVUcw7H-Knjzw,2062
@@ -25,6 +25,16 @@ scale_gp_beta/_utils/_transform.py,sha256=n7kskEWz6o__aoNvhFoGVyDoalNe6mJwp-g7BW
 scale_gp_beta/_utils/_typing.py,sha256=D0DbbNu8GnYQTSICnTSHDGsYXj8TcAKyhejb0XcnjtY,4602
 scale_gp_beta/_utils/_utils.py,sha256=ts4CiiuNpFiGB6YMdkQRh2SZvYvsl7mAF-JWHCcLDf4,12312
 scale_gp_beta/lib/.keep,sha256=wuNrz-5SXo3jJaJOJgz4vFHM41YH_g20F5cRQo0vLes,224
+scale_gp_beta/lib/tracing/__init__.py,sha256=UgyExbqAA2ljDEF4X4YFhtbBZuoQJ2IF4hkGs_xQEc0,226
+scale_gp_beta/lib/tracing/exceptions.py,sha256=vL2_GAfWEy8EfLhrBkDClLYTasOLnL-5zUpdCQnSzcs,107
+scale_gp_beta/lib/tracing/scope.py,sha256=kHrd0his8L2K_KXn2E6J9d565PliEdFoKRQ1d5ALTyk,3901
+scale_gp_beta/lib/tracing/span.py,sha256=qWeqiIuEDtCAPApMq5q_UWe9of204EStguja7j5eaXc,8143
+scale_gp_beta/lib/tracing/trace.py,sha256=jDZeT5-cE70FcsjC18CEGBdbDqYX36LqlnLIavK98RI,3382
+scale_gp_beta/lib/tracing/trace_exporter.py,sha256=dr8DfbHhUaJLcywuzfmoD_WQ30mXuKzDUHw5QWNWLfg,3032
+scale_gp_beta/lib/tracing/trace_queue_manager.py,sha256=DCR239JuOTr76Lo5jBPDhi8ZCqxCdH9GklXKuOyzUqE,5880
+scale_gp_beta/lib/tracing/tracing.py,sha256=xht7oAzxbgIE3fgH_botEDxBDKBr_8r5jrtFwtaVmiQ,7193
+scale_gp_beta/lib/tracing/types.py,sha256=CKdOHYIEZGQXWFSrn10iC-cyRw6_Lido9jC8pdSc3P0,1010
+scale_gp_beta/lib/tracing/util.py,sha256=8Oq4wLXRNOzh3CC1zRaBEr0h_WdXLrk536BUNKRddVE,1527
 scale_gp_beta/resources/__init__.py,sha256=Fyo05_2_pc5orfyTSIpxa3btmBTd45VasgibwSqbbKo,4942
 scale_gp_beta/resources/completions.py,sha256=4esj9lGTJAxt6wFvON126DvEGkMIChRZ6uZBOf56Aac,31868
 scale_gp_beta/resources/dataset_items.py,sha256=2d7O5zmqVEafJTxVwgbRz9yq-4T81dPPfFuPDRAaWqU,22510
@@ -107,7 +117,7 @@ scale_gp_beta/types/chat/completion_models_params.py,sha256=ETxafJIUx4tTvkiR-ZCr
 scale_gp_beta/types/chat/completion_models_response.py,sha256=Ctgj6o-QWPSdjBKzG9J4Id0-DjXu4UGGw1NR6-840Ec,403
 scale_gp_beta/types/chat/model_definition.py,sha256=NNgopTm900GD0Zs2YHkcvoW67uKaWUKVyPbhKBHvKdQ,817
 scale_gp_beta/types/files/__init__.py,sha256=OKfJYcKb4NObdiRObqJV_dOyDQ8feXekDUge2o_4pXQ,122
-scale_gp_beta-0.1.0a17.dist-info/METADATA,sha256=ZXduP_l4D_AXtqcwASCXrowWY5Af3_hdJFzmjysq9hk,16881
-scale_gp_beta-0.1.0a17.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
-scale_gp_beta-0.1.0a17.dist-info/licenses/LICENSE,sha256=x49Bj8r_ZpqfzThbmfHyZ_bE88XvHdIMI_ANyLHFFRE,11338
-scale_gp_beta-0.1.0a17.dist-info/RECORD,,
+scale_gp_beta-0.1.0a19.dist-info/METADATA,sha256=pOAEZo76uYkcEFTkgSza0Di3tD0gvC9vZCX3qWn4ySA,16881
+scale_gp_beta-0.1.0a19.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
+scale_gp_beta-0.1.0a19.dist-info/licenses/LICENSE,sha256=x49Bj8r_ZpqfzThbmfHyZ_bE88XvHdIMI_ANyLHFFRE,11338
+scale_gp_beta-0.1.0a19.dist-info/RECORD,,