wandb 0.18.2__py3-none-musllinux_1_2_x86_64.whl

wandb/sdk/data_types/trace_tree.py

@@ -0,0 +1,438 @@
+"""This module contains the `WBTraceTree` media type, and the supporting dataclasses.
+
+A `WBTraceTree` is a media object containing a root span and an
+arbitrary model dump as a serializable dictionary. Logging such media type will
+result in a W&B Trace Debugger panel being created in the workspace UI.
+"""
+
+import dataclasses
+import hashlib
+import json
+import typing
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+import wandb
+from wandb.sdk.data_types import _dtypes
+from wandb.sdk.data_types.base_types.media import Media
+from wandb.sdk.data_types.utils import _json_helper
+
+if TYPE_CHECKING:  # pragma: no cover
+    from wandb.sdk.artifacts.artifact import Artifact
+
+    from ..wandb_run import Run as LocalRun
+
+
+class StatusCode(str, Enum):
+    SUCCESS = "SUCCESS"
+    ERROR = "ERROR"
+
+    def __str__(self) -> str:
+        return str(self.value)
+
+
+class SpanKind(str, Enum):
+    LLM = "LLM"
+    CHAIN = "CHAIN"
+    AGENT = "AGENT"
+    TOOL = "TOOL"
+
+    def __str__(self) -> str:
+        return str(self.value)
+
+
+@dataclass()
+class Result:
+    inputs: Optional[Dict[str, Any]] = field(default=None)
+    outputs: Optional[Dict[str, Any]] = field(default=None)
+
+
+@dataclass()
+class Span:
+    span_id: Optional[str] = field(default=None)
+    name: Optional[str] = field(default=None)
+    start_time_ms: Optional[int] = field(default=None)
+    end_time_ms: Optional[int] = field(default=None)
+    status_code: Optional[StatusCode] = field(default=None)
+    status_message: Optional[str] = field(default=None)
+    attributes: Optional[Dict[str, Any]] = field(default=None)
+    results: Optional[List[Result]] = field(default=None)
+    child_spans: Optional[List["Span"]] = field(default=None)
+    span_kind: Optional[SpanKind] = field(default=None)
+
+    def add_attribute(self, key: str, value: Any) -> None:
+        if self.attributes is None:
+            self.attributes = {}
+        self.attributes[key] = value
+
+    def add_named_result(self, inputs: Dict[str, Any], outputs: Dict[str, Any]) -> None:
+        if self.results is None:
+            self.results = []
+        self.results.append(Result(inputs, outputs))
+
+    def add_child_span(self, span: "Span") -> None:
+        if self.child_spans is None:
+            self.child_spans = []
+        self.child_spans.append(span)
+
+
+class WBTraceTree(Media):
+    """Media object for trace tree data.
+
+    Arguments:
+        root_span (Span): The root span of the trace tree.
+        model_dict (dict, optional): A dictionary containing the model dump.
+            NOTE: model_dict is a completely-user-defined dict. The UI will render
+            a JSON viewer for this dict, giving special treatment to dictionaries
+            with a `_kind` key. This is because model vendors have such different
+            serialization formats that we need to be flexible here.
+    """
+
+    _log_type = "wb_trace_tree"
+
+    def __init__(
+        self,
+        root_span: Span,
+        model_dict: typing.Optional[dict] = None,
+    ):
+        super().__init__()
+        self._root_span = root_span
+        self._model_dict = model_dict
+
+    @classmethod
+    def get_media_subdir(cls) -> str:
+        return "media/wb_trace_tree"
+
+    def to_json(self, run: Optional[Union["LocalRun", "Artifact"]]) -> dict:
+        res = {"_type": self._log_type}
+        # Here we use `dumps` to put things into string format. This is because
+        # the complex data structures create problems for gorilla history to parquet.
+        if self._model_dict is not None:
+            model_dump_str = _safe_serialize(self._model_dict)
+            res["model_hash"] = _hash_id(model_dump_str)
+            res["model_dict_dumps"] = model_dump_str
+        res["root_span_dumps"] = _safe_serialize(dataclasses.asdict(self._root_span))
+        return res
+
+    def is_bound(self) -> bool:
+        return True
+
+
+class _WBTraceTreeFileType(_dtypes.Type):
+    name = "wb_trace_tree"
+    types = [WBTraceTree]
+
+
+_dtypes.TypeRegistry.add(_WBTraceTreeFileType)
+
+
+# generate a deterministic 16 character id based on input string
+def _hash_id(s: str) -> str:
+    return hashlib.md5(s.encode("utf-8")).hexdigest()[:16]
+
+
+def _fallback_serialize(obj: Any) -> str:
+    try:
+        return f"<<non-serializable: {type(obj).__qualname__}>>"
+    except Exception:
+        return "<<non-serializable>>"
+
+
+def _safe_serialize(obj: dict) -> str:
+    try:
+        return json.dumps(
+            _json_helper(obj, None),
+            skipkeys=True,
+            default=_fallback_serialize,
+        )
+    except Exception:
+        return "{}"
+
+
+class TraceAttribute:
+    """Descriptor for accessing and setting attributes of the `Trace` class."""
+
+    def __set_name__(self, owner: type, name: str) -> None:
+        self.name = name
+
+    def __get__(self, instance: "Trace", owner: type) -> Any:
+        return getattr(instance._span, self.name)
+
+    def __set__(self, instance: "Trace", value: Any) -> None:
+        setattr(instance._span, self.name, value)
+
+
+class Trace:
+    """A simplification of WBTraceTree and Span to manage a trace - a collection of spans, their metadata and hierarchy.
+
+    Args:
+        name: (str) The name of the root span.
+        kind: (str, optional) The kind of the root span.
+        status_code: (str, optional) The status of the root span, either "error" or "success".
+        status_message: (str, optional) Any status message associated with the root span.
+        metadata: (dict, optional) Any additional metadata for the root span.
+        start_time_ms: (int, optional) The start time of the root span in milliseconds.
+        end_time_ms: (int, optional) The end time of the root span in milliseconds.
+        inputs: (dict, optional) The named inputs of the root span.
+        outputs: (dict, optional) The named outputs of the root span.
+        model_dict: (dict, optional) A json serializable dictionary containing the model architecture details.
+
+    Example:
+        .. code-block:: python
+        ```
+        trace = Trace(
+            name="My awesome Model",
+            kind="LLM",
+            status_code= "SUCCESS",
+            metadata={"attr_1": 1, "attr_2": 2,},
+            start_time_ms=int(round(time.time() * 1000)),
+            end_time_ms=int(round(time.time() * 1000))+1000,
+            inputs={"user": "How old is google?"},
+            outputs={"assistant": "25 years old"},
+            model_dict={"_kind": "openai", "api_type": "azure"}
+              )
+        run = wandb.init(project=<my_awesome_project>,)
+        trace.log("my_trace")
+        wandb.finish()
+        ```
+    """
+
+    name = TraceAttribute()
+    status_code = TraceAttribute()
+    status_message = TraceAttribute()
+    start_time_ms = TraceAttribute()
+    end_time_ms = TraceAttribute()
+
+    def __init__(
+        self,
+        name: str,
+        kind: Optional[str] = None,
+        status_code: Optional[str] = None,
+        status_message: Optional[str] = None,
+        metadata: Optional[dict] = None,
+        start_time_ms: Optional[int] = None,
+        end_time_ms: Optional[int] = None,
+        inputs: Optional[dict] = None,
+        outputs: Optional[dict] = None,
+        model_dict: Optional[dict] = None,
+    ):
+        self._span = self._assert_and_create_span(
+            name=name,
+            kind=kind,
+            status_code=status_code,
+            status_message=status_message,
+            metadata=metadata,
+            start_time_ms=start_time_ms,
+            end_time_ms=end_time_ms,
+            inputs=inputs,
+            outputs=outputs,
+        )
+        if model_dict is not None:
+            assert isinstance(model_dict, dict), "Model dict must be a dictionary"
+        self._model_dict = model_dict
+
+    def _assert_and_create_span(
+        self,
+        name: str,
+        kind: Optional[str] = None,
+        status_code: Optional[str] = None,
+        status_message: Optional[str] = None,
+        metadata: Optional[dict] = None,
+        start_time_ms: Optional[int] = None,
+        end_time_ms: Optional[int] = None,
+        inputs: Optional[dict] = None,
+        outputs: Optional[dict] = None,
+    ) -> Span:
+        """Utility to assert the validity of the span parameters and create a span object.
+
+        Args:
+            name: The name of the span.
+            kind: The kind of the span.
+            status_code: The status code of the span.
+            status_message: The status message of the span.
+            metadata: Dictionary of metadata to be logged with the span.
+            start_time_ms: Start time of the span in milliseconds.
+            end_time_ms: End time of the span in milliseconds.
+            inputs: Dictionary of inputs to be logged with the span.
+            outputs: Dictionary of outputs to be logged with the span.
+
+        Returns:
+            A Span object.
+        """
+        if kind is not None:
+            assert (
+                kind.upper() in SpanKind.__members__
+            ), "Invalid span kind, can be one of 'LLM', 'AGENT', 'CHAIN', 'TOOL'"
+            kind = SpanKind(kind.upper())
+        if status_code is not None:
+            assert (
+                status_code.upper() in StatusCode.__members__
+            ), "Invalid status code, can be one of 'SUCCESS' or 'ERROR'"
+            status_code = StatusCode(status_code.upper())
+        if inputs is not None:
+            assert isinstance(inputs, dict), "Inputs must be a dictionary"
+        if outputs is not None:
+            assert isinstance(outputs, dict), "Outputs must be a dictionary"
+        if inputs or outputs:
+            result = Result(inputs=inputs, outputs=outputs)
+        else:
+            result = None
+
+        if metadata is not None:
+            assert isinstance(metadata, dict), "Metadata must be a dictionary"
+
+        return Span(
+            name=name,
+            span_kind=kind,
+            status_code=status_code,
+            status_message=status_message,
+            attributes=metadata,
+            start_time_ms=start_time_ms,
+            end_time_ms=end_time_ms,
+            results=[result] if result else None,
+        )
+
+    def add_child(
+        self,
+        child: "Trace",
+    ) -> "Trace":
+        """Utility to add a child span to the current span of the trace.
+
+        Args:
+            child: The child span to be added to the current span of the trace.
+
+        Returns:
+            The current trace object with the child span added to it.
+        """
+        self._span.add_child_span(child._span)
+        if self._model_dict is not None and child._model_dict is not None:
+            self._model_dict.update({child._span.name: child._model_dict})
+        return self
+
+    def add_inputs_and_outputs(self, inputs: dict, outputs: dict) -> "Trace":
+        """Add a result to the span of the current trace.
+
+        Args:
+            inputs: Dictionary of inputs to be logged with the span.
+            outputs: Dictionary of outputs to be logged with the span.
+
+        Returns:
+            The current trace object with the result added to it.
+        """
+        if self._span.results is None:
+            result = Result(inputs=inputs, outputs=outputs)
+            self._span.results = [result]
+        else:
+            result = Result(inputs=inputs, outputs=outputs)
+            self._span.results.append(result)
+        return self
+
+    def add_metadata(self, metadata: dict) -> "Trace":
+        """Add metadata to the span of the current trace."""
+        if self._span.attributes is None:
+            self._span.attributes = metadata
+        else:
+            self._span.attributes.update(metadata)
+        return self
+
+    @property
+    def metadata(self) -> Optional[Dict[str, str]]:
+        """Get the metadata of the trace.
+
+        Returns:
+            Dictionary of metadata.
+        """
+        return self._span.attributes
+
+    @metadata.setter
+    def metadata(self, value: Dict[str, str]) -> None:
+        """Set the metadata of the trace.
+
+        Args:
+            value: Dictionary of metadata to be set.
+        """
+        if self._span.attributes is None:
+            self._span.attributes = value
+        else:
+            self._span.attributes.update(value)
+
+    @property
+    def inputs(self) -> Optional[Dict[str, str]]:
+        """Get the inputs of the trace.
+
+        Returns:
+            Dictionary of inputs.
+        """
+        return self._span.results[-1].inputs if self._span.results else None
+
+    @inputs.setter
+    def inputs(self, value: Dict[str, str]) -> None:
+        """Set the inputs of the trace.
+
+        Args:
+            value: Dictionary of inputs to be set.
+        """
+        if self._span.results is None:
+            result = Result(inputs=value, outputs={})
+            self._span.results = [result]
+        else:
+            result = Result(inputs=value, outputs=self._span.results[-1].outputs)
+            self._span.results.append(result)
+
+    @property
+    def outputs(self) -> Optional[Dict[str, str]]:
+        """Get the outputs of the trace.
+
+        Returns:
+            Dictionary of outputs.
+        """
+        return self._span.results[-1].outputs if self._span.results else None
+
+    @outputs.setter
+    def outputs(self, value: Dict[str, str]) -> None:
+        """Set the outputs of the trace.
+
+        Args:
+            value: Dictionary of outputs to be set.
+        """
+        if self._span.results is None:
+            result = Result(inputs={}, outputs=value)
+            self._span.results = [result]
+        else:
+            result = Result(inputs=self._span.results[-1].inputs, outputs=value)
+            self._span.results.append(result)
+
+    @property
+    def kind(self) -> Optional[str]:
+        """Get the kind of the trace.
+
+        Returns:
+            The kind of the trace.
+        """
+        return self._span.span_kind.value if self._span.span_kind else None
+
+    @kind.setter
+    def kind(self, value: str) -> None:
+        """Set the kind of the trace.
+
+        Args:
+            value: The kind of the trace to be set.
+        """
+        assert (
+            value.upper() in SpanKind.__members__
+        ), "Invalid span kind, can be one of 'LLM', 'AGENT', 'CHAIN', 'TOOL'"
+        self._span.span_kind = SpanKind(value.upper())
+
+    def log(self, name: str) -> None:
+        """Log the trace to a wandb run.
+
+        Args:
+            name: The name of the trace to be logged
+        """
+        trace_tree = WBTraceTree(self._span, self._model_dict)
+        assert (
+            wandb.run is not None
+        ), "You must call wandb.init() before logging a trace"
+        assert len(name.strip()) > 0, "You must provide a valid name to log the trace"
+        wandb.run.log({name: trace_tree})

wandb/sdk/data_types/utils.py

@@ -0,0 +1,229 @@
+import datetime
+import logging
+import os
+import re
+from decimal import Decimal
+from typing import TYPE_CHECKING, Optional, Sequence, Union, cast
+
+import wandb
+from wandb import util
+
+from .base_types.media import BatchableMedia, Media
+from .base_types.wb_value import WBValue
+from .image import _server_accepts_image_filenames
+from .plotly import Plotly
+
+if TYPE_CHECKING:  # pragma: no cover
+    import matplotlib  # type: ignore
+    import pandas as pd
+    import plotly  # type: ignore
+
+    from ..wandb_run import Run as LocalRun
+
+    ValToJsonType = Union[
+        dict,
+        "WBValue",
+        Sequence["WBValue"],
+        "plotly.Figure",
+        "matplotlib.artist.Artist",
+        "pd.DataFrame",
+        object,
+    ]
+
+
+def history_dict_to_json(
+    run: Optional["LocalRun"],
+    payload: dict,
+    step: Optional[int] = None,
+    ignore_copy_err: Optional[bool] = None,
+) -> dict:
+    # Converts a History row dict's elements so they're friendly for JSON serialization.
+
+    if step is None:
+        # We should be at the top level of the History row; assume this key is set.
+        step = payload["_step"]
+
+    # We use list here because we were still seeing cases of RuntimeError dict changed size
+    for key in list(payload):
+        val = payload[key]
+        if isinstance(val, dict):
+            payload[key] = history_dict_to_json(
+                run, val, step=step, ignore_copy_err=ignore_copy_err
+            )
+        else:
+            payload[key] = val_to_json(
+                run, key, val, namespace=step, ignore_copy_err=ignore_copy_err
+            )
+
+    return payload
+
+
+# TODO: refine this
+def val_to_json(
+    run: Optional["LocalRun"],
+    key: str,
+    val: "ValToJsonType",
+    namespace: Optional[Union[str, int]] = None,
+    ignore_copy_err: Optional[bool] = None,
+) -> Union[Sequence, dict]:
+    # Converts a wandb datatype to its JSON representation.
+    if namespace is None:
+        raise ValueError(
+            "val_to_json must be called with a namespace(a step number, or 'summary') argument"
+        )
+
+    converted = val
+
+    if isinstance(val, (int, float, str, bool)):
+        # These are already JSON-serializable,
+        # no need to do the expensive checks below.
+        return converted  # type: ignore[return-value]
+
+    typename = util.get_full_typename(val)
+
+    if util.is_pandas_data_frame(val):
+        val = wandb.Table(dataframe=val)
+
+    elif util.is_matplotlib_typename(typename) or util.is_plotly_typename(typename):
+        val = Plotly.make_plot_media(val)
+    elif isinstance(val, (list, tuple, range)) and all(
+        isinstance(v, WBValue) for v in val
+    ):
+        assert run
+        # This check will break down if Image/Audio/... have child classes.
+        if (
+            len(val)
+            and isinstance(val[0], BatchableMedia)
+            and all(isinstance(v, type(val[0])) for v in val)
+        ):
+            if TYPE_CHECKING:
+                val = cast(Sequence["BatchableMedia"], val)
+
+            items = _prune_max_seq(val)
+
+            if _server_accepts_image_filenames():
+                for item in items:
+                    item.bind_to_run(
+                        run=run,
+                        key=key,
+                        step=namespace,
+                        ignore_copy_err=ignore_copy_err,
+                    )
+            else:
+                for i, item in enumerate(items):
+                    item.bind_to_run(
+                        run=run,
+                        key=key,
+                        step=namespace,
+                        id_=i,
+                        ignore_copy_err=ignore_copy_err,
+                    )
+                if run._attach_id and run._init_pid != os.getpid():
+                    wandb.termwarn(
+                        f"Attempting to log a sequence of {items[0].__class__.__name__} objects from multiple processes might result in data loss. Please upgrade your wandb server",
+                        repeat=False,
+                    )
+
+            return items[0].seq_to_json(items, run, key, namespace)
+        else:
+            # TODO(adrian): Good idea to pass on the same key here? Maybe include
+            # the array index?
+            # There is a bug here: if this array contains two arrays of the same type of
+            # anonymous media objects, their eventual names will collide.
+            # This used to happen. The frontend doesn't handle heterogeneous arrays
+            # raise ValueError(
+            #    "Mixed media types in the same list aren't supported")
+            return [
+                val_to_json(
+                    run, key, v, namespace=namespace, ignore_copy_err=ignore_copy_err
+                )
+                for v in val
+            ]
+
+    if isinstance(val, WBValue):
+        assert run
+        if isinstance(val, Media) and not val.is_bound():
+            if hasattr(val, "_log_type") and val._log_type in [
+                "table",
+                "partitioned-table",
+                "joined-table",
+            ]:
+                # Special conditional to log tables as artifact entries as well.
+                # I suspect we will generalize this as we transition to storing all
+                # files in an artifact
+                # we sanitize the key to meet the constraints
+                # in this case, leaving only alphanumerics or underscores.
+                sanitized_key = re.sub(r"[^a-zA-Z0-9_]+", "", key)
+                art = wandb.Artifact(f"run-{run.id}-{sanitized_key}", "run_table")
+                art.add(val, key)
+                run.log_artifact(art)
+
+            # Partitioned tables and joined tables do not support being bound to runs.
+            if not (
+                hasattr(val, "_log_type")
+                and val._log_type in ["partitioned-table", "joined-table"]
+            ):
+                val.bind_to_run(run, key, namespace)
+
+        return val.to_json(run)
+
+    return converted  # type: ignore
+
+
+def _prune_max_seq(seq: Sequence["BatchableMedia"]) -> Sequence["BatchableMedia"]:
+    # If media type has a max respect it
+    items = seq
+    if hasattr(seq[0], "MAX_ITEMS") and seq[0].MAX_ITEMS < len(seq):
+        logging.warning(
+            "Only %i %s will be uploaded."
+            % (seq[0].MAX_ITEMS, seq[0].__class__.__name__)
+        )
+        items = seq[: seq[0].MAX_ITEMS]
+    return items
+
+
+def _json_helper(val, artifact):
+    if isinstance(val, WBValue):
+        return val.to_json(artifact)
+    elif val.__class__ is dict:
+        res = {}
+        for key in val:
+            res[key] = _json_helper(val[key], artifact)
+        return res
+
+    if hasattr(val, "tolist"):
+        py_val = val.tolist()
+        if val.__class__.__name__ == "datetime64" and isinstance(py_val, int):
+            # when numpy datetime64 .tolist() returns an int, it is nanoseconds.
+            # need to convert to milliseconds
+            return _json_helper(py_val / int(1e6), artifact)
+        return _json_helper(py_val, artifact)
+    elif hasattr(val, "item"):
+        return _json_helper(val.item(), artifact)
+
+    if isinstance(val, datetime.datetime):
+        if val.tzinfo is None:
+            val = datetime.datetime(
+                val.year,
+                val.month,
+                val.day,
+                val.hour,
+                val.minute,
+                val.second,
+                val.microsecond,
+                tzinfo=datetime.timezone.utc,
+            )
+        return int(val.timestamp() * 1000)
+    elif isinstance(val, datetime.date):
+        return int(
+            datetime.datetime(
+                val.year, val.month, val.day, tzinfo=datetime.timezone.utc
+            ).timestamp()
+            * 1000
+        )
+    elif isinstance(val, (list, tuple)):
+        return [_json_helper(i, artifact) for i in val]
+    elif isinstance(val, Decimal):
+        return float(val)
+    else:
+        return util.json_friendly(val)[0]